llmgw/internal/control/agentsvc/heartbeat_test.go

package agentsvc

import (
	"log/slog"
	"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{},
		log: slog.New(slog.DiscardHandler)}
}

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

func TestAHeartbeatCarriesMemoryAddedToARunningNode(t *testing.T) {
	// Read once at registration, a node that was given 32 GB went on looking
	// like the 7 GB machine it had been, and every model that needed the
	// difference stayed unplaced until somebody restarted its agent.
	s := newTestService()
	nodeWith(s, inventory.Node{ID: "n1", State: inventory.Ready, MemoryTotalMB: 7 * 1024})

	s.applyHeartbeat("n1", &llmgwv1.Heartbeat{NodeId: "n1", MemoryTotalMb: 32 * 1024})

	if n, _ := s.inv.Get("n1"); n.MemoryTotalMB != 32*1024 {
		t.Errorf("memory is %d MB, expected the 32 GB the node now has", n.MemoryTotalMB)
	}

	// Zero is an agent that could not read it, not a node that lost its RAM.
	s.applyHeartbeat("n1", &llmgwv1.Heartbeat{NodeId: "n1"})
	if n, _ := s.inv.Get("n1"); n.MemoryTotalMB != 32*1024 {
		t.Errorf("a heartbeat without memory erased it: %d", n.MemoryTotalMB)
	}
}

func TestAHeartbeatCarriesAnEngineInstalledAfterRegistration(t *testing.T) {
	s := newTestService()
	nodeWith(s, inventory.Node{ID: "n1", State: inventory.Ready,
		EngineVersions: map[string]string{"llamacpp": "unknown"}})

	// Unchanged engines are not resent, and silence must not clear them.
	s.applyHeartbeat("n1", &llmgwv1.Heartbeat{NodeId: "n1", MemoryTotalMb: 1024})
	if n, _ := s.inv.Get("n1"); len(n.EngineVersions) != 1 {
		t.Fatalf("engines changed without being reported: %v", n.EngineVersions)
	}

	s.applyHeartbeat("n1", &llmgwv1.Heartbeat{
		NodeId: "n1", EnginesChanged: true,
		EngineVersions: map[string]string{"llamacpp": "unknown", "vllm": "unknown"},
	})
	n, _ := s.inv.Get("n1")
	if _, ok := n.EngineVersions["vllm"]; !ok {
		t.Errorf("the newly installed engine is not known: %v", n.EngineVersions)
	}
}

func TestALogRequestIsAnsweredByTheNodeItWasSentTo(t *testing.T) {
	s := newTestService()
	c := &conn{nodeID: "n1", send: make(chan *llmgwv1.ControlMessage, 4), closed: make(chan struct{})}
	s.conns["n1"] = c

	type result struct {
		lines, sources []string
		err            error
	}
	done := make(chan result, 1)
	go func() {
		l, src, err := s.Logs(t.Context(), "n1", "agent", 10)
		done <- result{l, src, err}
	}()

	req := (<-c.send).GetLogs()
	if req.GetSource() != "agent" || req.GetLines() != 10 {
		t.Fatalf("sent %v", req)
	}
	s.handle("n1", &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_Logs{Logs: &llmgwv1.LogsReply{
		RequestId: req.GetRequestId(), Lines: []string{"hello"}, Sources: []string{"agent", "shard-1"},
	}}})

	r := <-done
	if r.err != nil || len(r.lines) != 1 || r.lines[0] != "hello" || len(r.sources) != 2 {
		t.Errorf("got %+v", r)
	}

	if _, _, err := s.Logs(t.Context(), "nobody", "agent", 10); err == nil {
		t.Error("a node with no stream answered")
	}
}

func TestAHeartbeatCarriesHowBusyTheNodeIs(t *testing.T) {
	s := newTestService()
	nodeWith(s, inventory.Node{ID: "n1", State: inventory.Ready})

	s.applyHeartbeat("n1", &llmgwv1.Heartbeat{NodeId: "n1", Load: &llmgwv1.NodeLoad{
		CpuPercent: 62.5, CpuCount: 8, Load1: 3.2, MemoryUsedBytes: 6 << 30, MemoryAvailableBytes: 2 << 30,
		Shards: []*llmgwv1.ShardLoad{{ShardId: "r1-0", CpuCores: 4.5, MemoryRssBytes: 5 << 30}},
	}})
	n, _ := s.inv.Get("n1")
	if n.Load == nil || n.Load.CPUPercent != 62.5 || n.Load.MemUsedBytes != 6<<30 ||
		n.Load.Shards["r1-0"].CPUCores != 4.5 {
		t.Fatalf("load = %+v", n.Load)
	}

	// A heartbeat that measured nothing — the agent's first — leaves the last
	// measurement standing rather than blanking it.
	s.applyHeartbeat("n1", &llmgwv1.Heartbeat{NodeId: "n1"})
	if n, _ := s.inv.Get("n1"); n.Load == nil {
		t.Error("a heartbeat without load erased it")
	}
}

func TestAPortGoingDownTakesItsFabricAway(t *testing.T) {
	// Checked only at registration, a port that went down left the node on
	// its fabric until the agent restarted, still offered multi-node work.
	s := newTestService()
	nodeWith(s, inventory.Node{ID: "n1", State: inventory.Ready,
		Fabrics: []inventory.FabricAttachment{{Fabric: "ib-a", HCAs: []string{"mlx5_0"}, RateGbps: 400}}})

	s.applyHeartbeat("n1", &llmgwv1.Heartbeat{NodeId: "n1", FabricsChanged: true,
		IbPorts: []*llmgwv1.IBPort{{Hca: "mlx5_0", Port: 1, State: "DOWN", PhysState: "Polling",
			Counters: map[string]uint64{"link_downed": 4}}}})
	n, _ := s.inv.Get("n1")
	if len(n.Fabrics) != 0 {
		t.Errorf("the fabric should be gone: %+v", n.Fabrics)
	}
	if len(n.IBPorts) != 1 || n.IBPorts[0].Active() || n.IBPorts[0].Counters["link_downed"] != 4 {
		t.Errorf("ports = %+v", n.IBPorts)
	}

	// Back up: the fabric returns.
	s.applyHeartbeat("n1", &llmgwv1.Heartbeat{NodeId: "n1", FabricsChanged: true,
		Fabrics: []*llmgwv1.FabricAttachment{{Fabric: "ib-a", Hcas: []string{"mlx5_0"}, RateGbps: 400}}})
	if n, _ := s.inv.Get("n1"); !n.HasFabric("ib-a") {
		t.Error("the fabric should be back")
	}

	// A heartbeat that does not say the fabrics changed leaves them alone.
	s.applyHeartbeat("n1", &llmgwv1.Heartbeat{NodeId: "n1", IbPorts: []*llmgwv1.IBPort{{Hca: "mlx5_0", Port: 1, State: "ACTIVE"}}})
	if n, _ := s.inv.Get("n1"); !n.HasFabric("ib-a") {
		t.Error("a heartbeat without a fabric change emptied the fabrics")
	}
}

func TestAGPUFallingOutOfItsDomainTakesTheNodeOut(t *testing.T) {
	// A fabric manager restarting drops cards out of the domain; the node has
	// to stop being offered work that spans it until they are back.
	s := newTestService()
	fab := &discovery.GPUFabric{State: "Completed", ClusterUUID: "76b1ad76-7c24", CliqueID: "1"}
	nodeWith(s, inventory.Node{ID: "t1", State: inventory.Ready, GPUs: []discovery.GPU{
		{Index: 0, Fabric: fab}, {Index: 1, Fabric: fab}}})
	if n, _ := s.inv.Get("t1"); n.NVLinkDomain() == "" {
		t.Fatal("the node should start in its domain")
	}
	s.applyHeartbeat("t1", &llmgwv1.Heartbeat{NodeId: "t1",
		GpuFabric: []*llmgwv1.GPUFabricState{{Index: 1, State: "In Progress"}},
		Imex:      &llmgwv1.IMEX{State: "active", Channels: 1}})
	n, _ := s.inv.Get("t1")
	if n.NVLinkDomain() != "" {
		t.Errorf("a node with a card out of the domain is still in it")
	}
	if n.IMEX.State != "active" || n.IMEX.Channels != 1 {
		t.Errorf("imex = %+v", n.IMEX)
	}
}

func TestAnUpgradeWaitsForTheAgentsAnswer(t *testing.T) {
	s := newTestService()
	c := &conn{nodeID: "n1", send: make(chan *llmgwv1.ControlMessage, 4), closed: make(chan struct{})}
	s.conns["n1"] = c
	done := make(chan error, 1)
	go func() { done <- s.Upgrade(t.Context(), "n1", "v1.2.3", "https://r") }()

	req := (<-c.send).GetUpgrade()
	if req.GetVersion() != "v1.2.3" || req.GetReleases() != "https://r" {
		t.Fatalf("sent %v", req)
	}
	s.handle("n1", &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_Upgrade{Upgrade: &llmgwv1.UpgradeReply{
		RequestId: req.GetRequestId(), Version: "v1.2.3", Error: "checksum mismatch",
	}}})
	if err := <-done; err == nil || err.Error() != "checksum mismatch" {
		t.Errorf("the agent's refusal was not passed on: %v", err)
	}
}

func TestAnotherTenantOnACardIsJudgedEveryHeartbeat(t *testing.T) {
	// Found on cuda01: ComfyUI held 5 GB when the agent started, so the card
	// was kept out of placement. ComfyUI let go of it and the card stayed
	// out, saying "5014 MiB already allocated by another process", until the
	// agent was restarted.
	s := newTestService()
	nodeWith(s, inventory.Node{
		ID: "n1", State: inventory.Ready,
		GPUs: []discovery.GPU{{Index: 0, Health: discovery.GPUHealth{Problems: []string{
			discovery.ForeignMemoryProblem(5014), "2 uncorrected ECC errors"}}}},
	})

	s.applyHeartbeat("n1", &llmgwv1.Heartbeat{NodeId: "n1", GpuForeignMemoryMb: []int32{330}})
	n, _ := s.inv.Get("n1")
	if got := n.GPUs[0].Health.Problems; len(got) != 1 || got[0] != "2 uncorrected ECC errors" {
		t.Fatalf("problems are %q: the tenant left, and only the ECC errors should remain", got)
	}

	// And the other way: a tenant arriving after registration.
	s.applyHeartbeat("n1", &llmgwv1.Heartbeat{NodeId: "n1", GpuForeignMemoryMb: []int32{12000}})
	n, _ = s.inv.Get("n1")
	if got := n.GPUs[0].Health.Problems; len(got) != 2 || got[1] != discovery.ForeignMemoryProblem(12000) {
		t.Fatalf("problems are %q, want the new tenant's 12000 MiB", got)
	}

	// An agent that sends no reading changes nothing.
	s.applyHeartbeat("n1", &llmgwv1.Heartbeat{NodeId: "n1", GpuMemoryUsedMb: []int32{100}})
	n, _ = s.inv.Get("n1")
	if len(n.GPUs[0].Health.Problems) != 2 {
		t.Errorf("a heartbeat with no foreign reading cleared the card: %q", n.GPUs[0].Health.Problems)
	}
}

func TestAProblemOnANodeWithoutGPUsDoesNotPanic(t *testing.T) {
	s := newTestService()
	nodeWith(s, inventory.Node{ID: "cpu1", State: inventory.Ready})
	s.applyHeartbeat("cpu1", &llmgwv1.Heartbeat{NodeId: "cpu1", NewProblems: []string{"disk read errors"}})
	if _, ok := s.inv.Get("cpu1"); !ok {
		t.Fatal("the node is gone")
	}
}