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"
	"maps"
	"runtime"
	"slices"
	"sync/atomic"
	"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(),
		Features:       r.GetFeatures(),
		EngineVersions: r.GetEngineVersions(),
		EnginesProbed:  r.GetEnginesProbed(),
		MemoryTotalMB:  int(r.GetMemoryTotalMb()),
		CPUCount:       int(r.GetCpuCount()),
	}
	if h := r.GetHost(); h != nil {
		n.Host = discovery.Host{
			OSID: h.GetOsId(), OSVersion: h.GetOsVersion(), OSCodename: h.GetOsCodename(),
			Arch: h.GetArch(), CPUFlags: h.GetCpuFlags(), NVIDIADriver: h.GetNvidiaDriver(),
			MemBandwidthGBs: h.GetMemBandwidthGbs(),
		}
	}
	for _, w := range r.GetWeights() {
		n.Weights = append(n.Weights, discovery.LocalWeights{
			Path: w.GetPath(), SizeBytes: w.GetSizeBytes(), Shared: w.GetShared(),
		})
	}
	for _, s := range r.GetStorage() {
		n.Storage = append(n.Storage, discovery.Storage{
			Path:       s.GetPath(),
			FreeBytes:  s.GetFreeBytes(),
			TotalBytes: s.GetTotalBytes(),
			Shared:     s.GetShared(),
		})
	}
	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(),
			UtilizationPercent: int(g.GetUtilizationPercent()),
			CUDAVersion:        g.GetCudaVersion(),
			NVLinkPeers:        intSlice(g.GetNvlinkPeers()),
			NICPaths:           maps.Clone(g.GetNicPaths()),
			Fabric: fabricFrom(g.GetFabricState(), g.GetFabricStatus(), g.GetClusterUuid(),
				g.GetCliqueId(), g.GetFabricBandwidth()),
			Health: discovery.GPUHealth{
				ECCErrors:    int(g.GetEccErrors()),
				TemperatureC: int(g.GetTemperatureC()),
				Problems:     g.GetProblems(),
			},
			MIG: migFrom(g.GetMig()),
			PCIe: discovery.PCIeLink{Gen: int(g.GetPcieGen()), GPUGen: int(g.GetPcieGpuGen()),
				Width: int(g.GetPcieWidth()), WidthMax: int(g.GetPcieWidthMax())},
		})
	}
	n.IMEX = discovery.IMEX{State: r.GetImex().GetState(), Channels: int(r.GetImex().GetChannels())}
	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),
		CpuCount:      int32(runtime.NumCPU()),
	}
	for _, w := range inv.Weights {
		r.Weights = append(r.Weights, &llmgwv1.LocalWeights{
			Path: w.Path, SizeBytes: w.SizeBytes, Shared: w.Shared,
		})
	}
	for _, s := range inv.Storage {
		r.Storage = append(r.Storage, &llmgwv1.Storage{
			Path:       s.Path,
			FreeBytes:  s.FreeBytes,
			TotalBytes: s.TotalBytes,
			Shared:     s.Shared,
		})
	}
	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,
			UtilizationPercent: int32(g.UtilizationPercent),
			CudaVersion:        g.CUDAVersion,
			NvlinkPeers:        int32Slice(g.NVLinkPeers),
			NicPaths:           maps.Clone(g.NICPaths),
			Problems:           g.Health.Problems,
			TemperatureC:       int32(g.Health.TemperatureC),
			EccErrors:          int32(g.Health.ECCErrors),
			PcieGen:            int32(g.PCIe.Gen),
			PcieGpuGen:         int32(g.PCIe.GPUGen),
			PcieWidth:          int32(g.PCIe.Width),
			PcieWidthMax:       int32(g.PCIe.WidthMax),
		})
		if f := g.Fabric; f != nil {
			last := r.Gpus[len(r.Gpus)-1]
			last.FabricState, last.FabricStatus, last.ClusterUuid = f.State, f.Status, f.ClusterUUID
			last.CliqueId, last.FabricBandwidth = f.CliqueID, f.Bandwidth
		}
		if m := g.MIG; m != nil {
			r.Gpus[len(r.Gpus)-1].Mig = &llmgwv1.MIGDevice{
				ParentIndex: int32(m.ParentIndex), ParentUuid: m.ParentUUID,
				ParentMemoryTotalMb: int32(m.ParentMemoryTotalMB), Device: int32(m.Device),
				Profile: m.Profile, GpuInstance: int32(m.GPUInstance), ComputeInstance: int32(m.ComputeInstance),
			}
		}
	}
	r.Imex = &llmgwv1.IMEX{State: inv.IMEX.State, Channels: int32(inv.IMEX.Channels)}
	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(),
			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())
		}
		out = append(out, a)
	}
	return out
}

// StartShardFromCommand converts a reconciler command into a wire message.
// sourceRewrite is where a model is fetched while its mirror is down: set
// when mirrors are kept in step, else nil and the catalog's address as it is.
var sourceRewrite atomic.Pointer[func(string) string]

// SetSourceRewrite has every shard started from now on fetch its model, and
// adapters, from the address f gives: one on a mirror that answers.
func SetSourceRewrite(f func(string) string) { sourceRewrite.Store(&f) }

func rewriteSource(src string) string {
	if f := sourceRewrite.Load(); f != nil && src != "" {
		return (*f)(src)
	}
	return src
}

func StartShardFromCommand(cmd reconciler.ShardCommand) *llmgwv1.StartShard {
	p := cmd.Spec.Profile
	p.ModelPath = rewriteSource(p.ModelPath)
	if len(p.Adapters) > 0 {
		ads := slices.Clone(p.Adapters)
		for i := range ads {
			ads[i].Path = rewriteSource(ads[i].Path)
		}
		p.Adapters = ads
	}
	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,
			Adapters:    adaptersToWire(p.Adapters),
			KvEvents:    p.KVEvents,
			AllowPickle: p.AllowPickle,
			ModelSha256: p.ModelSHA256,
		},
		StartupTimeoutMs: cmd.StartupTimeout.Milliseconds(),
		Pool:             cmd.Pool,
		Fabric:           cmd.Fabric,
		ProfileFp:        cmd.ProfileFP,
		ApiKey:           cmd.Spec.APIKey,
		CpuThreads:       int32(cmd.Spec.Threads),
	}
}

// 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(),
		ProfileFP:     m.GetProfileFp(),
		APIKey:        m.GetApiKey(),
		Threads:       int(m.GetCpuThreads()),
		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(),
			Adapters:    adaptersFromWire(p.GetAdapters()),
			KVEvents:    p.GetKvEvents(),
			AllowPickle: p.GetAllowPickle(),
			ModelSHA256: p.GetModelSha256(),
		},
	}
	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
}

// migFrom is a MIG instance's parentage as the wire carries it, nil for a
// whole card and from every agent older than MIG support.
func migFrom(m *llmgwv1.MIGDevice) *discovery.MIGDevice {
	if m == nil {
		return nil
	}
	return &discovery.MIGDevice{
		ParentIndex: int(m.GetParentIndex()), ParentUUID: m.GetParentUuid(),
		ParentMemoryTotalMB: int(m.GetParentMemoryTotalMb()), Device: int(m.GetDevice()),
		Profile: m.GetProfile(), GPUInstance: int(m.GetGpuInstance()), ComputeInstance: int(m.GetComputeInstance()),
	}
}

// fabricFrom is a GPU's NVLink fabric as the wire carries it, nil when the
// driver reported none.
func fabricFrom(state, status, cluster, clique, bandwidth string) *discovery.GPUFabric {
	if state == "" {
		return nil
	}
	return &discovery.GPUFabric{State: state, Status: status, ClusterUUID: cluster, CliqueID: clique, Bandwidth: bandwidth}
}

func adaptersToWire(as []engines.Adapter) []*llmgwv1.LoraAdapter {
	var out []*llmgwv1.LoraAdapter
	for _, a := range as {
		out = append(out, &llmgwv1.LoraAdapter{Name: a.Name, Source: a.Path})
	}
	return out
}

func adaptersFromWire(as []*llmgwv1.LoraAdapter) []engines.Adapter {
	var out []engines.Adapter
	for _, a := range as {
		out = append(out, engines.Adapter{Name: a.GetName(), Path: a.GetSource()})
	}
	return out
}