AGENTS.md 9.5 KB

AGENTS.md — vllm-spark-launcher

Rust CLI tool for launching vLLM/SGLang containers across a DGX Spark cluster. SSH-based provisioning over russh. Supports two runtimes: Docker (default) and tmux (alternative).

Project Layout

  • Cargo.toml — binary crate vllm-spark-launcher. Edition 2024.
  • src/main.rs — CLI parsing, Docker orchestration, Doctor mode, health check, log streaming.
  • src/ssh.rs — SSH client session management, key loading, host key verification, retry logic.
  • src/tmux.rs — tmux session manager: start/stop/status/logs/attach across a cluster.
  • src/health.rs — tmux health-check helpers (unused, see below).
  • config.toml — runtime config: SSH key, cluster definitions, per-cluster NCCL settings.
  • recipes/<name>/recipe.toml — recipe files.

Commands

Docker runtime

cargo run serve --recipe glm51                             # Auto-cluster selection
cargo run serve --recipe glm51 --cluster 4nodes-switch     # Explicit cluster
cargo run serve --recipe glm51 --skip-before-serve         # Bypass pre-serve hooks
cargo run -- -d serve --recipe glm51                       # Dry-run (print only)
cargo run stop --recipe glm51                              # Stop containers on all nodes
cargo run health                                           # Health check all clusters
cargo run health --cluster 2nodes-dac                      # Single cluster
cargo run health --fingerprint                             # Show host keys for config.toml
cargo run init-host-keys                                   # Auto-write SSH host_keys
cargo run exec --cluster 4nodes -- "sudo nvidia-smi -lgc 300,2300"   # Run on all nodes
cargo run exec --cluster 4nodes --node dgx1 -- "uptime"            # Run on one node
cargo run exec --recipe glm51 -- "cat /tmp/foo"                     # Use recipe's cluster
cargo run logs --recipe glm51                              # Last 100 lines per node
cargo run logs --recipe glm51 --follow                     # Follow head node logs
cargo run doctor --recipe glm51                            # Check node prerequisites
cargo run doctor --recipe glm51 --repair                   # Auto-fix (build repaired image)

### SGLang recipe example

cargo run serve --recipe minimax-m3 # SGLang, 4 nodes, TP=4 cargo run -- -d serve --recipe minimax-m3 # Dry-run (see the commands) cargo run stop --recipe minimax-m3 # Stop SGLang containers cargo run logs --recipe minimax-m3 # Check SGLang logs


### Tmux runtime

cargo run tmux serve --recipe glm51 # Launch via tmux (no Docker) cargo run tmux stop --recipe glm51 # Stop tmux sessions cargo run tmux status --recipe glm51 # Detailed health per node cargo run tmux logs --recipe glm51 # Capture pane output cargo run tmux logs --recipe glm51 --follow # Stream pane output cargo run tmux attach --recipe glm51 # Interactive attach (head node)


All subprocess calls support `-d` / `--dry-run` as a global flag.

## config.toml

User-specific. Has hardcoded SSH key paths — won't work without editing. Structure:

toml ssh_user = "username" ssh_key_path = "/absolute/path/to/key"

[[clusters]] name = "my-cluster" master_ib_ip = "1.2.3.4" nodes = [

# host_key = "SHA256:..."  # run 'cargo run health --fingerprint' to discover
{ name = "dgx1.lan", ssh_ip = "10.0.0.1", ib_ip = "1.2.3.4", rank = 0 },
{ name = "dgx2.lan", ssh_ip = "10.0.0.2", ib_ip = "1.2.3.5", rank = 1 },

]

[nccl."my-cluster"] ib_hca = "rocep1s0f0,roceP2p1s0f0" socket_ifname = "enp1s0f0np0" ib_gid_index = 3 ib_disable = 0


### Recipe-level cluster override

A recipe can pin itself to a specific cluster (bypasses `--cluster` / auto-pick):

toml cluster = "4nodes-switch"


### Pre-serve hooks (`[before_serve."<cluster>"]`)

Run shell commands on every node of a cluster **before** `serve` / `tmux serve`
launches containers or tmux sessions.  Defined in `config.toml`:

toml [before_serve."4nodes-switch"] commands = [

# Default: abort the serve if this fails
"sudo nvidia-smi -lgc 300,2300",

# Explicit on_failure:
#   "abort"    — fail-fast (default)
#   "warn"     — log warning, continue
#   "continue" — stay silent, continue
{ command = "sudo nvidia-smi -pm 1", on_failure = "warn" },

]


Recipes can also declare pre-serve commands that **append** to the cluster
config (so recipe-level commands run after cluster-level):

toml before_serve = [

{ command = "echo recipe-level", on_failure = "abort" },

]


Commands are run sequentially on each node (so order matters) but nodes
run in parallel.  Use `--skip-before-serve` to bypass for a one-off
serve.

### `exec` aliases (`[exec_aliases]`)

Map short names to full shell commands for use with
`cargo run exec --alias`.  Defined in `config.toml`:

toml [exec_aliases] gpu-lock = "sudo nvidia-smi -lgc 300,2300" shutdown = "sudo shutdown now" reboot = "sudo reboot"


Then:

bash cargo run exec --cluster 4nodes --alias gpu-lock cargo run exec --cluster 4nodes --alias shutdown


`--alias <name>` and a trailing command are mutually exclusive.  If the
alias name is unknown, the error message lists the configured aliases.

## recipes/

TOML format. Directory name is intentionally misspelled. Each recipe lives in its own subdirectory:

toml engine = "vllm" # or "sglang" (default: "vllm") name = "Human Readable Name" docker_image = "my-image:tag" model = "org/model-name" port = 8000 required_nodes = 4 # optional: 0 = any cluster size, omitted = no check pipeline_parallel_size = 0 # 0 = auto (equals nnodes), only for vLLM tensor_parallel_size = 1 # for vLLM: per-node TP (= 1 for single-GPU nodes);

                                   # for SGLang: total TP across all GPUs (= nnodes * GPUs_per_node)

distributed_backend = "mp" # "mp" (multiprocessing) or "ray" cluster = "4nodes-switch" # optional: pin to a specific cluster

args = [

"--kv-cache-dtype", "fp8",
"--max-model-len", "32768",

]

Optional: override docker entrypoint (appended after image name)

entrypoint = ["python", "-m", "vllm.entrypoints.openai.api_server"]

[env] SOME_VAR = "value"


### Engine support

Both `vllm` and `sglang` engines are supported. The engine field controls:

**Host IP env var:**
- `vllm` → `-e VLLM_HOST_IP=<node.ib_ip>`
- `sglang` → `-e SGLANG_WORKER_IP=<node.ib_ip>`

**Model argument:**
- `vllm` → bare model name (e.g. `org/model`)
- `sglang` → `--model-path org/model`

**Distributed flags:**
- `vllm` → `--tensor-parallel-size 1`, `--pipeline-parallel-size <pp_size>`, `--distributed-executor-backend mp|ray`
- `sglang` → `--tp <tensor_parallel_size>`, no `--pipeline-parallel-size`, no `--distributed-executor-backend`

Both engines use `--nnodes`, `--node-rank`, and `--master-addr`.

## How Docker launch works

`build_docker_command` in `src/docker.rs` constructs the `docker run` command for each node. It is engine-aware:

### vLLM engine (default)

Appends these flags (cannot be overridden by recipe):

--tensor-parallel-size 1 --pipeline-parallel-size --nnodes --node-rank --master-addr --distributed-executor-backend mp|ray


### SGLang engine

Appends:

--tp # from recipe (e.g. 4 for 4 × single-GPU nodes) --nnodes --node-rank --master-addr


The host IP env var is `SGLANG_WORKER_IP` and the model is passed as `--model-path <model>`.

### Ray backend

For vLLM Ray backend, additionally appends:

-e RAY_ADDRESS=:6379 -e RAY_PORT= ```

(For SGLang with Ray backend, the container's entrypoint must handle Ray setup; SGLang's native multi-node uses mp backend with NCCL.)

Common to all runs

NCCL env vars come from the cluster's [nccl."<name>"] section in config.toml. Port mapping (-p) is applied on head node only.

Containers are labelled with:

  • sm121.recipe=<name>
  • sm121.rank=<rank>

--cluster auto picks the first cluster in config.toml that has an NCCL section.

How tmux launch works

tmux serve builds sessions on each node with a 3-window layout:

Window Content
0: model Main inference process (vLLM/SGLang)
1: gpu nvidia-smi -l 5 GPU monitor
2: shell Debug shell

For Ray backend, nodes start sequentially: head node first (polled for API readiness up to 150s), then workers. For MP backend, all nodes start in parallel.

Doctor mode

cargo run doctor --recipe <name> checks each node for:

  • Docker availability
  • NVIDIA GPU (nvidia-smi)
  • NCCL libraries
  • Engine version inside the image (vllm version / sglang version)
  • Ray availability (if backend is "ray")
  • Container entrypoint

With --repair, generates a Dockerfile that installs missing packages and rebuilds the image with a .repaired tag. The repaired image is auto-detected on subsequent serve runs.

Dependencies

  • clap 4.6 + derive — derive-based CLI
  • russh 0.60 + russh-keys 0.49 — SSH client
  • tokio full — async runtime
  • toml 1 + serde 1 — config/recipe deserialization
  • tracing + tracing-subscriber — logging
  • anyhow — error handling
  • futures 0.3 — join_all for concurrent SSH
  • sha2 0.10 — SHA-256 for host key fingerprints
  • base64 0.22 — Base64 encoding for fingerprints