What a shard computes on survives the agent's and the control plane's restarts

It was sent once, with READY, and kept only in the leader's memory: a
failover or an agent restart lost it. The agent now keeps it with the
shard's state on disk (ShardSpec.Compute, set by Supervisor.SetCompute) and
reports it with each running shard when it registers (RunningShard.compute),
which adoption carries onto the replica.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

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

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 {

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

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.

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
 }

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

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

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

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

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.

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