llmgw/internal/agent/agent.go

// Package agent is the node side of the control plane stream (plan §5.6).
package agent

import (
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"math/rand/v2"
	"net"
	"net/http"
	"sync"
	"time"

	"golang.org/x/net/http2"

	llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
	"gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1/llmgwv1connect"
	"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/agentsvc"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
)

// Config configures the agent.
type Config struct {
	NodeID   string            `koanf:"node_id"`
	Server   string            `koanf:"server"`
	Hostname string            `koanf:"hostname"`
	MgmtAddr string            `koanf:"mgmt_addr"`
	Labels   map[string]string `koanf:"labels"`

	// MgmtInterface is the ethernet interface NCCL bootstraps over.
	MgmtInterface string `koanf:"mgmt_interface"`
	// Fabrics labels this node's HCAs. It is explicit on purpose: OpenSM's
	// default subnet prefix is the same on every fabric, so nothing can be
	// inferred from the wire (plan §5.2).
	Fabrics []inventory.FabricAttachment `koanf:"fabrics"`

	HeartbeatInterval time.Duration `koanf:"heartbeat_interval"`
	// ReconnectBase is the first reconnect delay; it backs off from there.
	ReconnectBase time.Duration `koanf:"reconnect_base"`
	ReconnectMax  time.Duration `koanf:"reconnect_max"`

	Discovery discovery.Options `koanf:"discovery"`
	Engines   engines.Config    `koanf:"engines"`
	Superv    supervisor.Config `koanf:"supervisor"`

	Version string
}

// DefaultConfig returns sensible defaults.
func DefaultConfig() Config {
	return Config{
		MgmtInterface:     "eth0",
		HeartbeatInterval: 5 * time.Second,
		ReconnectBase:     time.Second,
		ReconnectMax:      30 * time.Second,
		Superv:            supervisor.DefaultConfig(),
		// NVML gives NVLink topology, so prefer it when this build has it.
		Discovery: discovery.Options{PreferNVML: true},
	}
}

func (c *Config) normalize() error {
	d := DefaultConfig()
	if c.NodeID == "" {
		return errors.New("node_id is required")
	}
	if c.Server == "" {
		return errors.New("server is required")
	}
	if c.MgmtInterface == "" {
		c.MgmtInterface = d.MgmtInterface
	}
	if c.HeartbeatInterval <= 0 {
		c.HeartbeatInterval = d.HeartbeatInterval
	}
	if c.ReconnectBase <= 0 {
		c.ReconnectBase = d.ReconnectBase
	}
	if c.ReconnectMax < c.ReconnectBase {
		c.ReconnectMax = max(c.ReconnectBase, d.ReconnectMax)
	}
	return nil
}

// Agent runs on a GPU node.
type Agent struct {
	cfg  Config
	log  *slog.Logger
	sup  *supervisor.Supervisor
	reg  *engines.Registry
	http connectHTTPClient

	mu  sync.Mutex
	inv discovery.Inventory
	// shardReplica remembers which replica a shard belongs to, so a status
	// report can name both. The control plane keys on the replica.
	shardReplica map[string]string
}

// connectHTTPClient is the HTTP client connect uses; an interface so tests can
// hand in an httptest server's client.
type connectHTTPClient interface {
	Do(*http.Request) (*http.Response, error)
}

// New builds an agent.
func New(cfg Config, httpClient connectHTTPClient, log *slog.Logger) (*Agent, error) {
	if err := cfg.normalize(); err != nil {
		return nil, err
	}
	if log == nil {
		log = slog.Default()
	}
	if httpClient == nil {
		httpClient = http.DefaultClient
	}

	health := supervisor.NewHTTPHealth(3 * time.Second)
	sup, err := supervisor.New(cfg.Superv, health, log)
	if err != nil {
		return nil, err
	}
	return &Agent{
		cfg: cfg, log: log, sup: sup,
		reg:          engines.NewRegistry(cfg.Engines),
		http:         httpClient,
		shardReplica: map[string]string{},
	}, nil
}

// Supervisor exposes the process supervisor, for tests and `ctl logs`.
func (a *Agent) Supervisor() *supervisor.Supervisor { return a.sup }

// Run connects to the control plane and stays connected until ctx ends.
//
// Reconnecting forever is the right behaviour for a node agent: a control plane
// restart or a network partition must not require touching every GPU node, and
// the engines keep serving traffic throughout because the router talks to them
// directly.
func (a *Agent) Run(ctx context.Context) error {
	if err := a.scan(ctx); err != nil {
		// Not fatal: a node whose GPUs cannot be read should still register, so
		// an operator can see it and why it is useless.
		a.log.Error("hardware discovery failed", "err", err)
	}

	attempt := 0
	for ctx.Err() == nil {
		err := a.session(ctx)
		if ctx.Err() != nil {
			break
		}
		attempt++
		delay := a.reconnectDelay(attempt)
		a.log.Warn("disconnected from the control plane, retrying",
			"err", err, "in", delay.Round(time.Millisecond), "attempt", attempt)

		select {
		case <-ctx.Done():
		case <-time.After(delay):
		}
		if err == nil {
			attempt = 0
		}
	}

	// Engines keep running when the agent stops on purpose: an agent restart
	// must not interrupt served traffic. Orphan adoption picks them up again.
	return ctx.Err()
}

func (a *Agent) reconnectDelay(attempt int) time.Duration {
	d := a.cfg.ReconnectBase << min(attempt-1, 16)
	if d > a.cfg.ReconnectMax || d <= 0 {
		d = a.cfg.ReconnectMax
	}
	// Jitter, so a fleet does not reconnect in lockstep after a control plane
	// restart and knock it over again.
	return d/2 + time.Duration(rand.Int64N(int64(d/2))) //nolint:gosec // jitter
}

// scan refreshes the node's hardware inventory.
func (a *Agent) scan(ctx context.Context) error {
	inv, err := discovery.Discover(ctx, a.cfg.Discovery, a.log)
	a.mu.Lock()
	a.inv = inv
	a.mu.Unlock()
	return err
}

// session runs one connection until it breaks.
func (a *Agent) session(ctx context.Context) error {
	client := llmgwv1connect.NewAgentServiceClient(a.http, a.cfg.Server)

	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	stream := client.Connect(ctx)
	defer func() { _ = stream.CloseRequest() }()

	a.mu.Lock()
	inv := a.inv
	a.mu.Unlock()

	if err := stream.Send(&llmgwv1.AgentMessage{
		Msg: &llmgwv1.AgentMessage_Register{
			Register: agentsvc.RegisterFromInventory(
				a.cfg.NodeID, a.cfg.Hostname, a.cfg.MgmtAddr, a.cfg.Version,
				a.cfg.Labels, inv, a.cfg.Fabrics),
		},
	}); err != nil {
		return fmt.Errorf("registering: %w", err)
	}

	// Sends are serialized through this channel: connect streams are not safe
	// for concurrent writes, and heartbeats race with status reports.
	out := make(chan *llmgwv1.AgentMessage, 32)
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		a.sendLoop(ctx, stream, out)
	}()
	defer wg.Wait()

	for {
		msg, err := stream.Receive()
		if err != nil {
			cancel()
			if errors.Is(err, io.EOF) || ctx.Err() != nil {
				return nil
			}
			return err
		}
		a.handle(ctx, msg, out)
	}
}

type sender interface {
	Send(*llmgwv1.AgentMessage) error
}

func (a *Agent) sendLoop(ctx context.Context, s sender, out <-chan *llmgwv1.AgentMessage) {
	t := time.NewTicker(a.cfg.HeartbeatInterval)
	defer t.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case msg := <-out:
			if err := s.Send(msg); err != nil {
				return
			}
		case <-t.C:
			if err := s.Send(a.heartbeat()); err != nil {
				return
			}
		}
	}
}

func (a *Agent) heartbeat() *llmgwv1.AgentMessage {
	a.mu.Lock()
	inv := a.inv
	a.mu.Unlock()

	hb := &llmgwv1.Heartbeat{NodeId: a.cfg.NodeID}
	for _, g := range inv.GPUs {
		hb.GpuMemoryUsedMb = append(hb.GpuMemoryUsedMb, int32(g.MemoryUsedMB))
	}
	return &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_Heartbeat{Heartbeat: hb}}
}

func (a *Agent) handle(ctx context.Context, msg *llmgwv1.ControlMessage, out chan<- *llmgwv1.AgentMessage) {
	switch m := msg.Msg.(type) {
	case *llmgwv1.ControlMessage_Start:
		go a.startShard(ctx, m.Start, out)
	case *llmgwv1.ControlMessage_Stop:
		go a.stopShard(ctx, m.Stop.GetShardId(), out)
	case *llmgwv1.ControlMessage_Drain:
		if err := a.sup.Drain(m.Drain.GetShardId()); err != nil {
			a.log.Warn("draining failed", "shard", m.Drain.GetShardId(), "err", err)
		}
	case *llmgwv1.ControlMessage_Ack:
		if !m.Ack.GetOk() {
			a.log.Warn("control plane rejected a message",
				"id", m.Ack.GetMessageId(), "err", m.Ack.GetError())
		}
	}
}

// startShard launches a shard and reports its progress.
//
// It runs in its own goroutine because model loading takes minutes, and the
// stream has to keep carrying heartbeats and other commands meanwhile.
func (a *Agent) startShard(ctx context.Context, m *llmgwv1.StartShard, out chan<- *llmgwv1.AgentMessage) {
	shardID, replicaID := m.GetShardId(), m.GetReplicaId()

	a.mu.Lock()
	a.shardReplica[shardID] = replicaID
	a.mu.Unlock()

	report := func(st supervisor.State, errMsg string, tail []string) {
		msg := &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_Shard{Shard: &llmgwv1.ShardStatus{
			NodeId: a.cfg.NodeID, ReplicaId: replicaID, ShardId: shardID,
			State: agentsvc.StateToWire(st), Error: errMsg, LastOutput: tail,
		}}}
		select {
		case out <- msg:
		case <-ctx.Done():
		}
	}

	spec, timeout := agentsvc.ShardSpecFromStart(m)

	eng, err := a.reg.Get(spec.Profile.Engine)
	if err != nil {
		report(supervisor.Failed, err.Error(), nil)
		return
	}

	a.mu.Lock()
	inv := a.inv
	a.mu.Unlock()

	// Validate before launching. An engine that refuses a profile does so after
	// minutes of loading; catching it here costs nothing.
	if err := eng.Validate(spec.Profile, gpuInfo(inv, spec.GPUIndices)); err != nil {
		report(supervisor.Failed, err.Error(), nil)
		return
	}

	cmd, err := eng.BuildCommand(spec)
	if err != nil {
		report(supervisor.Failed, err.Error(), nil)
		return
	}

	sh, err := a.sup.Start(ctx, shardID, spec, cmd)
	if err != nil {
		if errors.Is(err, supervisor.ErrAlreadyRunning) {
			// Commands are idempotent: the control plane retries, and a repeat
			// must not produce a second process (plan §5.6).
			report(sh.Status().State, "", nil)
			return
		}
		report(supervisor.Failed, err.Error(), sh.Logs(10))
		return
	}
	report(supervisor.Starting, "", nil)

	base := fmt.Sprintf("http://127.0.0.1:%d", spec.HTTPPort)
	if err := a.sup.WaitReady(ctx, sh, base, spec.Profile.ServedModelName, timeout); err != nil {
		report(supervisor.Failed, err.Error(), sh.Logs(20))
		return
	}
	report(supervisor.Ready, "", nil)

	// Watch for the engine dying later. Nothing else would notice: the router
	// would simply see failing requests.
	go a.watch(ctx, sh, replicaID, shardID, out)
}

// watch reports a shard that exits after it was ready.
func (a *Agent) watch(ctx context.Context, sh *supervisor.Shard, replicaID, shardID string, out chan<- *llmgwv1.AgentMessage) {
	t := time.NewTicker(time.Second)
	defer t.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-t.C:
			st := sh.Status()
			if !st.State.Terminal() {
				continue
			}
			msg := &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_Shard{Shard: &llmgwv1.ShardStatus{
				NodeId: a.cfg.NodeID, ReplicaId: replicaID, ShardId: shardID,
				State: agentsvc.StateToWire(st.State), Error: st.Err,
				ExitCode: int32(st.ExitCode), LastOutput: sh.Logs(20),
			}}}
			select {
			case out <- msg:
			case <-ctx.Done():
			}
			return
		}
	}
}

func (a *Agent) stopShard(ctx context.Context, shardID string, out chan<- *llmgwv1.AgentMessage) {
	a.mu.Lock()
	replicaID := a.shardReplica[shardID]
	delete(a.shardReplica, shardID)
	a.mu.Unlock()

	if err := a.sup.Stop(ctx, shardID); err != nil {
		a.log.Warn("stopping a shard failed", "shard", shardID, "err", err)
	}
	a.sup.Remove(shardID)

	msg := &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_Shard{Shard: &llmgwv1.ShardStatus{
		NodeId: a.cfg.NodeID, ReplicaId: replicaID, ShardId: shardID,
		State: agentsvc.StateToWire(supervisor.Stopped),
	}}}
	select {
	case out <- msg:
	case <-ctx.Done():
	}
}

// gpuInfo narrows the inventory to the GPUs a shard will use.
func gpuInfo(inv discovery.Inventory, indices []int) []engines.GPUInfo {
	want := map[int]struct{}{}
	for _, i := range indices {
		want[i] = struct{}{}
	}
	var out []engines.GPUInfo
	for _, g := range inv.GPUs {
		if _, ok := want[g.Index]; !ok {
			continue
		}
		out = append(out, engines.GPUInfo{
			Index: g.Index, Name: g.Name,
			ComputeCapability: g.ComputeCapability, MemoryTotalMB: g.MemoryTotalMB,
		})
	}
	return out
}

// HTTP2Client returns the client the agent should use to reach the control
// plane over plaintext HTTP/2.
//
// connect's bidirectional streams require HTTP/2, and an agent talking to a
// control plane over a trusted network has no TLS to negotiate it with. Go's
// default transport will not speak HTTP/2 without TLS, so it has to be asked
// explicitly.
//
// mTLS is plan §5.6's target and replaces this; until then the stream must not
// cross an untrusted network.
func HTTP2Client(timeout time.Duration) *http.Client {
	return &http.Client{
		Transport: &http2.Transport{
			AllowHTTP: true,
			DialTLSContext: func(ctx context.Context, network, addr string, _ *tls.Config) (net.Conn, error) {
				return (&net.Dialer{Timeout: timeout}).DialContext(ctx, network, addr)
			},
		},
	}
}