Files
llmgw/internal/agent/supervisor/supervisor.go
T
Ilker ManapandClaude Opus 5.5 2a1e0e6c4c LoRA adapters on vLLM, loaded on demand with replica affinity: max_loaded_adapters
A vLLM model can carry PEFT adapters, registered at start, or with
max_loaded_adapters held at most N per replica: the router loads an adapter
through vLLM's runtime LoRA API when a request first wants it, sends the
adapter's requests to a replica that holds it while that one is not
saturated, loads each cold adapter once, and unloads the least recently used
idle one to make room. llama.cpp keeps loading every adapter at start.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
2026-10-05 01:14:38 +02:00

893 lines
26 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"
"go.opentelemetry.io/otel/attribute"
"go.opentelemetry.io/otel/trace"
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
"gitea.mynodes.xyz/ilker/llmgw/internal/obs"
)
// tracerName names the spans the supervisor opens. There is no restart span
// here because the supervisor never restarts an engine on its own: a restart
// is the control plane's decision, traced there, and reaches a node as a stop
// and a start.
const tracerName = "gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
// 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"`
// RunAs is the user engines run as when the agent is root; empty is
// llmgw-engine, "root" keeps them root.
RunAs string `koanf:"run_as"`
}
// 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,
RunAs: DefaultRunAs,
}
}
func (c *Config) normalize() {
d := DefaultConfig()
if c.WorkDir == "" {
c.WorkDir = d.WorkDir
}
if c.RunAs == "" {
c.RunAs = d.RunAs
}
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, apiKey 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
// as is who engines run as; nil is as the agent.
as *runAs
}
// 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)
}
s := &Supervisor{cfg: cfg, log: log, health: health, shards: map[string]*Shard{}}
as, why := resolveRunAs(cfg.RunAs, cfg.WorkDir)
if why != "" {
log.Warn(why)
}
if as != nil {
log.Info("engines run as their own user", "user", as.name)
}
s.as = as
return s, 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).
//
// Traced as "launch engine": from the first process to the engine's being
// started, which for a multi-node worker includes waiting for its leader's
// rendezvous point. Loading the model is the readiness wait's, not this.
func (s *Supervisor) Start(ctx context.Context, id string, spec engines.ShardSpec, cmds []engines.Command) (_ *Shard, err error) {
ctx, span := obs.Tracer(tracerName).Start(ctx, "launch engine", trace.WithAttributes(
attribute.String("llmgw.shard", id),
attribute.String("llmgw.engine", spec.Profile.Engine),
attribute.Int("llmgw.processes", len(cmds)),
))
defer func() {
if errors.Is(err, ErrAlreadyRunning) {
// A repeated command, which is not this start failing.
span.SetAttributes(attribute.Bool("llmgw.already_running", true))
span.End()
return
}
obs.End(span, err)
}()
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}
s.dropTo(c)
// 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()
trace.SpanFromContext(ctx).SetAttributes(attribute.Int("llmgw.pid", c.Process.Pid))
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}
s.dropTo(proc)
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.
//
// Traced: a worker's start that took ten minutes because its leader was
// still loading reads as exactly that.
func waitForPort(ctx context.Context, addr string, timeout time.Duration) (err error) {
ctx, span := obs.Tracer(tracerName).Start(ctx, "wait for rendezvous",
trace.WithAttributes(attribute.String("llmgw.rendezvous", addr)))
defer func() { obs.End(span, err) }()
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.
//
// Traced as "wait for engine": the model loading, which is usually most of a
// start. The span says how many times the engine was asked and what it last
// answered, so one that never came up says why.
func (s *Supervisor) WaitReady(ctx context.Context, sh *Shard, baseURL, model string, timeout time.Duration) (err error) {
ctx, span := obs.Tracer(tracerName).Start(ctx, "wait for engine", trace.WithAttributes(
attribute.String("llmgw.shard", sh.ID),
attribute.String("llmgw.readiness", "health check"),
attribute.String("llmgw.timeout", timeout.String()),
))
polls := 0
var answer error
defer func() {
span.SetAttributes(attribute.Int("llmgw.polls", polls))
if answer != nil && err != nil {
span.SetAttributes(attribute.String("llmgw.last_answer", answer.Error()))
}
obs.End(span, err)
}()
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:
polls++
if answer = s.health.Ready(ctx, baseURL, model, sh.Spec.APIKey); answer != 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) (err error) {
_, span := obs.Tracer(tracerName).Start(ctx, "wait for engine", trace.WithAttributes(
attribute.String("llmgw.shard", sh.ID),
attribute.String("llmgw.readiness", "stays alive"),
attribute.String("llmgw.timeout", s.cfg.AliveGrace.String()),
))
defer func() { obs.End(span, err) }()
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.
//
// Traced as "stop engine", saying whether it had to be killed: an engine
// that ignores SIGTERM for a minute is a minute its replacement waits for
// the GPU.
func (s *Supervisor) Stop(ctx context.Context, id string) (err error) {
ctx, span := obs.Tracer(tracerName).Start(ctx, "stop engine",
trace.WithAttributes(attribute.String("llmgw.shard", id)))
defer func() { obs.End(span, err) }()
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() {
span.SetAttributes(attribute.Bool("llmgw.killed", true))
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
}
// EngineDir is the working directory engines run in: the engines' user's
// home when they run as one, the work directory otherwise. A relative path
// an engine is given resolves against it.
func (s *Supervisor) EngineDir() string {
if s.as != nil {
return s.as.home
}
return s.cfg.WorkDir
}
// dropTo makes a command run as the engines' user, when there is one: its
// credential, its home as the working directory, and its caches there.
func (s *Supervisor) dropTo(c *exec.Cmd) {
if s.as == nil {
return
}
c.SysProcAttr.Credential = s.as.cred
c.Dir = s.as.home
c.Env = append(c.Env, s.as.env()...)
}