From 25b3e129febb9da713da59872a3c153f6017e0b1 Mon Sep 17 00:00:00 2001 From: Ilker Manap Date: Wed, 7 Oct 2026 13:51:40 +0200 Subject: [PATCH] What a shard computes on survives the agent's and the control plane's restarts It was sent once, with READY, and kept only in the leader's memory: a failover or an agent restart lost it. The agent now keeps it with the shard's state on disk (ShardSpec.Compute, set by Supervisor.SetCompute) and reports it with each running shard when it registers (RunningShard.compute), which adoption carries onto the replica. Co-Authored-By: Claude Opus 5.5 (1M context) --- api/gen/llmgw/v1/agent.pb.go | 79 +++++++++++-------- api/proto/llmgw/v1/agent.proto | 5 ++ internal/agent/agent.go | 10 ++- internal/agent/engines/engine.go | 4 + internal/agent/gpucheck.go | 27 ++++--- internal/agent/gpucheck_test.go | 14 ++-- internal/agent/supervisor/supervisor.go | 14 ++++ internal/agent/supervisor/supervisor_test.go | 25 ++++++ internal/control/agentsvc/convert.go | 4 + internal/control/reconciler/adopt.go | 11 ++- .../control/reconciler/reconciler_test.go | 13 +++ 11 files changed, 152 insertions(+), 54 deletions(-) diff --git a/api/gen/llmgw/v1/agent.pb.go b/api/gen/llmgw/v1/agent.pb.go index ed765c8..e66d7c0 100644 --- a/api/gen/llmgw/v1/agent.pb.go +++ b/api/gen/llmgw/v1/agent.pb.go @@ -1197,7 +1197,11 @@ type RunningShard struct { // replica across its own restart, so without this a shard adopted after a // restart could not be compared with the catalog — which is exactly when the // catalog is most likely to have been edited. - ProfileFp string `protobuf:"bytes,15,opt,name=profile_fp,json=profileFp,proto3" json:"profile_fp,omitempty"` + ProfileFp string `protobuf:"bytes,15,opt,name=profile_fp,json=profileFp,proto3" json:"profile_fp,omitempty"` + // compute is where its engine said it put the model when it became ready, + // kept by the agent across its own restarts: a control plane that takes + // over, or restarts, learns it here. + Compute *Compute `protobuf:"bytes,16,opt,name=compute,proto3" json:"compute,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -1337,6 +1341,13 @@ func (x *RunningShard) GetProfileFp() string { return "" } +func (x *RunningShard) GetCompute() *Compute { + if x != nil { + return x.Compute + } + return nil +} + type GPU struct { state protoimpl.MessageState `protogen:"open.v1"` Index int32 `protobuf:"varint,1,opt,name=index,proto3" json:"index,omitempty"` @@ -4656,7 +4667,7 @@ const file_llmgw_v1_agent_proto_rawDesc = "" + "free_bytes\x18\x02 \x01(\x03R\tfreeBytes\x12\x1f\n" + "\vtotal_bytes\x18\x03 \x01(\x03R\n" + "totalBytes\x12\x16\n" + - "\x06shared\x18\x04 \x01(\bR\x06shared\"\xb6\x03\n" + + "\x06shared\x18\x04 \x01(\bR\x06shared\"\xe3\x03\n" + "\fRunningShard\x12\x1d\n" + "\n" + "replica_id\x18\x01 \x01(\tR\treplicaId\x12\x19\n" + @@ -4678,7 +4689,8 @@ const file_llmgw_v1_agent_proto_rawDesc = "" + "\vleader_port\x18\x0e \x01(\x05R\n" + "leaderPort\x12\x1d\n" + "\n" + - "profile_fp\x18\x0f \x01(\tR\tprofileFp\"\xca\x06\n" + + "profile_fp\x18\x0f \x01(\tR\tprofileFp\x12+\n" + + "\acompute\x18\x10 \x01(\v2\x11.llmgw.v1.ComputeR\acompute\"\xca\x06\n" + "\x03GPU\x12\x14\n" + "\x05index\x18\x01 \x01(\x05R\x05index\x12\x12\n" + "\x04uuid\x18\x02 \x01(\tR\x04uuid\x12\x12\n" + @@ -5108,36 +5120,37 @@ var file_llmgw_v1_agent_proto_depIdxs = []int32{ 7, // 28: llmgw.v1.Register.host:type_name -> llmgw.v1.Host 19, // 29: llmgw.v1.Register.imex:type_name -> llmgw.v1.IMEX 0, // 30: llmgw.v1.RunningShard.state:type_name -> llmgw.v1.ShardStatus.State - 46, // 31: llmgw.v1.GPU.nic_paths:type_name -> llmgw.v1.GPU.NicPathsEntry - 11, // 32: llmgw.v1.GPU.mig:type_name -> llmgw.v1.MIGDevice - 16, // 33: llmgw.v1.NodeLoad.shards:type_name -> llmgw.v1.ShardLoad - 15, // 34: llmgw.v1.NodeLoad.gpus:type_name -> llmgw.v1.GPUSensor - 47, // 35: llmgw.v1.IBPort.counters:type_name -> llmgw.v1.IBPort.CountersEntry - 18, // 36: llmgw.v1.IBPort.subnet_managers:type_name -> llmgw.v1.SubnetManager - 8, // 37: llmgw.v1.Heartbeat.storage:type_name -> llmgw.v1.Storage - 13, // 38: llmgw.v1.Heartbeat.weights:type_name -> llmgw.v1.LocalWeights - 48, // 39: llmgw.v1.Heartbeat.engine_versions:type_name -> llmgw.v1.Heartbeat.EngineVersionsEntry - 24, // 40: llmgw.v1.Heartbeat.provision:type_name -> llmgw.v1.ProvisionStatus - 14, // 41: llmgw.v1.Heartbeat.load:type_name -> llmgw.v1.NodeLoad - 17, // 42: llmgw.v1.Heartbeat.ib_ports:type_name -> llmgw.v1.IBPort - 12, // 43: llmgw.v1.Heartbeat.fabrics:type_name -> llmgw.v1.FabricAttachment - 20, // 44: llmgw.v1.Heartbeat.gpu_fabric:type_name -> llmgw.v1.GPUFabricState - 19, // 45: llmgw.v1.Heartbeat.imex:type_name -> llmgw.v1.IMEX - 23, // 46: llmgw.v1.Provision.engines:type_name -> llmgw.v1.EngineWant - 49, // 47: llmgw.v1.EngineWant.offline_sha256:type_name -> llmgw.v1.EngineWant.OfflineSha256Entry - 0, // 48: llmgw.v1.ShardStatus.state:type_name -> llmgw.v1.ShardStatus.State - 27, // 49: llmgw.v1.ShardStatus.fetch:type_name -> llmgw.v1.FetchProgress - 25, // 50: llmgw.v1.ShardStatus.compute:type_name -> llmgw.v1.Compute - 29, // 51: llmgw.v1.StartShard.profile:type_name -> llmgw.v1.EngineProfile - 50, // 52: llmgw.v1.EngineProfile.env:type_name -> llmgw.v1.EngineProfile.EnvEntry - 42, // 53: llmgw.v1.EngineProfile.adapters:type_name -> llmgw.v1.LoraAdapter - 1, // 54: llmgw.v1.AgentService.Connect:input_type -> llmgw.v1.AgentMessage - 3, // 55: llmgw.v1.AgentService.Connect:output_type -> llmgw.v1.ControlMessage - 55, // [55:56] is the sub-list for method output_type - 54, // [54:55] is the sub-list for method input_type - 54, // [54:54] is the sub-list for extension type_name - 54, // [54:54] is the sub-list for extension extendee - 0, // [0:54] is the sub-list for field type_name + 25, // 31: llmgw.v1.RunningShard.compute:type_name -> llmgw.v1.Compute + 46, // 32: llmgw.v1.GPU.nic_paths:type_name -> llmgw.v1.GPU.NicPathsEntry + 11, // 33: llmgw.v1.GPU.mig:type_name -> llmgw.v1.MIGDevice + 16, // 34: llmgw.v1.NodeLoad.shards:type_name -> llmgw.v1.ShardLoad + 15, // 35: llmgw.v1.NodeLoad.gpus:type_name -> llmgw.v1.GPUSensor + 47, // 36: llmgw.v1.IBPort.counters:type_name -> llmgw.v1.IBPort.CountersEntry + 18, // 37: llmgw.v1.IBPort.subnet_managers:type_name -> llmgw.v1.SubnetManager + 8, // 38: llmgw.v1.Heartbeat.storage:type_name -> llmgw.v1.Storage + 13, // 39: llmgw.v1.Heartbeat.weights:type_name -> llmgw.v1.LocalWeights + 48, // 40: llmgw.v1.Heartbeat.engine_versions:type_name -> llmgw.v1.Heartbeat.EngineVersionsEntry + 24, // 41: llmgw.v1.Heartbeat.provision:type_name -> llmgw.v1.ProvisionStatus + 14, // 42: llmgw.v1.Heartbeat.load:type_name -> llmgw.v1.NodeLoad + 17, // 43: llmgw.v1.Heartbeat.ib_ports:type_name -> llmgw.v1.IBPort + 12, // 44: llmgw.v1.Heartbeat.fabrics:type_name -> llmgw.v1.FabricAttachment + 20, // 45: llmgw.v1.Heartbeat.gpu_fabric:type_name -> llmgw.v1.GPUFabricState + 19, // 46: llmgw.v1.Heartbeat.imex:type_name -> llmgw.v1.IMEX + 23, // 47: llmgw.v1.Provision.engines:type_name -> llmgw.v1.EngineWant + 49, // 48: llmgw.v1.EngineWant.offline_sha256:type_name -> llmgw.v1.EngineWant.OfflineSha256Entry + 0, // 49: llmgw.v1.ShardStatus.state:type_name -> llmgw.v1.ShardStatus.State + 27, // 50: llmgw.v1.ShardStatus.fetch:type_name -> llmgw.v1.FetchProgress + 25, // 51: llmgw.v1.ShardStatus.compute:type_name -> llmgw.v1.Compute + 29, // 52: llmgw.v1.StartShard.profile:type_name -> llmgw.v1.EngineProfile + 50, // 53: llmgw.v1.EngineProfile.env:type_name -> llmgw.v1.EngineProfile.EnvEntry + 42, // 54: llmgw.v1.EngineProfile.adapters:type_name -> llmgw.v1.LoraAdapter + 1, // 55: llmgw.v1.AgentService.Connect:input_type -> llmgw.v1.AgentMessage + 3, // 56: llmgw.v1.AgentService.Connect:output_type -> llmgw.v1.ControlMessage + 56, // [56:57] is the sub-list for method output_type + 55, // [55:56] is the sub-list for method input_type + 55, // [55:55] is the sub-list for extension type_name + 55, // [55:55] is the sub-list for extension extendee + 0, // [0:55] is the sub-list for field type_name } func init() { file_llmgw_v1_agent_proto_init() } diff --git a/api/proto/llmgw/v1/agent.proto b/api/proto/llmgw/v1/agent.proto index c95af4b..82a6cd9 100644 --- a/api/proto/llmgw/v1/agent.proto +++ b/api/proto/llmgw/v1/agent.proto @@ -266,6 +266,11 @@ message RunningShard { // restart could not be compared with the catalog — which is exactly when the // catalog is most likely to have been edited. string profile_fp = 15; + + // compute is where its engine said it put the model when it became ready, + // kept by the agent across its own restarts: a control plane that takes + // over, or restarts, learns it here. + Compute compute = 16; } message GPU { diff --git a/internal/agent/agent.go b/internal/agent/agent.go index 1c6d354..45018f7 100644 --- a/internal/agent/agent.go +++ b/internal/agent/agent.go @@ -436,6 +436,7 @@ func (a *Agent) runningShards() []*llmgwv1.RunningShard { Pool: spec.Pool, Fabric: spec.Fabric, ProfileFp: spec.ProfileFP, + Compute: computeToWire(spec.Compute), LeaderAddr: a.cfg.MgmtAddr, LeaderPort: int32(spec.LeaderPort), }) @@ -1133,7 +1134,7 @@ func (a *Agent) startShard(ctx context.Context, m *llmgwv1.StartShard, out chan< // Where it put the model, by its own account: said with the readiness, // and a GPU placement it computes on the CPU refused even where // nvidia-smi cannot be asked. - compute, cerr := engineCompute(eng, spec.GPUIndices, sh.Logs(512)) + compute, said, cerr := engineCompute(eng, spec.GPUIndices, sh.Logs(512)) if err == nil && cerr != nil { err = cerr _ = a.sup.Stop(ctx, shardID) @@ -1150,8 +1151,11 @@ func (a *Agent) startShard(ctx context.Context, m *llmgwv1.StartShard, out chan< report(supervisor.Failed, err.Error(), sh.Logs(20)) return } - if compute != nil { - send(&llmgwv1.ShardStatus{State: agentsvc.StateToWire(supervisor.Ready), Compute: compute}) + if said { + // Kept with the shard, so an agent restarted, or a control plane + // taking over, still learns it. + a.sup.SetCompute(shardID, compute) + send(&llmgwv1.ShardStatus{State: agentsvc.StateToWire(supervisor.Ready), Compute: computeToWire(&compute)}) } else { report(supervisor.Ready, "", nil) } diff --git a/internal/agent/engines/engine.go b/internal/agent/engines/engine.go index eb80178..9b9d0fc 100644 --- a/internal/agent/engines/engine.go +++ b/internal/agent/engines/engine.go @@ -114,6 +114,10 @@ type ShardSpec struct { // lets the control plane tell, after its own restart, whether a shard that // is still running matches the profile the catalog holds now. ProfileFP string + + // Compute is where the engine said it put the model as it became ready, + // kept with the shard's state so an agent restarted later still knows. + Compute *Compute `json:",omitempty"` } // Host is where the engine listens. diff --git a/internal/agent/gpucheck.go b/internal/agent/gpucheck.go index 3bed7fc..a2b649e 100644 --- a/internal/agent/gpucheck.go +++ b/internal/agent/gpucheck.go @@ -156,26 +156,33 @@ func joinIndices(is []int) string { } // engineCompute is where an engine said it put its model, from its output -// so far — llama.cpp's, which says — nil for one that said nothing. A shard -// given GPUs that its engine says it computes on the CPU is an error, found -// this way even where nvidia-smi cannot be asked. -func engineCompute(eng engines.Engine, gpus []int, lines []string) (*llmgwv1.Compute, error) { +// so far — llama.cpp's, which says — and false for one that said nothing. A +// shard given GPUs that its engine says it computes on the CPU is an error, +// found this way even where nvidia-smi cannot be asked. +func engineCompute(eng engines.Engine, gpus []int, lines []string) (engines.Compute, bool, error) { if _, ok := eng.(engines.CPUFallback); !ok { - return nil, nil + return engines.Compute{}, false, nil } read, ok := engines.ReadLlamaCompute(lines) if !ok { - return nil, nil + return engines.Compute{}, false, nil } - c := &llmgwv1.Compute{Backend: read.Backend, Devices: read.Devices, - LayersOffloaded: int32(read.Offloaded), LayersTotal: int32(read.Total)} //nolint:gosec // layer counts if len(gpus) > 0 && read.OnTheCPU() { why := "its GPU backend saw no device it could use" if read.Total > 0 && !read.NoGPU { why = fmt.Sprintf("it offloaded %d of %d layers", read.Offloaded, read.Total) } - return c, fmt.Errorf("placed on GPU %v but computing on the CPU: %s, by llama.cpp's own account "+ + return read, true, fmt.Errorf("placed on GPU %v but computing on the CPU: %s, by llama.cpp's own account "+ "(llmgw ctl engine-check on this node says why)", gpus, why) } - return c, nil + return read, true, nil +} + +// computeToWire is a compute as it travels to the control plane. +func computeToWire(c *engines.Compute) *llmgwv1.Compute { + if c == nil { + return nil + } + return &llmgwv1.Compute{Backend: c.Backend, Devices: c.Devices, + LayersOffloaded: int32(c.Offloaded), LayersTotal: int32(c.Total)} //nolint:gosec // layer counts } diff --git a/internal/agent/gpucheck_test.go b/internal/agent/gpucheck_test.go index f73dc37..8ae73da 100644 --- a/internal/agent/gpucheck_test.go +++ b/internal/agent/gpucheck_test.go @@ -126,24 +126,24 @@ func TestForeignGPUMemoryLeavesOutTheAgentsOwnShards(t *testing.T) { func TestLlamaCppsOwnWordOnTheCPUIsEnoughToRefuse(t *testing.T) { eng := fallbackEngine{} // Without nvidia-smi to ask, the engine's own output says it. - c, err := engineCompute(eng, []int{0}, []string{ + c, said, err := engineCompute(eng, []int{0}, []string{ "warning: no usable GPU found, --gpu-layers option will be ignored", "load_tensors: offloaded 0/49 layers to GPU"}) - if err == nil || !strings.Contains(err.Error(), "computing on the CPU") || c.GetBackend() != "CPU" { + if err == nil || !said || !strings.Contains(err.Error(), "computing on the CPU") || c.Backend != "CPU" { t.Fatalf("a GPU shard llama.cpp put on the CPU: %v %+v", err, c) } - c, err = engineCompute(eng, []int{0}, []string{ + c, said, err = engineCompute(eng, []int{0}, []string{ "llama_model_load_from_file_impl: using device CUDA0 (Quadro RTX 6000) - 23800 MiB free", "load_tensors: offloaded 49/49 layers to GPU"}) - if err != nil || c.GetBackend() != "CUDA" || c.GetLayersOffloaded() != 49 { + if err != nil || !said || c.Backend != "CUDA" || c.Offloaded != 49 || computeToWire(&c).GetLayersOffloaded() != 49 { t.Errorf("on its GPU: %v %+v", err, c) } // A CPU shard is not refused for being on the CPU, and says so. - if c, err := engineCompute(eng, nil, []string{"load_tensors: offloaded 0/49 layers to GPU"}); err != nil || c.GetBackend() != "CPU" { + if c, _, err := engineCompute(eng, nil, []string{"load_tensors: offloaded 0/49 layers to GPU"}); err != nil || c.Backend != "CPU" { t.Errorf("a CPU shard: %v %+v", err, c) } // An engine that said nothing says nothing. - if c, err := engineCompute(eng, []int{0}, []string{"listening"}); err != nil || c != nil { - t.Errorf("nothing said: %v %+v", err, c) + if _, said, err := engineCompute(eng, []int{0}, []string{"listening"}); err != nil || said { + t.Errorf("nothing said: %v", err) } } diff --git a/internal/agent/supervisor/supervisor.go b/internal/agent/supervisor/supervisor.go index 259cd26..eed3c5c 100644 --- a/internal/agent/supervisor/supervisor.go +++ b/internal/agent/supervisor/supervisor.go @@ -644,6 +644,20 @@ func (s *Supervisor) Get(id string) (*Shard, error) { return sh, nil } +// SetCompute records where a shard's engine said it put the model, with the +// shard's state on disk. +func (s *Supervisor) SetCompute(id string, c engines.Compute) { + s.mu.Lock() + sh, ok := s.shards[id] + if ok { + sh.Spec.Compute = &c + } + s.mu.Unlock() + if ok && sh.pid > 0 { + _ = s.writeState(sh) + } +} + // Remove forgets a terminated shard. func (s *Supervisor) Remove(id string) { s.mu.Lock() diff --git a/internal/agent/supervisor/supervisor_test.go b/internal/agent/supervisor/supervisor_test.go index 3039a7a..5b542ef 100644 --- a/internal/agent/supervisor/supervisor_test.go +++ b/internal/agent/supervisor/supervisor_test.go @@ -848,3 +848,28 @@ func eventually(t *testing.T, limit time.Duration, msg string, cond func() bool) time.Sleep(10 * time.Millisecond) } } + +func TestWhatAShardComputesOnSurvivesTheAgent(t *testing.T) { + s, _ := newSupervisor(t) + spec := engines.ShardSpec{ReplicaID: "rep-1", ShardID: "a", HTTPPort: 20100} + sh, err := s.Start(t.Context(), "a", spec, shellCommand("sleep 30")) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = s.Stop(context.Background(), "a") }) + _ = sh + s.SetCompute("a", engines.Compute{Backend: "CUDA", Offloaded: 49, Total: 49}) + + s2, err := New(Config{WorkDir: s.cfg.WorkDir, StopTimeout: 2 * time.Second}, + &scriptedHealth{}, slog.New(slog.DiscardHandler)) + if err != nil { + t.Fatal(err) + } + orphans, err := s2.FindOrphans() + if err != nil || len(orphans) != 1 { + t.Fatalf("%v %d", err, len(orphans)) + } + if c := orphans[0].Spec.Compute; c == nil || c.Backend != "CUDA" || c.Offloaded != 49 { + t.Errorf("not kept: %+v", c) + } +} diff --git a/internal/control/agentsvc/convert.go b/internal/control/agentsvc/convert.go index 8e1cfae..d386cc5 100644 --- a/internal/control/agentsvc/convert.go +++ b/internal/control/agentsvc/convert.go @@ -169,6 +169,10 @@ func AdoptedFromRegister(nodeID string, shards []*llmgwv1.RunningShard) []reconc Fabric: s.GetFabric(), ProfileFP: s.GetProfileFp(), } + if c := s.GetCompute(); c != nil { + a.Compute = &reconciler.Compute{Backend: c.GetBackend(), Devices: c.GetDevices(), + Offloaded: int(c.GetLayersOffloaded()), Total: int(c.GetLayersTotal())} + } if a.Rank == 0 && s.GetLeaderAddr() != "" { a.Endpoint = fmt.Sprintf("http://%s:%d", s.GetLeaderAddr(), s.GetHttpPort()) } diff --git a/internal/control/reconciler/adopt.go b/internal/control/reconciler/adopt.go index b74722c..e3f4782 100644 --- a/internal/control/reconciler/adopt.go +++ b/internal/control/reconciler/adopt.go @@ -130,13 +130,19 @@ func (r *Reconciler) mergeAdopted(rep *Replica, group []AdoptedShard) { } for _, s := range group { if _, dup := have[s.ShardID]; dup { + // Known already, but what it computes on may not be. + for i := range rep.Shards { + if rep.Shards[i].ID == s.ShardID && rep.Shards[i].Compute == nil { + rep.Shards[i].Compute = s.Compute + } + } continue } have[s.ShardID] = struct{}{} rep.Shards = append(rep.Shards, Shard{ ID: s.ShardID, NodeID: s.NodeID, Rank: s.Rank, GPUIndices: s.GPUIndices, HCAs: s.HCAs, - HTTPPort: s.HTTPPort, State: s.State, + HTTPPort: s.HTTPPort, State: s.State, Compute: s.Compute, }) // Reclaim the GPUs and the memory, or the next placement hands them // out twice. The memory comes from the catalog rather than the agent: @@ -254,6 +260,9 @@ type AdoptedShard struct { ProfileFP string // Endpoint is the leader's base URL, set on the rank 0 shard. Endpoint string + // Compute is where the engine said it put the model, when the agent + // kept that. + Compute *Compute } // adoptedCores is what an adopted shard is taken to hold of its node's CPUs. diff --git a/internal/control/reconciler/reconciler_test.go b/internal/control/reconciler/reconciler_test.go index e7b9765..bcd1190 100644 --- a/internal/control/reconciler/reconciler_test.go +++ b/internal/control/reconciler/reconciler_test.go @@ -2120,3 +2120,16 @@ func TestMaxUnavailableStopsThatManyFirst(t *testing.T) { t.Errorf("one stopped first and one replaced beside itself: %d stopped, %d beside", stopped, beside) } } + +func TestAnAdoptedShardKeepsWhatItComputesOn(t *testing.T) { + // A control plane that took over, or restarted, learns where a replica + // computes from what its agent kept, not from a READY it never saw. + h := newHarness(t, testModel(1), gpuNode("1", 8)) + a := adopted("llama-99", "llama-99-r0", "1", 0, []int{0, 1, 2, 3}) + a.Compute = &Compute{Backend: "CUDA", Devices: []string{"CUDA0: H100"}, Offloaded: 81, Total: 81} + h.r.Adopt("1", []AdoptedShard{a}) + reps := h.r.Replicas() + if len(reps) != 1 || reps[0].Shards[0].Compute == nil || reps[0].Shards[0].Compute.Offloaded != 81 { + t.Fatalf("%+v", reps) + } +}