diff --git a/deploy/bootstrap/onboard-node.sh b/deploy/bootstrap/onboard-node.sh new file mode 100755 index 0000000..2078e54 --- /dev/null +++ b/deploy/bootstrap/onboard-node.sh @@ -0,0 +1,174 @@ +#!/usr/bin/env bash +# +# Give a machine an identity in an llmgw cluster. +# +# Run on the host that holds the certificate authority, because that is the one +# thing that cannot be delegated: this script issues the credential by which the +# control plane will believe the node is who it says it is. +# +# What it does, in the order it matters: +# 1. issue a client certificate whose common name is the node id +# 2. copy the CA root, the certificate and its key to the machine +# 3. trust the root system-wide, so other tooling there stops complaining +# 4. install the binary, the configuration and the unit +# 5. start the agent +# +# The node id, the VM name, the hostname and the certificate common name are all +# the same string. They do not have to be, and every time they are not, someone +# spends an afternoon working out which of the four a message is talking about. +set -euo pipefail + +usage() { + cat >&2 <&2; exit 1; } +[ -f "$PKI/ca-key.pem" ] || { echo "no authority in $PKI; run 'llmgw ctl cert init-ca' first" >&2; exit 1; } + +SSH="ssh -o BatchMode=yes -o ConnectTimeout=10 $SSH_USER@$ADDR" +say() { printf ' %s\n' "$*"; } + +echo "onboarding $NODE at $ADDR" +$SSH true || { echo "cannot reach $SSH_USER@$ADDR over ssh" >&2; exit 1; } + +# 1. The certificate. Issued fresh every time: re-onboarding a machine should +# not depend on finding an old key, and the old one keeps working until it +# is revoked, which is a separate decision. +say "issuing a certificate" +TMP="$(mktemp -d)" +trap 'rm -rf "$TMP"' EXIT +"$BINARY" ctl cert issue-client "$NODE" --dir "$PKI" --out "$TMP" >/dev/null + +# 2 and 3. Identity, and trusting the root so future tooling on that machine +# does not treat our certificates as strangers. +say "installing the certificate and trusting the root" +scp -q "$PKI/ca.pem" "$TMP/$NODE.pem" "$TMP/$NODE-key.pem" "$SSH_USER@$ADDR:/tmp/" +$SSH "set -e + sudo install -d -m 0755 /etc/llmgw + sudo install -m 0644 /tmp/ca.pem /etc/llmgw/ca.pem + sudo install -m 0644 /tmp/$NODE.pem /etc/llmgw/agent.pem + sudo install -m 0600 /tmp/$NODE-key.pem /etc/llmgw/agent-key.pem + sudo chown root:root /etc/llmgw/* + sudo install -m 0644 /tmp/ca.pem /usr/local/share/ca-certificates/llmgw-ca.crt + sudo update-ca-certificates >/dev/null 2>&1 || true + rm -f /tmp/ca.pem /tmp/$NODE.pem /tmp/$NODE-key.pem" + +# 4. The binary, the configuration and the unit. +say "installing the binary" +scp -q "$BINARY" "$SSH_USER@$ADDR:/tmp/llmgw" +$SSH 'sudo install -m 0755 /tmp/llmgw /usr/local/bin/llmgw && rm -f /tmp/llmgw' + +say "writing the configuration" +FABRICS="fabrics: []" +if [ -n "$FABRIC_NAME" ]; then + FABRICS="fabrics: + - name: $FABRIC_NAME + hcas: [$(echo "$FABRIC_HCAS" | sed 's/,/, /g')]" +fi +cat > "$TMP/agent.yaml" < "$TMP/llmgw-agent.service" <<'UNIT' +[Unit] +Description=llmgw node agent +After=network-online.target +Wants=network-online.target + +[Service] +ExecStart=/usr/local/bin/llmgw agent --config /etc/llmgw/agent.yaml +# Root, because the agent launches engines and needs device access. Engines +# outlive a restart of this unit and are adopted again on the next start, so +# KillMode must not take the whole process group with it. +KillMode=process +Restart=on-failure +RestartSec=2s +StateDirectory=llmgw-agent + +[Install] +WantedBy=multi-user.target +UNIT +scp -q "$TMP/llmgw-agent.service" "$SSH_USER@$ADDR:/tmp/" + +say "starting the agent" +$SSH 'set -e + sudo install -m 0644 -o root -g root /tmp/agent.yaml /etc/llmgw/agent.yaml + sudo install -m 0644 /tmp/llmgw-agent.service /etc/systemd/system/llmgw-agent.service + rm -f /tmp/agent.yaml /tmp/llmgw-agent.service + sudo systemctl daemon-reload + sudo systemctl enable --now llmgw-agent >/dev/null + sudo systemctl restart llmgw-agent' + +echo +echo "$NODE onboarded. It should appear in 'llmgw ctl nodes' within a few seconds." +echo "If it does not, the control plane's log says why — a certificate whose" +echo "common name does not match the claimed node id is refused by design." diff --git a/deploy/bootstrap/provision-node.sh b/deploy/bootstrap/provision-node.sh new file mode 100755 index 0000000..c980580 --- /dev/null +++ b/deploy/bootstrap/provision-node.sh @@ -0,0 +1,83 @@ +#!/usr/bin/env bash +# +# Create a VM on Proxmox for a new llmgw node. +# +# Run on the Proxmox host. It does the half of onboarding that is about +# hardware; onboard-node.sh does the half that is about identity, and the two +# are separate because they are done by different people with different access. +# +# Idempotent only in the sense that it refuses to touch an existing VMID. A node +# that already exists is not something a provisioning script should have an +# opinion about. +set -euo pipefail + +usage() { + cat >&2 </dev/null 2>&1; then + echo "VM $VMID already exists. Pick a free id; this script will not adopt one." >&2 + exit 1 +fi +[ -f "$IMAGE" ] || { echo "no cloud image at $IMAGE" >&2; exit 1; } + +echo "creating $NAME as VM $VMID at $IP" +qm create "$VMID" \ + --name "$NAME" --cores "$CORES" --memory "$MEMORY" \ + --net0 "virtio,bridge=$BRIDGE" --ostype l26 --agent enabled=1 + +# Imported rather than installed: a cloud image plus cloud-init is the whole +# reason a node takes a minute to build instead of an afternoon. +qm importdisk "$VMID" "$IMAGE" "$STORAGE" >/dev/null +qm set "$VMID" --scsihw virtio-scsi-pci --scsi0 "$STORAGE:vm-$VMID-disk-0,discard=on,ssd=1" >/dev/null +qm disk resize "$VMID" scsi0 "${DISK}G" >/dev/null +qm set "$VMID" --ide2 "$STORAGE:cloudinit" --boot order=scsi0 >/dev/null +qm set "$VMID" \ + --ipconfig0 "ip=$IP,gw=$GW" \ + --nameserver "$GW 8.8.8.8" \ + --ciuser "$USER" \ + --sshkeys "$SSHKEYS" >/dev/null + +qm start "$VMID" +echo "started. cloud-init takes a minute; then run onboard-node.sh from the CA host." diff --git a/deploy/examples/server.yaml b/deploy/examples/server.yaml index 6998421..ce87346 100644 --- a/deploy/examples/server.yaml +++ b/deploy/examples/server.yaml @@ -206,15 +206,51 @@ control: port_range: [20000, 20999] mgmt_interface: eth0 +# Each model needs a source — a weights path or a HuggingFace id — and at least +# one profile if the control plane is to place it. Without profiles the model is +# catalogued but never started, which is the router-only mode below. models: - name: llama-3.3-70b + source: meta-llama/Llama-3.3-70B-Instruct aliases: [llama] fallbacks: [qwen2.5-32b] routing_profile: chat + scaling: { min: 1, max: 4 } + profiles: + # Two nodes of four H100s each. A multi-node profile requires a fabric, + # and a replica never spans two of them. + - name: tp8 + engine: vllm + gpu_class: h100-80g + tp: 8 + nodes: 2 + gpus_per_node: 4 + priority: 100 + weight: 8 + startup_timeout: 30m + # Where a whole node is free, one node is cheaper to schedule and faster + # to start than two. + - name: tp4-a100 + engine: vllm + gpu_class: a100-40g + tp: 4 + priority: 50 + weight: 3 + startup_timeout: 20m - name: qwen2.5-32b + source: Qwen/Qwen2.5-32B-Instruct aliases: [qwen] routing_profile: batch + scaling: { min: 1, max: 2 } + profiles: + - name: tp2 + engine: sglang + gpu_class: h100-80g + tp: 2 + priority: 100 + weight: 2 + startup_timeout: 15m # A static replica list, used only when no model declares a profile. That is the # router-only mode for testing against sim-engines; with the control plane diff --git a/internal/agent/engines/engine.go b/internal/agent/engines/engine.go index cdc3929..7e17f60 100644 --- a/internal/agent/engines/engine.go +++ b/internal/agent/engines/engine.go @@ -142,6 +142,7 @@ func NewRegistry(cfg Config) *Registry { r.Add(&VLLM{cfg: cfg.VLLM}) r.Add(&SGLang{cfg: cfg.SGLang}) r.Add(&LlamaCpp{cfg: cfg.LlamaCpp}) + r.Add(&Sim{cfg: cfg.Sim}) return r } @@ -172,6 +173,7 @@ type Config struct { VLLM VLLMConfig `koanf:"vllm"` SGLang SGLangConfig `koanf:"sglang"` LlamaCpp LlamaCppConfig `koanf:"llamacpp"` + Sim SimConfig `koanf:"sim"` } // --- shared helpers --- diff --git a/internal/agent/engines/engines_test.go b/internal/agent/engines/engines_test.go index f3cd948..ef8cb89 100644 --- a/internal/agent/engines/engines_test.go +++ b/internal/agent/engines/engines_test.go @@ -397,14 +397,14 @@ func TestValidateRequiresModelPath(t *testing.T) { func TestRegistry(t *testing.T) { r := NewRegistry(Config{}) - if got := r.Names(); !slices.Equal(got, []string{"llamacpp", "sglang", "vllm"}) { + if got := r.Names(); !slices.Equal(got, []string{"llamacpp", "sglang", "sim", "vllm"}) { t.Errorf("registry contents: %v", got) } if _, err := r.Get("vllm"); err != nil { t.Errorf("vllm missing: %v", err) } err := func() error { _, err := r.Get("nope"); return err }() - if err == nil || !strings.Contains(err.Error(), "llamacpp, sglang, vllm") { + if err == nil || !strings.Contains(err.Error(), "llamacpp, sglang, sim, vllm") { t.Errorf("an unknown engine should list what is available: %v", err) } } diff --git a/internal/agent/engines/sim.go b/internal/agent/engines/sim.go new file mode 100644 index 0000000..d8d608b --- /dev/null +++ b/internal/agent/engines/sim.go @@ -0,0 +1,85 @@ +package engines + +import ( + "fmt" + "os" + "strconv" +) + +// SimConfig locates the simulated engine. +type SimConfig struct { + // Bin is the llmgw binary itself; the engine is one of its subcommands. + // Empty means whatever started this agent, which is right almost always. + Bin string `koanf:"bin"` +} + +// Sim launches llmgw's own simulated engine. +// +// It is in the production registry on purpose. A cluster can be brought up and +// exercised end to end — placement, fabrics, multi-node sequencing, failure and +// rebuild — on machines that have no GPUs, which is the only way most of that +// path gets tested at all. The multi-node launch sequence went four phases +// without ever running, because the only things exercising it were stub +// processes that exercised the command line and nothing else. +// +// Nothing about it pretends to be real: the model name it serves is whatever it +// is told, the tokens are invented, and it says so in its own logs. +type Sim struct{ cfg SimConfig } + +// Name identifies the engine in the catalog. +func (*Sim) Name() string { return "sim" } + +// HealthPath is the simulated engine's readiness endpoint. +func (*Sim) HealthPath() string { return "/health" } + +// MetricsPath is its Prometheus endpoint, in vLLM's vocabulary. +func (*Sim) MetricsPath() string { return "/metrics" } + +// ServesHTTP reports whether this rank answers the API. +// +// Only the leader, exactly as vLLM and SGLang behave — which is the behaviour +// this engine exists to let us test. +func (*Sim) ServesHTTP(s ShardSpec) bool { return s.Leader() } + +// Validate accepts anything the arithmetic allows. +// +// No compute-capability checks and no quantization rules: there is no hardware +// to be wrong about, and inventing refusals would make this a worse stand-in +// for the engines it stands in for. +func (*Sim) Validate(p Profile, gpus []GPUInfo) error { + return validateCommon(p, gpus) +} + +// Commands returns the single process that makes up one shard. +func (s *Sim) Commands(spec ShardSpec) ([]Command, error) { + bin := s.cfg.Bin + if bin == "" { + exe, err := os.Executable() + if err != nil { + return nil, fmt.Errorf("locating the llmgw binary for the sim engine: %w", err) + } + bin = exe + } + p := spec.Profile + + args := []string{"sim-engine", + "--listen", "0.0.0.0:" + strconv.Itoa(spec.HTTPPort), + "--model", p.ServedModelName, + } + if spec.NNodes > 1 { + args = append(args, + "--nnodes", strconv.Itoa(spec.NNodes), + "--node-rank", strconv.Itoa(spec.Rank), + "--dist-init-addr", spec.LeaderAddr+":"+strconv.Itoa(spec.LeaderPort), + ) + } + args = append(args, p.Args...) + + return []Command{{ + Name: "sim-engine", Path: bin, Args: args, + Env: materializeEnv(baseEnv(spec)), + // Workers dial the leader, so the leader has to be listening first. + // Both ranks express the same wait because both know the same address. + WaitForLeaderPort: spec.NNodes > 1 && !spec.Leader(), + }}, nil +} diff --git a/internal/agent/engines/sim_test.go b/internal/agent/engines/sim_test.go new file mode 100644 index 0000000..15fb770 --- /dev/null +++ b/internal/agent/engines/sim_test.go @@ -0,0 +1,88 @@ +package engines + +import ( + "slices" + "testing" +) + +func simSpec(rank, nnodes int) ShardSpec { + return ShardSpec{ + ReplicaID: "rep-1", ShardID: "sh-1", Rank: rank, NNodes: nnodes, + LeaderAddr: "10.0.0.5", LeaderPort: 20001, + GPUIndices: []int{0, 1}, HTTPPort: 20000, MgmtInterface: "eth0", + HCAs: []string{"mlx5_0", "mlx5_1"}, + Profile: Profile{ + Name: "tp4", Engine: "sim", TP: 4, GPUsPerNode: 2, Nodes: nnodes, + ModelPath: "/models/x", ServedModelName: "big-model", + }, + } +} + +func TestOnlyTheLeaderServes(t *testing.T) { + // The behaviour that made replicas fail while they were coming up perfectly + // well: health checking a worker's port, which never opens. + if !(&Sim{}).ServesHTTP(simSpec(0, 2)) { + t.Error("the leader does not serve") + } + if (&Sim{}).ServesHTTP(simSpec(1, 2)) { + t.Error("a worker claims to serve") + } +} + +func TestAMultiNodeShardCarriesItsRendezvous(t *testing.T) { + cmds, err := (&Sim{cfg: SimConfig{Bin: "/usr/local/bin/llmgw"}}).Commands(simSpec(1, 2)) + if err != nil { + t.Fatal(err) + } + if len(cmds) != 1 { + t.Fatalf("%d commands", len(cmds)) + } + c := cmds[0] + + if c.Args[0] != "sim-engine" { + t.Errorf("the engine is a subcommand of llmgw: %v", c.Args) + } + for flag, want := range map[string]string{ + "--nnodes": "2", "--node-rank": "1", "--dist-init-addr": "10.0.0.5:20001", + } { + if got, _ := flagValue(c.Args, flag); got != want { + t.Errorf("%s = %q, want %q", flag, got, want) + } + } + // A worker must wait for the leader's port, or the launch is a race. + if !c.WaitForLeaderPort { + t.Error("a worker does not wait for the leader") + } +} + +func TestTheLeaderDoesNotWaitForItself(t *testing.T) { + cmds, _ := (&Sim{cfg: SimConfig{Bin: "x"}}).Commands(simSpec(0, 2)) + if cmds[0].WaitForLeaderPort { + t.Error("the leader waits for a port it is about to open") + } +} + +func TestASingleNodeShardCarriesNoRendezvous(t *testing.T) { + cmds, _ := (&Sim{cfg: SimConfig{Bin: "x"}}).Commands(simSpec(0, 1)) + for _, flag := range []string{"--nnodes", "--node-rank", "--dist-init-addr"} { + if slices.Contains(cmds[0].Args, flag) { + t.Errorf("a single-node shard was given %s", flag) + } + } +} + +func TestTheShardIsStillPinnedToItsFabric(t *testing.T) { + // The simulated engine does not use NCCL, but the environment is generated + // by the same code that generates it for the ones that do — so a cluster + // exercised with this engine is exercising that too. + cmds, _ := (&Sim{cfg: SimConfig{Bin: "x"}}).Commands(simSpec(1, 2)) + env := cmds[0].Env + for _, want := range []string{ + "NCCL_IB_HCA=mlx5_0,mlx5_1", "NCCL_SOCKET_IFNAME=eth0", + "CUDA_VISIBLE_DEVICES=0,1", "NCCL_IB_DISABLE=0", + } { + if !slices.Contains(env, want) { + t.Errorf("the environment is missing %s", want) + } + } +} diff --git a/internal/cli/simengine.go b/internal/cli/simengine.go index 77a2245..85266d1 100644 --- a/internal/cli/simengine.go +++ b/internal/cli/simengine.go @@ -17,15 +17,39 @@ import ( func newSimEngineCmd() *cobra.Command { cfg := sim.Default() var listen string + var rv sim.Rendezvous cmd := &cobra.Command{ Use: "sim-engine", Short: "OpenAI-compatible fake inference engine (for testing)", Long: "Simulates the engine behaviours that matter to the router:\n" + "continuous batching, KV block accounting, prefix caching, and vLLM-shaped\n" + - "/metrics. Lets the router be developed and tested without GPUs.", + "/metrics. Lets the router be developed and tested without GPUs.\n\n" + + "With --nnodes above 1 it also simulates a multi-node replica: rank 0\n" + + "listens at --dist-init-addr and does not serve until every other rank has\n" + + "joined, and only rank 0 answers the API. That is the part of the\n" + + "multi-node path stub processes cannot exercise — the sequencing, the\n" + + "waiting, and one dead shard taking the replica with it.", RunE: func(cmd *cobra.Command, _ []string) error { log := logger() + + ctx, stop := signal.NotifyContext(cmd.Context(), syscall.SIGINT, syscall.SIGTERM) + defer stop() + + // Before anything binds the API port: a worker must never look + // ready, and a leader must not serve until the replica exists. + if err := rv.Join(ctx, log); err != nil { + return err + } + if rv.Multi() && !rv.Leader() { + // A worker contributes GPUs and answers nothing. Health + // checking its port is what used to fail replicas that were + // coming up perfectly well. + log.Info("worker rank running; this rank serves no API", "rank", rv.Rank) + <-ctx.Done() + return nil + } + eng := sim.New(cfg) defer eng.Close() @@ -47,9 +71,6 @@ func newSimEngineCmd() *cobra.Command { "prefix_cache_blocks", cfg.PrefixCacheBlocks, ) - ctx, stop := signal.NotifyContext(cmd.Context(), syscall.SIGINT, syscall.SIGTERM) - defer stop() - errCh := make(chan error, 1) go func() { errCh <- srv.Serve(ln) }() @@ -70,6 +91,11 @@ func newSimEngineCmd() *cobra.Command { f := cmd.Flags() f.StringVar(&listen, "listen", ":8000", "address to listen on") + f.IntVar(&rv.NNodes, "nnodes", 1, "ranks in this replica; above 1 simulates a multi-node launch") + f.IntVar(&rv.Rank, "node-rank", 0, "this rank; 0 is the leader and the only one that serves") + f.StringVar(&rv.Addr, "dist-init-addr", "", "the leader's rendezvous address") + f.DurationVar(&rv.Timeout, "rendezvous-timeout", 5*time.Minute, + "how long to wait for the replica to assemble before failing") f.StringVar(&cfg.Model, "model", cfg.Model, "model name to serve") f.IntVar(&cfg.MaxNumSeqs, "max-num-seqs", cfg.MaxNumSeqs, "concurrent sequences (continuous batch capacity)") f.IntVar(&cfg.BlockSize, "block-size", cfg.BlockSize, "tokens per KV cache block") diff --git a/internal/config/example_test.go b/internal/config/example_test.go new file mode 100644 index 0000000..11ae62b --- /dev/null +++ b/internal/config/example_test.go @@ -0,0 +1,62 @@ +package config + +import ( + "os" + "path/filepath" + "testing" +) + +// TestTheDocumentedAgentExampleLoads parses the file shipped in deploy/examples. +// +// It exists because the example was wrong for the entire life of the project +// and nothing noticed. The fabric struct carried no koanf tags, so the `name:` +// key the example documents bound to nothing, and an agent configured exactly +// as instructed refused to start with "fabrics[0]: name is required" — naming +// the key the operator had just written. +// +// Every fabric until then had been configured in a Go test, where the tags do +// not matter. A published example that has never been parsed is a guess. +func TestTheDocumentedAgentExampleLoads(t *testing.T) { + path := filepath.Join("..", "..", "deploy", "examples", "agent.yaml") + if _, err := os.Stat(path); err != nil { + t.Skipf("no example to check: %v", err) + } + + cfg, err := LoadAgent(path) + if err != nil { + t.Fatalf("the documented example does not load: %v", err) + } + if len(cfg.Fabrics) == 0 { + t.Fatal("the example documents a fabric and none was parsed") + } + f := cfg.Fabrics[0] + if f.Fabric == "" { + t.Error("the fabric name did not bind") + } + if len(f.HCAs) == 0 { + t.Error("the fabric's adapters did not bind") + } + if f.RateGbps == 0 { + t.Error("the fabric's rate did not bind") + } + if cfg.NodeID == "" || cfg.Server == "" { + t.Errorf("node_id=%q server=%q", cfg.NodeID, cfg.Server) + } +} + +func TestTheDocumentedServerExampleLoads(t *testing.T) { + path := filepath.Join("..", "..", "deploy", "examples", "server.yaml") + if _, err := os.Stat(path); err != nil { + t.Skipf("no example to check: %v", err) + } + cfg, err := Load(path) + if err != nil { + t.Fatalf("the documented example does not load: %v", err) + } + if err := cfg.Validate(); err != nil { + t.Fatalf("the documented example does not validate: %v", err) + } + if cfg.Listen == "" || len(cfg.Models) == 0 { + t.Errorf("listen=%q models=%d", cfg.Listen, len(cfg.Models)) + } +} diff --git a/internal/control/inventory/inventory.go b/internal/control/inventory/inventory.go index d8e7598..7c5ecee 100644 --- a/internal/control/inventory/inventory.go +++ b/internal/control/inventory/inventory.go @@ -50,16 +50,22 @@ func (s NodeState) String() string { func (s NodeState) Schedulable() bool { return s == Ready } // FabricAttachment is a node's connection to one fabric. +// The koanf tags are not decoration. Without them the fields bind from their Go +// names, so the documented `name:` key silently produced an empty fabric and +// `rate_gbps` was never read at all — and the agent then refused to start with +// "fabrics[0]: name is required", naming the key the operator had just written. +// Every fabric in this project was configured in a test until the first one was +// configured from a file. type FabricAttachment struct { // Fabric is the operator-assigned name. It is not derived from the GID // prefix: OpenSM's default is fe80:: everywhere, so a matching prefix // proves nothing (plan §5.2). - Fabric string `json:"fabric"` + Fabric string `json:"fabric" koanf:"name"` // HCAs are the adapters attached to this fabric. - HCAs []string `json:"hcas"` + HCAs []string `json:"hcas" koanf:"hcas"` // RateGbps is the slowest active port's rate, since a collective runs at // the speed of its slowest link. - RateGbps int `json:"rate_gbps"` + RateGbps int `json:"rate_gbps" koanf:"rate_gbps"` } // Node is one machine in the cluster. diff --git a/internal/sim/cluster.go b/internal/sim/cluster.go new file mode 100644 index 0000000..b591e5c --- /dev/null +++ b/internal/sim/cluster.go @@ -0,0 +1,116 @@ +package sim + +import ( + "context" + "errors" + "fmt" + "log/slog" + "net" + "time" +) + +// Rendezvous makes a sim-engine behave like a multi-node one. +// +// It exists because the multi-node launch path — leader first, workers joining, +// only rank 0 serving, the whole replica failing when one shard dies — can only +// be tested against something that actually rendezvouses. Stub processes +// exercise the command construction and nothing else, which is how that path +// stayed unverified through four phases. +// +// The protocol is the smallest thing that has the same shape as vLLM's Ray head +// or SGLang's dist-init: rank 0 listens, every other rank dials, and rank 0 +// does not serve until all of them have. +type Rendezvous struct { + // Addr is the leader's dist-init endpoint. Empty means single node. + Addr string + // Rank 0 is the leader. NNodes counts every rank including it. + Rank int + NNodes int + // Timeout bounds the wait. A rendezvous that never completes must fail the + // shard rather than hang it, which is the difference between a replica that + // gets rebuilt and one that occupies its GPUs forever. + Timeout time.Duration +} + +// Multi reports whether this shard is part of a multi-node replica. +func (r Rendezvous) Multi() bool { return r.NNodes > 1 && r.Addr != "" } + +// Leader reports whether this rank serves the API. +func (r Rendezvous) Leader() bool { return r.Rank == 0 } + +// Join blocks until the replica is assembled. +// +// The leader waits for every worker to dial in; a worker dials the leader and +// holds the connection open, so the leader learns immediately when one dies. +func (r Rendezvous) Join(ctx context.Context, log *slog.Logger) error { + if !r.Multi() { + return nil + } + if r.Timeout <= 0 { + r.Timeout = 5 * time.Minute + } + ctx, cancel := context.WithTimeout(ctx, r.Timeout) + defer cancel() + + if r.Leader() { + return r.host(ctx, log) + } + return r.dial(ctx, log) +} + +func (r Rendezvous) host(ctx context.Context, log *slog.Logger) error { + ln, err := net.Listen("tcp", r.Addr) + if err != nil { + return fmt.Errorf("binding the rendezvous address %s: %w", r.Addr, err) + } + // Deliberately not closed: holding it open is how a worker that restarts + // can reconnect, and how the leader notices one going away. + go func() { + <-ctx.Done() + _ = ln.Close() + }() + + want := r.NNodes - 1 + log.Info("waiting for workers", "rank", r.Rank, "want", want, "addr", r.Addr) + + joined := 0 + for joined < want { + conn, err := ln.Accept() + if err != nil { + if ctx.Err() != nil { + return fmt.Errorf("only %d of %d workers joined before the deadline", joined, want) + } + return err + } + joined++ + log.Info("a worker joined", "joined", joined, "want", want, "from", conn.RemoteAddr()) + // Held open rather than closed: the connection is the worker's + // liveness, exactly as a real distributed runtime would treat it. + go func() { <-ctx.Done(); _ = conn.Close() }() + } + log.Info("every worker joined; serving", "nnodes", r.NNodes) + return nil +} + +func (r Rendezvous) dial(ctx context.Context, log *slog.Logger) error { + log.Info("joining the leader", "rank", r.Rank, "addr", r.Addr) + + // Retried, because the leader is started first but "started" and "listening" + // are not the same instant, and a worker that gives up on the first refusal + // would make the whole replica a race. + var lastErr error + for ctx.Err() == nil { + conn, err := (&net.Dialer{Timeout: 5 * time.Second}).DialContext(ctx, "tcp", r.Addr) + if err == nil { + log.Info("joined", "rank", r.Rank) + go func() { <-ctx.Done(); _ = conn.Close() }() + return nil + } + lastErr = err + select { + case <-ctx.Done(): + case <-time.After(250 * time.Millisecond): + } + } + return fmt.Errorf("could not reach the leader at %s: %w", r.Addr, errors.Join(lastErr, ctx.Err())) +}