diff --git a/internal/agent/agent.go b/internal/agent/agent.go index 7e22889..7565252 100644 --- a/internal/agent/agent.go +++ b/internal/agent/agent.go @@ -393,8 +393,11 @@ func (a *Agent) settleAdopted(sh *supervisor.Shard, o supervisor.Orphan) { _ = a.sup.Stop(ctx, o.ShardID) return } - // What it computes on, for one started before the agent kept that. - if a.sup.ComputeOf(o.ShardID) == nil { + // What it computes on, for one started before the agent kept that — or + // kept as the CPU for a GPU placement by an agent that did not follow + // an engine into a process group of its own, which the check above has + // just found on its GPUs. + if c := a.sup.ComputeOf(o.ShardID); c == nil || (len(o.Spec.GPUIndices) > 0 && c.Backend == "CPU") { if c, ok := a.gpuCompute(ctx, o.Spec.GPUIndices, sh.Status().PID); ok { a.sup.SetCompute(o.ShardID, c) } @@ -1511,18 +1514,25 @@ func ResolveManagement(c Config) (addr, iface string) { // shard computes on as soon as that is known: an adopted engine's is found // after registration, which carried what was known then. func (a *Agent) reportCompute(ctx context.Context, out chan<- *llmgwv1.AgentMessage) { - sent := map[string]bool{} + // What was sent for each shard, so a compute that changes is sent again: + // an adopted shard's kept one is sent at once, and adopting it may then + // find it was kept wrong. + sent := map[string]string{} t := time.NewTicker(5 * time.Second) defer t.Stop() for range 60 { for _, st := range a.sup.Statuses() { - if st.State.Terminal() || sent[st.ID] { + if st.State.Terminal() { continue } c := a.sup.ComputeOf(st.ID) if c == nil { continue } + said := fmt.Sprint(*c) + if sent[st.ID] == said { + continue + } sh, err := a.sup.Get(st.ID) if err != nil { continue @@ -1533,7 +1543,7 @@ func (a *Agent) reportCompute(ctx context.Context, out chan<- *llmgwv1.AgentMess }}} select { case out <- msg: - sent[st.ID] = true + sent[st.ID] = said case <-ctx.Done(): return } diff --git a/internal/agent/engines/strata.go b/internal/agent/engines/strata.go index ff98817..0a4c6ad 100644 --- a/internal/agent/engines/strata.go +++ b/internal/agent/engines/strata.go @@ -38,8 +38,8 @@ func (*Strata) Name() string { return "strata" } func (*Strata) HealthPath() string { return "/health" } // MetricsPath is Strata's Prometheus text, in vLLM's names; its /metrics is -// JSON for its own monitor. -func (*Strata) MetricsPath() string { return "/metrics/prometheus" } +// JSON for its own monitor unless asked for this format. +func (*Strata) MetricsPath() string { return "/metrics?format=prometheus" } // ServesHTTP is true: a Strata replica is one process on one node. func (*Strata) ServesHTTP(ShardSpec) bool { return true } diff --git a/internal/agent/gpucheck.go b/internal/agent/gpucheck.go index d2588c6..0692d4e 100644 --- a/internal/agent/gpucheck.go +++ b/internal/agent/gpucheck.go @@ -20,6 +20,7 @@ var ( gpuProcesses = discovery.GPUProcesses gpuMemory = discovery.GPUMemoryByProcess processGroup = pgidOf + parentOf = ppidOf // gpuCheckWait is how long a shard has, once it answers, to show up on // its GPUs: the weights are on them by then, but nvidia-smi can lag. gpuCheckWait = 15 * time.Second @@ -116,19 +117,44 @@ func (a *Agent) foreignGPUMemory(ctx context.Context, shards []int) (foreign, us return foreign, used } -// holds reports whether any of pids is the shard's process or in its group: -// an engine that computes in a child process is still the shard. +// holds reports whether any of pids is the shard's process, in its group, +// or a process it started: an engine that computes in a child process is +// still the shard, and Strata's server starts its engine in a group of the +// engine's own. func holds(pids []int, pid, group int) bool { for _, p := range pids { - if p == pid || (group > 0 && processGroup(p) == group) { + if p == pid || (group > 0 && processGroup(p) == group) || descends(p, pid) { return true } } return false } +// descends reports whether process p was started, at any remove, by pid. +func descends(p, pid int) bool { + // Bounded: a chain this deep is no engine's, and a /proc read racing an + // exit must not loop. + for range 16 { + p = parentOf(p) + switch { + case p == pid: + return true + case p <= 1: + return false + } + } + return false +} + // pgidOf reads a process's group from /proc; zero when it cannot. -func pgidOf(pid int) int { +func pgidOf(pid int) int { return statField(pid, 2) } + +// ppidOf reads a process's parent from /proc; zero when it cannot. +func ppidOf(pid int) int { return statField(pid, 1) } + +// statField is one of the numbers after the command name in /proc's stat +// line: 1 is the parent, 2 the group. +func statField(pid, field int) int { b, err := os.ReadFile("/proc/" + strconv.Itoa(pid) + "/stat") if err != nil { return 0 @@ -141,10 +167,10 @@ func pgidOf(pid int) int { return 0 } f := strings.Fields(s[i+1:]) - if len(f) < 3 { + if len(f) <= field { return 0 } - n, _ := strconv.Atoi(f[2]) + n, _ := strconv.Atoi(f[field]) return n } diff --git a/internal/agent/gpucheck_test.go b/internal/agent/gpucheck_test.go index d2dd37e..45931bc 100644 --- a/internal/agent/gpucheck_test.go +++ b/internal/agent/gpucheck_test.go @@ -123,6 +123,30 @@ func TestForeignGPUMemoryLeavesOutTheAgentsOwnShards(t *testing.T) { } } +func TestAnEngineStartedInAGroupOfItsOwnIsStillTheShards(t *testing.T) { + // cuda01: Strata's server (259709) starts its engine (261072) in a + // group of the engine's own, and the engine holds the card. It read as + // a shard on the CPU, and as 22 GB of somebody else's memory. + oldMem, oldGroup, oldParent := gpuMemory, processGroup, parentOf + t.Cleanup(func() { gpuMemory, processGroup, parentOf = oldMem, oldGroup, oldParent }) + processGroup = func(pid int) int { return pid } + parentOf = func(pid int) int { + return map[int]int{261072: 259709, 259709: 258734, 258734: 1, 999: 1}[pid] + } + gpuMemory = func(context.Context, string) ([]discovery.GPUMemory, error) { + return []discovery.GPUMemory{{Index: 0, UsedMB: 23000, + Procs: []discovery.GPUProcess{{PID: 261072, UsedMB: 22226}, {PID: 999, UsedMB: 700}}}}, nil + } + a := &Agent{log: slog.Default()} + if c, ok := a.gpuCompute(context.Background(), []int{0}, 259709); !ok || c.Backend != "CUDA" || + len(c.Devices) != 1 || c.Devices[0] != "GPU 0: 21.7 GB" { + t.Errorf("the engine's child on GPU 0: %+v %v", c, ok) + } + if foreign, _ := a.foreignGPUMemory(context.Background(), []int{259709}); !slices.Equal(foreign, []int32{774}) { + t.Errorf("foreign %v; want only what the other process and nobody holds, [774]", foreign) + } +} + func TestLlamaCppsOwnWordOnTheCPUIsEnoughToRefuse(t *testing.T) { eng := fallbackEngine{} // Without nvidia-smi to ask, the engine's own output says it. diff --git a/internal/router/scrape/parse.go b/internal/router/scrape/parse.go index 144a001..9335031 100644 --- a/internal/router/scrape/parse.go +++ b/internal/router/scrape/parse.go @@ -51,10 +51,11 @@ func MappingFor(engine string) Mapping { case "sglang": return SGLangMapping() case "strata": - // vLLM's names, under a path of its own: Strata's /metrics is the - // JSON its monitor reads. + // vLLM's names, KV cache usage included, when its /metrics is asked + // for Prometheus' text; plain, it is the JSON its monitor reads, and + // /metrics/prometheus is a shorter set in llama.cpp's names. m := VLLMMapping() - m.Path = "/metrics/prometheus" + m.Path = "/metrics?format=prometheus" return m default: return VLLMMapping() diff --git a/internal/router/scrape/strata_test.go b/internal/router/scrape/strata_test.go new file mode 100644 index 0000000..7583e49 --- /dev/null +++ b/internal/router/scrape/strata_test.go @@ -0,0 +1,38 @@ +package scrape + +import ( + "strings" + "testing" +) + +// strataSample is Strata 0.2.7's /metrics?format=prometheus, cut down, as +// cuda01 served it: vLLM's names beside its own. +const strataSample = `# TYPE vllm:num_requests_running gauge +vllm:num_requests_running{model_name="qwen3.8-flash-next"} 1 +# TYPE vllm:num_requests_waiting gauge +vllm:num_requests_waiting{model_name="qwen3.8-flash-next"} 3 +# TYPE vllm:kv_cache_usage_perc gauge +vllm:kv_cache_usage_perc{model_name="qwen3.8-flash-next"} 0.25 +# TYPE vllm:prompt_tokens_total counter +vllm:prompt_tokens_total{model_name="qwen3.8-flash-next"} 4826 +# TYPE vllm:generation_tokens_total counter +vllm:generation_tokens_total{model_name="qwen3.8-flash-next"} 1979 +# TYPE strata:live_tok_s gauge +strata:live_tok_s 0 +` + +func TestStrataIsReadInVLLMsNamesFromItsPrometheusFormat(t *testing.T) { + m := MappingFor("strata") + // Plain /metrics is JSON, and /metrics/prometheus has llama.cpp's + // names without the KV cache: read there, Strata was unmeasured. + if m.Path != "/metrics?format=prometheus" { + t.Errorf("read at %q", m.Path) + } + s, err := Parse(strings.NewReader(strataSample), m) + if err != nil { + t.Fatal(err) + } + if s.Running != 1 || s.Waiting != 3 || s.KVCacheUsage != 0.25 { + t.Errorf("stats %+v", s) + } +}