diff --git a/api/gen/llmgw/v1/agent.pb.go b/api/gen/llmgw/v1/agent.pb.go index 30b1fd5..32acd6a 100644 --- a/api/gen/llmgw/v1/agent.pb.go +++ b/api/gen/llmgw/v1/agent.pb.go @@ -2953,8 +2953,12 @@ type Heartbeat struct { // is not handed out. Empty from an agent that does not measure it, or from // one that could not this time. GpuForeignMemoryMb []int32 `protobuf:"varint,19,rep,packed,name=gpu_foreign_memory_mb,json=gpuForeignMemoryMb,proto3" json:"gpu_foreign_memory_mb,omitempty"` - unknownFields protoimpl.UnknownFields - sizeCache protoimpl.SizeCache + // mem_bandwidth_gbs is how fast the node's memory reads, in GB/s, once + // the agent has measured it — in the background after it starts, so a + // registration usually has none yet. 0 means not measured. + MemBandwidthGbs float64 `protobuf:"fixed64,20,opt,name=mem_bandwidth_gbs,json=memBandwidthGbs,proto3" json:"mem_bandwidth_gbs,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache } func (x *Heartbeat) Reset() { @@ -3120,6 +3124,13 @@ func (x *Heartbeat) GetGpuForeignMemoryMb() []int32 { return nil } +func (x *Heartbeat) GetMemBandwidthGbs() float64 { + if x != nil { + return x.MemBandwidthGbs + } + return 0 +} + // Provision is the software a node should have: which engines, at which // versions. Nothing more. // @@ -5382,7 +5393,7 @@ const file_llmgw_v1_agent_proto_rawDesc = "" + "\x06status\x18\x03 \x01(\tR\x06status\x12!\n" + "\fcluster_uuid\x18\x04 \x01(\tR\vclusterUuid\x12\x1b\n" + "\tclique_id\x18\x05 \x01(\tR\bcliqueId\x12\x1c\n" + - "\tbandwidth\x18\x06 \x01(\tR\tbandwidth\"\xa9\a\n" + + "\tbandwidth\x18\x06 \x01(\tR\tbandwidth\"\xd5\a\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" + @@ -5404,7 +5415,8 @@ const file_llmgw_v1_agent_proto_rawDesc = "" + "gpu_fabric\x18\x10 \x03(\v2\x18.llmgw.v1.GPUFabricStateR\tgpuFabric\x12\"\n" + "\x04imex\x18\x11 \x01(\v2\x0e.llmgw.v1.IMEXR\x04imex\x12\x1b\n" + "\tcpu_count\x18\x12 \x01(\x05R\bcpuCount\x121\n" + - "\x15gpu_foreign_memory_mb\x18\x13 \x03(\x05R\x12gpuForeignMemoryMb\x1aA\n" + + "\x15gpu_foreign_memory_mb\x18\x13 \x03(\x05R\x12gpuForeignMemoryMb\x12*\n" + + "\x11mem_bandwidth_gbs\x18\x14 \x01(\x01R\x0fmemBandwidthGbs\x1aA\n" + "\x13EngineVersionsEntry\x12\x10\n" + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\";\n" + diff --git a/api/proto/llmgw/v1/agent.proto b/api/proto/llmgw/v1/agent.proto index cf7ac19..c0c5814 100644 --- a/api/proto/llmgw/v1/agent.proto +++ b/api/proto/llmgw/v1/agent.proto @@ -591,6 +591,11 @@ message Heartbeat { // is not handed out. Empty from an agent that does not measure it, or from // one that could not this time. repeated int32 gpu_foreign_memory_mb = 19; + + // mem_bandwidth_gbs is how fast the node's memory reads, in GB/s, once + // the agent has measured it — in the background after it starts, so a + // registration usually has none yet. 0 means not measured. + double mem_bandwidth_gbs = 20; } // Provision is the software a node should have: which engines, at which diff --git a/internal/agent/agent.go b/internal/agent/agent.go index 25759e0..31d311a 100644 --- a/internal/agent/agent.go +++ b/internal/agent/agent.go @@ -514,14 +514,13 @@ func (a *Agent) scan(ctx context.Context) error { inv, err := discovery.Discover(ctx, opts, a.log) host := discovery.DiscoverHost() host.MinComputeCapability = discovery.MinComputeCapability(inv.GPUs) - host.MemBandwidthGBs = discovery.MeasureMemoryBandwidth() + host.MemBandwidthGBs = a.memBandwidth() a.mu.Lock() a.inv = inv a.prov.Host = host a.mu.Unlock() a.log.Info("this node", "os", host.OSID+" "+host.OSVersion, "arch", host.Arch, - "cpu_flags", host.CPUFlags, "nvidia_driver", host.NVIDIADriver, - "memory_gbs", math.Round(host.MemBandwidthGBs*10)/10) + "cpu_flags", host.CPUFlags, "nvidia_driver", host.NVIDIADriver) return err } @@ -670,6 +669,7 @@ func (a *Agent) heartbeat() *llmgwv1.AgentMessage { // without anybody restarting anything. hb.MemoryTotalMb = int64(discovery.MemoryTotalMB()) hb.CpuCount = int32(runtime.NumCPU()) + hb.MemBandwidthGbs = a.memBandwidth() a.reprobeEngines() a.mu.Lock() @@ -1543,3 +1543,29 @@ func (a *Agent) reportCompute(ctx context.Context, out chan<- *llmgwv1.AgentMess } } } + +// memBandwidth is the node's memory read rate in GB/s, 0 until measured. +// Measured once a process, in the background, the first time it is asked +// for: it reads half a gigabyte, which must not hold up an agent reaching +// its control plane, and the heartbeat carries it when it is known. Once a +// process rather than an agent, so tests running many agents in one read it +// once. +func (a *Agent) memBandwidth() float64 { + bwOnce.Do(func() { + go func() { + bw := measureBandwidth() + bwBits.Store(math.Float64bits(bw)) + if bw > 0 { + a.log.Info("this node's memory", "read_gbs", math.Round(bw*10)/10) + } + }() + }) + return math.Float64frombits(bwBits.Load()) +} + +var ( + bwOnce sync.Once + bwBits atomic.Uint64 + // measureBandwidth is replaced in tests, which have no use for it. + measureBandwidth = discovery.MeasureMemoryBandwidth +) diff --git a/internal/agent/main_test.go b/internal/agent/main_test.go new file mode 100644 index 0000000..76842a1 --- /dev/null +++ b/internal/agent/main_test.go @@ -0,0 +1,13 @@ +package agent + +import ( + "os" + "testing" +) + +func TestMain(m *testing.M) { + // Half a gigabyte read per test binary, and contending with the tests + // for the CPUs: these never look at the figure. + measureBandwidth = func() float64 { return 0 } + os.Exit(m.Run()) +} diff --git a/internal/control/agentsvc/heartbeat_test.go b/internal/control/agentsvc/heartbeat_test.go index 9aee27a..a21821c 100644 --- a/internal/control/agentsvc/heartbeat_test.go +++ b/internal/control/agentsvc/heartbeat_test.go @@ -325,3 +325,18 @@ func TestAProblemOnANodeWithoutGPUsDoesNotPanic(t *testing.T) { t.Fatal("the node is gone") } } + +func TestAHeartbeatCarriesTheMemorysMeasuredSpeed(t *testing.T) { + // Measured in the background after the agent starts, so it arrives in a + // heartbeat, not in the registration. + s := newTestService() + nodeWith(s, inventory.Node{ID: "n1", State: inventory.Ready, MemoryTotalMB: 7 * 1024}) + s.applyHeartbeat("n1", &llmgwv1.Heartbeat{NodeId: "n1", MemBandwidthGbs: 41.9}) + if n, _ := s.inv.Get("n1"); n.Host.MemBandwidthGBs != 41.9 { + t.Errorf("bandwidth %v, want 41.9", n.Host.MemBandwidthGBs) + } + s.applyHeartbeat("n1", &llmgwv1.Heartbeat{NodeId: "n1"}) + if n, _ := s.inv.Get("n1"); n.Host.MemBandwidthGBs != 41.9 { + t.Errorf("a heartbeat without it erased it: %v", n.Host.MemBandwidthGBs) + } +} diff --git a/internal/control/agentsvc/service.go b/internal/control/agentsvc/service.go index 7bfeb3f..d3dafe7 100644 --- a/internal/control/agentsvc/service.go +++ b/internal/control/agentsvc/service.go @@ -547,10 +547,11 @@ func (s *Service) applyHeartbeat(nodeID string, hb *llmgwv1.Heartbeat) { storage := hb.GetStorage() memMB := int(hb.GetMemoryTotalMb()) cpus := int(hb.GetCpuCount()) + bw := hb.GetMemBandwidthGbs() if len(used) == 0 && len(foreign) == 0 && len(problems) == 0 && len(storage) == 0 && !hb.GetWeightsChanged() && memMB == 0 && cpus == 0 && !hb.GetEnginesChanged() && !hb.GetProvisionChanged() && hb.GetLoad() == nil && len(hb.GetIbPorts()) == 0 && - !hb.GetFabricsChanged() && len(hb.GetGpuFabric()) == 0 && hb.GetImex() == nil { + !hb.GetFabricsChanged() && len(hb.GetGpuFabric()) == 0 && hb.GetImex() == nil && bw == 0 { return } @@ -622,6 +623,9 @@ func (s *Service) applyHeartbeat(nodeID string, hb *llmgwv1.Heartbeat) { "from_mb", n.MemoryTotalMB, "to_mb", memMB) n.MemoryTotalMB = memMB } + if bw > 0 { + n.Host.MemBandwidthGBs = bw + } if cpus > 0 && cpus != n.CPUCount { if n.CPUCount > 0 { s.log.Info("node CPUs changed", "node", nodeID, "from", n.CPUCount, "to", cpus)