Files
llmgw/internal/agent/supervisor/supervisor.go
T
Ilker ManapandClaude Opus 5.5 9ee21dc814 Fix what a multi-node launch would have got wrong on real hardware
vLLM's mp launch gave workers no rank and the leader's address as their
own VLLM_HOST_IP; every node now runs serve with --nnodes, --node-rank,
--master-addr and --master-port, and the workers --headless, as vLLM
documents it.

NCCL_SOCKET_IFNAME and GLOO_SOCKET_IFNAME were the control plane's one
interface name for every node; the agent already worked out its own -
the one holding its management address - and now uses it.

On RoCE, NCCL's default GID index is a v1 link-local entry on most hosts,
which does not route. The agent finds the RoCE v2 entry holding an IPv4
address and passes it as NCCL_IB_GID_INDEX; InfiniBand gets none.

A worker whose only process needs the leader - a Ray join - started
without waiting for it, since the wait was honoured for sidecars alone;
it found no head, exited and failed the replica. It now waits.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
2026-09-29 00:17:31 +02:00

782 lines
22 KiB
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]
// The main process can need the rendezvous point too: a Ray worker's
// join is its only process, and started before the head listens it
// exits, failing the shard and with it the replica.
if cmd.WaitForLeaderPort && !sh.Spec.Leader() {
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, cmd.Name, err)
}
}
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
}