AGENTS.md 6.3 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 -- -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 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)

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:

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

cluster = "4nodes-switch"

recipes/

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

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)
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. Engine name is used at runtime for host IP variables:

  • vllm → VLLM_HOST_IP
  • sglang → SGLANG_WORKER_IP

How Docker launch works

build_docker_command in main.rs constructs the docker run command for each node. Flavour depends on distributed_backend:

MP backend (default)

Always appends these flags (cannot be overridden by recipe):

--tensor-parallel-size 1
--pipeline-parallel-size <computed>
--nnodes <total>
--node-rank <rank>
--master-addr <cluster.master_ib_ip>
--distributed-executor-backend mp

Ray backend

Appends:

--tensor-parallel-size <nnodes>
--node-rank <rank>
--master-addr <cluster.master_ib_ip>

Plus -e RAY_ADDRESS=<master_ib_ip>:6379 and -e RAY_PORT=<port>.

Common to all runs

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

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