diff --git a/CHANGELOG.md b/CHANGELOG.md index 36e4d219..e7d5830f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -170,6 +170,43 @@ tracking hijacked connections. The agent never reconnected and the assertion — on the connection map, which still held a stale entry — passed anyway. Tests now sever connections through their own listener and count opened streams. +### Multi-node launch + +Placement and environment generation had landed early, which made multi-node +look further along than it was: a multi-node replica could not come up at all. +Three things were missing, and all three came down to two assumptions that do +not hold once a replica spans nodes. + +**Not every rank serves HTTP.** The agent health-checked every shard at its own +port, but only the leader answers. Workers timed out, reported Failed, and took +the replica with them. `Engine.ServesHTTP(spec)` now says which ranks serve; a +worker is judged by staying alive instead. That is weaker evidence on purpose — +the replica's real readiness is the leader answering, which only happens once +every worker has joined. + +**A shard is not always one process.** `Engine.Commands(spec)` returns a +sequence: every command but the last is a sidecar kept alive for the shard's +lifetime and torn down with it, and the last one's exit is the shard's exit. +vLLM under Ray is why — the leader runs a Ray head and then the server against +it, and a worker runs only the join, with no vLLM process at all. `RayCommand` +had been written and tested but never called, so `launch_mode: ray` would have +run `vllm serve` against a cluster that did not exist. + +A command can also require the rendezvous point to be listening before the next +one starts, bounded by a timeout so a head that never appears fails the shard +rather than hanging it. Sidecars are killed with the shard: a Ray head left +behind holds the port the next replica needs. + +**Start order was backwards.** Workers were started first, on the theory that +the head should wait for them. Both engines put the rendezvous point on rank 0, +so a worker started first has nothing to connect to and spends its startup +budget retrying. vLLM's server blocks until the Ray cluster has the GPUs it +needs, so the leader can start immediately and waits on its own. + +None of this is verified against real engines. It is tested with stub processes; +the Ray sequencing and which launch mode a given vLLM version wants can only be +confirmed on hardware. + ### Still missing for phase 2 - mTLS. The stream is plaintext HTTP/2 today, so it must stay on a trusted diff --git a/internal/agent/agent.go b/internal/agent/agent.go index 116f31c5..2b671bd7 100644 --- a/internal/agent/agent.go +++ b/internal/agent/agent.go @@ -429,13 +429,13 @@ func (a *Agent) startShard(ctx context.Context, m *llmgwv1.StartShard, out chan< return } - cmd, err := eng.BuildCommand(spec) + cmds, err := eng.Commands(spec) if err != nil { report(supervisor.Failed, err.Error(), nil) return } - sh, err := a.sup.Start(ctx, shardID, spec, cmd) + sh, err := a.sup.Start(ctx, shardID, spec, cmds) if err != nil { if errors.Is(err, supervisor.ErrAlreadyRunning) { // Commands are idempotent: the control plane retries, and a repeat @@ -448,8 +448,16 @@ func (a *Agent) startShard(ctx context.Context, m *llmgwv1.StartShard, out chan< } report(supervisor.Starting, "", nil) - base := fmt.Sprintf("http://127.0.0.1:%d", spec.HTTPPort) - if err := a.sup.WaitReady(ctx, sh, base, spec.Profile.ServedModelName, timeout); err != nil { + // Only a rank that serves HTTP can be health-checked. A worker is judged by + // staying alive, and the replica's real readiness is the leader answering, + // which only happens once every worker has joined (plan §11, phase 4). + if eng.ServesHTTP(spec) { + base := fmt.Sprintf("http://127.0.0.1:%d", spec.HTTPPort) + err = a.sup.WaitReady(ctx, sh, base, spec.Profile.ServedModelName, timeout) + } else { + err = a.sup.WaitAlive(ctx, sh) + } + if err != nil { report(supervisor.Failed, err.Error(), sh.Logs(20)) return } diff --git a/internal/agent/engines/engine.go b/internal/agent/engines/engine.go index d9c211ec..c20cb4f6 100644 --- a/internal/agent/engines/engine.go +++ b/internal/agent/engines/engine.go @@ -83,6 +83,14 @@ type Command struct { Args []string // Env is the complete environment, not a delta. Env []string + // Name distinguishes the processes of a multi-process shard in logs. + Name string + // WaitForLeaderPort makes the supervisor hold the next command until + // LeaderAddr:LeaderPort accepts a connection. + // + // It exists for Ray: the head has to be listening before a worker can join, + // and both sides express the same wait because both dial the same address. + WaitForLeaderPort bool } // Engine builds launch commands for one inference engine. @@ -92,8 +100,24 @@ type Engine interface { // Validate rejects a profile that cannot work on these GPUs, before // anything is started. Validate(p Profile, gpus []GPUInfo) error - // BuildCommand produces the launch command for a shard. - BuildCommand(s ShardSpec) (Command, error) + + // Commands returns the processes that make up one shard, in start order. + // + // Most shards are a single process. vLLM under Ray is not: the leader runs + // a Ray head and then the server against it, and a worker runs only the Ray + // join. Returning a sequence rather than one command is what lets the + // supervisor treat both the same way — every command but the last is a + // sidecar kept alive for the shard's lifetime, and the last one's exit is + // the shard's exit. + Commands(s ShardSpec) ([]Command, error) + + // ServesHTTP reports whether this rank answers the OpenAI API. + // + // In a multi-node replica only the leader does. Health-checking a worker's + // port would time out and fail a replica that is coming up perfectly well, + // which is exactly what happened before this existed. + ServesHTTP(s ShardSpec) bool + // HealthPath and MetricsPath are appended to the shard's base URL. HealthPath() string MetricsPath() string diff --git a/internal/agent/engines/engines_test.go b/internal/agent/engines/engines_test.go index 349ecab2..ede59c85 100644 --- a/internal/agent/engines/engines_test.go +++ b/internal/agent/engines/engines_test.go @@ -1,11 +1,56 @@ package engines import ( + "fmt" "slices" "strings" "testing" ) +// buildOne returns the single command a shard produces, failing if it produced +// a different number. +func buildOne(t *testing.T, e Engine, s ShardSpec) Command { + t.Helper() + cmds := build(t, e, s) + if len(cmds) != 1 { + t.Fatalf("expected 1 command, got %d: %s", len(cmds), summarize(cmds)) + } + return cmds[0] +} + +// buildEngine returns the last command, which is the engine itself. The ones +// before it are sidecars. +func buildEngine(t *testing.T, e Engine, s ShardSpec) Command { + t.Helper() + cmds := build(t, e, s) + return cmds[len(cmds)-1] +} + +func build(t *testing.T, e Engine, s ShardSpec) []Command { + t.Helper() + cmds, err := e.Commands(s) + if err != nil { + t.Fatal(err) + } + if len(cmds) == 0 { + t.Fatal("no commands produced") + } + return cmds +} + +// summarize keeps failure output readable: a Command carries the whole +// environment, and printing it buries the actual problem. +func summarize(cmds []Command) string { + var b strings.Builder + for i, c := range cmds { + if i > 0 { + b.WriteString(" | ") + } + fmt.Fprintf(&b, "%s: %s %s", c.Name, c.Path, strings.Join(c.Args, " ")) + } + return b.String() +} + // envMap turns exec-form environment back into a map for assertions. func envMap(env []string) map[string]string { m := map[string]string{} @@ -61,10 +106,7 @@ func multiNodeSpec(rank int) ShardSpec { // --- environment generation: the part that breaks multi-node runs --- func TestSingleNodeDisablesIB(t *testing.T) { - cmd, err := (&VLLM{}).BuildCommand(singleNodeSpec()) - if err != nil { - t.Fatal(err) - } + cmd := buildOne(t, &VLLM{}, singleNodeSpec()) env := envMap(cmd.Env) if env["NCCL_IB_DISABLE"] != "1" { @@ -79,10 +121,7 @@ func TestSingleNodeDisablesIB(t *testing.T) { } func TestMultiNodePinsFabric(t *testing.T) { - cmd, err := (&VLLM{}).BuildCommand(multiNodeSpec(0)) - if err != nil { - t.Fatal(err) - } + cmd := buildEngine(t, &VLLM{}, multiNodeSpec(0)) env := envMap(cmd.Env) // Only the chosen fabric's adapters. Leaving this open lets NCCL pick an @@ -108,10 +147,7 @@ func TestProfileEnvOverridesGenerated(t *testing.T) { "NCCL_IB_HCA": "mlx5_9", "NCCL_NET_GDR_LEVEL": "SYS", } - cmd, err := (&VLLM{}).BuildCommand(s) - if err != nil { - t.Fatal(err) - } + cmd := buildEngine(t, &VLLM{}, s) env := envMap(cmd.Env) if env["NCCL_IB_HCA"] != "mlx5_9" { @@ -126,17 +162,14 @@ func TestProcessEnvironmentIsInherited(t *testing.T) { // Engines need PATH, HF_HOME, LD_LIBRARY_PATH and whatever else the unit // file set up. t.Setenv("LLMGW_TEST_MARKER", "present") - cmd, err := (&VLLM{}).BuildCommand(singleNodeSpec()) - if err != nil { - t.Fatal(err) - } + cmd := buildOne(t, &VLLM{}, singleNodeSpec()) if envMap(cmd.Env)["LLMGW_TEST_MARKER"] != "present" { t.Error("the process environment was not inherited") } } func TestEnvIsSorted(t *testing.T) { - cmd, _ := (&VLLM{}).BuildCommand(singleNodeSpec()) + cmd := buildOne(t, &VLLM{}, singleNodeSpec()) if !slices.IsSorted(cmd.Env) { t.Error("environment is not sorted; launches would not be reproducible in logs") } @@ -145,10 +178,7 @@ func TestEnvIsSorted(t *testing.T) { // --- vLLM --- func TestVLLMSingleNodeArgs(t *testing.T) { - cmd, err := (&VLLM{cfg: VLLMConfig{Bin: "/opt/vllm/bin/vllm"}}).BuildCommand(singleNodeSpec()) - if err != nil { - t.Fatal(err) - } + cmd := buildOne(t, &VLLM{cfg: VLLMConfig{Bin: "/opt/vllm/bin/vllm"}}, singleNodeSpec()) if cmd.Path != "/opt/vllm/bin/vllm" { t.Errorf("Path=%q", cmd.Path) } @@ -167,26 +197,54 @@ func TestVLLMSingleNodeArgs(t *testing.T) { } } -func TestVLLMRayMode(t *testing.T) { +func TestVLLMRayLeaderRunsTwoProcesses(t *testing.T) { + // The leader starts a Ray head and then the server against it. A single + // command would mean running vLLM against a cluster that does not exist. v := &VLLM{cfg: VLLMConfig{LaunchMode: "ray"}} - cmd, err := v.BuildCommand(multiNodeSpec(1)) + cmds, err := v.Commands(multiNodeSpec(0)) if err != nil { t.Fatal(err) } - if argValue(cmd.Args, "--distributed-executor-backend") != "ray" { - t.Errorf("backend not set: %v", cmd.Args) + if len(cmds) != 2 { + t.Fatalf("expected 2 commands for the Ray leader, got %d", len(cmds)) } - if got := envMap(cmd.Env)["RAY_ADDRESS"]; got != "10.0.0.10:20001" { + if cmds[0].Name != "ray" || !slices.Contains(cmds[0].Args, "--head") { + t.Errorf("first command is not the Ray head: %+v", cmds[0]) + } + if !cmds[0].WaitForLeaderPort { + t.Error("the server would start before the head is listening") + } + if cmds[1].Name != "vllm" || argValue(cmds[1].Args, "--distributed-executor-backend") != "ray" { + t.Errorf("second command is not the vLLM server: %+v", cmds[1]) + } + if got := envMap(cmds[1].Env)["RAY_ADDRESS"]; got != "10.0.0.10:20001" { t.Errorf("RAY_ADDRESS=%q", got) } } +func TestVLLMRayWorkerRunsOnlyRay(t *testing.T) { + // A Ray worker runs no vLLM process: the cluster hands its GPUs to the + // leader. Starting a second server there would bind a port nobody uses and + // compete for the same cards. + v := &VLLM{cfg: VLLMConfig{LaunchMode: "ray"}} + cmds, err := v.Commands(multiNodeSpec(1)) + if err != nil { + t.Fatal(err) + } + if len(cmds) != 1 { + t.Fatalf("expected 1 command for a Ray worker, got %d: %+v", len(cmds), cmds) + } + if cmds[0].Name != "ray" || slices.Contains(cmds[0].Args, "--head") { + t.Errorf("worker command wrong: %+v", cmds[0]) + } + if argValue(cmds[0].Args, "--address") != "10.0.0.10:20001" { + t.Errorf("worker is not joining the leader: %v", cmds[0].Args) + } +} + func TestVLLMMPMode(t *testing.T) { v := &VLLM{cfg: VLLMConfig{LaunchMode: "mp"}} - cmd, err := v.BuildCommand(multiNodeSpec(0)) - if err != nil { - t.Fatal(err) - } + cmd := buildOne(t, v, multiNodeSpec(0)) if argValue(cmd.Args, "--distributed-executor-backend") != "mp" { t.Errorf("backend not set: %v", cmd.Args) } @@ -198,51 +256,42 @@ func TestVLLMMPMode(t *testing.T) { func TestVLLMUnknownLaunchMode(t *testing.T) { v := &VLLM{cfg: VLLMConfig{LaunchMode: "sorcery"}} - if _, err := v.BuildCommand(multiNodeSpec(0)); err == nil { + if _, err := v.Commands(multiNodeSpec(0)); err == nil { t.Error("an unknown launch mode was accepted") } } -func TestVLLMRayCommand(t *testing.T) { - v := &VLLM{cfg: VLLMConfig{LaunchMode: "ray"}} +func TestOnlyLeaderServesHTTP(t *testing.T) { + // The fix for the bug that stopped multi-node replicas coming up at all: + // health-checking a worker's port times out and fails a healthy replica. + for _, e := range []Engine{&VLLM{}, &SGLang{}} { + if !e.ServesHTTP(multiNodeSpec(0)) { + t.Errorf("%s: the leader must serve HTTP", e.Name()) + } + if e.ServesHTTP(multiNodeSpec(1)) { + t.Errorf("%s: a worker must not be expected to serve HTTP", e.Name()) + } + if !e.ServesHTTP(singleNodeSpec()) { + t.Errorf("%s: a single-node shard must serve HTTP", e.Name()) + } + } +} - head, ok := v.RayCommand(multiNodeSpec(0)) - if !ok { - t.Fatal("no ray command for the leader") +func TestNoRayForSingleNodeOrMPMode(t *testing.T) { + ray := &VLLM{cfg: VLLMConfig{LaunchMode: "ray"}} + if cmds := buildOne(t, ray, singleNodeSpec()); cmds.Name != "vllm" { + t.Errorf("a single-node profile was given a Ray process: %+v", cmds) } - if !slices.Contains(head.Args, "--head") { - t.Errorf("rank 0 must start the head: %v", head.Args) - } - - worker, ok := v.RayCommand(multiNodeSpec(1)) - if !ok { - t.Fatal("no ray command for the worker") - } - if slices.Contains(worker.Args, "--head") { - t.Errorf("rank 1 must not start a head: %v", worker.Args) - } - if argValue(worker.Args, "--address") != "10.0.0.10:20001" { - t.Errorf("worker address wrong: %v", worker.Args) - } - - // Single node needs no Ray at all. - if _, ok := v.RayCommand(singleNodeSpec()); ok { - t.Error("a single-node profile was given a Ray command") - } - // Neither does mp mode. mp := &VLLM{cfg: VLLMConfig{LaunchMode: "mp"}} - if _, ok := mp.RayCommand(multiNodeSpec(0)); ok { - t.Error("mp mode was given a Ray command") + if cmds := buildOne(t, mp, multiNodeSpec(0)); cmds.Name != "vllm" { + t.Errorf("mp mode was given a Ray process: %+v", cmds) } } // --- SGLang --- func TestSGLangMultiNodeArgs(t *testing.T) { - cmd, err := (&SGLang{cfg: SGLangConfig{Python: "/opt/sglang/bin/python"}}).BuildCommand(multiNodeSpec(1)) - if err != nil { - t.Fatal(err) - } + cmd := buildOne(t, &SGLang{cfg: SGLangConfig{Python: "/opt/sglang/bin/python"}}, multiNodeSpec(1)) if cmd.Path != "/opt/sglang/bin/python" { t.Errorf("Path=%q", cmd.Path) } @@ -258,7 +307,7 @@ func TestSGLangMultiNodeArgs(t *testing.T) { } func TestSGLangSingleNodeOmitsMultiNodeFlags(t *testing.T) { - cmd, _ := (&SGLang{}).BuildCommand(singleNodeSpec()) + cmd := buildOne(t, &SGLang{}, singleNodeSpec()) for _, flag := range []string{"--nnodes", "--node-rank", "--dist-init-addr"} { if slices.Contains(cmd.Args, flag) { t.Errorf("%s was passed for a single-node profile", flag) @@ -270,7 +319,7 @@ func TestProfileArgsComeLast(t *testing.T) { // Operators must be able to override anything the adapter generated. s := singleNodeSpec() s.Profile.Args = []string{"--max-model-len", "32768"} - cmd, _ := (&VLLM{}).BuildCommand(s) + cmd := buildEngine(t, &VLLM{}, s) i := slices.Index(cmd.Args, "--max-model-len") j := slices.Index(cmd.Args, "--tensor-parallel-size") diff --git a/internal/agent/engines/sglang.go b/internal/agent/engines/sglang.go index b4c08786..40061bf9 100644 --- a/internal/agent/engines/sglang.go +++ b/internal/agent/engines/sglang.go @@ -35,11 +35,26 @@ func (*SGLang) Validate(p Profile, gpus []GPUInfo) error { return nil } -// BuildCommand produces the launch command for one shard. +// ServesHTTP reports whether this rank answers the API. // -// SGLang supports multi-node natively, so unlike vLLM there is no second -// process to sequence: every rank runs the same command with its own rank number. -func (g *SGLang) BuildCommand(s ShardSpec) (Command, error) { +// Every SGLang rank runs the same launcher, but only node-rank 0 binds the HTTP +// port; the others exist to contribute GPUs to the tensor-parallel group. +func (*SGLang) ServesHTTP(s ShardSpec) bool { return s.Leader() } + +// Commands returns the processes for one shard. +// +// SGLang supports multi-node natively, so there is only ever one process: every +// rank runs the same launcher with its own rank number and they rendezvous at +// the leader's dist-init address. +func (g *SGLang) Commands(s ShardSpec) ([]Command, error) { + c, err := g.serverCommand(s) + if err != nil { + return nil, err + } + return []Command{c}, nil +} + +func (g *SGLang) serverCommand(s ShardSpec) (Command, error) { python := g.cfg.Python if python == "" { python = "python3" @@ -64,5 +79,5 @@ func (g *SGLang) BuildCommand(s ShardSpec) (Command, error) { } args = append(args, p.Args...) - return Command{Path: python, Args: args, Env: materializeEnv(baseEnv(s))}, nil + return Command{Name: "sglang", Path: python, Args: args, Env: materializeEnv(baseEnv(s))}, nil } diff --git a/internal/agent/engines/vllm.go b/internal/agent/engines/vllm.go index 5e3b50aa..9eb8eef0 100644 --- a/internal/agent/engines/vllm.go +++ b/internal/agent/engines/vllm.go @@ -46,8 +46,38 @@ func (*VLLM) Validate(p Profile, gpus []GPUInfo) error { return nil } -// BuildCommand produces the launch command for one shard. -func (v *VLLM) BuildCommand(s ShardSpec) (Command, error) { +// ServesHTTP reports whether this rank answers the API. +// +// Under Ray a worker runs nothing but the Ray join, and under multiprocessing +// only the driver binds the port. Either way it is the leader alone. +func (*VLLM) ServesHTTP(s ShardSpec) bool { return s.Leader() } + +// Commands returns the processes for one shard, in start order. +// +// Single node, or multi-node under multiprocessing: one server process. Under +// Ray: the leader starts a head and then the server, a worker starts only the +// join. A Ray worker running no vLLM process at all is correct — the cluster +// gives the leader its GPUs. +func (v *VLLM) Commands(s ShardSpec) ([]Command, error) { + ray, useRay := v.rayCommand(s) + if useRay && !s.Leader() { + // A worker is the Ray join and nothing else. It blocks, so its exit is + // the shard's exit, which is what the supervisor needs. + return []Command{ray}, nil + } + + server, err := v.serverCommand(s) + if err != nil { + return nil, err + } + if useRay { + return []Command{ray, server}, nil + } + return []Command{server}, nil +} + +// serverCommand produces the vLLM server process. +func (v *VLLM) serverCommand(s ShardSpec) (Command, error) { bin := v.cfg.Bin if bin == "" { bin = "vllm" @@ -88,15 +118,12 @@ func (v *VLLM) BuildCommand(s ShardSpec) (Command, error) { // Profile arguments last, so an operator can override anything above. args = append(args, p.Args...) - return Command{Path: bin, Args: args, Env: materializeEnv(env)}, nil + return Command{Name: "vllm", Path: bin, Args: args, Env: materializeEnv(env)}, nil } -// RayCommand returns the command that joins this node to the Ray cluster, or +// rayCommand returns the command that starts or joins the Ray cluster, or // ok=false when the launch mode does not use Ray. -// -// It is separate from BuildCommand because it has to run first and finish -// before vLLM starts: the supervisor sequences the two. -func (v *VLLM) RayCommand(s ShardSpec) (Command, bool) { +func (v *VLLM) rayCommand(s ShardSpec) (Command, bool) { if s.NNodes <= 1 || (v.cfg.LaunchMode != "" && v.cfg.LaunchMode != "ray") { return Command{}, false } @@ -113,5 +140,10 @@ func (v *VLLM) RayCommand(s ShardSpec) (Command, bool) { } args = append(args, "--num-gpus", strconv.Itoa(len(s.GPUIndices))) - return Command{Path: bin, Args: args, Env: materializeEnv(baseEnv(s))}, true + return Command{ + Name: "ray", Path: bin, Args: args, Env: materializeEnv(baseEnv(s)), + // The head must be listening before anything else happens: a worker + // cannot join, and the leader's server cannot claim GPUs, until it is. + WaitForLeaderPort: true, + }, true } diff --git a/internal/agent/supervisor/supervisor.go b/internal/agent/supervisor/supervisor.go index 5b9975bd..f64d8158 100644 --- a/internal/agent/supervisor/supervisor.go +++ b/internal/agent/supervisor/supervisor.go @@ -8,9 +8,11 @@ import ( "fmt" "io" "log/slog" + "net" "os" "os/exec" "path/filepath" + "slices" "strings" "sync" "syscall" @@ -30,15 +32,24 @@ type Config struct { 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, + WorkDir: "/var/lib/llmgw-agent", + LogLines: 512, + StopTimeout: 60 * time.Second, + HealthInterval: 2 * time.Second, + RendezvousTimeout: 5 * time.Minute, + AliveGrace: 5 * time.Second, } } @@ -56,6 +67,12 @@ func (c *Config) normalize() { 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 @@ -78,10 +95,12 @@ type Shard struct { readyAt time.Time exitCode int - logs *ringBuffer - cmd *exec.Cmd - cancel context.CancelFunc - done chan struct{} + 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 @@ -157,7 +176,10 @@ var ErrAlreadyRunning = errors.New("shard is already running") // 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, cmd engines.Command) (*Shard, error) { +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() @@ -180,7 +202,7 @@ func (s *Supervisor) Start(ctx context.Context, id string, spec engines.ShardSpe s.shards[id] = sh s.mu.Unlock() - if err := s.launch(ctx, sh, cmd); err != nil { + if err := s.launch(ctx, sh, cmds); err != nil { sh.setState(Failed, err) close(sh.done) return sh, err @@ -188,10 +210,40 @@ func (s *Supervisor) Start(ctx context.Context, id string, spec engines.ShardSpe return sh, nil } -func (s *Supervisor) launch(ctx context.Context, sh *Shard, cmd engines.Command) error { +// 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 @@ -225,13 +277,56 @@ func (s *Supervisor) launch(ctx context.Context, sh *Shard, cmd engines.Command) "shard", sh.ID, "err", err) } - go s.pump(sh, stdout) + go s.pump(sh, stdout, cmd.Name) go s.wait(sh) return nil } -// pump copies engine output into the ring buffer and the log. -func (s *Supervisor) pump(sh *Shard, r io.Reader) { +// 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. @@ -239,6 +334,11 @@ func (s *Supervisor) pump(sh *Shard, r io.Reader) { 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) } @@ -310,6 +410,32 @@ func (s *Supervisor) WaitReady(ctx context.Context, sh *Shard, baseURL, model st } } +// 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 { @@ -334,11 +460,26 @@ func (s *Supervisor) Stop(ctx context.Context, id string) error { 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 { diff --git a/internal/agent/supervisor/supervisor_test.go b/internal/agent/supervisor/supervisor_test.go index 573f3e42..af1785f4 100644 --- a/internal/agent/supervisor/supervisor_test.go +++ b/internal/agent/supervisor/supervisor_test.go @@ -5,6 +5,7 @@ import ( "errors" "fmt" "log/slog" + "net" "os" "path/filepath" "strings" @@ -42,9 +43,9 @@ func newSupervisor(t *testing.T) (*Supervisor, *scriptedHealth) { return s, h } -// shellCommand builds a Command running a shell script. -func shellCommand(script string) engines.Command { - return engines.Command{Path: "/bin/sh", Args: []string{"-c", script}, Env: os.Environ()} +// shellCommand builds the single-command sequence for a shell script. +func shellCommand(script string) []engines.Command { + return []engines.Command{{Path: "/bin/sh", Args: []string{"-c", script}, Env: os.Environ()}} } func waitState(t *testing.T, sh *Shard, want State, limit time.Duration) { @@ -532,3 +533,197 @@ func TestCorruptStateFileIsCleanedUp(t *testing.T) { t.Error("the corrupt state file was left behind") } } + +// --- multi-process shards (plan §11, phase 4) --- + +func TestSidecarRunsBeforeEngine(t *testing.T) { + // The Ray shape: a supporting process starts first and keeps running, then + // the engine starts against it. + s, _ := newSupervisor(t) + dir := t.TempDir() + order := filepath.Join(dir, "order") + + // Each writes to the order file and to stdout: the file records the + // sequence, stdout exercises per-process log attribution. + cmds := []engines.Command{ + {Name: "sidecar", Path: "/bin/sh", Args: []string{"-c", + fmt.Sprintf("echo sidecar >> %s; echo hello-from-sidecar; sleep 30", order)}, Env: os.Environ()}, + {Name: "engine", Path: "/bin/sh", Args: []string{"-c", + fmt.Sprintf("echo engine >> %s; echo hello-from-engine; sleep 30", order)}, Env: os.Environ()}, + } + + sh, err := s.Start(t.Context(), "a", engines.ShardSpec{}, cmds) + if err != nil { + t.Fatal(err) + } + + deadline := time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + if b, err := os.ReadFile(order); err == nil && strings.Count(string(b), "\n") == 2 { + if got := strings.Fields(string(b)); got[0] != "sidecar" || got[1] != "engine" { + t.Fatalf("start order was %v", got) + } + break + } + time.Sleep(20 * time.Millisecond) + } + + // The shard's pid is the engine's, not the sidecar's: the engine's exit is + // what ends the shard. + if sh.Status().PID == 0 { + t.Fatal("no engine pid recorded") + } + // Both processes appear in one log tail, distinguishable by prefix. + deadline = time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + logs := strings.Join(sh.Logs(20), "\n") + if strings.Contains(logs, "[sidecar] hello-from-sidecar") && + strings.Contains(logs, "[engine] hello-from-engine") { + return + } + time.Sleep(20 * time.Millisecond) + } + t.Errorf("output is not attributed per process:\n%s", strings.Join(sh.Logs(20), "\n")) +} + +func TestStopKillsSidecar(t *testing.T) { + // A Ray head left behind holds the port the next replica needs, and would + // let a later shard join a cluster that serves nothing. + s, _ := newSupervisor(t) + dir := t.TempDir() + marker := filepath.Join(dir, "sidecar.pid") + + cmds := []engines.Command{ + {Name: "sidecar", Path: "/bin/sh", Args: []string{"-c", + fmt.Sprintf("echo $$ > %s; sleep 30", marker)}, Env: os.Environ()}, + {Name: "engine", Path: "/bin/sh", Args: []string{"-c", "sleep 30"}, Env: os.Environ()}, + } + if _, err := s.Start(t.Context(), "a", engines.ShardSpec{}, cmds); err != nil { + t.Fatal(err) + } + + var sidePID int + deadline := time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + if b, err := os.ReadFile(marker); err == nil { + if _, err := fmt.Sscanf(strings.TrimSpace(string(b)), "%d", &sidePID); err == nil && sidePID > 0 { + break + } + } + time.Sleep(10 * time.Millisecond) + } + if sidePID == 0 { + t.Fatal("the sidecar never started") + } + + if err := s.Stop(t.Context(), "a"); err != nil { + t.Fatal(err) + } + deadline = time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + if !processAlive(sidePID) { + return + } + time.Sleep(20 * time.Millisecond) + } + t.Errorf("the sidecar (%d) survived the stop and still holds its port", sidePID) +} + +func TestSidecarFailureFailsTheShard(t *testing.T) { + s, _ := newSupervisor(t) + cmds := []engines.Command{ + {Name: "sidecar", Path: "/nonexistent/binary", Env: os.Environ()}, + {Name: "engine", Path: "/bin/sh", Args: []string{"-c", "sleep 30"}, Env: os.Environ()}, + } + sh, err := s.Start(t.Context(), "a", engines.ShardSpec{}, cmds) + if err == nil { + t.Fatal("a shard whose sidecar could not start was reported healthy") + } + if !strings.Contains(err.Error(), "sidecar") { + t.Errorf("the error does not identify the sidecar: %v", err) + } + if sh.Status().State != Failed { + t.Errorf("state=%s", sh.Status().State) + } +} + +func TestWaitForLeaderPortBlocksUntilListening(t *testing.T) { + // 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. + s, _ := newSupervisor(t) + s.cfg.RendezvousTimeout = 5 * time.Second + + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + port := ln.Addr().(*net.TCPAddr).Port + _ = ln.Close() // nothing listening yet + + spec := engines.ShardSpec{LeaderAddr: "127.0.0.1", LeaderPort: port} + cmds := []engines.Command{ + {Name: "sidecar", Path: "/bin/sh", Args: []string{"-c", + fmt.Sprintf("sleep 0.4; nc -l %d", port)}, Env: os.Environ(), + WaitForLeaderPort: true}, + {Name: "engine", Path: "/bin/sh", Args: []string{"-c", "sleep 30"}, Env: os.Environ()}, + } + + start := time.Now() + if _, err := s.Start(t.Context(), "a", spec, cmds); err != nil { + t.Skipf("nc is unavailable for the rendezvous stub: %v", err) + } + if d := time.Since(start); d < 300*time.Millisecond { + t.Errorf("the engine started after %v, before the rendezvous point was listening", d) + } +} + +func TestRendezvousTimeout(t *testing.T) { + // Nothing ever listens, so the shard must fail rather than hang forever. + s, _ := newSupervisor(t) + s.cfg.RendezvousTimeout = 300 * time.Millisecond + + spec := engines.ShardSpec{LeaderAddr: "127.0.0.1", LeaderPort: 1} + cmds := []engines.Command{ + {Name: "sidecar", Path: "/bin/sh", Args: []string{"-c", "sleep 30"}, + Env: os.Environ(), WaitForLeaderPort: true}, + {Name: "engine", Path: "/bin/sh", Args: []string{"-c", "sleep 30"}, Env: os.Environ()}, + } + if _, err := s.Start(t.Context(), "a", spec, cmds); err == nil { + t.Fatal("the shard did not fail when the rendezvous point never appeared") + } +} + +func TestWaitAliveForNonServingRank(t *testing.T) { + // A worker has no endpoint to poll, so staying up is the only evidence + // available. This is what stopped multi-node replicas coming up at all. + s, _ := newSupervisor(t) + s.cfg.AliveGrace = 150 * time.Millisecond + + sh, err := s.Start(t.Context(), "a", engines.ShardSpec{Rank: 1}, shellCommand("sleep 30")) + if err != nil { + t.Fatal(err) + } + if err := s.WaitAlive(t.Context(), sh); err != nil { + t.Fatal(err) + } + if got := sh.Status().State; got != Ready { + t.Errorf("state=%s", got) + } +} + +func TestWaitAliveFailsOnEarlyExit(t *testing.T) { + s, _ := newSupervisor(t) + s.cfg.AliveGrace = 2 * time.Second + + sh, err := s.Start(t.Context(), "a", engines.ShardSpec{Rank: 1}, shellCommand("exit 1")) + if err != nil { + t.Fatal(err) + } + start := time.Now() + if err := s.WaitAlive(t.Context(), sh); err == nil { + t.Fatal("a worker that died immediately was reported ready") + } + if d := time.Since(start); d > 1500*time.Millisecond { + t.Errorf("waited %v for a process that had already exited", d) + } +} diff --git a/internal/control/reconciler/reconciler.go b/internal/control/reconciler/reconciler.go index 43360918..68629d2d 100644 --- a/internal/control/reconciler/reconciler.go +++ b/internal/control/reconciler/reconciler.go @@ -411,11 +411,17 @@ func (r *Reconciler) startReplica(ctx context.Context, m catalog.Model, pl *plac // agent status report can mutate it concurrently, and the race detector is // right to object. // - // Workers first, then the leader. SGLang's ranks rendezvous, and vLLM under - // Ray needs its workers attached before the head starts serving; starting - // rank 0 first just makes it wait. + // The leader first, then workers ascending. Both engines put the rendezvous + // point on rank 0 — SGLang's dist-init address and vLLM's Ray head — so a + // worker started first has nothing to connect to and spends its startup + // budget retrying. + // + // An earlier version started workers first, on the theory that the head + // should wait for them. That is backwards: vLLM's server blocks until the + // Ray cluster has the GPUs it needs, so the leader can be started + // immediately and will wait on its own. order := slices.Clone(rep.Shards) - slices.SortFunc(order, func(a, b Shard) int { return b.Rank - a.Rank }) + slices.SortFunc(order, func(a, b Shard) int { return a.Rank - b.Rank }) nodeIDs := rep.NodeIDs() r.mu.Lock() diff --git a/internal/control/reconciler/reconciler_test.go b/internal/control/reconciler/reconciler_test.go index 6b73b0ce..dab20b90 100644 --- a/internal/control/reconciler/reconciler_test.go +++ b/internal/control/reconciler/reconciler_test.go @@ -262,9 +262,15 @@ func TestPacksWhenSpreadingIsImpossible(t *testing.T) { } } -func TestShardOrderWorkersFirst(t *testing.T) { - // SGLang ranks rendezvous and vLLM under Ray needs its workers attached - // before the head serves, so rank 0 goes last. +func TestShardOrderLeaderFirst(t *testing.T) { + // Both engines put the rendezvous point on rank 0 — SGLang's dist-init + // address and vLLM's Ray head — so a worker started first has nothing to + // connect to and spends its startup budget retrying. + // + // This test previously asserted the opposite, on the theory that the head + // should wait for its workers. That is backwards: vLLM's server blocks + // until the Ray cluster has the GPUs it needs, so the leader can start + // immediately and waits on its own. h := newHarness(t, testModel(1, catalog.Profile{ Name: "2node", GPUClass: "h100-80g", TP: 16, Nodes: 2, GPUsPerNode: 8, }), gpuNode("1", 8, "ib-a"), gpuNode("2", 8, "ib-a")) @@ -275,8 +281,9 @@ func TestShardOrderWorkersFirst(t *testing.T) { if len(started) != 2 { t.Fatalf("expected 2 shards, got %d", len(started)) } - if started[0].Spec.Rank != 1 || started[1].Spec.Rank != 0 { - t.Errorf("start order was rank %d then %d; workers must start first", + if started[0].Spec.Rank != 0 || started[1].Spec.Rank != 1 { + t.Errorf("start order was rank %d then %d; the leader hosts the rendezvous "+ + "point and must start first", started[0].Spec.Rank, started[1].Spec.Rank) } // Every rank must be given the same rendezvous point. diff --git a/llmgw-plan.md b/llmgw-plan.md index 0b629e10..60c9bdcd 100644 --- a/llmgw-plan.md +++ b/llmgw-plan.md @@ -1285,22 +1285,30 @@ work_dir: /var/lib/llmgw-agent - [x] NCCL environment generation, HCA selection - [x] Replica-level failure handling (one dead shard kills the replica) -Placement and environment generation landed early, with phase 3's work. **The -launch path has not, and a multi-node replica cannot currently come up.** Three -specific things are missing: +The launch path is now built too: -- [ ] **Workers never become Ready.** The agent health-checks every shard at - `127.0.0.1:`, but in a multi-node replica only rank 0 serves HTTP. - Workers time out, report Failed, and take the replica with them. The fix - needs the engine adapter to say which ranks serve — `ServesHTTP(spec)` — - so a worker is judged by liveness and the replica's readiness is rank 0's. -- [ ] **Ray is never started.** `VLLM.RayCommand` is written and tested but - nothing calls it, so `launch_mode: ray` runs `vllm serve` against a - cluster that does not exist. The supervisor assumes one process per - shard; Ray needs two, sequenced. -- [ ] **Startup is fire-and-forget.** Workers are *sent* before rank 0 but - nothing waits for them. SGLang tolerates this because its ranks - rendezvous; Ray does not. +- [x] **Only the leader is health-checked.** `Engine.ServesHTTP(spec)` says which + ranks answer the API; a worker is judged by staying alive, and the + replica's real readiness is the leader answering, which only happens once + every worker has joined. Health-checking a worker's port used to time out + and fail a replica that was coming up perfectly well. +- [x] **A shard can be several processes.** `Engine.Commands(spec)` returns them + in start order: every command but the last is a sidecar kept alive for the + shard's lifetime and torn down with it, and the last one's exit is the + shard's. vLLM under Ray is the reason — the leader runs a Ray head and + then the server, a worker runs only the join and no vLLM process at all. +- [x] **The leader starts first.** Both engines put the rendezvous point on rank + 0, so a worker started first has nothing to connect to. An earlier version + had this backwards; vLLM's server blocks until the Ray cluster has the + GPUs it needs, so the leader can start immediately and waits on its own. +- [x] **Rendezvous is waited for, not assumed.** A command can require that + `LeaderAddr:LeaderPort` accepts a connection before the next one starts, + bounded by `rendezvous_timeout` so a head that never appears fails the + shard instead of hanging it. + +- [ ] **Not verified against real engines.** Everything above is tested with + stub processes. The Ray sequencing, and whether the deployed vLLM version + wants `ray` or `mp`, can only be confirmed on real hardware. - **Exit criterion:** a 2-node replica on each fabric; killing a node rebuilds the replica on the same fabric and never spreads it across two. - **Exit criterion:** a 2-node replica on each fabric; killing a node rebuilds the replica on the same fabric and never spreads it to the other one.