From 45e913bd10a57e304484b7ca72f2ddc4e9baf8e2 Mon Sep 17 00:00:00 2001 From: Ilker Manap Date: Sun, 20 Sep 2026 09:57:25 +0200 Subject: [PATCH] Refresh disk and weights while the agent runs, not once when it starts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Found by answering "where can we see this live". `ctl weights` was offering a file for deletion that I had deleted two days earlier, and the node's disk figure was the one it had when its agent started. Discover() is called at registration and nowhere else, so both were frozen. Disk is a fact about a moment, not about a machine. Sampled once, the figure is wrong on exactly the node worth watching — the one busy filling up — and it stays wrong until that agent restarts. The heartbeat already refreshes GPU memory for the same reason, with a comment about noticing a foreign process "after registration"; disk belonged beside it and I did not put it there. The stale list cut both ways. A model fetched after registration was invisible until a restart, so placement could not prefer the node that had just spent ten minutes downloading it, and the disk check could not see the space it used. Storage goes in every heartbeat, being a few numbers per filesystem. The file list goes only when a fingerprint of paths and sizes changes, with a flag saying the list is meant — without that, an empty list would be ambiguous between "this node now holds nothing" and "nothing to report", and the second reading would empty a node's inventory five seconds after it registered. Co-Authored-By: Claude Opus 5 (1M context) --- api/gen/llmgw/v1/agent.pb.go | 68 +++++++++--- api/proto/llmgw/v1/agent.proto | 16 +++ internal/agent/agent.go | 47 ++++++++ internal/control/agentsvc/heartbeat_test.go | 114 ++++++++++++++++++++ internal/control/agentsvc/service.go | 31 +++++- 5 files changed, 261 insertions(+), 15 deletions(-) create mode 100644 internal/control/agentsvc/heartbeat_test.go diff --git a/api/gen/llmgw/v1/agent.pb.go b/api/gen/llmgw/v1/agent.pb.go index 3ff8981d..45bafdfe 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 f4792bf9..ecf97f61 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 c5f36d2c..d1f1c2ae 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 00000000..255f5db2 --- /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 3309028a..f68d8931 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) }