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>
782 lines
22 KiB
Go
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
|
|
}
|