|
|
@@ -5,11 +5,25 @@ Rust CLI tool for launching vLLM/SGLang containers across a DGX Spark cluster. S
|
|
|
## 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/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/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.
|
|
|
+- `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
|
|
|
@@ -25,14 +39,24 @@ cargo run health # Health check all cl
|
|
|
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 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
|
|
|
@@ -66,9 +90,10 @@ 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 },
|
|
|
+ { 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"
|
|
|
@@ -77,6 +102,27 @@ 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):
|
|
|
@@ -148,21 +194,24 @@ 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
|
|
|
+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",
|
|
|
]
|
|
|
|
|
|
-# Optional: override docker entrypoint (appended after image name)
|
|
|
-entrypoint = ["python", "-m", "vllm.entrypoints.openai.api_server"]
|
|
|
+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"
|
|
|
@@ -192,30 +241,35 @@ Both engines use `--nnodes`, `--node-rank`, and `--master-addr`.
|
|
|
|
|
|
### vLLM engine (default)
|
|
|
|
|
|
-Appends these flags (cannot be overridden by recipe):
|
|
|
+Appends these flags:
|
|
|
|
|
|
```
|
|
|
---tensor-parallel-size 1
|
|
|
+--tensor-parallel-size <recipe.tensor_parallel_size>
|
|
|
--pipeline-parallel-size <computed>
|
|
|
--nnodes <total>
|
|
|
--node-rank <rank>
|
|
|
--master-addr <cluster.master_ib_ip>
|
|
|
---distributed-executor-backend mp|ray
|
|
|
+--master-port <config.init_port> # default 25000
|
|
|
+--distributed-executor-backend ray # only for ray backend (mp is auto-detected)
|
|
|
+--host 0.0.0.0
|
|
|
+--port <recipe.port>
|
|
|
```
|
|
|
|
|
|
+Non-head nodes also receive `--headless`.
|
|
|
+
|
|
|
### SGLang engine
|
|
|
|
|
|
Appends:
|
|
|
|
|
|
```
|
|
|
---tp <tensor_parallel_size> # from recipe (e.g. 4 for 4 × single-GPU nodes)
|
|
|
+--tp <tensor_parallel_size>
|
|
|
--nnodes <total>
|
|
|
--node-rank <rank>
|
|
|
---master-addr <cluster.master_ib_ip>
|
|
|
+--dist-init-addr <master_ib_ip>:<init_port>
|
|
|
+--host 0.0.0.0
|
|
|
+--port <recipe.port>
|
|
|
```
|
|
|
|
|
|
-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:
|
|
|
@@ -223,19 +277,38 @@ For vLLM Ray backend, additionally appends:
|
|
|
```
|
|
|
-e RAY_ADDRESS=<master_ib_ip>:6379
|
|
|
-e RAY_PORT=<port>
|
|
|
+-e RAY_NODE_IP_ADDRESS=<node.ib_ip>
|
|
|
+-e RAY_OVERRIDE_NODE_IP_ADDRESS=<node.ib_ip>
|
|
|
```
|
|
|
|
|
|
-(For SGLang with Ray backend, the container's entrypoint must handle Ray setup; SGLang's native multi-node uses `mp` backend with NCCL.)
|
|
|
+### 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 in `config.toml`. Port mapping (`-p`) is applied on **head node only**.
|
|
|
+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 with:
|
|
|
+Containers are labelled:
|
|
|
- `sm121.recipe=<name>`
|
|
|
- `sm121.rank=<rank>`
|
|
|
|
|
|
-`--cluster auto` picks the first cluster in `config.toml` that has an NCCL section.
|
|
|
+`--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
|
|
|
|
|
|
@@ -264,11 +337,15 @@ With `--repair`, generates a Dockerfile that installs missing packages and rebui
|
|
|
## Dependencies
|
|
|
|
|
|
- `clap` 4.6 + `derive` — derive-based CLI
|
|
|
-- `russh` 0.60 + `russh-keys` 0.49 — SSH client
|
|
|
+- `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.10 — SHA-256 for host key fingerprints
|
|
|
+- `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/`)
|