llmgw/internal/control/agentsvc/service.go

package agentsvc

import (
	"context"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"slices"
	"sync"

	"connectrpc.com/connect"

	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/supervisor"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
)

// StatusSink receives shard transitions. The reconciler implements it.
type StatusSink interface {
	OnShardStatus(replicaID, shardID string, state supervisor.State, errMsg string)
}

// Service is the control plane's AgentService implementation, and at the same
// time the reconciler's Dispatcher: it holds the open streams, so it is the
// only thing that can deliver a command to a node.
type Service struct {
	inv    *inventory.Store
	status StatusSink
	log    *slog.Logger

	mu    sync.RWMutex
	conns map[string]*conn
}

// conn is one agent's open stream.
type conn struct {
	nodeID string
	// send is serialized through a channel rather than a mutex around the
	// stream: connect's streams are not safe for concurrent writes, and the
	// reconciler dispatches to many nodes at once.
	send   chan *llmgwv1.ControlMessage
	closed chan struct{}
	once   sync.Once
}

func (c *conn) close() { c.once.Do(func() { close(c.closed) }) }

// New builds the service.
func New(inv *inventory.Store, status StatusSink, log *slog.Logger) *Service {
	if log == nil {
		log = slog.Default()
	}
	return &Service{inv: inv, status: status, log: log, conns: map[string]*conn{}}
}

// Handler returns the HTTP path and handler to mount.
func (s *Service) Handler() (string, http.Handler) {
	return llmgwv1connect.NewAgentServiceHandler(s)
}

// Connect serves one agent's stream for the lifetime of that agent.
func (s *Service) Connect(
	ctx context.Context,
	stream *connect.BidiStream[llmgwv1.AgentMessage, llmgwv1.ControlMessage],
) error {
	first, err := stream.Receive()
	if err != nil {
		return connect.NewError(connect.CodeInvalidArgument, fmt.Errorf("no opening message: %w", err))
	}
	reg := first.GetRegister()
	if reg == nil {
		// Registration first, always: the control plane has to know which node
		// it is talking to before it can route anything to it.
		return connect.NewError(connect.CodeInvalidArgument,
			errors.New("the first message must be a Register"))
	}
	nodeID := reg.GetNodeId()
	if nodeID == "" {
		return connect.NewError(connect.CodeInvalidArgument, errors.New("node_id is required"))
	}

	node := NodeFromRegister(reg)
	s.inv.Upsert(node)
	s.log.Info("agent registered",
		"node", nodeID, "gpus", len(node.GPUs), "fabrics", node.FabricNames(),
		"agent_version", reg.GetAgentVersion())

	c := s.attach(ctx, nodeID)
	defer s.detach(nodeID, c)

	if err := stream.Send(&llmgwv1.ControlMessage{
		Msg: &llmgwv1.ControlMessage_Ack{Ack: &llmgwv1.Ack{MessageId: "register", Ok: true}},
	}); err != nil {
		return err
	}

	// One goroutine writes, this one reads. Splitting them is what lets a
	// command be dispatched while the agent is mid-heartbeat.
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		s.writeLoop(ctx, stream, c)
	}()
	defer wg.Wait()

	for {
		msg, err := stream.Receive()
		if err != nil {
			c.close()
			if ctx.Err() != nil || errors.Is(err, io.EOF) {
				s.log.Info("agent disconnected", "node", nodeID)
				return nil
			}
			s.log.Warn("agent stream failed", "node", nodeID, "err", err)
			return nil
		}
		s.handle(nodeID, msg)
	}
}

func (s *Service) writeLoop(ctx context.Context,
	stream *connect.BidiStream[llmgwv1.AgentMessage, llmgwv1.ControlMessage], c *conn,
) {
	for {
		select {
		case <-ctx.Done():
			return
		case <-c.closed:
			return
		case msg := <-c.send:
			if err := stream.Send(msg); err != nil {
				s.log.Warn("sending to agent failed", "node", c.nodeID, "err", err)
				c.close()
				return
			}
		}
	}
}

func (s *Service) handle(nodeID string, msg *llmgwv1.AgentMessage) {
	switch m := msg.Msg.(type) {
	case *llmgwv1.AgentMessage_Heartbeat:
		s.inv.Touch(nodeID)
		s.applyHeartbeat(nodeID, m.Heartbeat)

	case *llmgwv1.AgentMessage_Shard:
		st := StateFromWire(m.Shard.GetState())
		s.status.OnShardStatus(m.Shard.GetReplicaId(), m.Shard.GetShardId(), st, m.Shard.GetError())
		if m.Shard.GetState() == llmgwv1.ShardStatus_STATE_FAILED {
			s.log.Error("shard failed",
				"node", nodeID, "shard", m.Shard.GetShardId(),
				"err", m.Shard.GetError(), "exit_code", m.Shard.GetExitCode(),
				"last_output", m.Shard.GetLastOutput())
		}

	case *llmgwv1.AgentMessage_Register:
		// A re-registration on a live stream means the agent's view of its own
		// hardware changed — a GPU fell out, say. Take it.
		s.inv.Upsert(NodeFromRegister(m.Register))
		s.log.Info("agent re-registered", "node", nodeID)
	}
}

// applyHeartbeat folds in the cheap facts a heartbeat carries, so a foreign
// process appearing on a card after registration is noticed without waiting for
// a reconnect.
func (s *Service) applyHeartbeat(nodeID string, hb *llmgwv1.Heartbeat) {
	used := hb.GetGpuMemoryUsedMb()
	problems := hb.GetNewProblems()
	if len(used) == 0 && len(problems) == 0 {
		return
	}

	n, ok := s.inv.Get(nodeID)
	if !ok {
		return
	}
	for i := range n.GPUs {
		if i < len(used) {
			n.GPUs[i].MemoryUsedMB = int(used[i])
		}
	}
	if len(problems) > 0 {
		// Problems are reported per node, not per GPU, so attach them to the
		// node's first GPU rather than guessing which one is affected. Any
		// problem is enough to keep that GPU out of placement, which is the
		// conservative direction.
		n.GPUs[0].Health.Problems = append(n.GPUs[0].Health.Problems, problems...)
	}
	s.inv.Upsert(n)
}

func (s *Service) attach(ctx context.Context, nodeID string) *conn {
	c := &conn{
		nodeID: nodeID,
		// Buffered so a burst of commands does not block the reconciler on a
		// slow agent; if it fills, the agent is too far behind to be useful.
		send:   make(chan *llmgwv1.ControlMessage, 64),
		closed: make(chan struct{}),
	}

	s.mu.Lock()
	if prev, ok := s.conns[nodeID]; ok {
		// A node reconnecting before the old stream was reaped. The new one
		// wins: the old is by definition the dead one.
		s.log.Warn("agent reconnected while an old stream was open", "node", nodeID)
		prev.close()
	}
	s.conns[nodeID] = c
	s.mu.Unlock()

	go func() {
		<-ctx.Done()
		c.close()
	}()
	return c
}

func (s *Service) detach(nodeID string, c *conn) {
	c.close()
	s.mu.Lock()
	if s.conns[nodeID] == c {
		delete(s.conns, nodeID)
	}
	s.mu.Unlock()

	// The node is not marked lost here. Losing a stream is not losing a node,
	// and the reconciler's two-stage expiry already handles silence with the
	// right delays (plan §6.3).
	s.log.Info("agent stream closed", "node", nodeID)
}

// ConnectedNodes lists the nodes with an open stream.
func (s *Service) ConnectedNodes() []string {
	s.mu.RLock()
	defer s.mu.RUnlock()
	out := make([]string, 0, len(s.conns))
	for id := range s.conns {
		out = append(out, id)
	}
	slices.Sort(out)
	return out
}

// --- Dispatcher ---

// ErrNotConnected means no agent stream is open for that node.
var ErrNotConnected = errors.New("no agent connected")

// StartShard implements reconciler.Dispatcher.
func (s *Service) StartShard(_ context.Context, cmd reconciler.ShardCommand) error {
	return s.send(cmd.NodeID, &llmgwv1.ControlMessage{
		Msg: &llmgwv1.ControlMessage_Start{Start: StartShardFromCommand(cmd)},
	})
}

// StopShard implements reconciler.Dispatcher.
func (s *Service) StopShard(_ context.Context, nodeID, shardID string) error {
	return s.send(nodeID, &llmgwv1.ControlMessage{
		Msg: &llmgwv1.ControlMessage_Stop{Stop: &llmgwv1.StopShard{ShardId: shardID}},
	})
}

// DrainShard implements reconciler.Dispatcher.
func (s *Service) DrainShard(_ context.Context, nodeID, shardID string) error {
	return s.send(nodeID, &llmgwv1.ControlMessage{
		Msg: &llmgwv1.ControlMessage_Drain{Drain: &llmgwv1.DrainShard{ShardId: shardID}},
	})
}

func (s *Service) send(nodeID string, msg *llmgwv1.ControlMessage) error {
	s.mu.RLock()
	c, ok := s.conns[nodeID]
	s.mu.RUnlock()
	if !ok {
		return fmt.Errorf("%w: node %s", ErrNotConnected, nodeID)
	}

	select {
	case c.send <- msg:
		return nil
	case <-c.closed:
		return fmt.Errorf("%w: node %s stream closed", ErrNotConnected, nodeID)
	default:
		// The queue is full, so this agent is far enough behind that queueing
		// more would only delay discovering it is broken.
		return fmt.Errorf("node %s is not keeping up with commands", nodeID)
	}
}