llmgw/internal/agent/supervisor/supervisor.go

package supervisor

import (
	"bufio"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"slices"
	"strings"
	"sync"
	"syscall"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
)

// Config configures the supervisor.
type Config struct {
	// WorkDir holds pid files and per-shard scratch space.
	WorkDir string `koanf:"work_dir"`
	// LogLines is how many lines of engine output to retain per shard.
	LogLines int `koanf:"log_lines"`
	// StopTimeout is how long a process gets to exit after SIGTERM before it is
	// killed. Engines need time to free GPU memory cleanly.
	StopTimeout time.Duration `koanf:"stop_timeout"`
	// HealthInterval is how often a starting shard is polled.
	HealthInterval time.Duration `koanf:"health_interval"`
	// RendezvousTimeout bounds the wait for a multi-node rendezvous point to
	// start listening.
	RendezvousTimeout time.Duration `koanf:"rendezvous_timeout"`
	// AliveGrace is how long a non-serving rank must stay up before it counts
	// as ready. There is no external signal for a worker, so staying alive is
	// the only evidence available.
	AliveGrace time.Duration `koanf:"alive_grace"`
}

// DefaultConfig returns sensible defaults.
func DefaultConfig() Config {
	return Config{
		WorkDir:           "/var/lib/llmgw-agent",
		LogLines:          512,
		StopTimeout:       60 * time.Second,
		HealthInterval:    2 * time.Second,
		RendezvousTimeout: 5 * time.Minute,
		AliveGrace:        5 * time.Second,
	}
}

func (c *Config) normalize() {
	d := DefaultConfig()
	if c.WorkDir == "" {
		c.WorkDir = d.WorkDir
	}
	if c.LogLines <= 0 {
		c.LogLines = d.LogLines
	}
	if c.StopTimeout <= 0 {
		c.StopTimeout = d.StopTimeout
	}
	if c.HealthInterval <= 0 {
		c.HealthInterval = d.HealthInterval
	}
	if c.RendezvousTimeout <= 0 {
		c.RendezvousTimeout = d.RendezvousTimeout
	}
	if c.AliveGrace <= 0 {
		c.AliveGrace = d.AliveGrace
	}
}

// HealthChecker reports whether a shard is serving. It is an interface so the
// supervisor can be tested without an HTTP engine.
type HealthChecker interface {
	// Ready reports whether the shard at baseURL serves model.
	Ready(ctx context.Context, baseURL, model string) error
}

// Shard is one supervised engine process.
type Shard struct {
	ID   string
	Spec engines.ShardSpec

	mu       sync.Mutex
	state    State
	pid      int
	err      error
	started  time.Time
	readyAt  time.Time
	exitCode int

	logs *ringBuffer
	cmd  *exec.Cmd
	// sidecars are the shard's supporting processes, torn down with it.
	sidecars []*exec.Cmd
	cancel   context.CancelFunc
	done     chan struct{}
	// adopted marks a shard inherited from a previous agent: it has no
	// exec.Cmd, so it is watched by polling and killed by signal alone.
	adopted bool
}

// Status is a snapshot of a shard, safe to send to the control plane.
type Status struct {
	ID       string
	State    State
	PID      int
	Err      string
	Started  time.Time
	ReadyAt  time.Time
	ExitCode int
}

// Status returns the shard's current status.
func (s *Shard) Status() Status {
	s.mu.Lock()
	defer s.mu.Unlock()
	st := Status{
		ID: s.ID, State: s.state, PID: s.pid,
		Started: s.started, ReadyAt: s.readyAt, ExitCode: s.exitCode,
	}
	if s.err != nil {
		st.Err = s.err.Error()
	}
	return st
}

// Logs returns the last n lines of engine output.
func (s *Shard) Logs(n int) []string { return s.logs.tail(n) }

func (s *Shard) setState(st State, err error) {
	s.mu.Lock()
	s.state = st
	if err != nil {
		s.err = err
	}
	if st == Ready && s.readyAt.IsZero() {
		s.readyAt = time.Now()
	}
	s.mu.Unlock()
}

// Supervisor starts, watches and stops engine processes on this node.
type Supervisor struct {
	cfg    Config
	log    *slog.Logger
	health HealthChecker

	mu     sync.Mutex
	shards map[string]*Shard
}

// New builds a supervisor.
func New(cfg Config, health HealthChecker, log *slog.Logger) (*Supervisor, error) {
	cfg.normalize()
	if log == nil {
		log = slog.Default()
	}
	if err := os.MkdirAll(cfg.WorkDir, 0o750); err != nil {
		return nil, fmt.Errorf("work dir: %w", err)
	}
	return &Supervisor{cfg: cfg, log: log, health: health, shards: map[string]*Shard{}}, nil
}

// ErrAlreadyRunning means a shard with that ID is already supervised.
var ErrAlreadyRunning = errors.New("shard is already running")

// Start launches a shard.
//
// It is idempotent by ID: a repeated StartShard for the same shard is a no-op
// rather than a second process, because the control plane retries commands
// (plan §5.6).
func (s *Supervisor) Start(ctx context.Context, id string, spec engines.ShardSpec, cmds []engines.Command) (*Shard, error) {
	if len(cmds) == 0 {
		return nil, fmt.Errorf("shard %s: no commands to run", id)
	}
	s.mu.Lock()
	if existing, ok := s.shards[id]; ok {
		s.mu.Unlock()
		if existing.Status().State.Terminal() {
			s.Remove(id)
		} else {
			return existing, ErrAlreadyRunning
		}
		s.mu.Lock()
	}

	sh := &Shard{
		ID:      id,
		Spec:    spec,
		state:   Pending,
		logs:    newRingBuffer(s.cfg.LogLines),
		started: time.Now(),
		done:    make(chan struct{}),
	}
	s.shards[id] = sh
	s.mu.Unlock()

	if err := s.launch(ctx, sh, cmds); err != nil {
		sh.setState(Failed, err)
		close(sh.done)
		return sh, err
	}
	return sh, nil
}

// launch starts a shard's processes in order.
//
// Every command but the last is a sidecar: it keeps running for the shard's
// lifetime and is torn down with it, but its exit is not what defines the
// shard's. The last command is the engine, and its exit is the shard's exit.
func (s *Supervisor) launch(ctx context.Context, sh *Shard, cmds []engines.Command) error {
	runCtx, cancel := context.WithCancel(context.WithoutCancel(ctx))
	sh.cancel = cancel

	for i, c := range cmds[:len(cmds)-1] {
		proc, err := s.spawn(runCtx, sh, c)
		if err != nil {
			cancel()
			return fmt.Errorf("sidecar %d (%s): %w", i, c.Name, err)
		}
		sh.mu.Lock()
		sh.sidecars = append(sh.sidecars, proc)
		sh.mu.Unlock()

		if c.WaitForLeaderPort {
			// The next process cannot start until the rendezvous point is
			// listening: a Ray worker cannot join a head that is not up, and
			// the leader's server cannot claim GPUs from a cluster that does
			// not exist yet.
			addr := fmt.Sprintf("%s:%d", sh.Spec.LeaderAddr, sh.Spec.LeaderPort)
			if err := waitForPort(runCtx, addr, s.cfg.RendezvousTimeout); err != nil {
				cancel()
				return fmt.Errorf("waiting for %s (%s): %w", addr, c.Name, err)
			}
		}
	}

	cmd := cmds[len(cmds)-1]

	c := exec.CommandContext(runCtx, cmd.Path, cmd.Args...)
	c.Env = cmd.Env
	c.Dir = s.cfg.WorkDir
	// Its own process group, so stopping kills the engine's children too.
	// Engines spawn workers; signalling only the parent leaves those workers
	// holding GPU memory.
	c.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	// CommandContext kills only the leader, so replace that with a group kill.
	c.Cancel = func() error { return killGroup(c.Process.Pid, syscall.SIGKILL) }

	stdout, err := c.StdoutPipe()
	if err != nil {
		cancel()
		return err
	}
	c.Stderr = c.Stdout.(io.Writer) // interleave, as the engine writes them

	if err := c.Start(); err != nil {
		cancel()
		return fmt.Errorf("starting %s: %w", cmd.Path, err)
	}

	sh.mu.Lock()
	sh.cmd = c
	sh.pid = c.Process.Pid
	sh.state = Starting
	sh.mu.Unlock()

	if err := s.writeState(sh); err != nil {
		s.log.Warn("could not write pid file; adoption after a restart will not work",
			"shard", sh.ID, "err", err)
	}

	go s.pump(sh, stdout, cmd.Name)
	go s.wait(sh)
	return nil
}

// spawn starts one process in its own group and pipes its output to the log.
func (s *Supervisor) spawn(ctx context.Context, sh *Shard, c engines.Command) (*exec.Cmd, error) {
	proc := exec.CommandContext(ctx, c.Path, c.Args...)
	proc.Env = c.Env
	proc.Dir = s.cfg.WorkDir
	proc.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	proc.Cancel = func() error { return killGroup(proc.Process.Pid, syscall.SIGKILL) }

	stdout, err := proc.StdoutPipe()
	if err != nil {
		return nil, err
	}
	proc.Stderr = proc.Stdout.(io.Writer)

	if err := proc.Start(); err != nil {
		return nil, fmt.Errorf("starting %s: %w", c.Path, err)
	}
	go s.pump(sh, stdout, c.Name)
	// Reaped so it does not become a zombie; its exit does not end the shard.
	go func() { _ = proc.Wait() }()
	return proc, nil
}

// waitForPort blocks until addr accepts a connection.
func waitForPort(ctx context.Context, addr string, timeout time.Duration) error {
	ctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	var d net.Dialer
	for {
		conn, err := d.DialContext(ctx, "tcp", addr)
		if err == nil {
			_ = conn.Close()
			return nil
		}
		select {
		case <-ctx.Done():
			return fmt.Errorf("not listening after %v: %w", timeout, ctx.Err())
		case <-time.After(200 * time.Millisecond):
		}
	}
}

// pump copies a process's output into the ring buffer and the log.
func (s *Supervisor) pump(sh *Shard, r io.Reader, name string) {
	sc := bufio.NewScanner(r)
	// Engine tracebacks and progress bars produce long lines; the default
	// 64 KiB limit truncates them exactly when they matter.
	sc.Buffer(make([]byte, 0, 64<<10), 1<<20)

	for sc.Scan() {
		line := sc.Text()
		if name != "" {
			// Prefixed so the two processes of a Ray shard stay distinguishable
			// in a single log tail.
			line = "[" + name + "] " + line
		}
		sh.logs.add(line)
		s.log.Debug("engine", "shard", sh.ID, "line", line)
	}
}

// wait reaps the process and records why it ended.
func (s *Supervisor) wait(sh *Shard) {
	defer close(sh.done)

	err := sh.cmd.Wait()

	sh.mu.Lock()
	prev := sh.state
	if ee := new(exec.ExitError); errors.As(err, &ee) {
		sh.exitCode = ee.ExitCode()
	}
	switch {
	case prev == Draining || prev == Stopped:
		// Expected: we asked it to stop.
		sh.state = Stopped
	case err == nil:
		// An engine exiting on its own is a failure, not a success: it is
		// supposed to serve until told otherwise.
		sh.state = Failed
		sh.err = errors.New("engine exited unexpectedly")
	default:
		sh.state = Failed
		sh.err = err
	}
	final, code := sh.state, sh.exitCode
	sh.mu.Unlock()

	s.removeState(sh)
	if final == Failed {
		s.log.Error("shard failed", "shard", sh.ID, "exit_code", code, "err", err,
			"last_output", strings.Join(sh.Logs(5), " | "))
	} else {
		s.log.Info("shard stopped", "shard", sh.ID, "exit_code", code)
	}
}

// WaitReady polls until the shard serves, it fails, or the timeout expires.
//
// The timeout comes from the profile because it is model-dependent: loading a
// 70B model across two nodes legitimately takes tens of minutes, while a small
// model failing to load for ten minutes means something is wrong.
func (s *Supervisor) WaitReady(ctx context.Context, sh *Shard, baseURL, model string, timeout time.Duration) error {
	ctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	t := time.NewTicker(s.cfg.HealthInterval)
	defer t.Stop()

	sh.setState(Loading, nil)
	for {
		select {
		case <-sh.done:
			return fmt.Errorf("shard %s exited during startup: %s", sh.ID, sh.Status().Err)
		case <-ctx.Done():
			last := strings.Join(sh.Logs(10), "\n")
			return fmt.Errorf("shard %s not ready within %v; last output:\n%s", sh.ID, timeout, last)
		case <-t.C:
			if err := s.health.Ready(ctx, baseURL, model); err != nil {
				continue
			}
			sh.setState(Ready, nil)
			return nil
		}
	}
}

// WaitAlive is readiness for a rank that serves no HTTP.
//
// A worker in a multi-node replica has no endpoint to poll, so the only
// evidence available is that it started and stayed started. That is weaker than
// a health check, and deliberately so: the replica's real readiness is the
// leader answering, which only happens once every worker has joined. Polling a
// worker's port instead would time out and fail a replica that is coming up
// perfectly well.
func (s *Supervisor) WaitAlive(ctx context.Context, sh *Shard) error {
	sh.setState(Loading, nil)

	select {
	case <-sh.done:
		return fmt.Errorf("shard %s exited during startup: %s", sh.ID, sh.Status().Err)
	case <-ctx.Done():
		return ctx.Err()
	case <-time.After(s.cfg.AliveGrace):
	}

	if st := sh.Status().State; st.Terminal() {
		return fmt.Errorf("shard %s is %s: %s", sh.ID, st, sh.Status().Err)
	}
	sh.setState(Ready, nil)
	return nil
}

// Drain marks a shard as taking no new work. The router stops routing to it;
// the process keeps running until Stop.
func (s *Supervisor) Drain(id string) error {
	sh, err := s.Get(id)
	if err != nil {
		return err
	}
	if st := sh.Status().State; !st.Live() {
		return fmt.Errorf("shard %s is %s, cannot drain", id, st)
	}
	sh.setState(Draining, nil)
	return nil
}

// Stop terminates a shard: SIGTERM to the whole group, then SIGKILL after
// StopTimeout.
func (s *Supervisor) Stop(ctx context.Context, id string) error {
	sh, err := s.Get(id)
	if err != nil {
		return err
	}

	sh.mu.Lock()
	pid := sh.pid
	sidecars := slices.Clone(sh.sidecars)
	if sh.state.Live() {
		sh.state = Draining // marks the exit as expected for wait() and poll()
	}
	sh.mu.Unlock()

	// Sidecars go too. A Ray head left behind holds the port the next replica
	// needs, and would let a later shard join a cluster that serves nothing.
	defer func() {
		for _, sc := range sidecars {
			if sc.Process == nil {
				continue
			}
			if err := killGroup(sc.Process.Pid, syscall.SIGTERM); err != nil &&
				!errors.Is(err, os.ErrProcessDone) {
				s.log.Warn("stopping a sidecar failed", "shard", id, "pid", sc.Process.Pid, "err", err)
			}
		}
	}()

	// A pid of zero means the process never started; an adopted shard has a pid
	// but no exec.Cmd, and must still be killed by signal.
	if pid == 0 {
		sh.setState(Stopped, nil)
		return nil
	}

	// Signal the group, not the process: engines fork workers that hold GPU
	// memory, and those outlive a signal sent only to the leader.
	if err := killGroup(pid, syscall.SIGTERM); err != nil && !errors.Is(err, os.ErrProcessDone) {
		s.log.Warn("SIGTERM failed", "shard", id, "err", err)
	}

	// force kills the group outright. exec.Cmd's cancel does this for shards we
	// started; an adopted shard has none, so signal its group directly.
	force := func() {
		if sh.cancel != nil {
			sh.cancel()
			return
		}
		if err := killGroup(pid, syscall.SIGKILL); err != nil && !errors.Is(err, os.ErrProcessDone) {
			s.log.Warn("SIGKILL failed", "shard", id, "err", err)
		}
	}

	select {
	case <-sh.done:
		return nil
	case <-time.After(s.cfg.StopTimeout):
		s.log.Warn("shard did not exit in time, killing", "shard", id, "timeout", s.cfg.StopTimeout)
		force()
	case <-ctx.Done():
		force()
		return ctx.Err()
	}

	select {
	case <-sh.done:
	case <-time.After(5 * time.Second):
		return fmt.Errorf("shard %s survived SIGKILL", id)
	}
	return nil
}

// StopAll stops every live shard.
func (s *Supervisor) StopAll(ctx context.Context) {
	for _, id := range s.IDs() {
		if err := s.Stop(ctx, id); err != nil {
			s.log.Warn("stopping shard failed", "shard", id, "err", err)
		}
	}
}

// Get returns a supervised shard.
func (s *Supervisor) Get(id string) (*Shard, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	sh, ok := s.shards[id]
	if !ok {
		return nil, fmt.Errorf("unknown shard %q", id)
	}
	return sh, nil
}

// Remove forgets a terminated shard.
func (s *Supervisor) Remove(id string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	delete(s.shards, id)
}

// IDs lists the supervised shards.
func (s *Supervisor) IDs() []string {
	s.mu.Lock()
	defer s.mu.Unlock()
	out := make([]string, 0, len(s.shards))
	for id := range s.shards {
		out = append(out, id)
	}
	return out
}

// Statuses returns every shard's status.
func (s *Supervisor) Statuses() []Status {
	s.mu.Lock()
	shards := make([]*Shard, 0, len(s.shards))
	for _, sh := range s.shards {
		shards = append(shards, sh)
	}
	s.mu.Unlock()

	out := make([]Status, 0, len(shards))
	for _, sh := range shards {
		out = append(out, sh.Status())
	}
	return out
}

// --- state files, for adoption after an agent restart ---

func (s *Supervisor) statePath(id string) string {
	return filepath.Join(s.cfg.WorkDir, "shard-"+id+".json")
}

// shardState is what survives an agent restart.
//
// The pid alone is not enough for two separate reasons. Pids are reused, so the
// command line is kept as a fingerprint — adopting an unrelated process would
// be worse than leaking the original. And the full spec is kept because a
// recovered shard has to be described back to the control plane: without it the
// control plane cannot tell which replica the process belongs to, and would
// place a duplicate while the original keeps holding the GPUs (plan §6.5).
type shardState struct {
	ShardID   string            `json:"shard_id"`
	ReplicaID string            `json:"replica_id"`
	PID       int               `json:"pid"`
	Cmdline   []string          `json:"cmdline"`
	StartedAt time.Time         `json:"started_at"`
	Spec      engines.ShardSpec `json:"spec"`
}

// writeState records everything needed to recognise and describe the process later.
func (s *Supervisor) writeState(sh *Shard) error {
	b, err := json.Marshal(shardState{
		ShardID:   sh.ID,
		ReplicaID: sh.Spec.ReplicaID,
		PID:       sh.pid,
		Cmdline:   sh.cmd.Args,
		StartedAt: sh.started,
		Spec:      sh.Spec,
	})
	if err != nil {
		return err
	}
	// Written atomically: a half-written state file after a crash would be
	// indistinguishable from a corrupt one, and the shard would be killed.
	tmp := s.statePath(sh.ID) + ".tmp"
	if err := os.WriteFile(tmp, b, 0o600); err != nil {
		return err
	}
	return os.Rename(tmp, s.statePath(sh.ID))
}

func (s *Supervisor) removeState(sh *Shard) {
	_ = os.Remove(s.statePath(sh.ID))
}

// Orphan is an engine process left behind by a previous agent.
type Orphan struct {
	ShardID   string
	ReplicaID string
	PID       int
	Cmdline   []string
	Spec      engines.ShardSpec
	StartedAt time.Time
	// Alive reports whether the process still exists.
	Alive bool
}

// FindOrphans scans the work directory for processes a previous agent started.
func (s *Supervisor) FindOrphans() ([]Orphan, error) {
	entries, err := os.ReadDir(s.cfg.WorkDir)
	if err != nil {
		return nil, err
	}

	var out []Orphan
	for _, e := range entries {
		name := e.Name()
		if !strings.HasPrefix(name, "shard-") || !strings.HasSuffix(name, ".json") {
			continue
		}

		b, err := os.ReadFile(filepath.Join(s.cfg.WorkDir, name))
		if err != nil {
			continue
		}
		var st shardState
		if err := json.Unmarshal(b, &st); err != nil || st.ShardID == "" || st.PID <= 0 {
			// A corrupt state file describes a process that cannot be
			// identified, so it cannot be adopted safely. Leave it for
			// KillOrphan to clean up by shard id from the filename.
			id := strings.TrimSuffix(strings.TrimPrefix(name, "shard-"), ".json")
			out = append(out, Orphan{ShardID: id})
			continue
		}

		out = append(out, Orphan{
			ShardID:   st.ShardID,
			ReplicaID: st.ReplicaID,
			PID:       st.PID,
			Cmdline:   st.Cmdline,
			Spec:      st.Spec,
			StartedAt: st.StartedAt,
			Alive:     processAlive(st.PID),
		})
	}
	return out, nil
}

// Adopt takes over a still-running engine from a previous agent.
//
// The process cannot be waited on — it is not this process's child — so it is
// watched by polling instead. That is the whole difference between an adopted
// shard and a started one; everything downstream treats them the same.
func (s *Supervisor) Adopt(o Orphan) (*Shard, error) {
	if !o.Alive {
		return nil, fmt.Errorf("shard %s: process %d is gone", o.ShardID, o.PID)
	}

	s.mu.Lock()
	if _, exists := s.shards[o.ShardID]; exists {
		s.mu.Unlock()
		return nil, ErrAlreadyRunning
	}
	sh := &Shard{
		ID:      o.ShardID,
		Spec:    o.Spec,
		state:   Starting,
		pid:     o.PID,
		logs:    newRingBuffer(s.cfg.LogLines),
		started: o.StartedAt,
		done:    make(chan struct{}),
		adopted: true,
	}
	s.shards[o.ShardID] = sh
	s.mu.Unlock()

	// Output is gone with the previous agent; say so rather than presenting an
	// empty log as if the engine had been silent.
	sh.logs.add("[adopted from a previous agent; earlier output is not available]")

	go s.poll(sh)
	return sh, nil
}

// poll watches an adopted process, which cannot be reaped with Wait.
func (s *Supervisor) poll(sh *Shard) {
	defer close(sh.done)

	// Polled rather than waited on, so the interval bounds how quickly a stop
	// is observed. Short enough that Stop does not sit out its whole timeout.
	t := time.NewTicker(200 * time.Millisecond)
	defer t.Stop()
	for range t.C {
		if processAlive(sh.pid) {
			continue
		}
		sh.mu.Lock()
		expected := sh.state == Draining || sh.state == Stopped
		if expected {
			sh.state = Stopped
		} else {
			sh.state = Failed
			sh.err = errors.New("adopted engine exited")
		}
		sh.mu.Unlock()

		s.removeState(sh)
		s.log.Info("adopted shard exited", "shard", sh.ID, "expected", expected)
		return
	}
}

// KillOrphan terminates a leftover process group and removes its pid file.
func (s *Supervisor) KillOrphan(o Orphan) error {
	if o.Alive {
		if err := killGroup(o.PID, syscall.SIGTERM); err != nil {
			return err
		}
		// Give it a moment to free GPU memory before insisting.
		deadline := time.Now().Add(s.cfg.StopTimeout)
		for time.Now().Before(deadline) {
			if !processAlive(o.PID) {
				break
			}
			time.Sleep(200 * time.Millisecond)
		}
		if processAlive(o.PID) {
			_ = killGroup(o.PID, syscall.SIGKILL)
		}
	}
	// An already-missing state file is success: the previous agent may have
	// reaped the process itself between the scan and this call.
	if err := os.Remove(s.statePath(o.ShardID)); err != nil && !errors.Is(err, os.ErrNotExist) {
		return err
	}
	return nil
}