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) + } +}