diff --git a/api/gen/llmgw/v1/agent.pb.go b/api/gen/llmgw/v1/agent.pb.go index 3ff8981..45bafdf 100644 --- a/api/gen/llmgw/v1/agent.pb.go +++ b/api/gen/llmgw/v1/agent.pb.go @@ -1033,8 +1033,22 @@ type Heartbeat struct { // a card after registration. GpuMemoryUsedMb []int32 `protobuf:"varint,2,rep,packed,name=gpu_memory_used_mb,json=gpuMemoryUsedMb,proto3" json:"gpu_memory_used_mb,omitempty"` NewProblems []string `protobuf:"bytes,3,rep,name=new_problems,json=newProblems,proto3" json:"new_problems,omitempty"` - unknownFields protoimpl.UnknownFields - sizeCache protoimpl.SizeCache + // storage for the same reason the GPU memory is here: it changes while the + // agent runs, and a figure sampled once at registration is a figure that is + // wrong on exactly the node worth watching — the one busy filling up. + // + // Always sent, because it is a few numbers per filesystem. + Storage []*Storage `protobuf:"bytes,4,rep,name=storage,proto3" json:"storage,omitempty"` + // weights only when they changed, because the list is long enough that + // sending it every few seconds would be paying for nothing most of the time. + // + // weights_changed is what says the list is meant: without it an empty list + // would be ambiguous between "this node now holds nothing" and "nothing to + // report", and the first would quietly empty a node's inventory. + WeightsChanged bool `protobuf:"varint,5,opt,name=weights_changed,json=weightsChanged,proto3" json:"weights_changed,omitempty"` + Weights []*LocalWeights `protobuf:"bytes,6,rep,name=weights,proto3" json:"weights,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache } func (x *Heartbeat) Reset() { @@ -1088,6 +1102,27 @@ func (x *Heartbeat) GetNewProblems() []string { return nil } +func (x *Heartbeat) GetStorage() []*Storage { + if x != nil { + return x.Storage + } + return nil +} + +func (x *Heartbeat) GetWeightsChanged() bool { + if x != nil { + return x.WeightsChanged + } + return false +} + +func (x *Heartbeat) GetWeights() []*LocalWeights { + if x != nil { + return x.Weights + } + return nil +} + // ShardStatus reports a shard's lifecycle transition. type ShardStatus struct { state protoimpl.MessageState `protogen:"open.v1"` @@ -1755,11 +1790,14 @@ const file_llmgw_v1_agent_proto_rawDesc = "" + "\x04path\x18\x01 \x01(\tR\x04path\x12\x1d\n" + "\n" + "size_bytes\x18\x02 \x01(\x03R\tsizeBytes\x12\x16\n" + - "\x06shared\x18\x03 \x01(\bR\x06shared\"t\n" + + "\x06shared\x18\x03 \x01(\bR\x06shared\"\xfc\x01\n" + "\tHeartbeat\x12\x17\n" + "\anode_id\x18\x01 \x01(\tR\x06nodeId\x12+\n" + "\x12gpu_memory_used_mb\x18\x02 \x03(\x05R\x0fgpuMemoryUsedMb\x12!\n" + - "\fnew_problems\x18\x03 \x03(\tR\vnewProblems\"\x8c\x03\n" + + "\fnew_problems\x18\x03 \x03(\tR\vnewProblems\x12+\n" + + "\astorage\x18\x04 \x03(\v2\x11.llmgw.v1.StorageR\astorage\x12'\n" + + "\x0fweights_changed\x18\x05 \x01(\bR\x0eweightsChanged\x120\n" + + "\aweights\x18\x06 \x03(\v2\x16.llmgw.v1.LocalWeightsR\aweights\"\x8c\x03\n" + "\vShardStatus\x12\x17\n" + "\anode_id\x18\x01 \x01(\tR\x06nodeId\x12\x1d\n" + "\n" + @@ -1887,16 +1925,18 @@ var file_llmgw_v1_agent_proto_depIdxs = []int32{ 8, // 13: llmgw.v1.Register.weights:type_name -> llmgw.v1.LocalWeights 4, // 14: llmgw.v1.Register.storage:type_name -> llmgw.v1.Storage 0, // 15: llmgw.v1.RunningShard.state:type_name -> llmgw.v1.ShardStatus.State - 0, // 16: llmgw.v1.ShardStatus.state:type_name -> llmgw.v1.ShardStatus.State - 12, // 17: llmgw.v1.StartShard.profile:type_name -> llmgw.v1.EngineProfile - 19, // 18: llmgw.v1.EngineProfile.env:type_name -> llmgw.v1.EngineProfile.EnvEntry - 1, // 19: llmgw.v1.AgentService.Connect:input_type -> llmgw.v1.AgentMessage - 2, // 20: llmgw.v1.AgentService.Connect:output_type -> llmgw.v1.ControlMessage - 20, // [20:21] is the sub-list for method output_type - 19, // [19:20] is the sub-list for method input_type - 19, // [19:19] is the sub-list for extension type_name - 19, // [19:19] is the sub-list for extension extendee - 0, // [0:19] is the sub-list for field type_name + 4, // 16: llmgw.v1.Heartbeat.storage:type_name -> llmgw.v1.Storage + 8, // 17: llmgw.v1.Heartbeat.weights:type_name -> llmgw.v1.LocalWeights + 0, // 18: llmgw.v1.ShardStatus.state:type_name -> llmgw.v1.ShardStatus.State + 12, // 19: llmgw.v1.StartShard.profile:type_name -> llmgw.v1.EngineProfile + 19, // 20: llmgw.v1.EngineProfile.env:type_name -> llmgw.v1.EngineProfile.EnvEntry + 1, // 21: llmgw.v1.AgentService.Connect:input_type -> llmgw.v1.AgentMessage + 2, // 22: llmgw.v1.AgentService.Connect:output_type -> llmgw.v1.ControlMessage + 22, // [22:23] is the sub-list for method output_type + 21, // [21:22] is the sub-list for method input_type + 21, // [21:21] is the sub-list for extension type_name + 21, // [21:21] is the sub-list for extension extendee + 0, // [0:21] 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 f4792bf..ecf97f6 100644 --- a/api/proto/llmgw/v1/agent.proto +++ b/api/proto/llmgw/v1/agent.proto @@ -199,6 +199,22 @@ message Heartbeat { // a card after registration. repeated int32 gpu_memory_used_mb = 2; repeated string new_problems = 3; + + // storage for the same reason the GPU memory is here: it changes while the + // agent runs, and a figure sampled once at registration is a figure that is + // wrong on exactly the node worth watching — the one busy filling up. + // + // Always sent, because it is a few numbers per filesystem. + repeated Storage storage = 4; + + // weights only when they changed, because the list is long enough that + // sending it every few seconds would be paying for nothing most of the time. + // + // weights_changed is what says the list is meant: without it an empty list + // would be ambiguous between "this node now holds nothing" and "nothing to + // report", and the first would quietly empty a node's inventory. + bool weights_changed = 5; + repeated LocalWeights weights = 6; } // ShardStatus reports a shard's lifecycle transition. diff --git a/internal/agent/agent.go b/internal/agent/agent.go index c5f36d2..d1f1c2a 100644 --- a/internal/agent/agent.go +++ b/internal/agent/agent.go @@ -3,7 +3,9 @@ package agent import ( "context" + "crypto/sha256" "crypto/tls" + "encoding/hex" "errors" "fmt" "io" @@ -172,6 +174,11 @@ type Agent struct { // report can name both. The control plane keys on the replica. shardReplica map[string]string fetcher *weights.Fetcher + + // lastWeights is the fingerprint of the file list most recently sent, so a + // heartbeat carries the list only when it changed. Touched only by the + // heartbeat goroutine, which is the only thing that reads or writes it. + lastWeights string } // connectHTTPClient is the HTTP client connect uses; an interface so tests can @@ -580,9 +587,49 @@ func (a *Agent) heartbeat() *llmgwv1.AgentMessage { for _, g := range inv.GPUs { hb.GpuMemoryUsedMb = append(hb.GpuMemoryUsedMb, int32(g.MemoryUsedMB)) } + + // Rescanned here rather than reused from registration. Disk is not a fact + // about a machine, it is a fact about a moment: a node that had 200 GB when + // its agent started can be full an hour later, and it is precisely the node + // busy filling up whose figure matters. + for _, st := range discovery.ScanStorage(a.cfg.ModelDirs) { + hb.Storage = append(hb.Storage, &llmgwv1.Storage{ + Path: st.Path, FreeBytes: st.FreeBytes, + TotalBytes: st.TotalBytes, Shared: st.Shared, + }) + } + + // The file list is longer and changes rarely, so it goes only when it has. + // The flag is what makes an empty list mean "this node now holds nothing" + // rather than "nothing to report" — without it, the first heartbeat after + // a model was deleted would be indistinguishable from silence. + weights := discovery.ScanWeights(a.cfg.ModelDirs) + if fp := weightsFingerprint(weights); fp != a.lastWeights { + a.lastWeights = fp + hb.WeightsChanged = true + for _, w := range weights { + hb.Weights = append(hb.Weights, &llmgwv1.LocalWeights{ + Path: w.Path, SizeBytes: w.SizeBytes, Shared: w.Shared, + }) + } + } return &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_Heartbeat{Heartbeat: hb}} } +// weightsFingerprint is enough to notice a model appearing, disappearing or +// finishing its download. +// +// Paths and sizes, because a fetch in progress has the same name as the file +// it will become — under a different name until it is renamed, but a shared +// filesystem can also grow a file in place. +func weightsFingerprint(ws []discovery.LocalWeights) string { + h := sha256.New() + for _, w := range ws { + fmt.Fprintf(h, "%s:%d\n", w.Path, w.SizeBytes) + } + return hex.EncodeToString(h.Sum(nil)) +} + // removeWeights deletes cached model files the control plane says are no // longer needed. // diff --git a/internal/control/agentsvc/heartbeat_test.go b/internal/control/agentsvc/heartbeat_test.go new file mode 100644 index 0000000..255f5db --- /dev/null +++ b/internal/control/agentsvc/heartbeat_test.go @@ -0,0 +1,114 @@ +package agentsvc + +import ( + "testing" + + llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1" + "gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery" + "gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory" +) + +func nodeWith(s *Service, n inventory.Node) { s.inv.Upsert(n) } + +func newTestService() *Service { + return &Service{inv: inventory.NewStore(), conns: map[string]*conn{}} +} + +func TestDiskIsRefreshedWhileTheAgentRuns(t *testing.T) { + // It was sampled once, at registration, and never again. Found on the live + // cluster: an agent running for two days still reported the disk it had + // when it started, and a file deleted in between was still being offered + // for deletion. Disk is a fact about a moment, and the node whose figure + // matters most is the one busy filling up. + s := newTestService() + nodeWith(s, inventory.Node{ + ID: "n1", + State: inventory.Ready, + Storage: []discovery.Storage{ + {Path: "/models", FreeBytes: 200 << 30, TotalBytes: 400 << 30}, + }, + }) + + s.applyHeartbeat("n1", &llmgwv1.Heartbeat{ + NodeId: "n1", + Storage: []*llmgwv1.Storage{ + {Path: "/models", FreeBytes: 5 << 30, TotalBytes: 400 << 30}, + }, + }) + + n, _ := s.inv.Get("n1") + free, total := n.LocalDisk() + if free != 5<<30 || total != 400<<30 { + t.Errorf("disk is %d of %d, want the figure the heartbeat carried", free, total) + } +} + +func TestAnEmptyWeightListOnlyCountsWhenTheAgentMeansIt(t *testing.T) { + // Most heartbeats carry no file list because nothing changed. Treating + // that as "this node holds nothing" would empty a node's inventory five + // seconds after it registered, and placement would stop preferring the + // node that actually has the model. + s := newTestService() + nodeWith(s, inventory.Node{ + ID: "n1", + State: inventory.Ready, + Weights: []discovery.LocalWeights{{Path: "/models/a.gguf", SizeBytes: 100}}, + }) + + // Silence about weights: the list must survive. + s.applyHeartbeat("n1", &llmgwv1.Heartbeat{ + NodeId: "n1", + Storage: []*llmgwv1.Storage{{Path: "/models", FreeBytes: 1, TotalBytes: 2}}, + }) + if n, _ := s.inv.Get("n1"); len(n.Weights) != 1 { + t.Fatalf("a heartbeat with no weight list emptied the inventory: %+v", n.Weights) + } + + // And when the agent says it means an empty list, it is taken. + s.applyHeartbeat("n1", &llmgwv1.Heartbeat{NodeId: "n1", WeightsChanged: true}) + if n, _ := s.inv.Get("n1"); len(n.Weights) != 0 { + t.Errorf("a deliberate empty list was ignored: %+v", n.Weights) + } +} + +func TestANewlyFetchedModelAppearsWithoutARestart(t *testing.T) { + // The other direction of the same bug: a model fetched after registration + // was invisible until the agent restarted, so placement could not prefer + // the node that had just spent ten minutes downloading it. + s := newTestService() + nodeWith(s, inventory.Node{ID: "n1", State: inventory.Ready}) + + s.applyHeartbeat("n1", &llmgwv1.Heartbeat{ + NodeId: "n1", + WeightsChanged: true, + Weights: []*llmgwv1.LocalWeights{ + {Path: "/models/new.gguf", SizeBytes: 4 << 30}, + }, + }) + + n, _ := s.inv.Get("n1") + if !n.HasWeights("/models/new.gguf") { + t.Errorf("a model fetched after registration did not appear: %+v", n.Weights) + } +} + +func TestAHeartbeatAboutNothingChangesNothing(t *testing.T) { + s := newTestService() + before := inventory.Node{ + ID: "n1", + State: inventory.Ready, + Weights: []discovery.LocalWeights{{Path: "/models/a.gguf", SizeBytes: 100}}, + Storage: []discovery.Storage{{Path: "/models", FreeBytes: 7, TotalBytes: 9}}, + } + nodeWith(s, before) + + s.applyHeartbeat("n1", &llmgwv1.Heartbeat{NodeId: "n1"}) + + n, _ := s.inv.Get("n1") + if len(n.Weights) != 1 { + t.Errorf("weights changed: %+v", n.Weights) + } + if free, _ := n.LocalDisk(); free != 7 { + t.Errorf("disk changed: %d", free) + } +} diff --git a/internal/control/agentsvc/service.go b/internal/control/agentsvc/service.go index 3309028..f68d893 100644 --- a/internal/control/agentsvc/service.go +++ b/internal/control/agentsvc/service.go @@ -14,6 +14,7 @@ import ( llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1" "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1/llmgwv1connect" + "gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery" "gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor" "gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory" "gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler" @@ -239,7 +240,8 @@ func (s *Service) handle(nodeID string, msg *llmgwv1.AgentMessage) { func (s *Service) applyHeartbeat(nodeID string, hb *llmgwv1.Heartbeat) { used := hb.GetGpuMemoryUsedMb() problems := hb.GetNewProblems() - if len(used) == 0 && len(problems) == 0 { + storage := hb.GetStorage() + if len(used) == 0 && len(problems) == 0 && len(storage) == 0 && !hb.GetWeightsChanged() { return } @@ -259,6 +261,33 @@ func (s *Service) applyHeartbeat(nodeID string, hb *llmgwv1.Heartbeat) { // conservative direction. n.GPUs[0].Health.Problems = append(n.GPUs[0].Health.Problems, problems...) } + + // Disk is a fact about a moment, not about a machine. Sampled only at + // registration it is wrong on exactly the node worth watching — the one + // busy filling up — and it stays wrong until that agent restarts. + if len(storage) > 0 { + n.Storage = n.Storage[:0] + for _, st := range storage { + n.Storage = append(n.Storage, discovery.Storage{ + Path: st.GetPath(), + FreeBytes: st.GetFreeBytes(), + TotalBytes: st.GetTotalBytes(), + Shared: st.GetShared(), + }) + } + } + + // Only when the agent says the list is meant. An empty list without the + // flag is silence, and treating it as "this node holds nothing" would + // empty a node's inventory on every heartbeat. + if hb.GetWeightsChanged() { + n.Weights = n.Weights[:0] + for _, w := range hb.GetWeights() { + n.Weights = append(n.Weights, discovery.LocalWeights{ + Path: w.GetPath(), SizeBytes: w.GetSizeBytes(), Shared: w.GetShared(), + }) + } + } s.inv.Upsert(n) }