llmgw/internal/control/agentsvc/convert.go

// Package agentsvc is the control plane's end of the agent stream (plan §5.6).
package agentsvc

import (
	"fmt"
	"time"

	llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
)

// NodeFromRegister converts a registration into an inventory node.
func NodeFromRegister(r *llmgwv1.Register) inventory.Node {
	n := inventory.Node{
		ID:             r.GetNodeId(),
		Hostname:       r.GetHostname(),
		MgmtAddr:       r.GetMgmtAddr(),
		Labels:         r.GetLabels(),
		State:          inventory.Ready,
		AgentVersion:   r.GetAgentVersion(),
		EngineVersions: r.GetEngineVersions(),
		MemoryTotalMB:  int(r.GetMemoryTotalMb()),
	}
	for _, g := range r.GetGpus() {
		n.GPUs = append(n.GPUs, discovery.GPU{
			Index:             int(g.GetIndex()),
			UUID:              g.GetUuid(),
			Name:              g.GetName(),
			MemoryTotalMB:     int(g.GetMemoryTotalMb()),
			MemoryUsedMB:      int(g.GetMemoryUsedMb()),
			ComputeCapability: g.GetComputeCapability(),
			NUMANode:          int(g.GetNumaNode()),
			PCIBusID:          g.GetPciBusId(),
			DriverVersion:     g.GetDriverVersion(),
			CUDAVersion:       g.GetCudaVersion(),
			NVLinkPeers:       intSlice(g.GetNvlinkPeers()),
			Health: discovery.GPUHealth{
				ECCErrors:    int(g.GetEccErrors()),
				TemperatureC: int(g.GetTemperatureC()),
				Problems:     g.GetProblems(),
			},
		})
	}
	for _, f := range r.GetFabrics() {
		n.Fabrics = append(n.Fabrics, inventory.FabricAttachment{
			Fabric:   f.GetFabric(),
			HCAs:     f.GetHcas(),
			RateGbps: int(f.GetRateGbps()),
		})
	}
	return n
}

// RegisterFromInventory builds the registration an agent sends.
func RegisterFromInventory(nodeID, hostname, mgmtAddr, agentVersion string,
	labels map[string]string, inv discovery.Inventory, fabrics []inventory.FabricAttachment,
) *llmgwv1.Register {
	r := &llmgwv1.Register{
		NodeId: nodeID, Hostname: hostname, MgmtAddr: mgmtAddr,
		Labels: labels, AgentVersion: agentVersion,
		MemoryTotalMb: int64(inv.MemoryTotalMB),
	}
	for _, g := range inv.GPUs {
		r.Gpus = append(r.Gpus, &llmgwv1.GPU{
			Index:             int32(g.Index),
			Uuid:              g.UUID,
			Name:              g.Name,
			MemoryTotalMb:     int32(g.MemoryTotalMB),
			MemoryUsedMb:      int32(g.MemoryUsedMB),
			ComputeCapability: g.ComputeCapability,
			NumaNode:          int32(g.NUMANode),
			PciBusId:          g.PCIBusID,
			DriverVersion:     g.DriverVersion,
			CudaVersion:       g.CUDAVersion,
			NvlinkPeers:       int32Slice(g.NVLinkPeers),
			Problems:          g.Health.Problems,
			TemperatureC:      int32(g.Health.TemperatureC),
			EccErrors:         int32(g.Health.ECCErrors),
		})
	}
	for _, f := range fabrics {
		r.Fabrics = append(r.Fabrics, &llmgwv1.FabricAttachment{
			Fabric: f.Fabric, Hcas: f.HCAs, RateGbps: int32(f.RateGbps),
		})
	}
	return r
}

// AdoptedFromRegister converts an agent's running-shard report into the
// reconciler's view.
func AdoptedFromRegister(nodeID string, shards []*llmgwv1.RunningShard) []reconciler.AdoptedShard {
	out := make([]reconciler.AdoptedShard, 0, len(shards))
	for _, s := range shards {
		a := reconciler.AdoptedShard{
			NodeID:     nodeID,
			ReplicaID:  s.GetReplicaId(),
			ShardID:    s.GetShardId(),
			Rank:       int(s.GetRank()),
			GPUIndices: intSlice(s.GetGpuIndices()),
			HCAs:       s.GetHcas(),
			HTTPPort:   int(s.GetHttpPort()),
			State:      StateFromWire(s.GetState()),
			Model:      s.GetModel(),
			Profile:    s.GetProfile(),
			Pool:       s.GetPool(),
			Fabric:     s.GetFabric(),
		}
		if a.Rank == 0 && s.GetLeaderAddr() != "" {
			a.Endpoint = fmt.Sprintf("http://%s:%d", s.GetLeaderAddr(), s.GetHttpPort())
		}
		out = append(out, a)
	}
	return out
}

// StartShardFromCommand converts a reconciler command into a wire message.
func StartShardFromCommand(cmd reconciler.ShardCommand) *llmgwv1.StartShard {
	p := cmd.Spec.Profile
	return &llmgwv1.StartShard{
		ReplicaId:     cmd.ReplicaID,
		ShardId:       cmd.ShardID,
		Rank:          int32(cmd.Spec.Rank),
		Nnodes:        int32(cmd.Spec.NNodes),
		LeaderAddr:    cmd.Spec.LeaderAddr,
		LeaderPort:    int32(cmd.Spec.LeaderPort),
		GpuIndices:    int32Slice(cmd.Spec.GPUIndices),
		Hcas:          cmd.Spec.HCAs,
		MgmtInterface: cmd.Spec.MgmtInterface,
		HttpPort:      int32(cmd.Spec.HTTPPort),
		Profile: &llmgwv1.EngineProfile{
			Name: p.Name, Engine: p.Engine, GpuClass: p.GPUClass,
			Tp: int32(p.TP), Pp: int32(p.PP),
			GpusPerNode: int32(p.GPUsPerNode), Nodes: int32(p.Nodes),
			Args: p.Args, Env: p.Env,
			ModelPath: p.ModelPath, ServedModelName: p.ServedModelName,
		},
		StartupTimeoutMs: cmd.StartupTimeout.Milliseconds(),
		Pool:             cmd.Pool,
		Fabric:           cmd.Fabric,
	}
}

// ShardSpecFromStart converts a wire message into an engine shard spec.
func ShardSpecFromStart(m *llmgwv1.StartShard) (engines.ShardSpec, time.Duration) {
	p := m.GetProfile()
	spec := engines.ShardSpec{
		ReplicaID:     m.GetReplicaId(),
		ShardID:       m.GetShardId(),
		Rank:          int(m.GetRank()),
		NNodes:        int(m.GetNnodes()),
		LeaderAddr:    m.GetLeaderAddr(),
		LeaderPort:    int(m.GetLeaderPort()),
		GPUIndices:    intSlice(m.GetGpuIndices()),
		HCAs:          m.GetHcas(),
		MgmtInterface: m.GetMgmtInterface(),
		HTTPPort:      int(m.GetHttpPort()),
		Pool:          m.GetPool(),
		Fabric:        m.GetFabric(),
		Profile: engines.Profile{
			Name: p.GetName(), Engine: p.GetEngine(), GPUClass: p.GetGpuClass(),
			TP: int(p.GetTp()), PP: int(p.GetPp()),
			GPUsPerNode: int(p.GetGpusPerNode()), Nodes: int(p.GetNodes()),
			Args: p.GetArgs(), Env: p.GetEnv(),
			ModelPath: p.GetModelPath(), ServedModelName: p.GetServedModelName(),
		},
	}
	return spec, time.Duration(m.GetStartupTimeoutMs()) * time.Millisecond
}

// stateWire maps supervisor states onto the wire enum.
var stateWire = map[supervisor.State]llmgwv1.ShardStatus_State{
	supervisor.Pending:  llmgwv1.ShardStatus_STATE_PENDING,
	supervisor.Starting: llmgwv1.ShardStatus_STATE_STARTING,
	supervisor.Loading:  llmgwv1.ShardStatus_STATE_LOADING,
	supervisor.Ready:    llmgwv1.ShardStatus_STATE_READY,
	supervisor.Draining: llmgwv1.ShardStatus_STATE_DRAINING,
	supervisor.Stopped:  llmgwv1.ShardStatus_STATE_STOPPED,
	supervisor.Failed:   llmgwv1.ShardStatus_STATE_FAILED,
}

// StateToWire converts a supervisor state for transmission.
func StateToWire(s supervisor.State) llmgwv1.ShardStatus_State {
	if w, ok := stateWire[s]; ok {
		return w
	}
	return llmgwv1.ShardStatus_STATE_UNSPECIFIED
}

// StateFromWire converts a received state.
//
// An unrecognised state is treated as Failed rather than guessed at: a newer
// agent reporting something this control plane does not understand must not be
// mistaken for a healthy shard.
func StateFromWire(w llmgwv1.ShardStatus_State) supervisor.State {
	for s, v := range stateWire {
		if v == w {
			return s
		}
	}
	return supervisor.Failed
}

func intSlice(v []int32) []int {
	if len(v) == 0 {
		return nil
	}
	out := make([]int, len(v))
	for i, x := range v {
		out[i] = int(x)
	}
	return out
}

func int32Slice(v []int) []int32 {
	if len(v) == 0 {
		return nil
	}
	out := make([]int32, len(v))
	for i, x := range v {
		out[i] = int32(x)
	}
	return out
}