AGENTS.md 13 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 entrypoint, command dispatch, i18n init.
  • src/cli.rs — clap derive-based CLI definitions (Commands, PowerAction, TmuxCommands).
  • src/config.rs — Config, Cluster, Node, Recipe, NcclConfig, HomeAssistantConfig deserialization; cluster resolution (exact + prefix match, auto-select smallest).
  • src/ssh.rs — SSH client session management, key loading, host key verification, retry logic.
  • src/docker/mod.rs — Docker command builder: serve, stop, health, logs, log streaming.
  • src/docker/tests.rs — Docker command builder tests (MP, SGLang, Ray, keep, port, env overrides).
  • src/tmux/mod.rs — tmux module root (session manager, serve, ops).
  • src/tmux/serve.rs — tmux serve logic (window/panel layout).
  • src/tmux/ops.rs — tmux stop/status/logs/attach operations.
  • src/health.rs — tmux health-check helpers.
  • src/power.rs — power on/off/switch flows (SSH shutdown + HA toggle).
  • src/home_assistant.rs — HA REST API client (turn_on/turn_off).
  • src/host_keys.rs — init-host-keys command (auto-write SSH host_keys).
  • src/exec.rs — exec subcommand: run commands on cluster/nodes, resolve aliases.
  • src/doctor.rs — Doctor mode: pre-flight checks and --repair (Dockerfile generation).
  • src/pre_serve.rs — BeforeServe/PreServeCommand types.
  • src/before_serve.rs — Pre-serve hook execution logic.
  • locales/{en,ru,es,it,fr,de}/main.ftl — Fluent i18n translation files.
  • config.toml — runtime config: SSH key, cluster definitions, per-cluster NCCL settings, HA, exec aliases.
  • 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 -- "..."                   # 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 exec --cluster 4nodes --alias gpu-lock           # Use exec alias
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)

### Power management

cargo run power off --cluster 4nodes-switch # SSH shutdown + HA turn_off cargo run power off --cluster 4 --node dgx1 # Prefix matching + single node cargo run power on --cluster 4nodes-switch # HA turn_on (wake nodes) cargo run power switch --from 4 --to 2 # Switch clusters (disable → enable) cargo run -- -d power off --cluster 4nodes-switch # Dry-run


### 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, entity = "switch.dgx1" },
{ name = "dgx2.lan", ssh_ip = "10.0.0.2", ib_ip = "1.2.3.5", rank = 1, entity = "switch.dgx2" },

] extra_entities = ["switch.network_switch"] # cluster-level switches (PDU, network)

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


### Home Assistant config

Optional `[home_assistant]` section for power management:

toml [home_assistant] url = "http://homeassistant.local:8123" token = "long-lived-access-token" shutdown_delay_seconds = 30


### Exec aliases

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


Use: `cargo run exec --cluster 4 --alias gpu-lock`

### 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 = 30000 # default: 30000 required_nodes = 4 # optional, omitted = no check pipeline_parallel_size = 0 # 0 = auto (equals nnodes), vLLM only tensor_parallel_size = 1 # per-node TP for vLLM, total TP for SGLang distributed_backend = "mp" # "mp" (multiprocessing) or "ray" shm_size = "64gb" # optional, only added when non-empty cluster = "4nodes-switch" # optional: pin to a specific cluster

args = [

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

]

entrypoint = ["vllm", "serve"] # optional, appended after image name

before_serve = [ # recipe-level pre-serve hooks (after cluster)

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

]

[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:

--tensor-parallel-size --pipeline-parallel-size --nnodes --node-rank --master-addr --master-port # default 25000 --distributed-executor-backend ray # only for ray backend (mp is auto-detected) --host 0.0.0.0 --port


Non-head nodes also receive `--headless`.

### SGLang engine

Appends:

--tp --nnodes --node-rank --dist-init-addr : --host 0.0.0.0 --port


### Ray backend

For vLLM Ray backend, additionally appends:

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

Env vars passed to every container

Cluster-driven (from [nccl."<name>"]): NCCL_IB_DISABLE, NCCL_IB_HCA, NCCL_SOCKET_IFNAME, NCCL_IB_GID_INDEX, GLOO_SOCKET_IFNAME, TP_SOCKET_IFNAME

Network interface (from nccl.socket_ifname): MN_IF_NAME, UCX_NET_DEVICES, OMPI_MCA_btl_tcp_if_include

Defaults: OMP_NUM_THREADS=4, NCCL_IGNORE_CPU_AFFINITY=1, NCCL_NET=IB, NCCL_DEBUG=WARN

Engine IP (from node.ib_ip): VLLM_HOST_IP (vLLM) or SGLANG_WORKER_IP (SGLang)

Head node only: -e VLLM_LOGGING_COLOR=1 -e TERM=xterm-256color (colored logs)

Recipe [env] entries are appended last and override any default.

Common to all runs

NCCL env vars come from the cluster's [nccl."<name>"] section. Port mapping -p <port>:<port> on the head node only. HF_HOME=/root/.cache/huggingface with volume mount.

Containers are labelled:

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

--cluster auto picks the smallest cluster (fewest nodes) that matches the recipe's required_nodes and has an NCCL config. Prefix matching works: --cluster 4 matches 4nodes-switch.

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.61 + russh-keys 0.50-beta — 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.11 — SHA-256 for host key fingerprints
  • base64 0.22 — Base64 encoding for fingerprints
  • dirs 6 — OS config dir for default config fallback
  • reqwest 0.13 — HTTP client for Home Assistant REST API
  • serde_json 1 — JSON for HA requests
  • fluent-i18n — i18n via Mozilla Fluent (locales in locales/)