-

2026-09-11 06:14:38 +02:00
parent bf64652300
commit e581949946
161 changed files with 31126 additions and 1 deletions
@@ -0,0 +1,180 @@
+package main
+
+import (
+	"context"
+	"encoding/json"
+	"flag"
+	"fmt"
+	"log"
+	"net"
+	"net/http"
+	"os"
+	"os/signal"
+	"strings"
+	"syscall"
+	"time"
+
+	"github.com/example/ollama-fair-gateway/internal/hoststats"
+)
+
+type telemetry struct {
+	MemoryUsedBytes   int64     `json:"memory_used_bytes,omitempty"`
+	MemoryTotalBytes  int64     `json:"memory_total_bytes,omitempty"`
+	VRAMUsedBytes     int64     `json:"vram_used_bytes,omitempty"`
+	VRAMTotalBytes    int64     `json:"vram_total_bytes,omitempty"`
+	GPUUtilizationPct float64   `json:"gpu_utilization_percent,omitempty"`
+	GPUTemperatureC   float64   `json:"gpu_temperature_c,omitempty"`
+	GPUPowerWatts     float64   `json:"gpu_power_watts,omitempty"`
+	Source            string    `json:"source,omitempty"`
+	UpdatedAt         time.Time `json:"updated_at,omitempty"`
+	Error             string    `json:"error,omitempty"`
+}
+
+type collector struct {
+	nvidia    bool
+	nvidiaGPU string
+	amd       bool
+	amdDevice string
+}
+
+func (c collector) collect(ctx context.Context) telemetry {
+	out := telemetry{UpdatedAt: time.Now().UTC()}
+	var errs []string
+	mctx, cancel := context.WithTimeout(ctx, 1500*time.Millisecond)
+	m, err := hoststats.ReadMemory(mctx)
+	cancel()
+	if err != nil {
+		errs = append(errs, "memory: "+err.Error())
+	} else {
+		out.MemoryTotalBytes, out.MemoryUsedBytes = m.TotalBytes, m.UsedBytes
+		out.Source = appendSource(out.Source, "host-memory")
+	}
+	if c.nvidia {
+		gctx, cancel := context.WithTimeout(ctx, 1500*time.Millisecond)
+		g, err := hoststats.ReadNVIDIA(gctx, c.nvidiaGPU)
+		cancel()
+		if err != nil {
+			errs = append(errs, "nvidia: "+err.Error())
+		} else {
+			out.VRAMUsedBytes, out.VRAMTotalBytes = g.MemoryUsedBytes, g.MemoryTotalBytes
+			out.GPUUtilizationPct, out.GPUTemperatureC, out.GPUPowerWatts = g.UtilizationPercent, g.TemperatureC, g.PowerWatts
+			out.Source = appendSource(out.Source, "nvidia-smi")
+		}
+	}
+	if c.amd {
+		gctx, cancel := context.WithTimeout(ctx, 1500*time.Millisecond)
+		g, err := hoststats.ReadAMD(gctx, c.amdDevice)
+		cancel()
+		if err != nil {
+			errs = append(errs, "amd: "+err.Error())
+		} else {
+			out.VRAMUsedBytes, out.VRAMTotalBytes = g.MemoryUsedBytes, g.MemoryTotalBytes
+			out.GPUUtilizationPct, out.GPUTemperatureC, out.GPUPowerWatts = g.UtilizationPercent, g.TemperatureC, g.PowerWatts
+			out.Source = appendSource(out.Source, "amdgpu-sysfs")
+		}
+	}
+	out.Error = strings.Join(errs, "; ")
+	return out
+}
+
+func appendSource(cur, next string) string {
+	if cur == "" {
+		return next
+	}
+	return cur + "+" + next
+}
+
+type cidrAllowlist struct{ nets []*net.IPNet }
+
+func parseCIDRs(raw string) (cidrAllowlist, error) {
+	var out cidrAllowlist
+	for _, part := range strings.Split(raw, ",") {
+		part = strings.TrimSpace(part)
+		if part == "" {
+			continue
+		}
+		_, n, err := net.ParseCIDR(part)
+		if err != nil {
+			return out, fmt.Errorf("invalid allow CIDR %q: %w", part, err)
+		}
+		out.nets = append(out.nets, n)
+	}
+	if len(out.nets) == 0 {
+		return out, fmt.Errorf("at least one allow CIDR is required")
+	}
+	return out, nil
+}
+
+func (a cidrAllowlist) allowed(remote string) bool {
+	host, _, err := net.SplitHostPort(remote)
+	if err != nil {
+		host = remote
+	}
+	ip := net.ParseIP(strings.Trim(host, "[]"))
+	if ip == nil {
+		return false
+	}
+	for _, n := range a.nets {
+		if n.Contains(ip) {
+			return true
+		}
+	}
+	return false
+}
+
+func main() {
+	listen := flag.String("listen", "127.0.0.1:11500", "listen address")
+	path := flag.String("path", "/telemetry", "telemetry path")
+	allow := flag.String("allow-cidrs", "127.0.0.1/32,::1/128", "comma-separated client CIDRs allowed to read telemetry")
+	nvidia := flag.Bool("nvidia-smi", false, "collect NVIDIA telemetry with nvidia-smi")
+	nvidiaGPU := flag.String("nvidia-gpu", "", "optional nvidia-smi GPU selector")
+	amd := flag.Bool("amd-sysfs", false, "collect Linux AMDGPU telemetry from sysfs")
+	amdDevice := flag.String("amd-device", "", "optional AMDGPU device path such as /sys/class/drm/card0/device; empty auto-detects")
+	once := flag.Bool("once", false, "print one telemetry sample as JSON and exit")
+	flag.Parse()
+	if !strings.HasPrefix(*path, "/") {
+		log.Fatal("-path must begin with /")
+	}
+	acl, err := parseCIDRs(*allow)
+	if err != nil {
+		log.Fatal(err)
+	}
+	c := collector{nvidia: *nvidia, nvidiaGPU: *nvidiaGPU, amd: *amd, amdDevice: *amdDevice}
+	if *once {
+		enc := json.NewEncoder(os.Stdout)
+		enc.SetIndent("", "  ")
+		if err := enc.Encode(c.collect(context.Background())); err != nil {
+			log.Fatal(err)
+		}
+		return
+	}
+	mux := http.NewServeMux()
+	mux.HandleFunc("/healthz", func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusNoContent) })
+	mux.HandleFunc(*path, func(w http.ResponseWriter, r *http.Request) {
+		if r.Method != http.MethodGet {
+			w.WriteHeader(http.StatusMethodNotAllowed)
+			return
+		}
+		if !acl.allowed(r.RemoteAddr) {
+			http.Error(w, "forbidden", http.StatusForbidden)
+			return
+		}
+		w.Header().Set("Content-Type", "application/json")
+		w.Header().Set("Cache-Control", "no-store")
+		w.Header().Set("X-Content-Type-Options", "nosniff")
+		_ = json.NewEncoder(w).Encode(c.collect(r.Context()))
+	})
+	srv := &http.Server{Addr: *listen, Handler: mux, ReadHeaderTimeout: 5 * time.Second, IdleTimeout: time.Minute}
+	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
+	defer stop()
+	go func() {
+		<-ctx.Done()
+		shutdown, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+		defer cancel()
+		_ = srv.Shutdown(shutdown)
+	}()
+	log.Printf("worker telemetry listening on http://%s%s allowed=%s nvidia=%t amd=%t", *listen, *path, *allow, *nvidia, *amd)
+	if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
+		log.Fatal(err)
+	}
+}
@@ -0,0 +1,19 @@
+package main
+
+import "testing"
+
+func TestCIDRAllowlist(t *testing.T) {
+	a, err := parseCIDRs("127.0.0.1/32,10.2.19.0/24")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !a.allowed("10.2.19.42:1234") || !a.allowed("127.0.0.1:1") || a.allowed("10.2.18.1:5") {
+		t.Fatal("unexpected allowlist result")
+	}
+}
+
+func TestAppendSource(t *testing.T) {
+	if got := appendSource("host-memory", "amdgpu-sysfs"); got != "host-memory+amdgpu-sysfs" {
+		t.Fatal(got)
+	}
+}