diff --git a/docs/get-started/dgx-station-preparation.mdx b/docs/get-started/dgx-station-preparation.mdx index 70b3d04f74..76c57aa172 100644 --- a/docs/get-started/dgx-station-preparation.mdx +++ b/docs/get-started/dgx-station-preparation.mdx @@ -162,6 +162,41 @@ The AI Developer Tools path can also leave the packaged CDI refresh units enable For the current support status and direct GPU policy boundaries, see [Platform Support](../../reference/platform-support). +## Prepare a Two-Station Pair + +Complete NVIDIA's [two-Station CX8 fabric playbook](https://build.nvidia.com/station/connect-two-stations/instructions) before running NemoClaw pair preparation. +The published [dual-Station Nemotron Ultra playbook](https://build.nvidia.com/station/nemoclaw/dual-nodes) is the runtime recipe tracked by this managed path. + +When no peer or model is selected explicitly, Station Express selects `nemotron-3-ultra-550b-a55b` and checks for one trusted reciprocal peer at the deterministic counterpart address on each of two configured private `/30` ConnectX-8 rails. +If both Stations pass reciprocal identity, GPU, route, neighbor, and jumbo-frame checks, NemoClaw uses the distributed Ultra recipe. +If no trusted pair qualifies, Express retains the existing single-Station Ultra recipe. +An explicit `NEMOCLAW_VLLM_MODEL` remains authoritative. +When `NEMOCLAW_DGX_STATION_PEER` is set, that exact peer must qualify or setup stops instead of falling back. + +Configure both Stations before installation with exactly two active 400 Gbit/s Ethernet rails, MTU 9000, and one usable private `/30` address per rail. +The reciprocal addresses must use direct-link routes, the expected peer MAC neighbors, and jumbo-frame connectivity in both directions. +The SSH target must already have usable host-key trust and non-interactive authentication, and the selected peer account must have passwordless `sudo` for remote preparation. +NemoClaw checks only the two deterministic `/30` counterpart addresses; it does not scan other addresses, configure the rails, enroll SSH trust, or accept a shared `/24` as equivalent evidence. + +Station preparation binds the preparing non-root account's UID in root-owned `/etc/nemoclaw/dual-station-controller-uid`. +Rebinding requires an administrator to remove that file before preparation is rerun as the replacement account. +If either Station requires a reboot, the installer stops with status `10` and names the host to reboot manually. +The initial local-host receipt preserves the accepted NemoClaw revision and Express selections. +After reciprocal peer qualification begins, owner-only pair state additionally binds the preparation helper, SSH host key, GPU identities, and reciprocal rails. +The installer never reboots either host automatically. + + +The dual-Station runtime uses unauthenticated Ray, NCCL, and vLLM coordination traffic, including the Ray head on TCP port `6379` and Ray's worker traffic. +Treat both Stations and every host that can reach either private rail as mutually trusted. +The operator owns physical rail isolation, host firewalling, SSH trust, and reboot control. +Do not use the pair path on a shared or routed network without separately reviewed isolation evidence. + + + +The dual-Station path is a Deferred evaluation. +It does not change the single-Station support status. + + ## Next Step After your host passes preparation or validation, continue with the [Quickstart](../quickstart). diff --git a/docs/inference/set-up-vllm.mdx b/docs/inference/set-up-vllm.mdx index 7704403c85..1a0be80ba8 100644 --- a/docs/inference/set-up-vllm.mdx +++ b/docs/inference/set-up-vllm.mdx @@ -19,10 +19,11 @@ The vLLM `/v1/responses` endpoint does not run the configured tool-call parser, Local vLLM does not require authentication by default. -NemoClaw needs port `8000` on host loopback for validation and on the OpenShell Docker bridge for sandbox traffic. +Existing-server and single-host managed-vLLM paths need port `8000` on host loopback for validation and on the OpenShell Docker bridge for sandbox traffic. Use a host firewall with default-deny inbound rules. Allow TCP port `8000` only from the OpenShell Docker subnet to its gateway address, keep loopback access, and deny the port on every other interface. Do not expose the port to your LAN or the internet. +The qualified dual-DGX Station path instead binds its bearer-protected `/v1` API to one selected private rail address and uses host networking on both runtime containers; follow the additional isolation guidance in the DGX Station section below. ## Use an Existing Server @@ -87,6 +88,7 @@ Managed profiles and model-specific recipes use immutable image digests: - DGX Spark and DGX Station models without a model-specific runtime use the `linux/arm64` digest `sha256:9204569b17ee4c0eff75194b8e6e458479c8aee18953b5ab9cf359fcdac659e2` with a compressed layer size of `9.60 GB` under `nvcr.io/nvidia/vllm:26.05.post1-py3`. - The DGX Station Nemotron 3 Ultra express recipe uses the multi-platform index digest `sha256:0fec7ec5f3e6bc168e54899935fb0557da908a4832a1dbc88e2debcf2f889416` under `vllm/vllm-openai:v0.22.0`; on DGX Station, that index selects a `linux/arm64` manifest with a compressed layer size of `10.67 GB`. It also pins Hugging Face revision `183968f87ae4cedce3039313cac1fd43d112c578` for the approximately `352.38 GB` model download. +- A qualified two-Station Ultra pair instead uses the ARM64 manifest digest `sha256:2cc49b81319f7a66a33dd8bd63a7bfddae079122b33ce51989b6828a1f038c37` under `vllm/vllm-openai:v0.25.1-aarch64`, with `10.24 GB` of compressed layers. - Generic Linux `arm64` hosts use `sha256:447995cbb57e6c7cf792cab95e9852e5f62b5fb6d2f39e030fa4eda9a54eadb4` with a compressed layer size of `9.28 GB` under `nvcr.io/nvidia/vllm:26.03.post1-py3`. - Generic Linux `amd64` hosts use `sha256:7be6c2f676c36059a494fe17254e69ae5c677535ba6191044e5fc8e42a91c773` with a compressed layer size of `8.93 GB` under `nvcr.io/nvidia/vllm:26.03.post1-py3`. @@ -181,7 +183,13 @@ When you start managed vLLM outside the installer express flow, NemoClaw uses th | DGX Station | `deepseek-ai/DeepSeek-V4-Flash` | | Linux with an NVIDIA GPU | `nvidia/NVIDIA-Nemotron-3-Nano-4B-FP8` | -On DGX Station, accepting the installer express prompt sets `NEMOCLAW_VLLM_MODEL=nemotron-3-ultra-550b-a55b` and overrides the profile default. +On DGX Station, accepting the installer express prompt selects `NEMOCLAW_VLLM_MODEL=nemotron-3-ultra-550b-a55b` and checks for a trusted reciprocal pair. +The installer prepares the local host with the pinned Station helper, derives exactly one counterpart from each of two configured private `/30` CX-8 rails, and consults existing SSH trust only for those two addresses. +At least one derived address must already be trusted; if both are trusted, their host keys must identify one coherent SSH host. +After the Stations pass reciprocal GPU, rail, MAC, route, neighbor, and jumbo-frame checks across both rails, the installer prepares the peer with the same reviewed helper and exports the qualified peer. +No managed-vLLM runtime image or model download starts before the pair-qualification or single-Station fallback decision completes. +If no trusted pair qualifies, Express retains the existing single-Station Ultra recipe. +Setting `NEMOCLAW_VLLM_MODEL` remains authoritative; setting `NEMOCLAW_DGX_STATION_PEER` requests that exact already-trusted peer and fails closed if it does not qualify. To select the existing `deepseek-v4-flash` recipe while retaining the same one-confirmation express flow, run: ```bash @@ -200,9 +208,22 @@ curl -fsSL https://www.nvidia.com/nemoclaw.sh | \ The flag requires an interactive terminal; in a `curl | bash` pipeline, `/dev/tty` must be available. Without terminal access, the installer stops before it installs Docker or build dependencies instead of silently continuing with another configuration. -The registered Ultra recipe tracks the [official DGX Station deployment guide](https://github.com/NVIDIA-NeMo/Nemotron/blob/287ae845639d2ce998998cb8fd1f70a3fa943c0b/usage-cookbook/Nemotron-3-Ultra/StationDeploymentGuide/README.md) and configures the pinned model revision, CPU offload, `16 GB` of shared memory, memory/stack ulimits, MTP speculative decoding, and the Nemotron reasoning and tool-call parsers. -NemoClaw intentionally keeps its existing bridge-networked managed-inference topology instead of importing the playbook's host-network setting. -The container publishes port `8000` through Docker, so apply the firewall guidance at the top of this page. +The registered single-Station Ultra recipe tracks the [official DGX Station deployment guide](https://github.com/NVIDIA-NeMo/Nemotron/blob/287ae845639d2ce998998cb8fd1f70a3fa943c0b/usage-cookbook/Nemotron-3-Ultra/StationDeploymentGuide/README.md) and configures the pinned model revision, CPU offload, `16 GB` of shared memory, memory/stack ulimits, MTP speculative decoding, and the Nemotron reasoning and tool-call parsers. +Configure the physical rails and SSH host-key and authentication trust before installation. +NemoClaw does not modify rail configuration, enroll SSH trust, or reboot either host automatically. +If the local host requires a reboot during initial preparation, the installer stops with status `10`; its owner-only receipt preserves the printed exact revision and Express selections. +After reciprocal peer qualification begins, the pair state additionally binds SSH, GPU, and rail identity and names either host that must be rebooted manually. +The registered two-Station Ultra recipe tracks the [latest NVIDIA dual-Station playbook](https://build.nvidia.com/station/nemoclaw/dual-nodes). +It pins vLLM `0.25.1` and Ray `2.56.0`, uses one tensor-parallel rank per Station with pipeline parallelism across the pair, serves the `nemotron-ultra` alias with a `262144`-token model limit, and retains the Nemotron reasoning and tool-call parsers. +The single-Station fallback keeps NemoClaw's bridge-networked managed-inference topology and publishes port `8000` through Docker. +The qualified dual-Station runtime intentionally uses Docker host networking on both containers so NCCL and RDMA can bind the validated direct-attach rails. +The head binds only the selected rank-0 address on a qualified private `/30` rail and requires the generated bearer key for `/v1`; `/health` remains unauthenticated for readiness. +The worker joins the Ray cluster and exposes no vLLM API. +Neither dual-Station container publishes a Docker port, so Docker bridge isolation and port-mapping rules do not protect this path. +As compensating controls, both containers run as the probed non-root UID and GID with a read-only root filesystem, all Linux capabilities dropped, `no-new-privileges`, only the selected GPU UUID and exact `uverbs` devices, and a read-only model cache; the worker does not receive the serving key. +Treat both Stations and their direct rails as one trusted runtime boundary. +Allow port `8000` from the OpenShell Docker subnet only to the selected rank-0 rail address, deny it on management and LAN interfaces, and restrict distributed-serving traffic to the reciprocal addresses on the two qualified private rails. +That distributed traffic includes the Ray head on TCP `6379` and Ray worker ports; do not expose either rail to an untrusted or routed network. Before managed vLLM setup on DGX Station, follow [Prepare DGX Station to Install NemoClaw](../../get-started/additional-setup/dgx-station-preparation). @@ -234,7 +255,9 @@ NEMOCLAW_PROVIDER=install-vllm \ On DGX Spark and DGX Station, `NEMOCLAW_PROVIDER=install-vllm` is sufficient for a non-interactive run. Add `NEMOCLAW_EXPERIMENTAL=1` on a generic Linux NVIDIA GPU host. Non-interactive runs use the profile default unless you set `NEMOCLAW_VLLM_MODEL`. -On DGX Station, a direct provider-only run therefore selects `deepseek-v4-flash`; the installer express flow sets the Nemotron 3 Ultra override for you. +The commands above invoke `nemoclaw onboard` directly, so a DGX Station run with no model or peer selects the `deepseek-v4-flash` profile default. +Supplying `NEMOCLAW_PROVIDER=install-vllm` to the shell installer enters the Station host-preparation boundary while retaining that profile default. +To request a non-interactive pair through this path, also set `NEMOCLAW_DGX_STATION_PEER`; the exact peer must qualify, and the installer selects the distributed Nemotron 3 Ultra recipe. For a headless DGX Station setup that selects DeepSeek V4 Flash explicitly, use the environment-variable path instead of `--station-deepseek`. diff --git a/docs/reference/commands.mdx b/docs/reference/commands.mdx index 261d5e4b8f..75590a9fbb 100644 --- a/docs/reference/commands.mdx +++ b/docs/reference/commands.mdx @@ -3562,7 +3562,9 @@ Set them before running `$$nemoclaw onboard`. | `SANDBOX_NAME` | sandbox name | Compatibility spelling used after `NEMOCLAW_SANDBOX_NAME` and `NEMOCLAW_SANDBOX`. | | `NEMOCLAW_INSTALL_REF` | git ref | For internal installer commands: the git ref to install from. A nonempty value takes precedence over `NEMOCLAW_INSTALL_TAG`. Overridden by the `--install-ref` flag. | | `NEMOCLAW_INSTALL_TAG` | release tag | For internal installer commands: the release tag to install when `NEMOCLAW_INSTALL_REF` is unset or empty. Defaults to the admin-promoted `lkg` tag when unset. Overridden by the `--install-tag` flag. | -| `NEMOCLAW_VLLM_MODEL` | registry slug or Hugging Face model id | Selects the model the managed-vLLM install path serves. Recognised slugs: `qwen3.6-27b`, `qwen3.6-35b-a3b-nvfp4`, `nemotron-3-nano-4b`, `deepseek-v4-flash`, `nemotron-3-ultra-550b-a55b`, `deepseek-r1-distill-70b`. Unset uses the per-platform profile default. The DGX Station express installer sets `nemotron-3-ultra-550b-a55b` explicitly. Gated models (e.g. `deepseek-r1-distill-70b`) require `HF_TOKEN` or `HUGGING_FACE_HUB_TOKEN`. | +| `NEMOCLAW_VLLM_MODEL` | registry slug or Hugging Face model id | Selects the model the managed-vLLM install path serves and remains authoritative during DGX Station installer setup. Recognised slugs: `qwen3.6-27b`, `qwen3.6-35b-a3b-nvfp4`, `nemotron-3-nano-4b`, `deepseek-v4-flash`, `nemotron-3-ultra-550b-a55b`, `deepseek-r1-distill-70b`. Station Express selects `nemotron-3-ultra-550b-a55b`; a qualified reciprocal pair uses the distributed topology, while no qualifying pair retains the single-Station Ultra topology. Outside Station Express, unset uses the per-platform profile default. Gated models (e.g. `deepseek-r1-distill-70b`) require `HF_TOKEN` or `HUGGING_FACE_HUB_TOKEN`. | +| `NEMOCLAW_DGX_STATION_PEER` | SSH host or `user@host` | Selects one exact, already-trusted DGX Station peer for Nemotron 3 Ultra pair qualification. The peer must match the reciprocal private `/30` rail and hardware checks; an explicit peer failure stops setup instead of falling back. NemoClaw does not enroll SSH trust or accept a port or SSH option in this value. When unset, DGX Station installer discovery checks only the two deterministic `/30` counterpart addresses. A peer cannot be combined with an explicit non-Ultra model; conflicting explicit selections fail before pair preparation. | +| `NEMOCLAW_DGX_STATION_SSH_BINDING` | opaque installer-managed token | Carries the qualified peer endpoint and host-key binding from DGX Station pair preparation into the current managed-vLLM install. The installer creates and clears this token; operators should not set or persist it. Missing, changed, or mismatched binding state fails before peer SSH or Docker work. | | `NEMOCLAW_VLLM_EXTRA_ARGS_JSON` | JSON array of non-blank strings | Appends advanced operator-owned tokens to the managed `vllm serve` command after NemoClaw's registry defaults. Example: `["--max-num-seqs","2"]`. Malformed JSON, non-string tokens, or blank tokens fail before Docker work starts. | | `NEMOCLAW_MINIMAL_BOOTSTRAP` | `1` to enable | Skips default OpenClaw workspace-template seeding for new pristine workspaces. Existing files are not deleted; refer to [Understand Runtime Changes](../manage-sandboxes/configure-sandboxes/understand-runtime-changes). | diff --git a/docs/resources/prompt-assets/dgx-station.md b/docs/resources/prompt-assets/dgx-station.md index 5f8a6cbfda..3e16fd7066 100644 --- a/docs/resources/prompt-assets/dgx-station.md +++ b/docs/resources/prompt-assets/dgx-station.md @@ -10,9 +10,10 @@ Use these instructions only after hardware detection confirms DGX Station. Use the selected maintained release's official installer as the authority for Station qualification, host preparation, model selection, consent, and reboot or login resume. Do not run the Station preparation helper separately or reproduce Express by pre-setting provider and model environment variables. -The installer provides these Station Express model choices: +The installer provides these Station Express choices: -1. The ordinary installer defaults to `nemotron-3-ultra-550b-a55b`, served as `nvidia/nemotron-3-ultra-550b-a55b`. +1. The ordinary installer selects `nemotron-3-ultra-550b-a55b` and checks for one already-trusted peer at the deterministic counterpart on each of two configured private `/30` ConnectX-8 rails. + A qualified pair uses the vLLM 0.25.1 and Ray 2.56.0 dual-Station recipe served as `nemotron-ultra`; otherwise it retains the single-Station Ultra recipe served as `nvidia/nemotron-3-ultra-550b-a55b`. 2. The explicit `--station-deepseek` flag selects `deepseek-v4-flash`, served as `deepseek-ai/DeepSeek-V4-Flash`. Both choices use the same Station detection, host-preparation, consent, suggested-policy, default-sandbox, and revision resume flow. @@ -22,21 +23,23 @@ Before asking for consent, explain all of these boundaries: - On generic Ubuntu, Station Express may install or change the pinned NVIDIA open driver, Docker with Buildx, NVIDIA Container Toolkit, and the reviewed factory `dkms` transition. On qualified factory images, the installer follows its bounded validation and repair path instead of replacing the factory stack. - Official Station preparation may add the trusted local account to the `docker` group, which grants root-equivalent control and is suitable only for a trusted single-user development host. - Official Station preparation may require an operator-controlled reboot and resumes only with the accepted NemoClaw revision. +- NemoClaw does not configure the two private rails, scan the network, enroll SSH trust, or reboot either Station. The operator owns physical isolation, firewalling, SSH trust, and manual reboots. +- The dual-Station runtime uses unauthenticated Ray, NCCL, and vLLM coordination traffic, including the Ray head on TCP port `6379` and Ray worker traffic. Both Stations and every host that can reach either rail must be mutually trusted; a shared `/24` is not equivalent to the required direct private `/30` rails. - Nemotron Ultra Express discloses an approximately `352 GB` model download. DeepSeek Express downloads its pinned vLLM container and model data. Both require enough space on the model-cache filesystem and Docker storage. - DGX Station is tested with limitations across qualified profiles on one physical DGX Station GB300. - Dual-Station configurations are not yet validated, and dedicated CI coverage is not available. -Ask: "Which DGX Station Express model would you like?" +Ask: "Which DGX Station Express option would you like?" Choices: -1. Nemotron 3 Ultra 550B, the ordinary installer default. +1. Automatic pair selection: use Nemotron 3 Ultra 550B with a qualified trusted pair when available; otherwise use the single-Station Ultra recipe. 2. DeepSeek V4 Flash, the explicit `--station-deepseek` override. 3. Neither, let me choose the runtime and model normally. If a Station Express model is selected: - Set `NEMOCLAW_AGENT` to the agent already selected in the starter prompt. -- For Nemotron Ultra, run the ordinary installer without `--station-deepseek`. +- For automatic pair selection, run the ordinary installer without `--station-deepseek`. Do not supply a peer unless the user already selected an exact, pretrusted peer; an explicit peer must qualify or setup stops rather than falling back. - For DeepSeek, pass `--station-deepseek` and no other model-selection override. - Do not set `NEMOCLAW_PROVIDER`, `NEMOCLAW_VLLM_MODEL`, `NEMOCLAW_MODEL`, `NEMOCLAW_NON_INTERACTIVE`, `NEMOCLAW_YES`, `NEMOCLAW_ACCEPT_THIRD_PARTY_SOFTWARE`, or `NEMOCLAW_NO_EXPRESS`. - Leave `NEMOCLAW_SANDBOX_NAME`, `NEMOCLAW_POLICY_TIER`, web-search settings, and messaging settings unset so the installer applies its Express defaults. @@ -45,6 +48,7 @@ If a Station Express model is selected: - Let the installer present its third-party-software notice and complete Express summary. Keep each official confirmation visible, wait for the user's response, and do not pre-answer or suppress it. - Do not pass `--force-station-install` unless the installer rejects release metadata on genuine Station GB300 hardware and the user separately chooses the documented temporary override. - Follow the command that the installer prints after a required reboot or login transition. +- Describe Ultra as the ordinary Express selection. Describe the distributed two-Station topology as selected only after the installer reports that reciprocal Station, GPU, rail, MAC, route, neighbor, and jumbo-frame checks qualified the pair. If Station Express is declined, continue with the normal provider selection. Offer existing vLLM when a ready server is detected, managed vLLM, supported local Ollama, and every hosted or compatible provider supported by the selected agent. diff --git a/docs/resources/starter-prompt.md b/docs/resources/starter-prompt.md index 9965d363c2..257641dd68 100644 --- a/docs/resources/starter-prompt.md +++ b/docs/resources/starter-prompt.md @@ -79,9 +79,9 @@ Set `NEMOCLAW_AGENT=langchain-deepagents-code` for Deep Agents, or use `nemo-dee After the readiness check, load exactly one matching instruction asset before provider selection: -- Confirmed DGX Spark: [DGX Spark Express instructions](https://raw.githubusercontent.com/NVIDIA/NemoClaw/ebf29b024c7a24e69ab8386b68e426d31a4c6821/docs/resources/prompt-assets/dgx-spark.md). -- Confirmed DGX Station: [DGX Station installation instructions](https://raw.githubusercontent.com/NVIDIA/NemoClaw/ebf29b024c7a24e69ab8386b68e426d31a4c6821/docs/resources/prompt-assets/dgx-station.md). -- Officially detected Windows WSL: [Windows WSL Express instructions](https://raw.githubusercontent.com/NVIDIA/NemoClaw/ebf29b024c7a24e69ab8386b68e426d31a4c6821/docs/resources/prompt-assets/windows-wsl.md). +- Confirmed DGX Spark: [DGX Spark Express instructions](https://raw.githubusercontent.com/NVIDIA/NemoClaw/bf46e62f901825f19e570c17f8c870a0eae04fbc/docs/resources/prompt-assets/dgx-spark.md). +- Confirmed DGX Station: [DGX Station installation instructions](https://raw.githubusercontent.com/NVIDIA/NemoClaw/bf46e62f901825f19e570c17f8c870a0eae04fbc/docs/resources/prompt-assets/dgx-station.md). +- Officially detected Windows WSL: [Windows WSL Express instructions](https://raw.githubusercontent.com/NVIDIA/NemoClaw/bf46e62f901825f19e570c17f8c870a0eae04fbc/docs/resources/prompt-assets/windows-wsl.md). Read the matching raw Markdown file completely and follow it before continuing. Do not load a platform asset for any other computer. diff --git a/scripts/checks/vitest-project-overlap.mts b/scripts/checks/vitest-project-overlap.mts index f668590f85..d7e21b9b37 100644 --- a/scripts/checks/vitest-project-overlap.mts +++ b/scripts/checks/vitest-project-overlap.mts @@ -48,6 +48,9 @@ const INSTALLER_INTEGRATION_TESTS = new Set([ "test/install-openshell-version-check.test.ts", "test/install-preflight-docker-bootstrap.test.ts", "test/install-preflight.test.ts", + "test/install-station-controller-binding.test.ts", + "test/install-station-pair-preparation.test.ts", + "test/install-station-resume-cleanup.test.ts", "test/install-station-dgx-os.test.ts", "test/install-station-docker-repository.test.ts", "test/install-station-host-preparation.test.ts", diff --git a/scripts/install.sh b/scripts/install.sh index 02882a6257..6e73907539 100755 --- a/scripts/install.sh +++ b/scripts/install.sh @@ -259,18 +259,30 @@ resolve_nemoclaw_gateway_port() { printf "%s" "$port" } +nemoclaw_state_root() { + local home="${HOME%/}" + [ -n "$home" ] || home="/" + if [ "$home" = "/" ]; then + printf "/.nemoclaw" + else + printf "%s/.nemoclaw" "$home" + fi +} + nemoclaw_state_dir() { - local port + local port root port="$(resolve_nemoclaw_gateway_port)" || return 1 + root="$(nemoclaw_state_root)" || return 1 if [ "$port" -eq 8080 ]; then - printf "%s/.nemoclaw" "$HOME" + printf "%s" "$root" else - printf "%s/.nemoclaw/gateways/%s" "$HOME" "$port" + printf "%s/gateways/%s" "$root" "$port" fi } assert_nemoclaw_state_path_safe() { - local target="$1" root="${HOME}/.nemoclaw" current relative component + local target="$1" root current relative component + root="$(nemoclaw_state_root)" || return 1 case "$target" in "$root" | "$root"/*) ;; *) error "Refusing NemoClaw state path outside ${root}: ${target}" ;; @@ -296,7 +308,7 @@ assert_nemoclaw_state_path_safe() { ensure_nemoclaw_state_dir() { local state_dir root gateways_dir state_dir="$(nemoclaw_state_dir)" || return 1 - root="${HOME}/.nemoclaw" + root="$(nemoclaw_state_root)" || return 1 gateways_dir="${root}/gateways" assert_nemoclaw_state_path_safe "$state_dir" (umask 077 && mkdir -p "$state_dir") || error "Could not create NemoClaw state directory: ${state_dir}" @@ -3353,10 +3365,16 @@ validate_express_platform_boundary() { STATION_ULTRA_VLLM_MODEL="nemotron-3-ultra-550b-a55b" STATION_ULTRA_SERVED_MODEL="nvidia/nemotron-3-ultra-550b-a55b" +STATION_ULTRA_DUAL_SERVED_MODEL="nemotron-ultra" +STATION_ULTRA_LEGACY_VLLM_IMAGE="vllm/vllm-openai@sha256:0fec7ec5f3e6bc168e54899935fb0557da908a4832a1dbc88e2debcf2f889416" STATION_DEEPSEEK_VLLM_MODEL="deepseek-v4-flash" STATION_DEEPSEEK_SERVED_MODEL="deepseek-ai/DeepSeek-V4-Flash" _SELECTED_EXPRESS_PLATFORM="" _STATION_EXPRESS_RESUME_REVISION="" +_STATION_EXPRESS_MODEL_WAS_EXPLICIT=0 +_STATION_EXPRESS_DEFERRED_MANAGED_PAIR=0 +_STATION_EXPRESS_MIGRATING_LEGACY_HEAD=0 +_STATION_INSTALL_MODE="" _STATION_EXPRESS_RESUME_LOADED="" _STATION_EXPRESS_RESUME_GENERATION="" _STATION_EXPRESS_RESUME_GATEWAY_PORT="" @@ -3479,18 +3497,32 @@ preflight_explicit_express_flags() { configure_station_express_model() { local selected_model selected_model="$(normalize_station_vllm_model "${NEMOCLAW_VLLM_MODEL:-}")" + _STATION_EXPRESS_MODEL_WAS_EXPLICIT=0 if [ "${STATION_DEEPSEEK:-}" = "1" ]; then + _STATION_EXPRESS_MODEL_WAS_EXPLICIT=1 NEMOCLAW_VLLM_MODEL="$STATION_DEEPSEEK_VLLM_MODEL" NEMOCLAW_MODEL="$STATION_DEEPSEEK_SERVED_MODEL" elif [ -z "$selected_model" ]; then - NEMOCLAW_VLLM_MODEL="$STATION_ULTRA_VLLM_MODEL" - NEMOCLAW_MODEL="$STATION_ULTRA_SERVED_MODEL" + if [ "${_STATION_INSTALL_MODE:-}" = "express" ] || [ -n "${NEMOCLAW_DGX_STATION_PEER:-}" ]; then + # Station Express keeps its single-host Ultra default. Pair qualification + # changes only the serving topology. An explicit peer also selects Ultra + # for a provider-driven, non-interactive pair setup. + NEMOCLAW_VLLM_MODEL="$STATION_ULTRA_VLLM_MODEL" + NEMOCLAW_MODEL="$STATION_ULTRA_SERVED_MODEL" + else + # A direct provider selection outside Express keeps the Station profile + # default unless the operator selects a model or peer. + unset NEMOCLAW_VLLM_MODEL + fi else + [ "$selected_model" != "auto" ] \ + || error "NEMOCLAW_VLLM_MODEL=auto is reserved for DGX Station reboot-resume state. Unset NEMOCLAW_VLLM_MODEL to request automatic selection." + _STATION_EXPRESS_MODEL_WAS_EXPLICIT=1 case "$selected_model" in "$STATION_ULTRA_VLLM_MODEL" | "nvidia/nvidia-nemotron-3-ultra-550b-a55b-nvfp4") NEMOCLAW_MODEL="$STATION_ULTRA_SERVED_MODEL" ;; - "$STATION_ULTRA_SERVED_MODEL") + "$STATION_ULTRA_SERVED_MODEL" | "$STATION_ULTRA_DUAL_SERVED_MODEL") # The served alias is useful in route output but is not a Hugging Face # repository ID. Normalize it to the registered model slug before the # existing managed-vLLM selector consumes it. @@ -3502,12 +3534,30 @@ configure_station_express_model() { ;; esac fi - export NEMOCLAW_VLLM_MODEL + if [ -n "${NEMOCLAW_VLLM_MODEL:-}" ]; then + export NEMOCLAW_VLLM_MODEL + fi if [ -n "${NEMOCLAW_MODEL:-}" ]; then export NEMOCLAW_MODEL fi } +station_dual_model_requested() { + local selected_model + selected_model="$(normalize_station_vllm_model "${NEMOCLAW_VLLM_MODEL:-}")" + case "$selected_model" in + "$STATION_ULTRA_VLLM_MODEL" | "$STATION_ULTRA_SERVED_MODEL" | "$STATION_ULTRA_DUAL_SERVED_MODEL" | "nvidia/nvidia-nemotron-3-ultra-550b-a55b-nvfp4") + return 0 + ;; + "") + [ "${_STATION_INSTALL_MODE:-}" = "express" ] || [ -n "${NEMOCLAW_DGX_STATION_PEER:-}" ] + ;; + *) + return 1 + ;; + esac +} + station_express_resume_file() { local state_dir state_dir="$(nemoclaw_state_dir)" || return 1 @@ -3543,6 +3593,10 @@ validate_station_express_resume_revision() { [[ "${1:-}" =~ ^[0-9a-f]{40}$ ]] } +validate_station_install_mode() { + [[ "${1:-}" == "express" || "${1:-}" == "provider" ]] +} + validate_station_express_resume_generation() { [[ "${1:-}" =~ ^[0-9a-f]{32}$ ]] } @@ -3581,11 +3635,20 @@ validate_station_express_resume_ports_distinct() { station_express_resume_generation() { local generation - [[ -r /proc/sys/kernel/random/uuid ]] \ - || error "Could not generate a DGX Station express resume receipt identity." - IFS= read -r generation "$temp_file"; then + "$gateway_port" "$dashboard_port" "$vllm_port" "$mode" >"$temp_file"; then rm -f "$temp_file" error "Could not write DGX Station express resume state under ${state_dir}." fi @@ -3793,11 +3894,16 @@ save_station_express_resume() { _STATION_EXPRESS_RESUME_GATEWAY_PORT="$gateway_port" _STATION_EXPRESS_RESUME_DASHBOARD_PORT="$dashboard_port" _STATION_EXPRESS_RESUME_VLLM_PORT="$vllm_port" + _STATION_INSTALL_MODE="$mode" } station_express_resume_command() { - printf 'curl -fsSL https://www.nvidia.com/nemoclaw.sh | NEMOCLAW_INSTALL_TAG=%s NEMOCLAW_AGENT=%s NEMOCLAW_SANDBOX_NAME=%s NEMOCLAW_POLICY_TIER=%s NEMOCLAW_GATEWAY_PORT=%s NEMOCLAW_DASHBOARD_PORT=%s NEMOCLAW_VLLM_PORT=%s bash' \ - "$_STATION_EXPRESS_RESUME_REVISION" "$_STATION_EXPRESS_RESUME_AGENT" \ + local provider_assignment="" + if [ "${_STATION_INSTALL_MODE:-express}" = "provider" ]; then + provider_assignment="NEMOCLAW_PROVIDER=install-vllm " + fi + printf 'curl -fsSL https://www.nvidia.com/nemoclaw.sh | %sNEMOCLAW_INSTALL_TAG=%s NEMOCLAW_AGENT=%s NEMOCLAW_SANDBOX_NAME=%s NEMOCLAW_POLICY_TIER=%s NEMOCLAW_GATEWAY_PORT=%s NEMOCLAW_DASHBOARD_PORT=%s NEMOCLAW_VLLM_PORT=%s bash' \ + "$provider_assignment" "$_STATION_EXPRESS_RESUME_REVISION" "$_STATION_EXPRESS_RESUME_AGENT" \ "$_STATION_EXPRESS_RESUME_SANDBOX" "$_STATION_EXPRESS_RESUME_POLICY_TIER" \ "$_STATION_EXPRESS_RESUME_GATEWAY_PORT" "$_STATION_EXPRESS_RESUME_DASHBOARD_PORT" \ "$_STATION_EXPRESS_RESUME_VLLM_PORT" @@ -3945,6 +4051,9 @@ express_wsl_can_use_windows_host_ollama() { activate_express_install() { local platform="$1" _SELECTED_EXPRESS_PLATFORM="$platform" + if [ "$platform" = "DGX Station" ]; then + _STATION_INSTALL_MODE="express" + fi NON_INTERACTIVE=1 export NEMOCLAW_NON_INTERACTIVE=1 export NEMOCLAW_NON_INTERACTIVE_SUDO_MODE=prompt @@ -3980,6 +4089,19 @@ activate_express_install() { esac } +resume_loaded_station_install() { + local platform="$1" + if [ "$_STATION_INSTALL_MODE" = "provider" ]; then + _SELECTED_EXPRESS_PLATFORM="$platform" + export NEMOCLAW_PROVIDER=install-vllm + configure_station_express_model + info "Detected DGX Station. Resuming the accepted managed-vLLM provider setup after host preparation." + else + info "Detected DGX Station. Resuming the accepted express install after host preparation." + activate_express_install "$platform" + fi +} + run_station_host_preparation() { # Public curl|bash starts in the root bootstrap, which clones the complete # selected ref before executing this payload. Keep the sibling lookup and @@ -4011,9 +4133,81 @@ filter_station_host_preparation_output() { done } +station_local_default_docker() { + ( + unset DOCKER_HOST DOCKER_CONTEXT + docker --context default "$@" + ) +} + +station_managed_dual_head_running() { + command_exists docker || return 1 + + local inspection name running managed role schema cluster launch_contract api_fingerprint transaction + inspection="$( + station_local_default_docker container inspect --format \ + '{{.Name}} {{.State.Running}} {{index .Config.Labels "com.nvidia.nemoclaw.managed-vllm"}} {{index .Config.Labels "com.nvidia.nemoclaw.vllm-role"}} {{index .Config.Labels "com.nvidia.nemoclaw.vllm-launch-schema"}} {{index .Config.Labels "com.nvidia.nemoclaw.vllm-cluster"}} {{index .Config.Labels "com.nvidia.nemoclaw.vllm-launch-contract"}} {{index .Config.Labels "com.nvidia.nemoclaw.vllm-api-key-fingerprint"}} {{index .Config.Labels "com.nvidia.nemoclaw.vllm-transaction"}}' \ + nemoclaw-vllm 2>/dev/null + )" || return 1 + read -r name running managed role schema cluster launch_contract api_fingerprint transaction <<<"$inspection" + [[ "$name" == "/nemoclaw-vllm" && + "$running" == "true" && + "$managed" == "true" && + "$role" == "head" && + "$schema" == "1" && + "$cluster" =~ ^[a-f0-9]{64}$ && + "$launch_contract" =~ ^[a-f0-9]{64}$ && + "$api_fingerprint" =~ ^[a-f0-9]{64}$ && + "$transaction" =~ ^[a-f0-9]{32}$ ]] +} + +station_migratable_legacy_single_head_running() { + command_exists docker || return 1 + + local inspection name running image managed role endpoint cluster gpu schema launch_contract api_fingerprint transaction + inspection="$( + station_local_default_docker container inspect --format \ + '{{.Name}}|{{.State.Running}}|{{.Config.Image}}|{{with index .Config.Labels "com.nvidia.nemoclaw.managed-vllm"}}{{.}}{{else}}-{{end}}|{{with index .Config.Labels "com.nvidia.nemoclaw.vllm-role"}}{{.}}{{else}}-{{end}}|{{with index .Config.Labels "com.nvidia.nemoclaw.vllm-endpoint"}}{{.}}{{else}}-{{end}}|{{with index .Config.Labels "com.nvidia.nemoclaw.vllm-cluster"}}{{.}}{{else}}-{{end}}|{{with index .Config.Labels "com.nvidia.nemoclaw.vllm-gpu"}}{{.}}{{else}}-{{end}}|{{with index .Config.Labels "com.nvidia.nemoclaw.vllm-launch-schema"}}{{.}}{{else}}-{{end}}|{{with index .Config.Labels "com.nvidia.nemoclaw.vllm-launch-contract"}}{{.}}{{else}}-{{end}}|{{with index .Config.Labels "com.nvidia.nemoclaw.vllm-api-key-fingerprint"}}{{.}}{{else}}-{{end}}|{{with index .Config.Labels "com.nvidia.nemoclaw.vllm-transaction"}}{{.}}{{else}}-{{end}}' \ + nemoclaw-vllm 2>/dev/null + )" || return 1 + IFS='|' read -r name running image managed role endpoint cluster gpu schema launch_contract api_fingerprint transaction <<<"$inspection" + [[ "$name" == "/nemoclaw-vllm" && + "$running" == "true" && + "$image" == "$STATION_ULTRA_LEGACY_VLLM_IMAGE" && + "$managed" == "true" && + "$role" == "-" && + "$endpoint" == "-" && + "$cluster" == "-" && + "$gpu" == "-" && + "$schema" == "-" && + "$launch_contract" == "-" && + "$api_fingerprint" == "-" && + "$transaction" == "-" ]] +} + ensure_station_express_host() { [[ "${_SELECTED_EXPRESS_PLATFORM:-}" == "DGX Station" ]] || return 0 + _STATION_EXPRESS_DEFERRED_MANAGED_PAIR=0 + _STATION_EXPRESS_MIGRATING_LEGACY_HEAD=0 + if station_dual_model_requested && station_managed_dual_head_running; then + # The host-preparation helper intentionally refuses active workloads. Only + # this complete dual-head ownership contract may defer local preparation; + # the post-Node coordinator revalidates the reciprocal physical pair before + # the existing lifecycle is allowed to reuse it. + _STATION_EXPRESS_DEFERRED_MANAGED_PAIR=1 + info "Found a complete running NemoClaw-managed dual-Station head candidate; deferring host preparation until reciprocal pair and lifecycle validation." + return 0 + fi + if station_dual_model_requested && station_migratable_legacy_single_head_running; then + # This exact frozen legacy head is the only single-host workload that the + # dual lifecycle can migrate. Bind its controller without running the + # workload-rejecting host probes, then prepare the newly qualified peer. + _STATION_EXPRESS_MIGRATING_LEGACY_HEAD=1 + info "Found the exact running legacy single-Station Ultra head; deferring workload-safe controller binding and reciprocal peer preparation." + return 0 + fi + # Publish the accepted, secret-free recipe before the helper can mutate the # host. The same receipt then binds reboot/login continuation and onboarding # recovery to this exact Express attempt. @@ -4066,8 +4260,222 @@ ensure_station_express_host() { esac } +station_dual_pair_resume_file() { + local state_dir + state_dir="$(nemoclaw_state_dir)" || return 1 + printf '%s/station-dual-pair-resume.json' "$state_dir" +} + +station_dual_pair_resume_pending() { + local state_file + state_file="$(station_dual_pair_resume_file)" || return 1 + [[ -e "$state_file" || -L "$state_file" ]] || return 1 + assert_nemoclaw_state_path_safe "$state_file" + return 0 +} + +validate_station_pair_selection() { + [[ "${_SELECTED_EXPRESS_PLATFORM:-}" == "DGX Station" ]] || return 0 + station_dual_model_requested && return 0 + [ -z "${NEMOCLAW_DGX_STATION_PEER:-}" ] \ + || error "NEMOCLAW_DGX_STATION_PEER requires the DGX Station dual-serving model. Unset NEMOCLAW_VLLM_MODEL or select ${STATION_ULTRA_VLLM_MODEL}; the explicit model override remains authoritative." + station_dual_pair_resume_pending \ + && error "A dual-DGX Station pair resume is pending; refusing to bypass exact pair revalidation with model '${NEMOCLAW_VLLM_MODEL:-}'." + return 0 +} + +parse_station_dual_pair_result() { + # The single-quoted payload is JavaScript, not shell interpolation. + # shellcheck disable=SC2016 + node -e ' + const fs = require("node:fs"); + const net = require("node:net"); + const path = require("node:path"); + const fail = () => process.exit(2); + try { + const raw = fs.readFileSync(0, "utf8"); + const value = JSON.parse(raw); + if (!value || typeof value !== "object" || Array.isArray(value)) fail(); + if (value.kind === "single-station") { + if (typeof value.reason !== "string" || value.reason.length === 0 || value.reason.length > 4096) fail(); + process.stdout.write("single-station\n"); + process.exit(0); + } + if (value.kind !== "ready" && value.kind !== "reboot-required") fail(); + const peer = value.peerTarget; + if (typeof peer !== "string" || peer.length === 0 || peer.length > 286 || peer !== peer.trim()) fail(); + const parts = peer.split("@"); + if (parts.length > 2) fail(); + const user = parts.length === 2 ? parts[0] : ""; + const host = parts[parts.length - 1]; + const safeUser = /^[A-Za-z_][A-Za-z0-9._-]*$/; + const safeHost = /^[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?(?:\.[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?)*$/; + const numericHost = /^[0-9.]+$/.test(host); + if ((user && !safeUser.test(user)) || (net.isIP(host) !== 4 && (numericHost || !safeHost.test(host)))) fail(); + const identity = value.identity; + if (!identity || typeof identity !== "object" || Array.isArray(identity)) fail(); + if (identity.peerTarget !== peer || !/^[a-f0-9]{64}$/.test(identity.hostKeyDigest)) fail(); + if (!/^GPU-[A-Za-z0-9-]+$/.test(identity.localGpuUuid) || !/^GPU-[A-Za-z0-9-]+$/.test(identity.peerGpuUuid)) fail(); + if (!Array.isArray(identity.rails) || identity.rails.length !== 2) fail(); + const mac = /^(?:[0-9a-f]{2}:){5}[0-9a-f]{2}$/; + for (const rail of identity.rails) { + if (!rail || typeof rail !== "object" || Array.isArray(rail)) fail(); + if (net.isIP(rail.localAddress) !== 4 || net.isIP(rail.peerAddress) !== 4) fail(); + if (!mac.test(rail.localMac) || !mac.test(rail.peerMac)) fail(); + } + const token = value.sshBinding; + if (typeof token !== "string" || token.length === 0 || token.length > 8192 || !/^[A-Za-z0-9_-]+$/.test(token)) fail(); + let handoff; + try { + const decoded = Buffer.from(token, "base64url"); + if (decoded.toString("base64url") !== token) fail(); + handoff = JSON.parse(decoded.toString("utf8")); + } catch { + fail(); + } + if (!handoff || typeof handoff !== "object" || Array.isArray(handoff)) fail(); + if (JSON.stringify(Object.keys(handoff).sort()) !== JSON.stringify(["bindingFile", "hostKeyDigest"])) fail(); + const bindingFile = handoff.bindingFile; + if (typeof bindingFile !== "string" || bindingFile.length === 0 || bindingFile.length > 4096 || bindingFile !== bindingFile.trim() || /[\u0000-\u001f\u007f]/.test(bindingFile)) fail(); + if (!path.isAbsolute(bindingFile) || path.normalize(bindingFile) !== bindingFile) fail(); + if (handoff.hostKeyDigest !== identity.hostKeyDigest) fail(); + process.stdout.write(`${value.kind}\n${peer}\n${token}\n`); + } catch { + fail(); + } + ' +} + +ensure_station_express_pair() { + [[ "${_SELECTED_EXPRESS_PLATFORM:-}" == "DGX Station" ]] || return 0 + validate_station_pair_selection + if ! station_dual_model_requested; then + return 0 + fi + + local coordinator="${SCRIPT_DIR}/prepare-dual-dgx-station.mts" + local helper="${SCRIPT_DIR}/prepare-dgx-station-host.sh" + [[ -f "$coordinator" ]] || error "Dual DGX Station preparation coordinator is missing: ${coordinator}" + [[ -f "$helper" ]] || error "DGX Station host preparation helper is missing: ${helper}" + + local state_dir state_file revision output parsed kind peer_target ssh_binding status=0 + state_dir="$(ensure_nemoclaw_state_dir)" || error "Could not prepare owner-only NemoClaw state for dual DGX Station discovery." + state_file="${state_dir}/station-dual-pair-resume.json" + assert_nemoclaw_state_path_safe "$state_file" + revision="$(station_installer_revision)" + + local -a pair_command=( + node --no-warnings --experimental-strip-types "$coordinator" + --helper "$helper" + --state "$state_file" + --revision "$revision" + ) + if [ -n "${NEMOCLAW_DGX_STATION_PEER:-}" ]; then + pair_command+=(--explicit-peer "$NEMOCLAW_DGX_STATION_PEER") + fi + if [ "${_STATION_EXPRESS_DEFERRED_MANAGED_PAIR:-0}" = "1" ]; then + pair_command+=(--reuse-existing-managed-pair) + fi + if [ "${_STATION_EXPRESS_MIGRATING_LEGACY_HEAD:-0}" = "1" ]; then + pair_command+=(--migrate-legacy-single-head) + fi + + info "Checking for one pretrusted reciprocal dual-DGX Station peer on the two direct /30 rail counterparts." + # Publish the companion mode/model state before the coordinator can persist + # pair identity and mutate the remote host. A crash, mismatch, or later + # onboarding failure must resume this exact accepted setup without returning + # to a prompt or silently falling back. + save_station_express_resume + output="$("${pair_command[@]}")" || status=$? + case "$status" in + 0 | 10) ;; + *) + if ! station_dual_pair_resume_pending; then + clear_station_express_resume + fi + error "Dual DGX Station preparation failed. No serving-model or vLLM-image pull was started; review the station-pair diagnostics above and rerun the same revision." + ;; + esac + if ! parsed="$(printf '%s' "$output" | parse_station_dual_pair_result)"; then + if ! station_dual_pair_resume_pending; then + clear_station_express_resume + fi + error "Dual DGX Station preparation returned invalid result data; refusing to continue." + fi + kind="$(printf '%s\n' "$parsed" | sed -n '1p')" + peer_target="$(printf '%s\n' "$parsed" | sed -n '2p')" + ssh_binding="$(printf '%s\n' "$parsed" | sed -n '3p')" + + case "$kind" in + single-station) + [ "$status" -eq 0 ] \ + || error "Dual DGX Station preparation returned an inconsistent reboot result; refusing to continue." + [ "${_STATION_EXPRESS_DEFERRED_MANAGED_PAIR:-0}" != "1" ] \ + || error "The running managed dual-Station head could not be matched to its trusted reciprocal peer; refusing single-Station fallback." + [ "${_STATION_EXPRESS_MIGRATING_LEGACY_HEAD:-0}" != "1" ] \ + || error "The running legacy single-Station head could not be matched to a trusted reciprocal peer; refusing migration and single-Station fallback." + [ -z "${NEMOCLAW_DGX_STATION_PEER:-}" ] \ + || error "The explicit DGX Station peer could not be qualified; refusing single-Station fallback." + station_dual_pair_resume_pending \ + && error "Dual DGX Station preparation returned a single-Station result while exact pair resume state is pending; refusing to discard it." + clear_station_express_resume + if [ "${_STATION_EXPRESS_MODEL_WAS_EXPLICIT:-0}" = "0" ]; then + NEMOCLAW_VLLM_MODEL="$STATION_ULTRA_VLLM_MODEL" + NEMOCLAW_MODEL="$STATION_ULTRA_SERVED_MODEL" + export NEMOCLAW_VLLM_MODEL NEMOCLAW_MODEL + fi + unset NEMOCLAW_DGX_STATION_SSH_BINDING + info "No trusted reciprocal dual-DGX Station pair was detected; using the existing single-Station Ultra recipe." + ;; + ready) + [ "$status" -eq 0 ] \ + || error "Dual DGX Station preparation returned an inconsistent ready result; refusing to continue." + station_dual_pair_resume_pending \ + || error "Dual DGX Station preparation returned ready without exact pair resume state; refusing to continue." + NEMOCLAW_DGX_STATION_PEER="$peer_target" + NEMOCLAW_DGX_STATION_SSH_BINDING="$ssh_binding" + export NEMOCLAW_DGX_STATION_PEER NEMOCLAW_DGX_STATION_SSH_BINDING + NEMOCLAW_MODEL="$STATION_ULTRA_DUAL_SERVED_MODEL" + export NEMOCLAW_MODEL + if [ "${_STATION_EXPRESS_MODEL_WAS_EXPLICIT:-0}" = "0" ]; then + NEMOCLAW_VLLM_MODEL="$STATION_ULTRA_VLLM_MODEL" + export NEMOCLAW_VLLM_MODEL + fi + ok "Trusted reciprocal dual-DGX Station pair is ready (${peer_target})" + ;; + reboot-required) + [ "$status" -eq 10 ] \ + || error "Dual DGX Station preparation returned an inconsistent reboot result; refusing to continue." + station_dual_pair_resume_pending \ + || error "Dual DGX Station preparation requested a reboot without exact pair resume state; refusing to continue." + save_station_express_resume + warn "Peer DGX Station ${peer_target} was prepared and requires a manual reboot." + info "On peer ${peer_target}, run: sudo reboot" + info "After the peer is back online, rerun the accepted revision on this Station:" + info "$(station_express_resume_command)" + exit 10 + ;; + *) + error "Dual DGX Station preparation returned an unsupported result; refusing to continue." + ;; + esac +} + +clear_station_dual_pair_resume() { + local state_file coordinator="${SCRIPT_DIR}/prepare-dual-dgx-station.mts" + state_file="$(station_dual_pair_resume_file)" || return 0 + assert_nemoclaw_state_path_safe "$state_file" + [[ -e "$state_file" || -L "$state_file" || -e "${state_file}.ssh-binding" || -L "${state_file}.ssh-binding" ]] || return 0 + [[ -f "$coordinator" ]] || error "Dual DGX Station preparation coordinator is missing: ${coordinator}" + node --no-warnings --experimental-strip-types "$coordinator" --state "$state_file" --clear-state >/dev/null \ + || error "Could not safely clear completed dual DGX Station resume state: ${state_file}" +} + prepare_installer_host() { maybe_offer_express_install + # Reject conflicting explicit Station selections and pending-pair bypasses + # before the local host-preparation helper can mutate packages or Docker. + validate_station_pair_selection if [[ "${_SELECTED_EXPRESS_PLATFORM:-}" == "DGX Station" ]]; then # Station qualification is deliberately scoped to the local factory # runtime. Normalize the Docker target before any preparation probe so an @@ -4120,8 +4528,8 @@ describe_express_install() { esac else show_hf_authentication="1" - inference_summary="managed local vLLM with NVIDIA Nemotron 3 Ultra 550B" - inference_disclosure="Managed vLLM pulls the pinned Station image and approximately 352 GB model, then runs a local inference container." + inference_summary="managed local vLLM with NVIDIA Nemotron 3 Ultra 550B and automatic Station topology selection" + inference_disclosure="A pretrusted reciprocal dual-Station pair selects distributed serving; otherwise NemoClaw uses the existing single-Station Ultra recipe. The pinned image and approximately 352 GB model are pulled only after topology qualification." fi case "$(classify_dgx_station_release)" in supported-dgx-os) @@ -4211,6 +4619,25 @@ maybe_offer_express_install() { validate_express_platform_boundary "$platform" validate_force_station_install_override "$platform" validate_station_deepseek_override "$platform" + + # Pair state is written before remote mutation. It therefore outranks every + # prompt/skip path and must recover the companion exact-revision mode/model + # state before any generic or single-Station setup can continue. + if station_dual_pair_resume_pending; then + [ "$platform" = "DGX Station" ] \ + || error "A dual-DGX Station pair resume is pending, but this host no longer satisfies the DGX Station preparation boundary. Refusing to continue without exact pair revalidation." + [ "${NEMOCLAW_NO_EXPRESS:-}" != "1" ] \ + || error "A dual-DGX Station pair resume is pending; finish exact pair revalidation before disabling Station setup." + case "${NEMOCLAW_PROVIDER:-}" in + "") ;; + install-vllm) _STATION_INSTALL_MODE="provider" ;; + *) error "A dual-DGX Station pair resume is pending; finish exact pair revalidation before changing providers." ;; + esac + load_station_express_resume \ + || error "Dual-DGX Station pair state exists without its required installer resume state. Refusing to prompt, fall back, or continue; restore the owner-only state from the accepted revision." + resume_loaded_station_install "$platform" + return 0 + fi # Not on a platform we have an express recipe for — say nothing. if [ -z "$platform" ]; then return 0 @@ -4218,18 +4645,36 @@ maybe_offer_express_install() { # On a supported platform but a skip condition applies — explain why so # the user understands they could have gotten express otherwise. if [ "${NEMOCLAW_NO_EXPRESS:-}" = "1" ]; then - if [ "$platform" = "DGX Station" ]; then clear_station_express_resume; fi + if [ "$platform" = "DGX Station" ]; then + station_dual_pair_resume_pending \ + && error "A dual-DGX Station pair resume is pending; finish exact pair revalidation before disabling Station setup." + clear_station_express_resume + fi info "Detected ${platform}. Skipping express prompt (NEMOCLAW_NO_EXPRESS=1)." return 0 fi if [ -n "${NEMOCLAW_PROVIDER:-}" ]; then - if [ "$platform" = "DGX Station" ]; then clear_station_express_resume; fi + if [ "$platform" = "DGX Station" ] && [ "$NEMOCLAW_PROVIDER" = "install-vllm" ]; then + # An explicit managed-vLLM provider selects the same Station host/pair + # preparation boundary without forcing the rest of the express policy. + # Honor an existing exact-revision reboot resume before configuration. + _STATION_INSTALL_MODE="provider" + load_station_express_resume || true + _SELECTED_EXPRESS_PLATFORM="$platform" + configure_station_express_model + info "Detected ${platform}. Using Station preparation for the explicitly selected managed-vLLM provider." + return 0 + fi + if [ "$platform" = "DGX Station" ]; then + station_dual_pair_resume_pending \ + && error "A dual-DGX Station pair resume is pending; finish exact pair revalidation before changing providers." + clear_station_express_resume + fi info "Detected ${platform}. Skipping express prompt (NEMOCLAW_PROVIDER=${NEMOCLAW_PROVIDER} already set)." return 0 fi if [ "$platform" = "DGX Station" ] && load_station_express_resume; then - info "Detected DGX Station. Resuming the accepted express install after host preparation." - activate_express_install "$platform" + resume_loaded_station_install "$platform" return 0 fi if [ "${NON_INTERACTIVE:-}" = "1" ]; then @@ -4392,6 +4837,7 @@ main() { step 1 "Node.js" install_nodejs ensure_supported_runtime + ensure_station_express_pair step 2 "${_CLI_DISPLAY} CLI" # Ollama and vLLM install/upgrade and model pulls are owned by @@ -4459,6 +4905,14 @@ main() { fi finalize_install + clear_station_resume_after_completed_onboarding +} + +clear_station_resume_after_completed_onboarding() { + [[ "${_SELECTED_EXPRESS_PLATFORM:-}" == "DGX Station" ]] || return 0 + [[ "${ONBOARD_RAN:-false}" == true ]] || return 0 + clear_station_dual_pair_resume + clear_station_express_resume } # Print the completion summary, then propagate a fatal/non-zero result when the diff --git a/scripts/lib/dgx-station-peer.mts b/scripts/lib/dgx-station-peer.mts new file mode 100644 index 0000000000..74575705ea --- /dev/null +++ b/scripts/lib/dgx-station-peer.mts @@ -0,0 +1,901 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import net from "node:net"; +import { isDgxStationGb300Product } from "../../src/lib/inference/dgx-station-identity.ts"; +import { stationKnownHostsDigest } from "../../src/lib/inference/vllm-station-ssh-binding.ts"; + +export const DUAL_STATION_RESUME_SCHEMA_VERSION = 1; +export const STATION_PREP_REBOOT_REQUIRED_EXIT = 10; +export const STATION_PREP_LOGIN_REQUIRED_EXIT = 11; + +const DIRECT_RAIL_PREFIX_LENGTH = 30; +const SAFE_TARGET_PATTERN = + /^[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?(?:\.[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?)*$/; +const SAFE_USERNAME_PATTERN = /^[A-Za-z_][A-Za-z0-9._-]*$/; +const SAFE_DEVICE_PATTERN = /^[A-Za-z0-9][A-Za-z0-9_.:-]{0,63}$/; +const GPU_UUID_PATTERN = /^GPU-[A-Za-z0-9-]+$/; +const HOST_KEY_DIGEST_PATTERN = /^[a-f0-9]{64}$/; +const HOST_KEY_FINGERPRINT_PATTERN = /^SHA256:[A-Za-z0-9+/]{16,86}={0,2}$/; +const MAC_PATTERN = /^(?:[0-9a-f]{2}:){5}[0-9a-f]{2}$/; +const MAX_KNOWN_HOSTS_BYTES = 64 * 1024; +const MAX_KNOWN_HOSTS_LINE_BYTES = 16 * 1024; + +export type StationPrepMode = "--check" | "--apply" | "--verify" | "--bind-controller"; + +export interface StationIpv4Address { + address: string; + prefixLength: number; +} + +export interface StationDiscoveryRail { + netdev: string; + macAddress: string; + pciAddress: string; + pciName: string; + state: string; + linkLayer: string; + speedMbps: number; + mtu: number; + ipv4Addresses: StationIpv4Address[]; +} + +export interface StationDiscoveryGpu { + index: number; + name: string; + uuid: string; +} + +export interface StationDiscoveryHost { + schemaVersion: 1; + hostname: string; + productName: string; + architecture: string; + gpus: StationDiscoveryGpu[]; + rails: StationDiscoveryRail[]; +} + +export interface PretrustedSshTarget { + requestedTarget: string; + sshTarget: string; + resolvedHost: string; + sshUser: string; + port: number; + lookupHost: string; + hostKeyDigest: string; + keyFingerprints: string[]; + knownHostsLines: string[]; +} + +export interface RailConnectivityRequest { + netdev: string; + sourceAddress: string; + peerAddress: string; + expectedPeerMac: string; +} + +export interface DualStationRailIdentity { + localAddress: string; + localMac: string; + peerAddress: string; + peerMac: string; +} + +export interface DualStationPairIdentity { + peerTarget: string; + hostKeyDigest: string; + localGpuUuid: string; + peerGpuUuid: string; + rails: DualStationRailIdentity[]; +} + +export interface DualStationResumeState extends DualStationPairIdentity { + schemaVersion: 1; + revision: string; + helperSha256: string; + phase: "remote-preparation" | "remote-reboot-required" | "ready"; +} + +export type DualStationPreparationResult = + | { kind: "single-station"; reason: string } + | { + kind: "ready"; + peerTarget: string; + identity: DualStationPairIdentity; + binding: PretrustedSshTarget; + } + | { + kind: "reboot-required"; + peerTarget: string; + identity: DualStationPairIdentity; + binding: PretrustedSshTarget; + }; + +export interface DualStationPreparationOptions { + revision: string; + helperSha256: string; + explicitPeer?: string; + reuseExistingManagedPair?: boolean; + migrateLegacySingleStationHead?: boolean; +} + +export interface DualStationPreparationDeps { + runLocalHelper(mode: StationPrepMode): number; + probeLocalHost(): StationDiscoveryHost; + inspectPretrustedTarget(target: string): PretrustedSshTarget | null; + probePeerHost(target: PretrustedSshTarget): StationDiscoveryHost; + probeLocalConnectivity(requests: readonly RailConnectivityRequest[]): boolean; + probePeerConnectivity( + target: PretrustedSshTarget, + requests: readonly RailConnectivityRequest[], + ): boolean; + runRemoteHelper(target: PretrustedSshTarget, mode: StationPrepMode): number; + readResumeState(): DualStationResumeState | null; + writeResumeState(state: DualStationResumeState): void; + clearResumeState(): void; + log(message: string): void; +} + +type QualifiedRail = { + rail: StationDiscoveryRail; + address: string; + peerAddress: string; + subnet: string; +}; + +type DiscoveryPlan = { + identity: DualStationPairIdentity; + localConnectivity: RailConnectivityRequest[]; + peerConnectivity: RailConnectivityRequest[]; +}; + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +function requireString(value: unknown, label: string, maxLength: number): string { + if ( + typeof value !== "string" || + value.length === 0 || + value.length > maxLength || + /[\u0000-\u001f\u007f]/.test(value) + ) { + throw new Error(`${label} must be a non-empty printable string`); + } + return value; +} + +function requireArray(value: unknown, label: string, maxLength: number): unknown[] { + if (!Array.isArray(value) || value.length > maxLength) { + throw new Error(`${label} must be an array with at most ${String(maxLength)} entries`); + } + return value; +} + +function requireInteger(value: unknown, label: string, min: number, max: number): number { + if (!Number.isInteger(value) || (value as number) < min || (value as number) > max) { + throw new Error(`${label} must be an integer between ${String(min)} and ${String(max)}`); + } + return value as number; +} + +function normalizeMac(value: unknown, label: string): string { + const mac = requireString(value, label, 17).toLowerCase(); + if (!MAC_PATTERN.test(mac) || mac === "00:00:00:00:00:00") { + throw new Error(`${label} must be a nonzero canonical MAC address`); + } + const firstOctet = Number.parseInt(mac.slice(0, 2), 16); + if ((firstOctet & 1) !== 0) throw new Error(`${label} must be a unicast MAC address`); + return mac; +} + +function requireIpv4(value: unknown, label: string): string { + const address = requireString(value, label, 15); + if (net.isIP(address) !== 4) throw new Error(`${label} must be IPv4`); + return address; +} + +function isSafeTargetHost(hostname: string): boolean { + return ( + net.isIP(hostname) === 4 || (!/^[0-9.]+$/.test(hostname) && SAFE_TARGET_PATTERN.test(hostname)) + ); +} + +export function validateStationPeerTarget(raw: string): string { + if (raw.length === 0 || raw !== raw.trim() || raw.length > 286) { + throw new Error("Station peer must be one canonical SSH host or user@host"); + } + if (/[/,:;`'"\\$(){}[\]<>|&!?*\s\u0000-\u001f\u007f]/.test(raw)) { + throw new Error("Station peer must be one canonical SSH host or user@host"); + } + const parts = raw.split("@"); + if (parts.length > 2) throw new Error("Station peer must be one canonical SSH host or user@host"); + const username = parts.length === 2 ? parts[0] : ""; + const hostname = parts.at(-1) ?? ""; + const validHost = isSafeTargetHost(hostname); + if ( + !validHost || + (parts.length === 2 && username.length === 0) || + (username.length > 0 && !SAFE_USERNAME_PATTERN.test(username)) + ) { + throw new Error("Station peer must be one canonical SSH host or user@host"); + } + return raw; +} + +function ipv4ToNumber(address: string): number { + return address + .split(".") + .map(Number) + .reduce((value, octet) => value * 256 + octet, 0); +} + +function numberToIpv4(value: number): string { + return [24, 16, 8, 0].map((shift) => Math.floor(value / 2 ** shift) % 256).join("."); +} + +function isPrivateIpv4(address: string): boolean { + const value = ipv4ToNumber(address); + return ( + (value >= ipv4ToNumber("10.0.0.0") && value <= ipv4ToNumber("10.255.255.255")) || + (value >= ipv4ToNumber("172.16.0.0") && value <= ipv4ToNumber("172.31.255.255")) || + (value >= ipv4ToNumber("192.168.0.0") && value <= ipv4ToNumber("192.168.255.255")) + ); +} + +export function deriveSlash30Counterpart(address: string, prefixLength = 30): string | null { + if (prefixLength !== DIRECT_RAIL_PREFIX_LENGTH || net.isIP(address) !== 4) return null; + if (!isPrivateIpv4(address)) return null; + const value = ipv4ToNumber(address); + const network = Math.floor(value / 4) * 4; + const host = value - network; + if (host === 1) return numberToIpv4(network + 2); + if (host === 2) return numberToIpv4(network + 1); + return null; +} + +function subnetOfSlash30(address: string): string { + return `${numberToIpv4(Math.floor(ipv4ToNumber(address) / 4) * 4)}/30`; +} + +export function parseStationDiscoveryHost(value: unknown): StationDiscoveryHost { + if (!isRecord(value) || value.schemaVersion !== 1) { + throw new Error("Station discovery probe schema is unsupported"); + } + const gpus = requireArray(value.gpus, "Station discovery GPUs", 16).map((entry, index) => { + if (!isRecord(entry)) throw new Error(`Station discovery GPU ${String(index)} is invalid`); + const uuid = requireString(entry.uuid, `Station discovery GPU ${String(index)} UUID`, 128); + if (!GPU_UUID_PATTERN.test(uuid)) { + throw new Error(`Station discovery GPU ${String(index)} UUID is invalid`); + } + return { + index: requireInteger(entry.index, `Station discovery GPU ${String(index)} index`, 0, 1024), + name: requireString(entry.name, `Station discovery GPU ${String(index)} name`, 256), + uuid, + }; + }); + const rails = requireArray(value.rails, "Station discovery rails", 16).map((entry, index) => { + if (!isRecord(entry)) throw new Error(`Station discovery rail ${String(index)} is invalid`); + const netdev = requireString( + entry.netdev, + `Station discovery rail ${String(index)} netdev`, + 64, + ); + if (!SAFE_DEVICE_PATTERN.test(netdev)) { + throw new Error(`Station discovery rail ${String(index)} netdev is unsafe`); + } + const pciAddress = requireString( + entry.pciAddress, + `Station discovery rail ${String(index)} PCI address`, + 32, + ); + if (!/^[0-9A-Fa-f]{4}:[0-9A-Fa-f]{2}:[0-9A-Fa-f]{2}\.[0-7]$/.test(pciAddress)) { + throw new Error(`Station discovery rail ${String(index)} PCI address is invalid`); + } + return { + netdev, + macAddress: normalizeMac(entry.macAddress, `Station discovery rail ${String(index)} MAC`), + pciAddress, + pciName: requireString( + entry.pciName, + `Station discovery rail ${String(index)} PCI name`, + 512, + ), + state: requireString(entry.state, `Station discovery rail ${String(index)} state`, 64), + linkLayer: requireString( + entry.linkLayer, + `Station discovery rail ${String(index)} link layer`, + 64, + ), + speedMbps: requireInteger( + entry.speedMbps, + `Station discovery rail ${String(index)} speed`, + -1, + 1_000_000, + ), + mtu: requireInteger(entry.mtu, `Station discovery rail ${String(index)} MTU`, -1, 1_000_000), + ipv4Addresses: requireArray( + entry.ipv4Addresses, + `Station discovery rail ${String(index)} IPv4 addresses`, + 16, + ).map((rawAddress, addressIndex) => { + if (!isRecord(rawAddress)) { + throw new Error( + `Station discovery rail ${String(index)} address ${String(addressIndex)} is invalid`, + ); + } + return { + address: requireIpv4( + rawAddress.address, + `Station discovery rail ${String(index)} address ${String(addressIndex)}`, + ), + prefixLength: requireInteger( + rawAddress.prefixLength, + `Station discovery rail ${String(index)} prefix ${String(addressIndex)}`, + 1, + 32, + ), + }; + }), + }; + }); + return { + schemaVersion: 1, + hostname: requireString(value.hostname, "Station discovery hostname", 256), + productName: requireString(value.productName, "Station discovery product", 512), + architecture: requireString(value.architecture, "Station discovery architecture", 64), + gpus, + rails, + }; +} + +function selectedGb300(host: StationDiscoveryHost, label: string): StationDiscoveryGpu { + const matches = host.gpus.filter((gpu) => /\bGB300\b/i.test(gpu.name)); + if (matches.length !== 1) throw new Error(`${label} must expose exactly one GB300 GPU`); + return matches[0]; +} + +function assertStationIdentity(host: StationDiscoveryHost, label: string): void { + if ( + !isDgxStationGb300Product(host.productName) || + !/^(?:aarch64|arm64)$/i.test(host.architecture) + ) { + throw new Error(`${label} is not a verified arm64 DGX Station GB300`); + } +} + +function qualifyRails(host: StationDiscoveryHost, label: string): QualifiedRail[] { + const cx8 = host.rails.filter((rail) => /ConnectX[- ]?8|\bCX-?8\b/i.test(rail.pciName)); + if (cx8.length !== 2) throw new Error(`${label} must expose exactly two CX-8 rails`); + const result = cx8.map((rail, index): QualifiedRail => { + if ( + !/\bACTIVE\b/i.test(rail.state) || + rail.linkLayer.toLowerCase() !== "ethernet" || + rail.speedMbps !== 400_000 || + rail.mtu !== 9000 + ) { + throw new Error(`${label} rail ${String(index + 1)} is not active 400G Ethernet MTU 9000`); + } + const plausible = rail.ipv4Addresses + .map((entry) => ({ + entry, + peer: deriveSlash30Counterpart(entry.address, entry.prefixLength), + })) + .filter((entry): entry is { entry: StationIpv4Address; peer: string } => entry.peer !== null); + if (plausible.length !== 1) { + throw new Error( + `${label} rail ${String(index + 1)} must have exactly one usable private /30 address`, + ); + } + return { + rail, + address: plausible[0].entry.address, + peerAddress: plausible[0].peer, + subnet: subnetOfSlash30(plausible[0].entry.address), + }; + }); + if ( + new Set(result.map((entry) => entry.rail.netdev)).size !== 2 || + new Set(result.map((entry) => entry.rail.macAddress)).size !== 2 || + new Set(result.map((entry) => entry.rail.pciAddress)).size !== 2 || + new Set(result.map((entry) => entry.subnet)).size !== 2 || + new Set(result.map((entry) => entry.peerAddress)).size !== 2 + ) { + throw new Error(`${label} CX-8 rail identity is ambiguous`); + } + return result.sort((left, right) => left.subnet.localeCompare(right.subnet)); +} + +export function deriveDiscoveryCandidates(host: StationDiscoveryHost): string[] { + assertStationIdentity(host, "Local host"); + selectedGb300(host, "Local host"); + return qualifyRails(host, "Local host").map((entry) => entry.peerAddress); +} + +function peerHostFromTarget(target: string): string { + return target.slice(target.lastIndexOf("@") + 1); +} + +function buildDiscoveryPlan( + binding: PretrustedSshTarget, + local: StationDiscoveryHost, + peer: StationDiscoveryHost, + automatic: boolean, +): DiscoveryPlan { + assertStationIdentity(local, "Local host"); + assertStationIdentity(peer, "Peer host"); + const localGpu = selectedGb300(local, "Local host"); + const peerGpu = selectedGb300(peer, "Peer host"); + if (localGpu.uuid === peerGpu.uuid) { + throw new Error("Peer SSH target resolved back to the local Station GPU"); + } + const localRails = qualifyRails(local, "Local host"); + const peerRails = qualifyRails(peer, "Peer host"); + const matched = localRails.map((localRail) => { + const peers = peerRails.filter( + (peerRail) => + peerRail.subnet === localRail.subnet && + peerRail.address === localRail.peerAddress && + peerRail.peerAddress === localRail.address, + ); + if (peers.length !== 1) { + throw new Error("Peer did not report one reciprocal address and MAC on each /30 rail"); + } + return { local: localRail, peer: peers[0] }; + }); + if (new Set(matched.map((entry) => entry.peer.rail.macAddress)).size !== 2) { + throw new Error("Peer rail MAC identity is ambiguous"); + } + if ( + automatic && + !matched.some((entry) => entry.peer.address === peerHostFromTarget(binding.requestedTarget)) + ) { + throw new Error("Pretrusted discovery target is not one of the reciprocal peer rail addresses"); + } + const rails = matched.map( + (entry): DualStationRailIdentity => ({ + localAddress: entry.local.address, + localMac: entry.local.rail.macAddress, + peerAddress: entry.peer.address, + peerMac: entry.peer.rail.macAddress, + }), + ); + return { + identity: { + peerTarget: binding.sshTarget, + hostKeyDigest: binding.hostKeyDigest, + localGpuUuid: localGpu.uuid, + peerGpuUuid: peerGpu.uuid, + rails, + }, + localConnectivity: matched.map((entry) => ({ + netdev: entry.local.rail.netdev, + sourceAddress: entry.local.address, + peerAddress: entry.peer.address, + expectedPeerMac: entry.peer.rail.macAddress, + })), + peerConnectivity: matched.map((entry) => ({ + netdev: entry.peer.rail.netdev, + sourceAddress: entry.peer.address, + peerAddress: entry.local.address, + expectedPeerMac: entry.local.rail.macAddress, + })), + }; +} + +function validateRailIdentity(value: unknown, label: string): DualStationRailIdentity { + if (!isRecord(value)) throw new Error(`${label} must be an object`); + const localAddress = requireIpv4(value.localAddress, `${label}.localAddress`); + const peerAddress = requireIpv4(value.peerAddress, `${label}.peerAddress`); + if ( + deriveSlash30Counterpart(localAddress) !== peerAddress || + deriveSlash30Counterpart(peerAddress) !== localAddress + ) { + throw new Error(`${label} must contain reciprocal private /30 addresses`); + } + return { + localAddress, + localMac: normalizeMac(value.localMac, `${label}.localMac`), + peerAddress, + peerMac: normalizeMac(value.peerMac, `${label}.peerMac`), + }; +} + +export function parseDualStationResumeState(value: unknown): DualStationResumeState { + if (!isRecord(value) || value.schemaVersion !== DUAL_STATION_RESUME_SCHEMA_VERSION) { + throw new Error("Dual-Station resume state schema is unsupported"); + } + const revision = requireString(value.revision, "Dual-Station resume revision", 40); + if (!/^[a-f0-9]{40}$/.test(revision)) { + throw new Error("Dual-Station resume revision is invalid"); + } + const helperSha256 = requireString(value.helperSha256, "Dual-Station resume helper SHA-256", 64); + if (!HOST_KEY_DIGEST_PATTERN.test(helperSha256)) { + throw new Error("Dual-Station resume helper SHA-256 is invalid"); + } + const peerTarget = validateStationPeerTarget( + requireString(value.peerTarget, "Dual-Station resume peer target", 286), + ); + const hostKeyDigest = requireString( + value.hostKeyDigest, + "Dual-Station resume host-key digest", + 64, + ); + if (!HOST_KEY_DIGEST_PATTERN.test(hostKeyDigest)) { + throw new Error("Dual-Station resume host-key digest is invalid"); + } + const localGpuUuid = requireString(value.localGpuUuid, "Dual-Station local GPU UUID", 128); + const peerGpuUuid = requireString(value.peerGpuUuid, "Dual-Station peer GPU UUID", 128); + if ( + !GPU_UUID_PATTERN.test(localGpuUuid) || + !GPU_UUID_PATTERN.test(peerGpuUuid) || + localGpuUuid === peerGpuUuid + ) { + throw new Error("Dual-Station resume GPU identity is invalid"); + } + const phase = value.phase; + if (phase !== "remote-preparation" && phase !== "remote-reboot-required" && phase !== "ready") { + throw new Error("Dual-Station resume phase is invalid"); + } + const rails = requireArray(value.rails, "Dual-Station resume rails", 2) + .map((entry, index) => validateRailIdentity(entry, `Dual-Station resume rail ${String(index)}`)) + .sort((left, right) => left.localAddress.localeCompare(right.localAddress)); + if ( + rails.length !== 2 || + new Set(rails.map((rail) => rail.localAddress)).size !== 2 || + new Set(rails.map((rail) => rail.peerAddress)).size !== 2 || + new Set(rails.map((rail) => rail.localMac)).size !== 2 || + new Set(rails.map((rail) => rail.peerMac)).size !== 2 + ) { + throw new Error("Dual-Station resume rail identity is ambiguous"); + } + return { + schemaVersion: 1, + revision, + helperSha256, + phase, + peerTarget, + hostKeyDigest, + localGpuUuid, + peerGpuUuid, + rails, + }; +} + +export function validateResumeFileMetadata( + metadata: { isFile: boolean; isSymbolicLink: boolean; uid: number; mode: number; size: number }, + expectedUid: number, +): void { + if (metadata.isSymbolicLink || !metadata.isFile) { + throw new Error("Dual-Station resume state must be a regular file, not a symlink"); + } + if (metadata.uid !== expectedUid) { + throw new Error("Dual-Station resume state is not owned by the current user"); + } + if ((metadata.mode & 0o777) !== 0o600) { + throw new Error("Dual-Station resume state must have mode 0600"); + } + if (metadata.size <= 0 || metadata.size > 16 * 1024) { + throw new Error("Dual-Station resume state size is invalid"); + } +} + +function canonicalPairIdentity(value: DualStationPairIdentity): DualStationPairIdentity { + return { + peerTarget: value.peerTarget, + hostKeyDigest: value.hostKeyDigest, + localGpuUuid: value.localGpuUuid, + peerGpuUuid: value.peerGpuUuid, + rails: [...value.rails].sort((left, right) => + left.localAddress.localeCompare(right.localAddress), + ), + }; +} + +function samePair(left: DualStationPairIdentity, right: DualStationPairIdentity): boolean { + return ( + JSON.stringify(canonicalPairIdentity(left)) === JSON.stringify(canonicalPairIdentity(right)) + ); +} + +function samePhysicalSshIdentity(left: PretrustedSshTarget, right: PretrustedSshTarget): boolean { + return ( + left.sshUser === right.sshUser && + left.port === right.port && + left.hostKeyDigest === right.hostKeyDigest + ); +} + +function validateKnownHostsLookupHost(binding: PretrustedSshTarget): void { + const expected = + binding.port === 22 + ? binding.resolvedHost + : `[${binding.resolvedHost}]:${String(binding.port)}`; + if (binding.lookupHost !== expected) { + throw new Error("Pretrusted SSH target has an invalid known-hosts lookup identity"); + } +} + +function validateBinding(binding: PretrustedSshTarget): void { + const requestedTarget = validateStationPeerTarget(binding.requestedTarget); + const sshTarget = validateStationPeerTarget(binding.sshTarget); + if (requestedTarget !== sshTarget) { + throw new Error("Pretrusted SSH target changed after configuration resolution"); + } + if (!isSafeTargetHost(binding.resolvedHost)) { + throw new Error("Pretrusted SSH target resolved to an unsafe host"); + } + const explicitUser = sshTarget.includes("@") ? sshTarget.slice(0, sshTarget.indexOf("@")) : null; + if ( + !SAFE_USERNAME_PATTERN.test(binding.sshUser) || + (explicitUser !== null && explicitUser !== binding.sshUser) || + !Number.isInteger(binding.port) || + binding.port < 1 || + binding.port > 65535 + ) { + throw new Error("Pretrusted SSH target has an unsafe user or port"); + } + validateKnownHostsLookupHost(binding); + if (!HOST_KEY_DIGEST_PATTERN.test(binding.hostKeyDigest)) { + throw new Error("Pretrusted SSH target has an invalid host-key digest"); + } + if ( + !Array.isArray(binding.keyFingerprints) || + binding.keyFingerprints.length === 0 || + binding.keyFingerprints.some( + (fingerprint) => + typeof fingerprint !== "string" || !HOST_KEY_FINGERPRINT_PATTERN.test(fingerprint), + ) || + !Array.isArray(binding.knownHostsLines) || + binding.knownHostsLines.length === 0 || + binding.knownHostsLines.some( + (line) => + typeof line !== "string" || + line.length === 0 || + Buffer.byteLength(line, "utf8") > MAX_KNOWN_HOSTS_LINE_BYTES || + line !== line.trim() || + line.startsWith("#") || + /[\u0000\r\n]/.test(line), + ) + ) { + throw new Error("Pretrusted SSH target has invalid known-hosts evidence"); + } + const knownHosts = `${[...new Set(binding.knownHostsLines)].sort().join("\n")}\n`; + if ( + Buffer.byteLength(knownHosts, "utf8") > MAX_KNOWN_HOSTS_BYTES || + stationKnownHostsDigest(knownHosts) !== binding.hostKeyDigest + ) { + throw new Error("Pretrusted SSH target known-hosts evidence does not match its digest"); + } +} + +function selectPretrustedTarget( + options: DualStationPreparationOptions, + local: StationDiscoveryHost, + resume: DualStationResumeState | null, + deps: DualStationPreparationDeps, +): { binding: PretrustedSshTarget; automatic: boolean } | DualStationPreparationResult { + const candidates = deriveDiscoveryCandidates(local); + const explicitPeer = options.explicitPeer?.trim() ?? ""; + if (explicitPeer) validateStationPeerTarget(explicitPeer); + + if (resume) { + if (resume.revision !== options.revision) { + throw new Error( + `Dual-Station resume requires NemoClaw revision ${resume.revision}; current revision is ${options.revision}`, + ); + } + if (resume.helperSha256 !== options.helperSha256) { + throw new Error("The reviewed Station host-preparation helper changed during reboot resume"); + } + if (explicitPeer && explicitPeer !== resume.peerTarget) { + throw new Error("Explicit Station peer does not match the reboot-resume pair"); + } + if (!explicitPeer && !candidates.includes(peerHostFromTarget(resume.peerTarget))) { + throw new Error("The reboot-resume peer is no longer a derived local /30 counterpart"); + } + const binding = deps.inspectPretrustedTarget(resume.peerTarget); + if (!binding) throw new Error("The reboot-resume peer is no longer pretrusted"); + validateBinding(binding); + if (binding.hostKeyDigest !== resume.hostKeyDigest) { + throw new Error("The reboot-resume peer host-key identity changed"); + } + return { binding, automatic: !explicitPeer }; + } + + if (explicitPeer) { + const binding = deps.inspectPretrustedTarget(explicitPeer); + if (!binding) + throw new Error("Explicit Station peer is not pretrusted; SSH trust was not changed"); + validateBinding(binding); + return { binding, automatic: false }; + } + + const trusted: PretrustedSshTarget[] = []; + for (const candidate of candidates) { + try { + const binding = deps.inspectPretrustedTarget(candidate); + if (!binding) continue; + validateBinding(binding); + trusted.push(binding); + } catch (error) { + deps.log( + `Ignoring derived peer ${candidate}: pre-existing SSH trust is unusable (${(error as Error).message})`, + ); + } + } + if (trusted.length === 0) { + return { + kind: "single-station", + reason: "No derived dual-rail peer address has pre-existing SSH host-key trust", + }; + } + if (trusted.length === 2 && !samePhysicalSshIdentity(trusted[0], trusted[1])) { + return { + kind: "single-station", + reason: "The two derived rail addresses map to different pretrusted SSH identities", + }; + } + const binding = [...trusted].sort((left, right) => + left.requestedTarget.localeCompare(right.requestedTarget), + )[0]; + return { binding, automatic: true }; +} + +function fallbackOrThrow( + strict: boolean, + reason: string, +): Extract { + if (strict) throw new Error(reason); + return { kind: "single-station", reason }; +} + +export function prepareDualStationPair( + options: DualStationPreparationOptions, + deps: DualStationPreparationDeps, +): DualStationPreparationResult { + if (!/^[a-f0-9]{40}$/.test(options.revision)) { + throw new Error("Exact NemoClaw revision is required for dual-Station preparation"); + } + if (!HOST_KEY_DIGEST_PATTERN.test(options.helperSha256)) { + throw new Error("Exact Station host-preparation helper SHA-256 is required"); + } + if (options.reuseExistingManagedPair && options.migrateLegacySingleStationHead) { + throw new Error("Managed-pair reuse and legacy single-head migration are mutually exclusive"); + } + + if (!options.reuseExistingManagedPair && !options.migrateLegacySingleStationHead) { + deps.log("Checking the local Station with the reviewed host-preparation helper"); + if (deps.runLocalHelper("--check") !== 0) { + throw new Error("Local DGX Station host-preparation check failed before peer contact"); + } + if (deps.runLocalHelper("--verify") !== 0) { + throw new Error("Local DGX Station verification failed before peer contact"); + } + } else if (options.reuseExistingManagedPair) { + deps.log("Revalidating the exact running managed pair without disrupting its workloads"); + } else { + deps.log("Revalidating the exact running legacy single-Station head before migration"); + } + + const resume = deps.readResumeState(); + let local: StationDiscoveryHost; + try { + local = deps.probeLocalHost(); + deriveDiscoveryCandidates(local); + } catch (error) { + if (resume || options.explicitPeer?.trim()) throw error; + return { + kind: "single-station", + reason: `Local direct-rail discovery is unavailable: ${(error as Error).message}`, + }; + } + + const selected = selectPretrustedTarget(options, local, resume, deps); + if ("kind" in selected) return selected; + const strict = Boolean( + resume || + options.explicitPeer?.trim() || + options.reuseExistingManagedPair || + options.migrateLegacySingleStationHead, + ); + const { binding, automatic } = selected; + + let peer: StationDiscoveryHost; + try { + peer = deps.probePeerHost(binding); + } catch (error) { + return fallbackOrThrow( + strict, + `Trusted peer identity probe failed: ${(error as Error).message}`, + ); + } + + let plan: DiscoveryPlan; + try { + plan = buildDiscoveryPlan(binding, local, peer, automatic); + } catch (error) { + return fallbackOrThrow(strict, `Trusted peer was not reciprocal: ${(error as Error).message}`); + } + if (resume && !samePair(resume, plan.identity)) { + throw new Error("The physical dual-Station pair changed during reboot resume"); + } + + let connectivityReady = false; + try { + connectivityReady = + deps.probeLocalConnectivity(plan.localConnectivity) && + deps.probePeerConnectivity(binding, plan.peerConnectivity); + } catch { + connectivityReady = false; + } + if (!connectivityReady) { + return fallbackOrThrow( + strict, + "Trusted peer failed direct-route, neighbor-MAC, or jumbo-frame checks", + ); + } + + const state: DualStationResumeState = { + schemaVersion: 1, + revision: options.revision, + helperSha256: options.helperSha256, + phase: "remote-preparation", + ...plan.identity, + }; + deps.writeResumeState(state); + deps.log("Binding the local Station controller account to the qualified pair"); + if (deps.runLocalHelper("--bind-controller") !== 0) { + throw new Error("Local DGX Station controller UID binding failed"); + } + if (options.reuseExistingManagedPair) { + deps.log("Binding the reciprocal peer controller account without disrupting managed inference"); + if (deps.runRemoteHelper(binding, "--bind-controller") !== 0) { + throw new Error("Peer DGX Station controller UID binding failed"); + } + deps.writeResumeState({ ...state, phase: "ready" }); + return { + kind: "ready", + peerTarget: binding.sshTarget, + identity: plan.identity, + binding, + }; + } + deps.log(`Preparing reciprocal peer ${binding.sshTarget} with the exact reviewed helper`); + + if (deps.runRemoteHelper(binding, "--check") !== 0) { + throw new Error( + "Peer DGX Station host-preparation check failed; the selected pair remains pinned", + ); + } + const applyStatus = deps.runRemoteHelper(binding, "--apply"); + if (applyStatus === STATION_PREP_REBOOT_REQUIRED_EXIT) { + deps.writeResumeState({ ...state, phase: "remote-reboot-required" }); + return { + kind: "reboot-required", + peerTarget: binding.sshTarget, + identity: plan.identity, + binding, + }; + } + if (applyStatus !== 0) { + if (applyStatus !== STATION_PREP_LOGIN_REQUIRED_EXIT) { + throw new Error("Peer DGX Station host preparation failed; refusing single-Station fallback"); + } + deps.log("Peer Docker access requires a new login; reopening SSH before verification"); + } + if (deps.runRemoteHelper(binding, "--bind-controller") !== 0) { + throw new Error("Peer DGX Station controller UID binding failed"); + } + if (deps.runRemoteHelper(binding, "--verify") !== 0) { + throw new Error("Peer DGX Station verification failed; refusing single-Station fallback"); + } + + deps.writeResumeState({ ...state, phase: "ready" }); + return { + kind: "ready", + peerTarget: binding.sshTarget, + identity: plan.identity, + binding, + }; +} diff --git a/scripts/lib/station-vllm-conflict.sh b/scripts/lib/station-vllm-conflict.sh index 82696b564f..0c91178d30 100644 --- a/scripts/lib/station-vllm-conflict.sh +++ b/scripts/lib/station-vllm-conflict.sh @@ -9,7 +9,9 @@ _STATION_LOCAL_VLLM_SELECTED="" station_local_vllm_resume_file() { - printf '%s/.nemoclaw/station-local-vllm-resume' "$HOME" + local state_root + state_root="$(nemoclaw_state_root)" || return 1 + printf '%s/station-local-vllm-resume' "$state_root" } assert_station_local_vllm_resume_file_safe() { diff --git a/scripts/prepare-dgx-station-host.sh b/scripts/prepare-dgx-station-host.sh index eeb9e8557e..a9ec9d506e 100755 --- a/scripts/prepare-dgx-station-host.sh +++ b/scripts/prepare-dgx-station-host.sh @@ -5,7 +5,7 @@ set -Eeuo pipefail umask 077 -readonly SCRIPT_VERSION="2026-07-18.1" +readonly SCRIPT_VERSION="2026-07-20.1" readonly REBOOT_REQUIRED_EXIT=10 readonly LOGIN_REQUIRED_EXIT=11 readonly MIN_FREE_KIB=$((20 * 1024 * 1024)) @@ -50,6 +50,9 @@ readonly -a RETAINED_DKMS_VERSIONS=( readonly ACCEPTANCE_IMAGE="docker.io/library/ubuntu@sha256:7f622ca8766bccb22f04242ecb6f19f770b2f08827dc4b8c707de5e78a6da7ab" readonly STATE_DIR="${HOME}/.local/state/station-bootstrap" readonly INSTALL_BOOT_MARKER="${STATE_DIR}/install-boot-id" +readonly NEMOCLAW_CONFIG_DIR="/etc/nemoclaw" +readonly DUAL_STATION_CONTROLLER_UID_FILE="${NEMOCLAW_CONFIG_DIR}/dual-station-controller-uid" +readonly DUAL_STATION_CONTROLLER_UID_MODE="0644" DOCKER_BASELINE_CAPTURED=0 DOCKER_CONTAINER_BASELINE="" DOCKER_CONTAINER_BASELINE_TOTAL=0 @@ -339,13 +342,27 @@ on_error() { exit "$rc" } +# Remote pair preparation must never fall back to an interactive sudo prompt. +# Keep the helper's local behavior unchanged unless its caller explicitly opts +# into this strict mode after proving `sudo -n true` succeeds. +sudo() { + if [[ "${NEMOCLAW_STATION_PREP_SUDO_NONINTERACTIVE:-0}" == "1" && "${1:-}" != "-n" ]]; then + command sudo -n "$@" + else + command sudo "$@" + fi +} + usage() { cat <<'EOF' -Usage: prepare-dgx-station-host.sh --check|--apply|--verify [--force-station-install] +Usage: prepare-dgx-station-host.sh --check|--apply|--verify|--bind-controller [--force-station-install] --check Read-only eligibility and current-state report. --apply Install exact prerequisites or finish post-reboot runtime setup. --verify Read-only host verification plus ephemeral GPU container tests. + --bind-controller + Bind only the current non-root controller UID without inspecting + or disrupting an already-running managed inference workload. --force-station-install Bypass only the DGX release-metadata allowlist. ARM64 Ubuntu 24.04, Station GB300 hardware, workload quiescence, and all factory-runtime @@ -366,7 +383,7 @@ parse_args() { FORCE_STATION_INSTALL=0 for arg in "$@"; do case "$arg" in - --check | --apply | --verify | --classify-dgx-release) + --check | --apply | --verify | --bind-controller | --classify-dgx-release) [[ -z "$MODE" ]] || return 1 MODE="$arg" ;; @@ -375,7 +392,8 @@ parse_args() { esac done [[ -n "$MODE" ]] || return 1 - [[ "$MODE" != "--classify-dgx-release" || "$FORCE_STATION_INSTALL" == "0" ]] + [[ "$MODE" != "--classify-dgx-release" || "$FORCE_STATION_INSTALL" == "0" ]] \ + && [[ "$MODE" != "--bind-controller" || "$FORCE_STATION_INSTALL" == "0" ]] } is_station_gb300_product() { @@ -1134,17 +1152,25 @@ host_docker_sudo() { } query_host_docker() { - local output + local output allow_sudo=0 DOCKER_QUERY_OUTPUT="" command -v docker >/dev/null 2>&1 || return 2 if output="$(host_docker "$@" 2>/dev/null)"; then DOCKER_QUERY_OUTPUT="$output" return 0 fi - if [[ "$MODE" == "--apply" ]] && output="$(host_docker_sudo "$@" 2>/dev/null)"; then + if [[ "$MODE" == "--apply" || + ("$MODE" == "--check" && "${NEMOCLAW_STATION_PREP_SUDO_NONINTERACTIVE:-0}" == "1") ]]; then + allow_sudo=1 + fi + if ((allow_sudo == 1)) && output="$(host_docker_sudo "$@" 2>/dev/null)"; then DOCKER_QUERY_OUTPUT="$output" if ((DOCKER_QUERY_USES_SUDO == 0)); then - info "docker_access=sudo_until_group_membership_is_active" + if [[ "$MODE" == "--check" ]]; then + info "docker_access=sudo_for_noninteractive_read_only_check" + else + info "docker_access=sudo_until_group_membership_is_active" + fi DOCKER_QUERY_USES_SUDO=1 fi return 0 @@ -1634,6 +1660,112 @@ install_exact_file_or_reuse() { info "${label}=installed path=${target}" } +preparation_controller_uid_for() { + local effective_uid=${1:-} sudo_uid=${2:-} uid + if [[ "$effective_uid" == "0" ]]; then + uid="$sudo_uid" + else + uid="$effective_uid" + fi + if ! [[ "$uid" =~ ^[1-9][0-9]*$ && ${#uid} -le 10 ]] || ((10#$uid > 4294967295)); then + fatal "Station preparation must be run by a non-root controller account" + fi + printf '%s\n' "$uid" +} + +preparation_controller_uid() { + preparation_controller_uid_for "$EUID" "${SUDO_UID:-}" +} + +root_directory_is_safe_unprivileged() { + local path=$1 expected_mode=${2:-} metadata uid gid mode + test ! -L "$path" || return 1 + test -d "$path" || return 1 + metadata="$(stat -c '%u %g %a' -- "$path")" || return 1 + read -r uid gid mode <<<"$metadata" + [[ "$uid" == "0" && "$gid" == "0" && "$mode" =~ ^[0-7]{3,4}$ ]] || return 1 + (((8#$mode & 0022) == 0)) || return 1 + [[ -z "$expected_mode" || "$mode" == "${expected_mode#0}" ]] +} + +root_regular_file_is_safe_unprivileged() { + local path=$1 expected_mode=$2 metadata uid gid mode + test ! -L "$path" || return 1 + test -f "$path" || return 1 + metadata="$(stat -c '%u %g %a' -- "$path")" || return 1 + read -r uid gid mode <<<"$metadata" + [[ "$uid" == "0" && "$gid" == "0" && "$mode" == "${expected_mode#0}" ]] +} + +verify_dual_station_controller_uid_binding() { + local expected_uid=${1:-} config_dir=${2:-$NEMOCLAW_CONFIG_DIR} + local binding_file=${3:-$DUAL_STATION_CONTROLLER_UID_FILE} + [[ -n "$expected_uid" ]] || expected_uid="$(preparation_controller_uid)" + root_directory_is_safe_unprivileged "$config_dir" 0755 \ + || fatal "NemoClaw configuration directory must be root-owned with mode 0755 so the prepared controller can read its binding: ${config_dir}" + root_regular_file_is_safe_unprivileged "$binding_file" "$DUAL_STATION_CONTROLLER_UID_MODE" \ + || fatal "Dual-Station controller UID binding must be a root-owned regular file with mode ${DUAL_STATION_CONTROLLER_UID_MODE}: ${binding_file}" + if ! printf '%s\n' "$expected_uid" | cmp -s - "$binding_file"; then + fatal "Dual-Station is already bound to a different controller UID; an administrator must remove ${binding_file} before rebinding" + fi + info "dual_station_controller_uid=verified uid=${expected_uid} path=${binding_file}" +} + +ensure_dual_station_controller_uid_binding() { + local config_dir=${1:-$NEMOCLAW_CONFIG_DIR} + local binding_file=${2:-$DUAL_STATION_CONTROLLER_UID_FILE} + local expected_uid parent tmp published=0 + expected_uid="$(preparation_controller_uid)" + parent="$(dirname "$config_dir")" + ensure_root_directory_safe \ + "$config_dir" "$parent" 0755 "NemoClaw configuration directory" + root_directory_is_safe_unprivileged "$config_dir" 0755 \ + || fatal "NemoClaw configuration directory must be root-owned with mode 0755 before binding the controller: ${config_dir}" + sudo test ! -L "$binding_file" \ + || fatal "Dual-Station controller UID binding must not be a symbolic link: ${binding_file}" + if sudo test -e "$binding_file"; then + verify_dual_station_controller_uid_binding "$expected_uid" "$config_dir" "$binding_file" + return 0 + fi + + tmp="$(sudo mktemp "${config_dir}/.dual-station-controller-uid.XXXXXXXXXX")" \ + || fatal "Could not create the root-owned Dual-Station controller UID candidate" + if [[ "$tmp" != "${config_dir}/.dual-station-controller-uid."* || "$tmp" == *$'\n'* ]]; then + sudo rm -f -- "$tmp" || true + fatal "Root-owned Dual-Station controller UID candidate path was invalid" + fi + if ! printf '%s\n' "$expected_uid" | sudo tee "$tmp" >/dev/null; then + sudo rm -f -- "$tmp" || true + fatal "Could not write the Dual-Station controller UID candidate" + fi + if ! sudo chown root:root "$tmp" || ! sudo chmod "$DUAL_STATION_CONTROLLER_UID_MODE" "$tmp"; then + sudo rm -f -- "$tmp" || true + fatal "Could not secure the Dual-Station controller UID candidate" + fi + if ! root_regular_file_is_safe "$tmp" "$DUAL_STATION_CONTROLLER_UID_MODE"; then + sudo rm -f -- "$tmp" || true + fatal "Dual-Station controller UID candidate metadata was unsafe" + fi + if ! printf '%s\n' "$expected_uid" | sudo cmp -s - "$tmp"; then + sudo rm -f -- "$tmp" || true + fatal "Dual-Station controller UID candidate content was invalid" + fi + if sudo ln -- "$tmp" "$binding_file" 2>/dev/null; then + published=1 + fi + sudo rm -f -- "$tmp" \ + || fatal "Could not remove the Dual-Station controller UID candidate" + if ((published == 0)); then + sudo test ! -L "$binding_file" \ + || fatal "Dual-Station controller UID binding must not be a symbolic link: ${binding_file}" + sudo test -e "$binding_file" \ + || fatal "Could not publish the Dual-Station controller UID binding without replacement" + else + info "dual_station_controller_uid=installed uid=${expected_uid} path=${binding_file}" + fi + verify_dual_station_controller_uid_binding "$expected_uid" "$config_dir" "$binding_file" +} + ensure_docker_repository_source() { local docker_asc=$1 docker_gpg=$2 docker_gpg_list=$3 docker_asc_list=$4 local source_target=/etc/apt/sources.list.d/docker.list @@ -2523,6 +2655,7 @@ run_apply() { run_verify() { common_preflight + require_command cmp require_command docker require_command nvidia-ctk require_command nvidia-smi @@ -2539,6 +2672,17 @@ run_verify() { verify_host } +run_bind_controller() { + require_command cmp + require_command stat + require_command sudo + check_platform + acquire_sudo + ensure_dual_station_controller_uid_binding \ + "$NEMOCLAW_CONFIG_DIR" "$DUAL_STATION_CONTROLLER_UID_FILE" + info "CONTROLLER_UID_BINDING_READY" +} + main() { if ! parse_args "$@"; then usage >&2 @@ -2550,6 +2694,8 @@ main() { fi if [[ "$MODE" == "--apply" ]]; then setup_log + elif [[ "$MODE" == "--bind-controller" ]]; then + info "version=${SCRIPT_VERSION} mode=${MODE} log=disabled_binding_only" else info "version=${SCRIPT_VERSION} mode=${MODE} log=disabled_read_only" fi @@ -2559,6 +2705,7 @@ main() { --check) run_check ;; --apply) run_apply ;; --verify) run_verify ;; + --bind-controller) run_bind_controller ;; esac } diff --git a/scripts/prepare-dual-dgx-station.mts b/scripts/prepare-dual-dgx-station.mts new file mode 100755 index 0000000000..4d5f888b54 --- /dev/null +++ b/scripts/prepare-dual-dgx-station.mts @@ -0,0 +1,1107 @@ +#!/usr/bin/env -S node --no-warnings --experimental-strip-types +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { spawnSync } from "node:child_process"; +import { createHash, randomBytes } from "node:crypto"; +import fs from "node:fs"; +import net from "node:net"; +import os from "node:os"; +import path from "node:path"; +import { fileURLToPath } from "node:url"; +import { + clearDualStationSshBinding, + encodeDualStationSshBindingHandoff, + stationKnownHostsDigest, + writeDualStationSshBinding, +} from "../src/lib/inference/vllm-station-ssh-binding.ts"; +import { + type DualStationPreparationDeps, + type DualStationResumeState, + type PretrustedSshTarget, + parseDualStationResumeState, + parseStationDiscoveryHost, + prepareDualStationPair, + type RailConnectivityRequest, + type StationDiscoveryHost, + type StationPrepMode, + validateResumeFileMetadata, + validateStationPeerTarget, +} from "./lib/dgx-station-peer.mts"; + +const COMMAND_TIMEOUT_MS = 60_000; +const HELPER_TIMEOUT_MS = 2 * 60 * 60_000; +const MAX_PROBE_OUTPUT_BYTES = 1024 * 1024; + +type CommandResult = { + status: number | null; + stdout: string; + stderr: string; + error?: string; +}; + +type CliOptions = { + helperPath: string; + statePath: string; + revision: string; + explicitPeer?: string; + reuseExistingManagedPair: boolean; + migrateLegacySingleStationHead: boolean; + clearState: boolean; +}; + +type SshConfig = Map; + +const SUBPROCESS_ENV_NAMES = new Set([ + "HOME", + "USER", + "LOGNAME", + "SHELL", + "PATH", + "TERM", + "HOSTNAME", + "NODE_ENV", + "TMPDIR", + "TMP", + "TEMP", + "LANG", + "HTTP_PROXY", + "HTTPS_PROXY", + "NO_PROXY", + "http_proxy", + "https_proxy", + "no_proxy", + "SSL_CERT_FILE", + "SSL_CERT_DIR", + "NODE_EXTRA_CA_CERTS", + "GIT_SSL_CAINFO", + "GIT_SSL_CAPATH", + "CURL_CA_BUNDLE", + "REQUESTS_CA_BUNDLE", + "SSH_AUTH_SOCK", +]); + +const STATION_DISCOVERY_PROBE = String.raw` +import csv +import json +from pathlib import Path +import platform +import re +import socket +import subprocess + +def read_text(path): + try: + return Path(path).read_text(encoding="utf-8").rstrip("\x00").strip() + except (OSError, UnicodeError): + return "" + +def run(argv, timeout=5): + try: + result = subprocess.run( + argv, + stdin=subprocess.DEVNULL, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + timeout=timeout, + check=False, + ) + return result.returncode, result.stdout.strip() + except (FileNotFoundError, OSError, subprocess.TimeoutExpired): + return 127, "" + +def product_name(): + for candidate in ( + "/sys/class/dmi/id/product_name", + "/sys/devices/virtual/dmi/id/product_name", + "/sys/firmware/devicetree/base/model", + ): + value = read_text(candidate) + if value: + return value + return "" + +def gpu_inventory(): + rc, output = run([ + "nvidia-smi", + "--query-gpu=index,name,uuid", + "--format=csv,noheader,nounits", + ]) + if rc != 0: + return [] + result = [] + for row in csv.reader(output.splitlines()): + if len(row) != 3: + continue + try: + index = int(row[0].strip()) + except ValueError: + continue + result.append({"index": index, "name": row[1].strip(), "uuid": row[2].strip()}) + return result + +def ipv4_addresses(netdev): + rc, output = run(["ip", "-j", "-4", "address", "show", "dev", netdev]) + if rc != 0: + return [] + try: + links = json.loads(output) + except json.JSONDecodeError: + return [] + result = [] + for link in links if isinstance(links, list) else []: + for address in link.get("addr_info", []): + if address.get("family") != "inet" or address.get("scope") == "host": + continue + local = address.get("local") + prefix = address.get("prefixlen") + if isinstance(local, str) and isinstance(prefix, int): + result.append({"address": local, "prefixLength": prefix}) + return result + +def rail_inventory(): + rc, output = run(["ibdev2netdev"]) + if rc != 0: + return [] + rails = [] + pattern = re.compile(r"^(\S+)\s+port\s+(\d+)\s+==>\s+(\S+)\s+\(([^)]*)\)") + for line in output.splitlines(): + match = pattern.match(line.strip()) + if not match: + continue + rdma_device, raw_port, netdev, _reported_state = match.groups() + port = int(raw_port) + device_path = Path("/sys/class/net") / netdev / "device" + try: + pci_address = device_path.resolve(strict=True).name + except OSError: + pci_address = "" + rc_lspci, pci_name = run(["lspci", "-D", "-s", pci_address]) if pci_address else (127, "") + if rc_lspci != 0: + pci_name = "" + try: + speed_mbps = int(read_text(Path("/sys/class/net") / netdev / "speed")) + except ValueError: + speed_mbps = -1 + try: + mtu = int(read_text(Path("/sys/class/net") / netdev / "mtu")) + except ValueError: + mtu = -1 + ib_port = Path("/sys/class/infiniband") / rdma_device / "ports" / str(port) + rails.append({ + "netdev": netdev, + "macAddress": read_text(Path("/sys/class/net") / netdev / "address").lower(), + "pciAddress": pci_address, + "pciName": pci_name, + "state": read_text(ib_port / "state"), + "linkLayer": read_text(ib_port / "link_layer"), + "speedMbps": speed_mbps, + "mtu": mtu, + "ipv4Addresses": ipv4_addresses(netdev), + }) + return rails + +print(json.dumps({ + "schemaVersion": 1, + "hostname": socket.gethostname(), + "productName": product_name(), + "architecture": platform.machine(), + "gpus": gpu_inventory(), + "rails": rail_inventory(), +}, separators=(",", ":"))) +`; + +export const CONNECTIVITY_PROBE = String.raw` +import ipaddress +import json +import subprocess +import sys + +def run(argv, timeout=5): + try: + result = subprocess.run( + argv, + stdin=subprocess.DEVNULL, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + timeout=timeout, + check=False, + ) + return result.returncode, result.stdout.strip() + except (FileNotFoundError, OSError, subprocess.TimeoutExpired): + return 127, "" + +if len(sys.argv[1:]) != 6: + raise SystemExit("expected two netdev/source/peer triples") + +checks = [] +for offset in (0, 3): + netdev, source, peer = sys.argv[1 + offset:4 + offset] + route_device = "" + route_source = "" + route_gateway = None + route_scope = "" + peer_mac = "" + peer_neighbor_state = "" + rc, output = run(["ip", "-j", "route", "get", peer, "from", source, "oif", netdev]) + if rc == 0: + try: + routes = json.loads(output) + route = routes[0] if isinstance(routes, list) and routes else {} + route_device = route.get("dev", "") if isinstance(route, dict) else "" + route_source = route.get("prefsrc", route.get("src", route.get("from", ""))) if isinstance(route, dict) else "" + route_gateway = route.get("gateway") if isinstance(route, dict) else None + except json.JSONDecodeError: + pass + network = str(ipaddress.ip_network(source + "/30", strict=False)) + link_rc, link_output = run(["ip", "-j", "route", "show", "exact", network]) + if link_rc == 0: + try: + link_routes = json.loads(link_output) + matching_routes = [ + route for route in link_routes if ( + isinstance(route, dict) + and route.get("dst") == network + and route.get("dev") == netdev + and route.get("gateway") is None + ) + ] if isinstance(link_routes, list) else [] + if len(matching_routes) == 1: + route_scope = matching_routes[0].get("scope", "") + except json.JSONDecodeError: + pass + ping_rc, _ = run([ + "ping", "-4", "-M", "do", "-s", "8972", "-c", "1", "-W", "2", "-I", source, peer, + ]) + neighbor_rc, neighbor_output = run(["ip", "-j", "neighbor", "show", "to", peer]) + if neighbor_rc == 0: + try: + neighbors = json.loads(neighbor_output) + matching_neighbors = [ + neighbor for neighbor in neighbors if ( + isinstance(neighbor, dict) + and neighbor.get("dst") == peer + and neighbor.get("dev") == netdev + ) + ] if isinstance(neighbors, list) else [] + if len(matching_neighbors) == 1: + peer_mac = str(matching_neighbors[0].get("lladdr", "")).lower() + raw_state = matching_neighbors[0].get("state", "") + peer_neighbor_state = ",".join(raw_state) if isinstance(raw_state, list) else str(raw_state) + except json.JSONDecodeError: + pass + checks.append({ + "netdev": netdev, + "sourceAddress": source, + "peerAddress": peer, + "routeDevice": route_device, + "routeSource": route_source, + "routeGateway": route_gateway, + "routeScope": route_scope, + "peerMac": peer_mac, + "peerNeighborState": peer_neighbor_state, + "jumboPing": ping_rc == 0, + }) + +print(json.dumps({"schemaVersion": 1, "checks": checks}, separators=(",", ":"))) +`; + +export function buildStationPrepSubprocessEnv( + source: NodeJS.ProcessEnv = process.env, +): NodeJS.ProcessEnv { + const env: NodeJS.ProcessEnv = {}; + for (const [name, value] of Object.entries(source)) { + if ( + value !== undefined && + (SUBPROCESS_ENV_NAMES.has(name) || name.startsWith("LC_") || name.startsWith("XDG_")) + ) { + env[name] = value; + } + } + env.LC_ALL = "C"; + env.LANG = "C"; + return env; +} + +function runCommand( + file: string, + args: readonly string[], + input = "", + timeout = COMMAND_TIMEOUT_MS, + maxBuffer = MAX_PROBE_OUTPUT_BYTES, +): CommandResult { + const result = spawnSync(file, [...args], { + encoding: "utf8", + input, + timeout, + maxBuffer, + killSignal: "SIGKILL", + windowsHide: true, + env: buildStationPrepSubprocessEnv(), + }); + return { + status: result.status, + stdout: result.stdout ?? "", + stderr: result.stderr ?? "", + error: result.error?.message, + }; +} + +function commandSucceeded(result: CommandResult, requireOutput = false): boolean { + return ( + result.status === 0 && !result.error && (!requireOutput || result.stdout.trim().length > 0) + ); +} + +function runStreamingCommand( + file: string, + args: readonly string[], + input: string, + timeout = HELPER_TIMEOUT_MS, +): number { + const result = spawnSync(file, [...args], { + input, + timeout, + killSignal: "SIGKILL", + windowsHide: true, + env: buildStationPrepSubprocessEnv(), + // Keep stdout machine-readable for the coordinator result while allowing + // long package and acceptance-image operations to stream without a + // bounded child-process buffer. + stdio: ["pipe", process.stderr.fd, process.stderr.fd], + }); + if (result.error) { + process.stderr.write(`[station-pair] ${result.error.message}\n`); + } + return result.status ?? 1; +} + +export function strictStationPrepSshTransportArgs(): string[] { + return [ + "-T", + "-o", + "BatchMode=yes", + "-o", + "StrictHostKeyChecking=yes", + "-o", + "VerifyHostKeyDNS=no", + "-o", + "NoHostAuthenticationForLocalhost=no", + "-o", + "NumberOfPasswordPrompts=0", + "-o", + "PasswordAuthentication=no", + "-o", + "KbdInteractiveAuthentication=no", + "-o", + "PreferredAuthentications=publickey", + "-o", + "ConnectTimeout=5", + "-o", + "ConnectionAttempts=1", + "-o", + "ServerAliveInterval=5", + "-o", + "ServerAliveCountMax=1", + "-o", + "ClearAllForwardings=yes", + "-o", + "ForwardAgent=no", + "-o", + "ForwardX11=no", + "-o", + "ForwardX11Trusted=no", + "-o", + "Tunnel=no", + "-o", + "UpdateHostKeys=no", + "-o", + "ControlMaster=no", + "-o", + "ControlPath=none", + "-o", + "PermitLocalCommand=no", + "-o", + "RemoteCommand=none", + "-o", + "ProxyCommand=none", + "-o", + "ProxyJump=none", + "-o", + "KnownHostsCommand=none", + "-o", + "LogLevel=ERROR", + ]; +} + +function parseSshConfig(stdout: string): SshConfig { + const values = new Map(); + for (const rawLine of stdout.split(/\r?\n/)) { + const line = rawLine.trim(); + if (!line) continue; + const separator = line.search(/\s/); + if (separator <= 0) throw new Error("ssh -G returned malformed effective configuration"); + const key = line.slice(0, separator).toLowerCase(); + const value = line.slice(separator).trim(); + values.set(key, [...(values.get(key) ?? []), value]); + } + return values; +} + +function assertStrictSshConfig(values: SshConfig): void { + const exactly = (key: string, allowed: readonly string[]): boolean => { + const observed = (values.get(key) ?? []).map((value) => value.toLowerCase()); + return observed.length === 1 && allowed.includes(observed[0]); + }; + const absentOrNone = (key: string): boolean => { + const observed = (values.get(key) ?? []).map((value) => value.toLowerCase()); + return observed.length === 0 || (observed.length === 1 && observed[0] === "none"); + }; + const sendEnv = (values.get("sendenv") ?? []).map((value) => value.toLowerCase()); + if ( + !exactly("batchmode", ["yes"]) || + !exactly("stricthostkeychecking", ["yes", "true"]) || + !exactly("verifyhostkeydns", ["false", "no"]) || + !exactly("nohostauthenticationforlocalhost", ["false", "no"]) || + !exactly("permitlocalcommand", ["no"]) || + !exactly("forwardagent", ["no"]) || + !exactly("forwardx11", ["no"]) || + !exactly("forwardx11trusted", ["no"]) || + !exactly("tunnel", ["false", "no"]) || + !exactly("updatehostkeys", ["false", "no"]) || + !exactly("controlmaster", ["false", "no"]) || + !absentOrNone("controlpath") || + !absentOrNone("remotecommand") || + !absentOrNone("proxycommand") || + !absentOrNone("proxyjump") || + !absentOrNone("localcommand") || + !absentOrNone("knownhostscommand") || + values.has("localforward") || + values.has("remoteforward") || + values.has("dynamicforward") || + values.has("setenv") || + !sendEnv.every((value) => value === "lang" || value === "lc_*") + ) { + throw new Error("Effective SSH configuration is unsafe for Station peer preparation"); + } +} + +function oneSshConfigValue(values: SshConfig, key: string): string { + const entries = values.get(key) ?? []; + if (entries.length !== 1 || entries[0].length === 0) { + throw new Error(`Effective SSH configuration must define exactly one ${key}`); + } + return entries[0]; +} + +function fingerprintKnownHostKey(keyType: string, keyData: string): string | null { + const result = runCommand( + "ssh-keygen", + ["-l", "-E", "sha256", "-f", "-"], + `${keyType} ${keyData}\n`, + ); + if (!commandSucceeded(result, true)) return null; + return result.stdout.match(/\b(SHA256:[A-Za-z0-9+/]{16,86}={0,2})\b/)?.[1] ?? null; +} + +function knownHostEvidence( + lookupHost: string, + files: readonly string[], +): { + lines: string[]; + fingerprints: string[]; + digest: string; +} | null { + const lines = new Set(); + const fingerprints = new Set(); + for (const file of files) { + if (!path.isAbsolute(file)) continue; + let metadata: fs.Stats; + try { + metadata = fs.lstatSync(file); + } catch { + continue; + } + const uid = process.getuid?.(); + if ( + !metadata.isFile() || + metadata.isSymbolicLink() || + uid === undefined || + (metadata.uid !== uid && metadata.uid !== 0) || + (metadata.mode & 0o022) !== 0 + ) { + continue; + } + const result = runCommand("ssh-keygen", ["-F", lookupHost, "-f", file]); + if (!commandSucceeded(result, true)) continue; + for (const rawLine of result.stdout.split(/\r?\n/)) { + const line = rawLine.trim(); + if (!line || line.startsWith("#") || /[\u0000\r\n]/.test(line)) continue; + const fields = line.split(/\s+/); + const marker = fields[0]?.startsWith("@") ? fields.shift() : ""; + if (fields.length < 3) continue; + const [_hosts, keyType, keyData] = fields; + if (!/^(?:ssh-|ecdsa-|sk-)[A-Za-z0-9@._+-]+$/.test(keyType)) continue; + if (!/^[A-Za-z0-9+/]+={0,3}$/.test(keyData)) continue; + const fingerprint = fingerprintKnownHostKey(keyType, keyData); + if (!fingerprint) continue; + lines.add(line); + // Preserve matching revocations in the private pinned file so the + // subsequent SSH connection cannot resurrect a key the operator or + // system administrator explicitly revoked. A revoked line alone is not + // positive trust evidence. + if (marker === "@revoked") continue; + fingerprints.add(fingerprint); + } + } + if (lines.size === 0 || fingerprints.size === 0) return null; + const retainedLines = [...lines].sort(); + return { + lines: retainedLines, + fingerprints: [...fingerprints].sort(), + digest: stationKnownHostsDigest(`${retainedLines.join("\n")}\n`), + }; +} + +export function inspectPretrustedSshTarget(target: string): PretrustedSshTarget | null { + validateStationPeerTarget(target); + const configResult = runCommand( + "ssh", + ["-G", ...strictStationPrepSshTransportArgs(), "--", target], + "", + ); + if (!commandSucceeded(configResult, true)) return null; + const config = parseSshConfig(configResult.stdout); + assertStrictSshConfig(config); + const resolvedHost = oneSshConfigValue(config, "hostname"); + const sshUser = oneSshConfigValue(config, "user"); + const portText = oneSshConfigValue(config, "port"); + const port = Number(portText); + if (!Number.isInteger(port) || port < 1 || port > 65535) { + throw new Error("Effective SSH port is invalid"); + } + const requestedHost = target.slice(target.lastIndexOf("@") + 1); + if (net.isIP(requestedHost) === 4 && resolvedHost !== requestedHost) { + throw new Error("Automatic rail target was remapped by SSH configuration"); + } + const alias = config.get("hostkeyalias")?.[0]; + const baseLookupHost = alias && alias.toLowerCase() !== "none" ? alias : resolvedHost; + validateStationPeerTarget(baseLookupHost); + const lookupHost = port === 22 ? baseLookupHost : `[${baseLookupHost}]:${String(port)}`; + const knownHostFiles = [ + ...(config.get("userknownhostsfile") ?? []), + ...(config.get("globalknownhostsfile") ?? []), + ].flatMap((entry) => entry.split(/\s+/).filter((value) => value && value !== "none")); + const evidence = knownHostEvidence(lookupHost, knownHostFiles); + if (!evidence) return null; + return { + requestedTarget: target, + sshTarget: target, + resolvedHost, + sshUser, + port, + lookupHost, + hostKeyDigest: evidence.digest, + keyFingerprints: evidence.fingerprints, + knownHostsLines: evidence.lines, + }; +} + +function parseHostResult(result: CommandResult, label: string): StationDiscoveryHost { + if (!commandSucceeded(result, true)) { + throw new Error(`${label} failed${result.error ? `: ${result.error}` : ""}`); + } + if (Buffer.byteLength(result.stdout, "utf8") > MAX_PROBE_OUTPUT_BYTES) { + throw new Error(`${label} output is too large`); + } + let value: unknown; + try { + value = JSON.parse(result.stdout); + } catch { + throw new Error(`${label} returned invalid JSON`); + } + return parseStationDiscoveryHost(value); +} + +function validateConnectivityArgs(requests: readonly RailConnectivityRequest[]): string[] { + if (requests.length !== 2) throw new Error("Station connectivity requires exactly two rails"); + const args: string[] = []; + for (const request of requests) { + if (!/^[A-Za-z0-9][A-Za-z0-9_.:-]{0,63}$/.test(request.netdev)) { + throw new Error("Station connectivity netdev is unsafe"); + } + if (net.isIP(request.sourceAddress) !== 4 || net.isIP(request.peerAddress) !== 4) { + throw new Error("Station connectivity addresses are invalid"); + } + if (!/^(?:[0-9a-f]{2}:){5}[0-9a-f]{2}$/.test(request.expectedPeerMac)) { + throw new Error("Station connectivity peer MAC is invalid"); + } + args.push(request.netdev, request.sourceAddress, request.peerAddress); + } + return args; +} + +function connectivityMatches( + result: CommandResult, + requests: readonly RailConnectivityRequest[], +): boolean { + if (!commandSucceeded(result, true)) return false; + let value: unknown; + try { + value = JSON.parse(result.stdout); + } catch { + return false; + } + if ( + typeof value !== "object" || + value === null || + !("schemaVersion" in value) || + value.schemaVersion !== 1 || + !("checks" in value) || + !Array.isArray(value.checks) || + value.checks.length !== requests.length + ) { + return false; + } + const checks = value.checks as unknown[]; + return requests.every((request) => + checks.some( + (check) => + typeof check === "object" && + check !== null && + "netdev" in check && + check.netdev === request.netdev && + "sourceAddress" in check && + check.sourceAddress === request.sourceAddress && + "peerAddress" in check && + check.peerAddress === request.peerAddress && + "routeDevice" in check && + check.routeDevice === request.netdev && + "routeSource" in check && + check.routeSource === request.sourceAddress && + "routeGateway" in check && + check.routeGateway === null && + "routeScope" in check && + typeof check.routeScope === "string" && + check.routeScope.toLowerCase() === "link" && + "peerMac" in check && + check.peerMac === request.expectedPeerMac && + "peerNeighborState" in check && + typeof check.peerNeighborState === "string" && + /^(?:REACHABLE|STALE|DELAY|PROBE|PERMANENT|NOARP)(?:,(?:REACHABLE|STALE|DELAY|PROBE|PERMANENT|NOARP))*$/i.test( + check.peerNeighborState, + ) && + "jumboPing" in check && + check.jumboPing === true, + ), + ); +} + +export function buildRemoteHelperCommand(helperSha256: string, mode: StationPrepMode): string { + if (!/^[a-f0-9]{64}$/.test(helperSha256)) throw new Error("Helper SHA-256 is invalid"); + if ( + mode !== "--check" && + mode !== "--apply" && + mode !== "--verify" && + mode !== "--bind-controller" + ) { + throw new Error("Helper mode is invalid"); + } + return [ + "set -eu", + "umask 077", + 'd=$(mktemp -d "${TMPDIR:-/tmp}/nemoclaw-station-prep.XXXXXX")', + "trap 'rm -rf -- \"$d\"' EXIT HUP INT TERM", + 'f="$d/prepare-dgx-station-host.sh"', + 'cat >"$f"', + `test "$(sha256sum "$f" | awk '{print $1}')" = "${helperSha256}"`, + 'chmod 0600 "$f"', + "sudo -n true", + `NEMOCLAW_STATION_PREP_SUDO_NONINTERACTIVE=1 bash "$f" ${mode}`, + ].join("; "); +} + +function assertSecureStateDirectory(directory: string): void { + let metadata: fs.Stats; + try { + metadata = fs.lstatSync(directory); + } catch (error) { + if ((error as NodeJS.ErrnoException).code === "ENOENT") { + throw new Error("Dual-Station resume directory must already exist"); + } + throw error; + } + const uid = process.getuid?.(); + if ( + !metadata.isDirectory() || + metadata.isSymbolicLink() || + uid === undefined || + metadata.uid !== uid || + (metadata.mode & 0o077) !== 0 + ) { + throw new Error("Dual-Station resume directory must be owner-only and symlink-free"); + } +} + +export function readDualStationResumeState(statePath: string): DualStationResumeState | null { + const directory = path.dirname(statePath); + assertSecureStateDirectory(directory); + const noFollow = fs.constants.O_NOFOLLOW; + if (typeof noFollow !== "number") throw new Error("O_NOFOLLOW is required for resume state"); + let fd: number; + try { + fd = fs.openSync(statePath, fs.constants.O_RDONLY | noFollow | fs.constants.O_NONBLOCK); + } catch (error) { + const code = (error as NodeJS.ErrnoException).code; + if (code === "ENOENT") return null; + if (code === "ELOOP") throw new Error("Dual-Station resume state must not be a symlink"); + throw error; + } + try { + const metadata = fs.fstatSync(fd); + const uid = process.getuid?.(); + if (uid === undefined) throw new Error("Current user identity is unavailable"); + validateResumeFileMetadata( + { + isFile: metadata.isFile(), + isSymbolicLink: metadata.isSymbolicLink(), + uid: metadata.uid, + mode: metadata.mode, + size: metadata.size, + }, + uid, + ); + const raw = fs.readFileSync(fd, "utf8"); + let value: unknown; + try { + value = JSON.parse(raw); + } catch { + throw new Error("Dual-Station resume state is malformed JSON"); + } + return parseDualStationResumeState(value); + } finally { + fs.closeSync(fd); + } +} + +export function writeDualStationResumeState( + statePath: string, + state: DualStationResumeState, +): void { + const validated = parseDualStationResumeState(state); + const directory = path.dirname(statePath); + assertSecureStateDirectory(directory); + const noFollow = fs.constants.O_NOFOLLOW; + if (typeof noFollow !== "number") throw new Error("O_NOFOLLOW is required for resume state"); + const temporary = `${statePath}.tmp.${randomBytes(12).toString("hex")}`; + let fd: number | null = null; + let failure: { error: unknown } | null = null; + try { + fd = fs.openSync( + temporary, + fs.constants.O_WRONLY | fs.constants.O_CREAT | fs.constants.O_EXCL | noFollow, + 0o600, + ); + fs.writeFileSync(fd, `${JSON.stringify(validated)}\n`, "utf8"); + fs.fsyncSync(fd); + fs.closeSync(fd); + fd = null; + fs.renameSync(temporary, statePath); + const directoryFd = fs.openSync(directory, fs.constants.O_RDONLY); + let directoryFailure: { error: unknown } | null = null; + try { + fs.fsyncSync(directoryFd); + } catch (error) { + directoryFailure = { error }; + } + try { + fs.closeSync(directoryFd); + } catch (error) { + directoryFailure ??= { error }; + } + if (directoryFailure) throw directoryFailure.error; + } catch (error) { + failure = { error }; + } + if (fd !== null) { + try { + fs.closeSync(fd); + } catch (error) { + failure ??= { error }; + } + } + try { + fs.unlinkSync(temporary); + } catch (error) { + if ((error as NodeJS.ErrnoException).code !== "ENOENT") failure ??= { error }; + } + if (failure) throw failure.error; +} + +export function clearDualStationResumeState(statePath: string): void { + const current = readDualStationResumeState(statePath); + clearDualStationSshBinding(statePath); + if (!current) return; + + fs.unlinkSync(statePath); + const directoryFd = fs.openSync(path.dirname(statePath), fs.constants.O_RDONLY); + let failure: { error: unknown } | null = null; + try { + fs.fsyncSync(directoryFd); + } catch (error) { + failure = { error }; + } + try { + fs.closeSync(directoryFd); + } catch (error) { + failure ??= { error }; + } + if (failure) throw failure.error; +} + +function parseCliOptions(args: readonly string[]): CliOptions { + const values = new Map(); + let reuseExistingManagedPair = false; + let migrateLegacySingleStationHead = false; + let clearState = false; + for (let index = 0; index < args.length; index += 1) { + const arg = args[index]; + if (arg === "--reuse-existing-managed-pair") { + reuseExistingManagedPair = true; + continue; + } + if (arg === "--migrate-legacy-single-head") { + migrateLegacySingleStationHead = true; + continue; + } + if (arg === "--clear-state") { + clearState = true; + continue; + } + if (!arg.startsWith("--") || index + 1 >= args.length) { + throw new Error(`Unexpected argument: ${arg}`); + } + if (values.has(arg)) throw new Error(`Duplicate argument: ${arg}`); + values.set(arg, args[(index += 1)]); + } + const helperPath = values.get("--helper") ?? ""; + const statePath = values.get("--state") ?? ""; + const revision = values.get("--revision") ?? ""; + if (!path.isAbsolute(statePath)) throw new Error("--state must be an absolute path"); + if (!clearState && !path.isAbsolute(helperPath)) { + throw new Error("--helper must be an absolute path"); + } + if (!clearState && !/^[a-f0-9]{40}$/.test(revision)) { + throw new Error("--revision must be an exact commit SHA"); + } + const explicitPeer = values.get("--explicit-peer"); + if (explicitPeer !== undefined) validateStationPeerTarget(explicitPeer); + const allowed = new Set(["--helper", "--state", "--revision", "--explicit-peer"]); + for (const key of values.keys()) { + if (!allowed.has(key)) throw new Error(`Unexpected argument: ${key}`); + } + return { + helperPath, + statePath, + revision, + explicitPeer, + reuseExistingManagedPair, + migrateLegacySingleStationHead, + clearState, + }; +} + +function assertHelperFile(helperPath: string): Buffer { + const noFollow = fs.constants.O_NOFOLLOW; + if (typeof noFollow !== "number") throw new Error("O_NOFOLLOW is required for the helper"); + let fd: number; + try { + fd = fs.openSync(helperPath, fs.constants.O_RDONLY | noFollow | fs.constants.O_NONBLOCK); + } catch (error) { + if ((error as NodeJS.ErrnoException).code === "ELOOP") { + throw new Error("Station host-preparation helper must not be a symlink"); + } + throw error; + } + try { + const metadata = fs.fstatSync(fd); + if (!metadata.isFile() || metadata.size <= 0 || metadata.size > 2 ** 20) { + throw new Error("Station host-preparation helper must be a bounded regular file"); + } + return fs.readFileSync(fd); + } finally { + fs.closeSync(fd); + } +} + +function sshArgs( + binding: PretrustedSshTarget, + pinnedKnownHostsPath: string, + remoteCommand: string, +): string[] { + return [ + ...strictStationPrepSshTransportArgs(), + "-o", + `UserKnownHostsFile=${pinnedKnownHostsPath}`, + "-o", + "GlobalKnownHostsFile=/dev/null", + "-o", + `HostKeyAlias=${binding.lookupHost}`, + "--", + binding.sshTarget, + remoteCommand, + ]; +} + +function createRuntimeDeps(options: CliOptions): { + deps: DualStationPreparationDeps; + helperSha256: string; + cleanup(): void; +} { + const helperBytes = assertHelperFile(options.helperPath); + const helperSha256 = createHash("sha256").update(helperBytes).digest("hex"); + const temporaryDirectory = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-station-pair-")); + fs.chmodSync(temporaryDirectory, 0o700); + const pinnedHelperPath = path.join(temporaryDirectory, "prepare-dgx-station-host.sh"); + fs.writeFileSync(pinnedHelperPath, helperBytes, { flag: "wx", mode: 0o600 }); + const pinnedFiles = new Map(); + + const pinnedKnownHosts = (binding: PretrustedSshTarget): string => { + const cached = pinnedFiles.get(binding.hostKeyDigest); + if (cached) return cached; + const file = path.join(temporaryDirectory, `known-hosts-${binding.hostKeyDigest}`); + fs.writeFileSync(file, `${binding.knownHostsLines.join("\n")}\n`, { + encoding: "utf8", + flag: "wx", + mode: 0o600, + }); + pinnedFiles.set(binding.hostKeyDigest, file); + return file; + }; + + const runLocalHelper = (mode: StationPrepMode): number => { + return runStreamingCommand("bash", [pinnedHelperPath, mode], ""); + }; + + const deps: DualStationPreparationDeps = { + runLocalHelper, + probeLocalHost: () => + parseHostResult( + runCommand("python3", ["-"], STATION_DISCOVERY_PROBE), + "Local Station identity probe", + ), + inspectPretrustedTarget: inspectPretrustedSshTarget, + probePeerHost: (binding) => + parseHostResult( + runCommand( + "ssh", + sshArgs(binding, pinnedKnownHosts(binding), "python3 -"), + STATION_DISCOVERY_PROBE, + ), + "Peer Station identity probe", + ), + probeLocalConnectivity: (requests) => { + const args = validateConnectivityArgs(requests); + return connectivityMatches( + runCommand("python3", ["-", ...args], CONNECTIVITY_PROBE), + requests, + ); + }, + probePeerConnectivity: (binding, requests) => { + const args = validateConnectivityArgs(requests); + return connectivityMatches( + runCommand( + "ssh", + sshArgs(binding, pinnedKnownHosts(binding), ["python3", "-", ...args].join(" ")), + CONNECTIVITY_PROBE, + ), + requests, + ); + }, + runRemoteHelper: (binding, mode) => { + return runStreamingCommand( + "ssh", + sshArgs(binding, pinnedKnownHosts(binding), buildRemoteHelperCommand(helperSha256, mode)), + helperBytes.toString("utf8"), + ); + }, + readResumeState: () => readDualStationResumeState(options.statePath), + writeResumeState: (state) => writeDualStationResumeState(options.statePath, state), + clearResumeState: () => clearDualStationResumeState(options.statePath), + log: (message) => process.stderr.write(`[station-pair] ${message}\n`), + }; + + return { + deps, + helperSha256, + cleanup: () => fs.rmSync(temporaryDirectory, { recursive: true, force: true }), + }; +} + +export function runCli(args: readonly string[]): number { + const options = parseCliOptions(args); + if (options.clearState) { + clearDualStationResumeState(options.statePath); + process.stdout.write(`${JSON.stringify({ kind: "cleared" })}\n`); + return 0; + } + const runtime = createRuntimeDeps(options); + try { + const result = prepareDualStationPair( + { + revision: options.revision, + helperSha256: runtime.helperSha256, + explicitPeer: options.explicitPeer, + reuseExistingManagedPair: options.reuseExistingManagedPair, + migrateLegacySingleStationHead: options.migrateLegacySingleStationHead, + }, + runtime.deps, + ); + if (result.kind === "single-station") { + runtime.deps.log(`Using the existing single-Station path: ${result.reason}`); + // A single-Station result is possible only without pair resume state. + // Remove any owner-only binding orphan left by an interrupted earlier + // cleanup so it cannot linger without the pair identity that owns it. + if (readDualStationResumeState(options.statePath)) { + throw new Error("Single-Station fallback cannot discard exact pair resume state"); + } + clearDualStationSshBinding(options.statePath); + process.stdout.write(`${JSON.stringify(result)}\n`); + return 0; + } + if ( + result.binding.sshTarget !== result.peerTarget || + result.binding.hostKeyDigest !== result.identity.hostKeyDigest + ) { + throw new Error("Qualified Station SSH binding does not match the prepared pair identity"); + } + const binding = writeDualStationSshBinding(options.statePath, result.binding); + process.stdout.write( + `${JSON.stringify({ + kind: result.kind, + peerTarget: result.peerTarget, + identity: result.identity, + sshBinding: encodeDualStationSshBindingHandoff(binding), + })}\n`, + ); + return result.kind === "reboot-required" ? 10 : 0; + } finally { + runtime.cleanup(); + } +} + +function isMainModule(): boolean { + const invoked = process.argv[1]; + return Boolean(invoked && path.resolve(invoked) === path.resolve(fileURLToPath(import.meta.url))); +} + +if (isMainModule()) { + try { + process.exitCode = runCli(process.argv.slice(2)); + } catch (error) { + process.stderr.write(`[station-pair] ERROR: ${(error as Error).message}\n`); + process.exitCode = 1; + } +} diff --git a/src/lib/actions/sandbox/rebuild-local-provider-recreate.test.ts b/src/lib/actions/sandbox/rebuild-local-provider-recreate.test.ts index 0eb5dcef86..6edc921d4f 100644 --- a/src/lib/actions/sandbox/rebuild-local-provider-recreate.test.ts +++ b/src/lib/actions/sandbox/rebuild-local-provider-recreate.test.ts @@ -115,6 +115,7 @@ const localProviderScenarios = [ applyLocalInferenceRoute, run: () => ({ status: 0 }), VLLM_LOCAL_CREDENTIAL_ENV: "NEMOCLAW_VLLM_LOCAL_TOKEN", + getManagedVllmProviderBinding: () => null, ...unusedCommonInferenceDeps, }, ), diff --git a/src/lib/inference/context-window.ts b/src/lib/inference/context-window.ts index 4f0aa6158e..433a091814 100644 --- a/src/lib/inference/context-window.ts +++ b/src/lib/inference/context-window.ts @@ -10,10 +10,12 @@ * window kept for a cloud model → silent under-utilization). */ -import { VLLM_PORT } from "../core/ports"; import { DEFAULT_CONTEXT_WINDOW } from "./config"; import { + getLocalProviderHealthEndpoint, + getManagedDualStationVllmProviderBinding, getOllamaWarmupCommand, + probeVllmModels, type RunCaptureFn, resolveOllamaRuntimeContextWindow, } from "./local"; @@ -44,10 +46,22 @@ const defaultContextWindowDeps: ContextWindowDeps = { probeVllmContextWindow: (model: string): number | null => { // Same source onboard uses: GET /v1/models on the host vLLM server and read // max_model_len (handles both NemoClaw-launched and bring-your-own vLLM). - const { runCapture } = require("../runner") as { runCapture: RunCaptureFn }; - const raw = runCapture(["curl", "-sf", `http://127.0.0.1:${VLLM_PORT}/v1/models`], { - ignoreError: true, - }); + let managedBinding: ReturnType; + try { + managedBinding = getManagedDualStationVllmProviderBinding(); + } catch { + return null; + } + let raw: string; + if (managedBinding) { + const result = probeVllmModels(managedBinding.baseUrl, managedBinding.apiKey); + raw = result.ok ? result.body : ""; + } else { + const { runCapture } = require("../runner") as { runCapture: RunCaptureFn }; + const endpoint = getLocalProviderHealthEndpoint("vllm-local"); + if (!endpoint) return null; + raw = runCapture(["curl", "-sf", endpoint], { ignoreError: true }); + } if (!raw) return null; let parsed: unknown; try { diff --git a/src/lib/inference/dgx-station-identity.ts b/src/lib/inference/dgx-station-identity.ts new file mode 100644 index 0000000000..4b3d111a75 --- /dev/null +++ b/src/lib/inference/dgx-station-identity.ts @@ -0,0 +1,11 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +/** Match the product identifiers reported by supported DGX Station GB300 firmware. */ +export function isDgxStationGb300Product(productName: string): boolean { + return ( + /(? string | null; + onManagedHeadObserved?: () => void; +} + +const lifecycle = vi.hoisted(() => ({ + baseUrl: vi.fn<(overrides?: ManagedBaseUrlOverrides) => string | null>(), +})); + +vi.mock("./vllm-station-cluster-lifecycle", () => ({ + getDualStationManagedVllmBaseUrl: lifecycle.baseUrl, +})); + +import { + CONTAINER_REACHABILITY_IMAGE, + getLocalProviderBaseUrl, + getLocalProviderContainerReachabilityCheck, + getLocalProviderHealthCheck, + getLocalProviderHealthEndpoint, + getManagedDualStationVllmProviderBinding, + getManagedDualStationVllmProviderState, + LOCAL_INFERENCE_SANDBOX_HOST_URL_ENV, + probeLocalProviderHealth, + probeVllmModels, + validateLocalProvider, +} from "./local"; + +const BASE_URL = "http://10.40.0.1:8000"; +const API_KEY = "e".repeat(64); +const OTHER_API_KEY = "f".repeat(64); +const MANAGED_BASE_URL_BY_API_KEY = new Map([ + [undefined, BASE_URL], + [API_KEY, BASE_URL], + [null, null], +]); + +let actualLifecycle: typeof import("./vllm-station-cluster-lifecycle"); + +beforeAll(async () => { + actualLifecycle = await vi.importActual("./vllm-station-cluster-lifecycle"); +}); + +function productionManagedBaseUrlResolver( + expectedApiKey = API_KEY, + apiKeyFingerprint = actualLifecycle.dualStationVllmApiKeyFingerprint(expectedApiKey), +) { + const row = [ + "a".repeat(64), + actualLifecycle.DUAL_STATION_VLLM_HEAD_CONTAINER_NAME, + "running", + DUAL_STATION_VLLM_RUNTIME.image, + "true", + "head", + BASE_URL, + "b".repeat(64), + "GPU-12345678", + "2", + "c".repeat(64), + apiKeyFingerprint, + "d".repeat(32), + ].join("\t"); + + return (overrides: ManagedBaseUrlOverrides = {}): string | null => + actualLifecycle.getDualStationManagedVllmBaseUrl({ + dockerCapture: () => row, + buildLocalDockerEnv: () => ({}), + loadApiKey: overrides.loadApiKey ?? (() => null), + onManagedHeadObserved: overrides.onManagedHeadObserved ?? (() => undefined), + localInterfaceAddresses: () => ["10.40.0.1"], + }); +} + +beforeEach(() => { + vi.stubEnv(LOCAL_INFERENCE_SANDBOX_HOST_URL_ENV, undefined); + lifecycle.baseUrl.mockReset(); + lifecycle.baseUrl.mockImplementation( + (overrides) => MANAGED_BASE_URL_BY_API_KEY.get(overrides?.loadApiKey?.()) ?? null, + ); +}); + +afterEach(() => vi.unstubAllEnvs()); + +describe("managed dual-Station vLLM authentication", () => { + it("keeps an explicit sandbox host override ahead of managed endpoint recovery", () => { + const loadApiKeyImpl = vi.fn(() => API_KEY); + expect(getLocalProviderBaseUrl("vllm-local", { hostUrl: "http://explicit-host" })).toBe( + "http://explicit-host:8000/v1", + ); + expect( + getManagedDualStationVllmProviderBinding({ + hostUrl: "http://explicit-host", + loadApiKeyImpl, + }), + ).toBeNull(); + expect(loadApiKeyImpl).not.toHaveBeenCalled(); + }); + + it("does not load a stale or unsafe key without a recovered managed endpoint", () => { + const loadApiKeyImpl = vi.fn(() => { + throw new Error(`unsafe ${API_KEY}`); + }); + lifecycle.baseUrl.mockReturnValue(null); + + expect(getManagedDualStationVllmProviderBinding({ loadApiKeyImpl })).toBeNull(); + expect(loadApiKeyImpl).not.toHaveBeenCalled(); + }); + + it("keeps legacy health unauthenticated even when stale key loading would fail", () => { + const loadVllmApiKeyImpl = vi.fn(() => { + throw new Error(`unsafe ${API_KEY}`); + }); + const runCurlProbeImpl = vi.fn(() => ({ + ok: true, + httpStatus: 200, + curlStatus: 0, + body: '{"data":[{"id":"served/model"}]}', + stderr: "", + message: "HTTP 200", + })); + lifecycle.baseUrl.mockReturnValue(null); + + const result = probeLocalProviderHealth("vllm-local", { + model: "served/model", + loadVllmApiKeyImpl, + runCurlProbeImpl, + }); + + expect(result?.ok).toBe(true); + expect(loadVllmApiKeyImpl).not.toHaveBeenCalled(); + expect(runCurlProbeImpl).toHaveBeenCalledWith([ + "-sS", + "--connect-timeout", + "3", + "--max-time", + "5", + "http://127.0.0.1:8000/v1/models", + ]); + }); + + it("returns one atomic provider endpoint and credential binding", () => { + expect(getManagedDualStationVllmProviderBinding({ loadApiKeyImpl: () => API_KEY })).toEqual({ + baseUrl: `${BASE_URL}/v1`, + apiKey: API_KEY, + }); + expect(getManagedDualStationVllmProviderState({ loadApiKeyImpl: () => API_KEY })).toEqual({ + kind: "ready", + baseUrl: `${BASE_URL}/v1`, + apiKey: API_KEY, + }); + }); + + it("uses /health only for unauthenticated availability checks", () => { + expect(getLocalProviderHealthEndpoint("vllm-local")).toBe(`${BASE_URL}/v1/models`); + expect(getLocalProviderHealthCheck("vllm-local")).toEqual([ + "curl", + "-sf", + "--connect-timeout", + "3", + "--max-time", + "5", + "--noproxy", + "*", + "--write-out", + "%{http_code}", + `${BASE_URL}/health`, + ]); + expect(getLocalProviderContainerReachabilityCheck("vllm-local")).toEqual([ + "docker", + "--context", + "default", + "run", + "--rm", + "--add-host", + "host.openshell.internal:host-gateway", + CONTAINER_REACHABILITY_IMAGE, + "--connect-timeout", + "5", + "--max-time", + "10", + "--noproxy", + "*", + "-sf", + "-w", + "%{http_code}", + `${BASE_URL}/health`, + ]); + }); + + it("pins reachability and diagnostics to the local daemon despite a persisted remote context", () => { + const capture = vi.fn((argv: readonly string[]) => (argv[0] === "curl" ? "200" : "")); + + const result = validateLocalProvider("vllm-local", capture, () => undefined); + const dockerCommands = capture.mock.calls + .map(([argv]) => argv) + .filter((argv) => argv[0] === "docker"); + + expect(result.ok).toBe(false); + expect(dockerCommands).toHaveLength(5); + expect( + dockerCommands.every((argv) => argv.slice(0, 3).join(" ") === "docker --context default"), + ).toBe(true); + }); + + it("passes bearer auth through a private curl config and cleans it up", () => { + let configPath = ""; + const result = probeVllmModels(`${BASE_URL}/v1`, API_KEY, { + runCurlProbeImpl: (argv, options) => { + expect(argv).not.toContain(API_KEY); + expect(argv.at(-1)).toBe(`${BASE_URL}/v1/models`); + const configIndex = argv.indexOf("--config"); + configPath = argv[configIndex + 1] ?? ""; + expect(options?.trustedConfigFiles).toEqual([configPath]); + expect(options?.pinnedAddresses).toEqual([]); + expect(fs.readFileSync(configPath, "utf8")).toContain(`Authorization: Bearer ${API_KEY}`); + return { + ok: true, + httpStatus: 200, + curlStatus: 0, + body: '{"data":[{"id":"served/model"}]}', + stderr: "", + message: "HTTP 200", + }; + }, + }); + + expect(result.ok).toBe(true); + expect(configPath).not.toBe(""); + expect(fs.existsSync(configPath)).toBe(false); + }); + + it("keeps authenticated model inventory authoritative for configured-model health", () => { + const runCurlProbeImpl = vi.fn((argv: string[]) => ({ + ok: true, + httpStatus: 200, + curlStatus: 0, + body: '{"data":[{"id":"different/model"}]}', + stderr: "", + message: "HTTP 200", + })); + const result = probeLocalProviderHealth("vllm-local", { + model: "required/model", + loadVllmApiKeyImpl: () => API_KEY, + runCurlProbeImpl, + }); + + expect(runCurlProbeImpl).toHaveBeenCalledOnce(); + const probeArgv = runCurlProbeImpl.mock.calls[0]?.[0] ?? []; + expect(probeArgv).not.toContain(API_KEY); + expect(result?.ok).toBe(false); + expect(result?.failureLabel).toBe("unhealthy"); + expect(result?.detail).toContain("required/model"); + expect(result?.detail).toContain("different/model"); + }); + + it("fails closed through production lifecycle recovery when the managed key is absent", () => { + const runCurlProbeImpl = vi.fn(); + const result = probeLocalProviderHealth("vllm-local", { + getManagedVllmBaseUrlImpl: productionManagedBaseUrlResolver(), + model: "required/model", + loadVllmApiKeyImpl: () => null, + runCurlProbeImpl, + }); + + expect(runCurlProbeImpl).not.toHaveBeenCalled(); + expect(result?.ok).toBe(false); + expect(result?.failureLabel).toBe("unauthorized"); + expect(result?.detail).not.toContain(API_KEY); + }); + + it("fails closed through production lifecycle recovery when managed key state is unsafe", () => { + const runCurlProbeImpl = vi.fn(); + const result = probeLocalProviderHealth("vllm-local", { + getManagedVllmBaseUrlImpl: productionManagedBaseUrlResolver(), + model: "required/model", + loadVllmApiKeyImpl: () => { + throw new Error(`unsafe ${API_KEY}`); + }, + runCurlProbeImpl, + }); + + expect(runCurlProbeImpl).not.toHaveBeenCalled(); + expect(result?.ok).toBe(false); + expect(result?.failureLabel).toBe("unhealthy"); + expect(result?.detail).not.toContain(API_KEY); + }); + + it("returns invalid-auth when the private key does not match the managed lifecycle", () => { + expect( + getManagedDualStationVllmProviderState({ + getManagedBaseUrlImpl: productionManagedBaseUrlResolver(OTHER_API_KEY), + loadApiKeyImpl: () => API_KEY, + }), + ).toEqual({ kind: "invalid-auth", reason: "mismatched" }); + }); + + it("fails closed when an owned managed head has invalid auth fingerprint metadata", () => { + const runCurlProbeImpl = vi.fn(); + const result = probeLocalProviderHealth("vllm-local", { + getManagedVllmBaseUrlImpl: productionManagedBaseUrlResolver(API_KEY, ""), + loadVllmApiKeyImpl: () => API_KEY, + runCurlProbeImpl, + }); + + expect(runCurlProbeImpl).not.toHaveBeenCalled(); + expect(result).toMatchObject({ + ok: false, + endpoint: "managed dual-Station vLLM", + failureLabel: "unhealthy", + }); + }); +}); diff --git a/src/lib/inference/local.test.ts b/src/lib/inference/local.test.ts index 8cb4b6e471..80fd021069 100644 --- a/src/lib/inference/local.test.ts +++ b/src/lib/inference/local.test.ts @@ -31,6 +31,7 @@ import { getOllamaModelOptions, getOllamaProbeCommand, getOllamaWarmupCommand, + isLocalProviderProbeOutputHealthy, isOllamaRunnerCrash, LOCAL_INFERENCE_SANDBOX_HOST_URL_ENV, parseOllamaList, @@ -175,6 +176,16 @@ describe("local inference helpers", () => { ]); }); + it("requires HTTP 200 for managed health output and rejects curl connection status 000", () => { + expect(isLocalProviderProbeOutputHealthy("http://10.40.0.1:8000/health", "200")).toBe(true); + expect(isLocalProviderProbeOutputHealthy("http://10.40.0.1:8000/health", "204")).toBe(false); + expect(isLocalProviderProbeOutputHealthy("http://10.40.0.1:8000/health", "000")).toBe(false); + expect(isLocalProviderProbeOutputHealthy("http://127.0.0.1:8000/v1/models", "000")).toBe(false); + expect( + isLocalProviderProbeOutputHealthy("http://127.0.0.1:8000/v1/models", '{"data":[]}'), + ).toBe(true); + }); + it("validates a reachable local provider", () => { let callCount = 0; const mockCapture = () => { diff --git a/src/lib/inference/local.ts b/src/lib/inference/local.ts index e2bcb0eab8..c9c29d970b 100644 --- a/src/lib/inference/local.ts +++ b/src/lib/inference/local.ts @@ -45,7 +45,9 @@ import { resetOllamaRuntimeContextWindowAutoState, resolveOllamaRuntimeContextWindow as resolveOllamaRuntimeContextWindowWithHost, } from "./ollama-runtime-context"; +import { loadDualStationVllmApiKey } from "./vllm-api-key"; import { applyVllmRuntimeContextWindow as applyVllmRuntimeContextWindowFromModels } from "./vllm-runtime-context"; +import { getDualStationManagedVllmBaseUrl } from "./vllm-station-cluster-lifecycle"; export type { OllamaRuntimeModelStatus } from "./ollama-runtime-context"; @@ -250,6 +252,10 @@ export interface LocalProviderHealthProbeOptions { * state root (written by inference/ollama/proxy.ts during onboard). */ loadOllamaProxyTokenImpl?: () => string | null; + /** Reads the managed dual-Station vLLM key. Injectable so tests stay deterministic. */ + loadVllmApiKeyImpl?: () => string | null; + /** Recovers an owned managed endpoint while validating the injected key in one lifecycle read. */ + getManagedVllmBaseUrlImpl?: ManagedDualStationVllmBaseUrlResolver; } function defaultLoadOllamaProxyToken(): string | null { @@ -272,6 +278,52 @@ function runLocalCurlProbe(argv: string[], opts: CurlProbeOptions = {}): CurlPro return runCurlProbe(argv, { ...opts, env: buildSubprocessEnv(), replaceEnv: true }); } +export interface VllmModelsProbeOptions { + runCurlProbeImpl?: (argv: string[], opts?: CurlProbeOptions) => CurlProbeResult; +} + +/** Query vLLM's authoritative model inventory without exposing its bearer in process argv. */ +export function probeVllmModels( + baseUrl: string, + apiKey: string, + options: VllmModelsProbeOptions = {}, +): CurlProbeResult { + const runCurlProbeImpl = options.runCurlProbeImpl ?? runLocalCurlProbe; + let authConfig: ReturnType | undefined; + try { + authConfig = createBearerAuthConfig(apiKey, { prefix: "nemoclaw-vllm-auth" }); + return runCurlProbeImpl( + [ + "-sS", + "--connect-timeout", + "3", + "--max-time", + "5", + ...authConfig.args, + `${baseUrl.replace(/\/+$/, "")}/models`, + ], + { + trustedConfigFiles: authConfig.trustedConfigFiles, + // Managed dual-Station endpoints are recovered from owned container + // labels and use a direct-attached RFC1918 rail. Never delegate that + // request through an ambient HTTP proxy. + pinnedAddresses: [], + }, + ); + } catch { + return { + ok: false, + httpStatus: 0, + curlStatus: 1, + body: "", + stderr: "", + message: "Could not prepare the authenticated vLLM model probe.", + }; + } finally { + authConfig?.cleanup(); + } +} + // A 200 response on `/api/tags` alone is not enough to call Ollama healthy — // a captive HTTP_PROXY, a stale listener, or a stub on the loopback port can // all answer with arbitrary 2xx bodies that look healthy at the curl-status @@ -349,22 +401,116 @@ function normalizeLocalInferenceHostUrl(raw: string | null | undefined): string return null; } -function getLocalInferenceSandboxHostUrl(): string { +function configuredLocalInferenceHostUrl(hostUrl?: string | null): string | null { return ( - normalizeLocalInferenceHostUrl(process.env[LOCAL_INFERENCE_SANDBOX_HOST_URL_ENV]) || - HOST_GATEWAY_URL + normalizeLocalInferenceHostUrl(hostUrl) || + normalizeLocalInferenceHostUrl(process.env[LOCAL_INFERENCE_SANDBOX_HOST_URL_ENV]) ); } +function recoveredManagedDualStationVllmBaseUrl(): string | null { + return configuredLocalInferenceHostUrl() ? null : getDualStationManagedVllmBaseUrl(); +} + +export interface ManagedDualStationVllmProviderBinding { + baseUrl: string; + apiKey: string; +} + +type ManagedDualStationVllmBaseUrlResolver = (overrides?: { + loadApiKey?: () => string | null; + onManagedHeadObserved?: () => void; +}) => string | null; + +export type ManagedDualStationVllmProviderState = + | { kind: "absent" } + | { kind: "invalid-auth"; reason: "missing" | "unsafe" | "mismatched" } + | ({ kind: "ready" } & ManagedDualStationVllmProviderBinding); + +export interface ManagedDualStationVllmProviderBindingOptions { + hostUrl?: string | null; + getManagedBaseUrlImpl?: ManagedDualStationVllmBaseUrlResolver; + loadApiKeyImpl?: () => string | null; +} + +/** Recover endpoint and credential as one lifecycle-validated state. */ +export function getManagedDualStationVllmProviderState( + options: ManagedDualStationVllmProviderBindingOptions = {}, +): ManagedDualStationVllmProviderState { + const configuredHostUrl = configuredLocalInferenceHostUrl(options.hostUrl); + if (configuredHostUrl) return { kind: "absent" }; + + const loadApiKey = options.loadApiKeyImpl ?? loadDualStationVllmApiKey; + let keyRead = false; + let managedHeadObserved = false; + let apiKey: string | null = null; + let authFailure: "missing" | "unsafe" | null = null; + let managedBaseUrl: string | null; + try { + managedBaseUrl = (options.getManagedBaseUrlImpl ?? getDualStationManagedVllmBaseUrl)({ + onManagedHeadObserved: () => { + managedHeadObserved = true; + }, + loadApiKey: () => { + keyRead = true; + try { + apiKey = loadApiKey(); + if (!apiKey) authFailure = "missing"; + return apiKey; + } catch { + authFailure = "unsafe"; + return null; + } + }, + }); + } catch (error) { + // A malformed key can make lifecycle fingerprint validation throw. Once + // key recovery began, treat every such failure as unsafe authentication; + // unrelated endpoint-inspection failures retain their existing behavior. + if (keyRead) return { kind: "invalid-auth", reason: "unsafe" }; + throw error; + } + + // Production recovery reports a structurally owned managed head before it + // validates fingerprint metadata. Test resolvers may still signal ownership + // by invoking the key reader, so only the absence of both signals permits the + // legacy single-host path. + if (!managedHeadObserved && !keyRead) return { kind: "absent" }; + if (!managedBaseUrl || !apiKey) { + return { kind: "invalid-auth", reason: authFailure ?? "mismatched" }; + } + return { kind: "ready", baseUrl: `${managedBaseUrl}/v1`, apiKey }; +} + +/** Compatibility binding for onboarding and context-window callers. */ +export function getManagedDualStationVllmProviderBinding( + options: ManagedDualStationVllmProviderBindingOptions = {}, +): ManagedDualStationVllmProviderBinding | null { + const state = getManagedDualStationVllmProviderState(options); + if (state.kind === "absent") return null; + if (state.kind === "invalid-auth") { + if (state.reason !== "missing") { + throw new Error("Managed dual-Station vLLM authentication is unsafe or mismatched."); + } + throw new Error("Managed dual-Station vLLM authentication is missing."); + } + return { baseUrl: state.baseUrl, apiKey: state.apiKey }; +} + export function getLocalProviderBaseUrl( provider: string, options: { hostUrl?: string | null } = {}, ): string | null { - const hostUrl = - normalizeLocalInferenceHostUrl(options.hostUrl) || getLocalInferenceSandboxHostUrl(); + const configuredHostUrl = configuredLocalInferenceHostUrl(options.hostUrl); + const hostUrl = configuredHostUrl || HOST_GATEWAY_URL; switch (provider) { - case "vllm-local": + case "vllm-local": { + if (!configuredHostUrl) { + const dualStationBaseUrl = recoveredManagedDualStationVllmBaseUrl(); + if (dualStationBaseUrl) return `${dualStationBaseUrl}/v1`; + } return `${hostUrl}:${VLLM_PORT}/v1`; + } case "ollama-local": // Containers reach Ollama through the auth proxy, not directly. return `${hostUrl}:${getOllamaContainerPort()}/v1`; @@ -375,8 +521,10 @@ export function getLocalProviderBaseUrl( export function getLocalProviderValidationBaseUrl(provider: string): string | null { switch (provider) { - case "vllm-local": - return `http://127.0.0.1:${VLLM_PORT}/v1`; + case "vllm-local": { + const dualStationBaseUrl = recoveredManagedDualStationVllmBaseUrl(); + return dualStationBaseUrl ? `${dualStationBaseUrl}/v1` : `http://127.0.0.1:${VLLM_PORT}/v1`; + } case "ollama-local": return `http://${getResolvedOllamaHost()}:${OLLAMA_PORT}/v1`; default: @@ -386,8 +534,12 @@ export function getLocalProviderValidationBaseUrl(provider: string): string | nu export function getLocalProviderHealthEndpoint(provider: string): string | null { switch (provider) { - case "vllm-local": - return `http://127.0.0.1:${VLLM_PORT}/v1/models`; + case "vllm-local": { + const dualStationBaseUrl = recoveredManagedDualStationVllmBaseUrl(); + return dualStationBaseUrl + ? `${dualStationBaseUrl}/v1/models` + : `http://127.0.0.1:${VLLM_PORT}/v1/models`; + } case "ollama-local": return `http://${getResolvedOllamaHost()}:${OLLAMA_PORT}/api/tags`; default: @@ -395,8 +547,38 @@ export function getLocalProviderHealthEndpoint(provider: string): string | null } } +/** Lightweight endpoint used only to prove that the local service is reachable. */ +export function getLocalProviderAvailabilityEndpoint(provider: string): string | null { + if (provider === "vllm-local") { + const dualStationBaseUrl = recoveredManagedDualStationVllmBaseUrl(); + if (dualStationBaseUrl) return `${dualStationBaseUrl}/health`; + } + return getLocalProviderHealthEndpoint(provider); +} + +export function isLocalProviderProbeOutputHealthy(endpoint: string, output: string): boolean { + const normalized = output.trim(); + if (!normalized || normalized === "000") return false; + return endpoint.endsWith("/health") ? normalized === "200" : true; +} + export function getLocalProviderHealthCheck(provider: string): string[] | null { - const endpoint = getLocalProviderHealthEndpoint(provider); + const endpoint = getLocalProviderAvailabilityEndpoint(provider); + if (provider === "vllm-local" && endpoint?.endsWith("/health")) { + return [ + "curl", + "-sf", + "--connect-timeout", + "3", + "--max-time", + "5", + "--noproxy", + "*", + "--write-out", + "%{http_code}", + endpoint, + ]; + } return endpoint ? ["curl", ...buildValidatedCurlCommandArgs(["-sf", endpoint])] : null; } @@ -417,7 +599,10 @@ export function isLocalProviderHostHealthy( const command = getLocalProviderHealthCheck(provider); if (!command) return false; const capture = runCaptureImpl ?? runCapture; - return Boolean(capture(command, { ignoreError: true })); + return isLocalProviderProbeOutputHealthy( + command.at(-1) ?? "", + capture(command, { ignoreError: true }), + ); } export function getLocalProviderLabel(provider: string): string | null { @@ -554,14 +739,56 @@ export function probeLocalProviderHealth( provider: string, options: LocalProviderHealthProbeOptions = {}, ): LocalProviderHealthStatus | null { - const endpoint = getLocalProviderHealthEndpoint(provider); const providerLabel = getLocalProviderLabel(provider); - if (!endpoint || !providerLabel) { - return null; + if (!providerLabel) return null; + + let managedState: ManagedDualStationVllmProviderState = { kind: "absent" }; + if (provider === "vllm-local") { + try { + managedState = getManagedDualStationVllmProviderState({ + getManagedBaseUrlImpl: options.getManagedVllmBaseUrlImpl, + loadApiKeyImpl: options.loadVllmApiKeyImpl, + }); + } catch { + return { + ok: false, + providerLabel, + endpoint: "managed dual-Station vLLM", + failureLabel: "unhealthy", + probeLabel: "vllm backend", + detail: + "Local vLLM authentication state could not be inspected safely. Re-run `nemoclaw onboard` to repair the managed dual-Station provider.", + }; + } + } + if (managedState.kind === "invalid-auth") { + const missingAuth = managedState.reason === "missing"; + return { + ok: false, + providerLabel, + endpoint: "managed dual-Station vLLM", + failureLabel: missingAuth ? "unauthorized" : "unhealthy", + probeLabel: "vllm backend", + detail: missingAuth + ? "Local vLLM requires its managed bearer credential, but no private key is available. Re-run `nemoclaw onboard` to repair the dual-Station provider." + : "Local vLLM authentication state is unsafe or does not match the managed service. Re-run `nemoclaw onboard` to repair the managed dual-Station provider.", + }; } + const managedBinding = managedState.kind === "ready" ? managedState : null; + const endpoint = managedBinding + ? `${managedBinding.baseUrl}/models` + : provider === "vllm-local" + ? `http://127.0.0.1:${VLLM_PORT}/v1/models` + : getLocalProviderHealthEndpoint(provider); + if (!endpoint) return null; const runCurlProbeImpl = options.runCurlProbeImpl ?? runLocalCurlProbe; - const result = runCurlProbeImpl(["-sS", "--connect-timeout", "3", "--max-time", "5", endpoint]); + let result: CurlProbeResult; + if (managedBinding) { + result = probeVllmModels(managedBinding.baseUrl, managedBinding.apiKey, { runCurlProbeImpl }); + } else { + result = runCurlProbeImpl(["-sS", "--connect-timeout", "3", "--max-time", "5", endpoint]); + } // Per #3265 the status line is renamed `Inference ():` for local // providers so the upcoming `Inference (auth proxy):` subprobe lines render @@ -661,9 +888,10 @@ export function probeLocalProviderHealth( export function getLocalProviderContainerReachabilityCheck(provider: string): string[] | null { switch (provider) { - case "vllm-local": + case "vllm-local": { + const dualStationBaseUrl = recoveredManagedDualStationVllmBaseUrl(); return [ - "docker", + ...(dualStationBaseUrl ? ["docker", "--context", "default"] : ["docker"]), "run", "--rm", "--add-host", @@ -673,9 +901,14 @@ export function getLocalProviderContainerReachabilityCheck(provider: string): st "5", "--max-time", "10", + ...(dualStationBaseUrl ? ["--noproxy", "*"] : []), "-sf", - `http://host.openshell.internal:${VLLM_PORT}/v1/models`, + ...(dualStationBaseUrl ? ["-w", "%{http_code}"] : []), + dualStationBaseUrl + ? `${dualStationBaseUrl}/health` + : `http://host.openshell.internal:${VLLM_PORT}/v1/models`, ]; + } case "ollama-local": // Check the auth proxy port, not Ollama directly. The proxy listens // on 0.0.0.0 and is reachable from containers; Ollama is on 127.0.0.1. @@ -730,12 +963,12 @@ export function validateLocalProvider( } const output = capture(command, { ignoreError: true }); - if (!output) { + if (!isLocalProviderProbeOutputHealthy(command.at(-1) ?? "", output)) { switch (provider) { case "vllm-local": return { ok: false, - message: `Local vLLM was selected, but nothing is responding on http://127.0.0.1:${VLLM_PORT}.`, + message: `Local vLLM was selected, but nothing is responding on ${getLocalProviderHealthEndpoint(provider) ?? "the configured endpoint"}.`, }; case "ollama-local": return { @@ -755,7 +988,7 @@ export function validateLocalProvider( // Retry container reachability check with backoff for (let attempt = 1; attempt <= CONTAINER_CHECK_MAX_ATTEMPTS; attempt++) { const containerOutput = capture(containerCommand, { ignoreError: true }); - if (containerOutput) { + if (isLocalProviderProbeOutputHealthy(containerCommand.at(-1) ?? "", containerOutput)) { return { ok: true }; } if (attempt < CONTAINER_CHECK_MAX_ATTEMPTS) { @@ -770,7 +1003,7 @@ export function validateLocalProvider( case "vllm-local": return { ok: false, - message: `Local vLLM is responding on 127.0.0.1, but the Docker container reachability check failed for http://host.openshell.internal:${VLLM_PORT}. This may be a Docker networking issue — the sandbox uses a different network path and may still work.`, + message: `Local vLLM is responding on the host, but the Docker container reachability check failed for ${getContainerCheckUrl(provider)}. This may be a Docker networking issue — the sandbox uses a different network path and may still work.`, diagnostic, }; case "ollama-local": @@ -790,8 +1023,12 @@ export function validateLocalProvider( function getContainerCheckUrl(provider: string): string { switch (provider) { - case "vllm-local": - return `http://host.openshell.internal:${VLLM_PORT}/v1/models`; + case "vllm-local": { + const dualStationBaseUrl = recoveredManagedDualStationVllmBaseUrl(); + return dualStationBaseUrl + ? `${dualStationBaseUrl}/health` + : `http://host.openshell.internal:${VLLM_PORT}/v1/models`; + } case "ollama-local": return `http://host.openshell.internal:${getOllamaContainerPort()}/api/tags`; default: @@ -801,11 +1038,15 @@ function getContainerCheckUrl(provider: string): string { function collectContainerDiagnostic(provider: string, capture: RunCaptureFn): string { const url = getContainerCheckUrl(provider); + const dockerCommand = + provider === "vllm-local" && recoveredManagedDualStationVllmBaseUrl() + ? ["docker", "--context", "default"] + : ["docker"]; try { // Get HTTP status code const httpStatus = capture( [ - "docker", + ...dockerCommand, "run", "--rm", "--add-host", @@ -828,7 +1069,7 @@ function collectContainerDiagnostic(provider: string, capture: RunCaptureFn): st // Get /etc/hosts to see host-gateway resolution const hostsOutput = capture( [ - "docker", + ...dockerCommand, "run", "--rm", "--add-host", diff --git a/src/lib/inference/nim.ts b/src/lib/inference/nim.ts index 3a29757950..852bd500da 100644 --- a/src/lib/inference/nim.ts +++ b/src/lib/inference/nim.ts @@ -23,6 +23,7 @@ import { VLLM_PORT } from "../core/ports"; import { sleepSeconds } from "../core/wait"; import { runCapture } from "../runner"; import { isSafeModelId } from "../validation"; +import { isDgxStationGb300Product } from "./dgx-station-identity"; import { type Arm64WslDockerDesktopGpuProver, isDenylistedNvidiaGpuName, @@ -297,13 +298,7 @@ function detectTegraHostGpu(): { name: string; platform: NvidiaPlatform } | null export function detectNvidiaPlatform(): NvidiaPlatform { const model = readPlatformModel(); if (/DGX[_\s-]+Spark/i.test(model)) return "spark"; - if ( - /(? { + vi.doUnmock("node:fs"); + vi.resetModules(); + for (const directory of temporaryDirectories.splice(0)) { + fs.rmSync(directory, { force: true, recursive: true }); + } +}); + +describe("dual-Station vLLM API key persistence", () => { + it("creates one private 256-bit key and reuses it", () => { + const stateDir = path.join(temporaryDirectory(), "state"); + const generated = ensureDualStationVllmApiKey({ + stateDir, + randomBytes: () => Buffer.alloc(32, 0xab), + }); + + expect(generated).toBe("ab".repeat(32)); + expect(loadDualStationVllmApiKey({ stateDir })).toBe(generated); + expect(ensureDualStationVllmApiKey({ stateDir, randomBytes: () => Buffer.alloc(32, 1) })).toBe( + generated, + ); + expect(fs.statSync(stateDir).mode & 0o777).toBe(0o700); + expect(fs.statSync(dualStationVllmApiKeyPath(stateDir)).mode & 0o777).toBe(0o600); + }); + + it("returns null when no key has been provisioned", () => { + expect(loadDualStationVllmApiKey({ stateDir: temporaryDirectory() })).toBeNull(); + }); + + it("rejects malformed or overly permissive key files", () => { + const stateDir = temporaryDirectory(); + const filePath = dualStationVllmApiKeyPath(stateDir); + fs.writeFileSync(filePath, `${"ab".repeat(32)}\n`, { mode: 0o644 }); + expect(() => loadDualStationVllmApiKey({ stateDir })).toThrow("group or others"); + + fs.chmodSync(filePath, 0o600); + fs.writeFileSync(filePath, "not-a-key\n"); + expect(() => loadDualStationVllmApiKey({ stateDir })).toThrow("malformed"); + }); + + it("refuses to follow a symbolic-link key path", () => { + const root = temporaryDirectory(); + const stateDir = path.join(root, "state"); + fs.mkdirSync(stateDir, { mode: 0o700 }); + const target = path.join(root, "target"); + fs.writeFileSync(target, `${"cd".repeat(32)}\n`, { mode: 0o600 }); + fs.symlinkSync(target, dualStationVllmApiKeyPath(stateDir)); + + expect(() => loadDualStationVllmApiKey({ stateDir })).toThrow("symbolic link"); + expect(() => ensureDualStationVllmApiKey({ stateDir })).toThrow("symbolic link"); + }); + + it("fails closed when secure no-follow opens are unavailable", async () => { + vi.resetModules(); + vi.doMock("node:fs", async (importOriginal) => { + const actual = await importOriginal(); + const constants = { ...actual.constants, O_NOFOLLOW: undefined }; + const actualDefault = (actual as { default?: typeof fs }).default ?? actual; + return { ...actual, constants, default: { ...actualDefault, constants } }; + }); + const unavailable = await import("./vllm-api-key"); + + expect(() => + unavailable.ensureDualStationVllmApiKey({ + stateDir: path.join(temporaryDirectory(), "state"), + }), + ).toThrow("Secure no-follow file opens are unavailable"); + }); +}); diff --git a/src/lib/inference/vllm-api-key.ts b/src/lib/inference/vllm-api-key.ts new file mode 100644 index 0000000000..0eda887c52 --- /dev/null +++ b/src/lib/inference/vllm-api-key.ts @@ -0,0 +1,129 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import crypto from "node:crypto"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; + +import { DEFAULT_GATEWAY_PORT } from "../core/ports"; +import { nemoclawStateRoot } from "../state/state-root"; +import { ensureLocalAdapterStateDir } from "./local-adapter-lifecycle"; + +export const DUAL_STATION_VLLM_API_KEY_FILE = "dual-station-vllm-api-key"; +export const DUAL_STATION_VLLM_API_KEY_PATTERN = /^[a-f0-9]{64}$/; + +export interface DualStationVllmApiKeyOptions { + stateDir?: string; + randomBytes?: (size: number) => Buffer; +} + +function defaultStateDir(): string { + // The managed vLLM service is host-global rather than gateway-scoped. Every + // gateway therefore reads the same key even when NEMOCLAW_GATEWAY_PORT is + // changed for a second sandbox. + return nemoclawStateRoot(os.homedir(), DEFAULT_GATEWAY_PORT); +} + +export function dualStationVllmApiKeyPath(stateDir = defaultStateDir()): string { + return path.join(stateDir, DUAL_STATION_VLLM_API_KEY_FILE); +} + +function assertPrivateRegularFile(stat: fs.Stats, filePath: string): void { + if (!stat.isFile()) { + throw new Error(`Refusing to read dual-Station vLLM API key from non-file path: ${filePath}`); + } + if ((stat.mode & 0o077) !== 0) { + throw new Error( + `Dual-Station vLLM API key file must not be accessible by group or others: ${filePath}`, + ); + } + if (typeof process.getuid === "function" && stat.uid !== process.getuid()) { + throw new Error(`Dual-Station vLLM API key file is not owned by the current user: ${filePath}`); + } +} + +/** Load the host-global managed endpoint key, failing closed on unsafe state. */ +export function loadDualStationVllmApiKey( + options: Pick = {}, +): string | null { + const filePath = dualStationVllmApiKeyPath(options.stateDir ?? defaultStateDir()); + const noFollow = fs.constants.O_NOFOLLOW; + if (typeof noFollow !== "number") { + throw new Error("Secure no-follow file opens are unavailable on this platform"); + } + const nonBlock = fs.constants.O_NONBLOCK ?? 0; + let fd: number | undefined; + try { + try { + fd = fs.openSync(filePath, fs.constants.O_RDONLY | noFollow | nonBlock); + } catch (error) { + const code = (error as NodeJS.ErrnoException).code; + if (code === "ENOENT") return null; + if (code === "ELOOP") { + throw new Error( + `Refusing to read dual-Station vLLM API key through a symbolic link: ${filePath}`, + ); + } + throw error; + } + const opened = fs.fstatSync(fd); + assertPrivateRegularFile(opened, filePath); + if (opened.size < 64 || opened.size > 65) { + throw new Error(`Dual-Station vLLM API key file is malformed: ${filePath}`); + } + const value = fs.readFileSync(fd, "utf8").trim(); + if (!DUAL_STATION_VLLM_API_KEY_PATTERN.test(value)) { + throw new Error(`Dual-Station vLLM API key file is malformed: ${filePath}`); + } + return value; + } finally { + if (fd !== undefined) fs.closeSync(fd); + } +} + +/** Create the managed endpoint key once, or reuse the existing private key. */ +export function ensureDualStationVllmApiKey(options: DualStationVllmApiKeyOptions = {}): string { + const stateDir = options.stateDir ?? defaultStateDir(); + ensureLocalAdapterStateDir(stateDir); + const existing = loadDualStationVllmApiKey({ stateDir }); + if (existing) return existing; + + const randomBytes = options.randomBytes ?? crypto.randomBytes; + const value = randomBytes(32).toString("hex"); + if (!DUAL_STATION_VLLM_API_KEY_PATTERN.test(value)) { + throw new Error("Could not generate a valid dual-Station vLLM API key"); + } + + const filePath = dualStationVllmApiKeyPath(stateDir); + const noFollow = fs.constants.O_NOFOLLOW; + if (typeof noFollow !== "number") { + throw new Error("Secure no-follow file opens are unavailable on this platform"); + } + let fd: number | undefined; + try { + fd = fs.openSync( + filePath, + fs.constants.O_WRONLY | fs.constants.O_CREAT | fs.constants.O_EXCL | noFollow, + 0o600, + ); + const opened = fs.fstatSync(fd); + assertPrivateRegularFile(opened, filePath); + fs.writeFileSync(fd, `${value}\n`, "utf8"); + fs.fsyncSync(fd); + } catch (error) { + if ((error as NodeJS.ErrnoException).code === "EEXIST") { + const raced = loadDualStationVllmApiKey({ stateDir }); + if (raced) return raced; + } + throw error; + } finally { + if (fd !== undefined) fs.closeSync(fd); + } + + const persisted = loadDualStationVllmApiKey({ stateDir }); + if (persisted !== value) { + throw new Error("Could not verify the persisted dual-Station vLLM API key"); + } + return value; +} diff --git a/src/lib/inference/vllm-contracts.test.ts b/src/lib/inference/vllm-contracts.test.ts new file mode 100644 index 0000000000..c63334e0e1 --- /dev/null +++ b/src/lib/inference/vllm-contracts.test.ts @@ -0,0 +1,97 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { beforeEach, describe, expect, it, vi } from "vitest"; + +const mocks = vi.hoisted(() => ({ + dockerPullWithProgressWatchdog: vi.fn(), +})); + +vi.mock("../adapters/docker", async (importOriginal) => ({ + ...(await importOriginal()), + dockerPullWithProgressWatchdog: mocks.dockerPullWithProgressWatchdog, +})); + +import { + assertVllmRegistryDigestRef, + detectVllmProfile, + pullImage, + resolveVllmServedModelId, + VLLM_IMAGES, +} from "./vllm"; +import { VLLM_MODELS } from "./vllm-models"; + +beforeEach(() => vi.clearAllMocks()); + +describe("vLLM served route identity", () => { + it("uses one safe served-model override and rejects ambiguous aliases (#6315)", () => { + expect(resolveVllmServedModelId("catalog/model", [])).toBe("catalog/model"); + expect(resolveVllmServedModelId("catalog/model", ["--served-model-name", "served/model"])).toBe( + "served/model", + ); + expect(() => + resolveVllmServedModelId("catalog/model", [ + "--served-model-name", + "served/one", + "served/two", + ]), + ).toThrow("exactly one safe model ID"); + }); +}); + +describe("managed vLLM image distribution boundary", () => { + const digest = `sha256:${"a".repeat(64)}`; + + it("accepts repository-qualified immutable registry digests", () => { + expect(() => assertVllmRegistryDigestRef(`vllm/vllm-openai@${digest}`)).not.toThrow(); + expect(() => + assertVllmRegistryDigestRef(`registry.example.test:5000/team/runtime@${digest}`), + ).not.toThrow(); + }); + + it.each([ + `sha256:${"a".repeat(64)}`, + "vllm/vllm-openai:latest", + `ubuntu@${digest}`, + `vllm/vllm-openai@sha256:${"A".repeat(64)}`, + `vllm/vllm-openai@${digest}suffix`, + ` vllm/vllm-openai@${digest}`, + `vllm/vllm-openai@${digest} `, + ])("rejects an unpullable or mutable product image reference %j", (image) => { + expect(() => assertVllmRegistryDigestRef(image)).toThrow( + /pullable immutable registry reference/, + ); + }); + + it("keeps every shipped managed-vLLM image on a registry digest", () => { + const platformRefs = Object.values(VLLM_IMAGES).flatMap((imageSet) => + Object.values(imageSet) + .map((value) => + typeof value === "object" && value !== null && "ref" in value ? String(value.ref) : null, + ) + .filter((ref): ref is string => ref !== null), + ); + const runtimeRefs = VLLM_MODELS.map((model) => model.runtime?.image).filter( + (ref): ref is string => typeof ref === "string", + ); + const refs = new Set([...platformRefs, ...runtimeRefs]); + + expect(refs.size).toBeGreaterThan(0); + for (const ref of refs) { + expect(() => assertVllmRegistryDigestRef(ref), ref).not.toThrow(); + } + }); + + it("refuses a local image ID before invoking Docker pull", async () => { + const profile = { + ...detectVllmProfile({ platform: "station", type: "nvidia" })!, + image: `sha256:${"a".repeat(64)}`, + }; + + await expect(pullImage(profile)).resolves.toEqual({ + ok: false, + reason: expect.stringContaining("Local image IDs"), + }); + expect(mocks.dockerPullWithProgressWatchdog).not.toHaveBeenCalled(); + }); +}); diff --git a/src/lib/inference/vllm-docker-env.test.ts b/src/lib/inference/vllm-docker-env.test.ts index af06687ff3..20497497a3 100644 --- a/src/lib/inference/vllm-docker-env.test.ts +++ b/src/lib/inference/vllm-docker-env.test.ts @@ -1,11 +1,34 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -import { afterEach, describe, expect, it, vi } from "vitest"; -import { buildVllmDockerEnv } from "./vllm-docker-env"; +import { spawnSync } from "node:child_process"; +import fs from "node:fs"; + +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { + buildLocalDualStationDockerEnv, + buildRemoteVllmDockerEnv, + buildVllmDockerEnv, +} from "./vllm-docker-env"; +import { + clearDualStationSshBinding, + stationKnownHostsDigest, + writeDualStationSshBinding, +} from "./vllm-station-ssh-binding"; +import { + createDualStationSshBindingFixture, + type DualStationSshBindingFixture, +} from "./vllm-station-ssh-binding.test-support"; + +let sshFixture: DualStationSshBindingFixture; + +beforeEach(() => { + sshFixture = createDualStationSshBindingFixture("station@dgx-peer.example.test"); +}); afterEach(() => { vi.unstubAllEnvs(); + sshFixture.cleanup(); }); describe("managed vLLM Docker client environment", () => { @@ -38,4 +61,100 @@ describe("managed vLLM Docker client environment", () => { expect(env.DOCKER_CONTEXT).toBe("requested-context"); expect(env.DOCKER_HOST).toBeUndefined(); }); + + it("pins a canonical SSH daemon and strips incompatible ambient Docker selectors", () => { + vi.stubEnv("DOCKER_API_VERSION", "1.48"); + vi.stubEnv("DOCKER_CERT_PATH", "/tmp/ambient-docker-certs"); + vi.stubEnv("DOCKER_CONFIG", "/tmp/nemoclaw-docker-config"); + vi.stubEnv("DOCKER_CONTEXT", "ambient-context"); + vi.stubEnv("DOCKER_HOST", "tcp://ambient.example.test:2376"); + vi.stubEnv("DOCKER_TLS", "1"); + vi.stubEnv("DOCKER_TLS_VERIFY", "1"); + vi.stubEnv("SSH_AUTH_SOCK", "/tmp/ssh-agent.sock"); + vi.stubEnv("OPENSHELL_GATEWAY_AUTH_TOKEN", "must-not-cross-ssh"); + vi.stubEnv("UNRELATED_SECRET", "do-not-forward"); + + const env = buildRemoteVllmDockerEnv(sshFixture.binding); + + expect(env).toEqual( + expect.objectContaining({ + DOCKER_HOST: "ssh://station@192.168.50.20", + SSH_AUTH_SOCK: "/tmp/ssh-agent.sock", + }), + ); + expect(env.PATH).toBe(sshFixture.binding.sshWrapperDirectory); + expect(env.DOCKER_API_VERSION).toBeUndefined(); + expect(env.DOCKER_CERT_PATH).toBeUndefined(); + expect(env.DOCKER_CONFIG).toBeUndefined(); + expect(env.DOCKER_CONTEXT).toBeUndefined(); + expect(env.DOCKER_TLS).toBeUndefined(); + expect(env.DOCKER_TLS_VERIFY).toBeUndefined(); + expect(env.UNRELATED_SECRET).toBeUndefined(); + expect(env.OPENSHELL_GATEWAY_AUTH_TOKEN).toBeUndefined(); + }); + + it("pins the dual-Station head to the physical host default Docker daemon", () => { + const env = buildLocalDualStationDockerEnv( + { SAFE_MARKER: "kept" }, + { + DOCKER_HOST: "ssh://wrong-daemon", + DOCKER_CONTEXT: "wrong-context", + DOCKER_CONFIG: "/tmp/wrong-config", + }, + ); + + expect(env.SAFE_MARKER).toBe("kept"); + expect(env.DOCKER_HOST).toBeUndefined(); + expect(env.DOCKER_CONTEXT).toBe("default"); + expect(env.DOCKER_CONFIG).toBeUndefined(); + }); + + it("rejects a changed qualified host-key pin before constructing a Docker environment", () => { + fs.appendFileSync(sshFixture.binding.knownHostsFile, "changed\n"); + + expect(() => buildRemoteVllmDockerEnv(sshFixture.binding)).toThrow( + "Station SSH known-hosts binding changed after qualification", + ); + }); + + it("keeps an existing environment pinned when a later qualification writes a new version", () => { + const first = sshFixture.binding; + const firstEnv = buildRemoteVllmDockerEnv(first); + const replacementHost = "192.168.50.21"; + const replacementLines = [ + `${replacementHost} ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIcmVwbGFjZW1lbnQ=`, + ]; + const second = writeDualStationSshBinding( + sshFixture.resumeStatePath, + { + ...sshFixture.identity, + resolvedHost: replacementHost, + lookupHost: replacementHost, + hostKeyDigest: stationKnownHostsDigest(`${replacementLines.join("\n")}\n`), + knownHostsLines: replacementLines, + }, + { dockerCliFile: sshFixture.dockerCliFile }, + ); + const secondEnv = buildRemoteVllmDockerEnv(second); + + expect(second.sshWrapperDirectory).not.toBe(first.sshWrapperDirectory); + expect(firstEnv.PATH).toBe(first.sshWrapperDirectory); + expect(secondEnv.PATH).toBe(second.sshWrapperDirectory); + expect(firstEnv.DOCKER_HOST).toBe("ssh://station@192.168.50.20"); + expect(secondEnv.DOCKER_HOST).toBe(`ssh://station@${replacementHost}`); + expect(buildRemoteVllmDockerEnv(first)).toEqual(firstEnv); + expect(fs.existsSync(first.dockerShimFile)).toBe(true); + expect(fs.existsSync(second.dockerShimFile)).toBe(true); + }); + + it("fails closed instead of falling through to ambient Docker after cleanup", () => { + const env = buildRemoteVllmDockerEnv(sshFixture.binding); + + clearDualStationSshBinding(sshFixture.resumeStatePath); + + const result = spawnSync("docker", ["version"], { env, encoding: "utf8" }); + expect(result.status).toBeNull(); + expect(result.error).toMatchObject({ code: "ENOENT" }); + expect(() => buildRemoteVllmDockerEnv(sshFixture.binding)).toThrow(); + }); }); diff --git a/src/lib/inference/vllm-docker-env.ts b/src/lib/inference/vllm-docker-env.ts index 986d2286e9..6dcf3f14f9 100644 --- a/src/lib/inference/vllm-docker-env.ts +++ b/src/lib/inference/vllm-docker-env.ts @@ -1,7 +1,14 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +import { isIP } from "node:net"; + import { buildSubprocessEnv } from "../subprocess-env"; +import { + assertDualStationSshBindingFiles, + type DualStationSshBinding, + dualStationDockerSshUri, +} from "./vllm-station-ssh-binding"; const DOCKER_CLIENT_ENV_NAMES = [ "DOCKER_API_VERSION", @@ -13,6 +20,75 @@ const DOCKER_CLIENT_ENV_NAMES = [ "DOCKER_TLS_VERIFY", ] as const; +const REMOTE_DOCKER_INCOMPATIBLE_ENV_NAMES = [ + "DOCKER_API_VERSION", + "DOCKER_CERT_PATH", + "DOCKER_CONFIG", + "DOCKER_CONTEXT", + "DOCKER_TLS", + "DOCKER_TLS_VERIFY", +] as const; +const SSH_TRANSPORT_ENV_NAMES = [ + "HOME", + "PATH", + "USER", + "LOGNAME", + "SHELL", + "SSH_AUTH_SOCK", + "LANG", + "LC_ALL", + "LC_CTYPE", + "TMPDIR", +] as const; +const CANONICAL_SSH_HOST_PATTERN = + /^[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?(?:\.[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?)*$/; +const CANONICAL_SSH_USERNAME_PATTERN = /^[A-Za-z_][A-Za-z0-9._-]*$/; + +function validateRemoteDockerSshUri(value: string): string { + const invalid = () => + new Error( + "Remote Docker host must be a canonical ssh://[user@]host[:port] URI without a password, path, query, or fragment", + ); + if (typeof value !== "string" || !value || value !== value.trim() || value.includes("\0")) { + throw invalid(); + } + + let parsed: URL; + try { + parsed = new URL(value); + } catch { + throw invalid(); + } + + if ( + parsed.protocol !== "ssh:" || + !parsed.hostname || + parsed.password || + parsed.pathname || + parsed.search || + parsed.hash || + (parsed.username && !CANONICAL_SSH_USERNAME_PATTERN.test(parsed.username)) + ) { + throw invalid(); + } + + const bracketedIpv6 = parsed.hostname.startsWith("[") && parsed.hostname.endsWith("]"); + const bareHostname = bracketedIpv6 ? parsed.hostname.slice(1, -1) : parsed.hostname; + const validHostname = bracketedIpv6 + ? bareHostname === bareHostname.toLowerCase() && isIP(bareHostname) === 6 + : isIP(bareHostname) === 4 || CANONICAL_SSH_HOST_PATTERN.test(bareHostname); + const port = parsed.port ? Number(parsed.port) : null; + if (!validHostname || (port !== null && (!Number.isInteger(port) || port < 1 || port > 65535))) { + throw invalid(); + } + + const canonical = `ssh://${parsed.username ? `${parsed.username}@` : ""}${parsed.hostname}${ + parsed.port ? `:${parsed.port}` : "" + }`; + if (value !== canonical) throw invalid(); + return canonical; +} + /** * Use one Docker client selection for every managed-vLLM subprocess while * retaining the repository's child-process environment sanitization. @@ -32,3 +108,45 @@ export function buildVllmDockerEnv( } return env; } + +/** Select the physical host's default Docker daemon, ignoring ambient client routing. */ +export function buildLocalDualStationDockerEnv( + extra: Record = {}, + source: NodeJS.ProcessEnv = process.env, +): Record { + const env = buildVllmDockerEnv(extra, source); + for (const name of DOCKER_CLIENT_ENV_NAMES) delete env[name]; + // Docker otherwise falls back to config.json's persisted currentContext, + // which may point at a remote daemon even with every selector env unset. + env.DOCKER_CONTEXT = "default"; + return env; +} + +/** Minimal environment shared by strict SSH probes and Docker's SSH helper. */ +export function buildVllmSshTransportEnv( + extra: Record = {}, + source: NodeJS.ProcessEnv = process.env, +): Record { + const env: Record = {}; + for (const name of SSH_TRANSPORT_ENV_NAMES) { + const value = source[name]; + if (value !== undefined) env[name] = value; + } + return { ...env, ...extra }; +} + +/** + * Select one explicitly configured Docker-over-SSH daemon without allowing an + * ambient context, client config, API pin, or TCP/TLS settings to influence it. + */ +export function buildRemoteVllmDockerEnv( + binding: DualStationSshBinding, + source: NodeJS.ProcessEnv = process.env, +): Record { + assertDualStationSshBindingFiles(binding); + const remoteHost = validateRemoteDockerSshUri(dualStationDockerSshUri(binding)); + const env = buildVllmSshTransportEnv({ DOCKER_HOST: remoteHost }, source); + for (const name of REMOTE_DOCKER_INCOMPATIBLE_ENV_NAMES) delete env[name]; + env.PATH = binding.sshWrapperDirectory; + return env; +} diff --git a/src/lib/inference/vllm-dual-station.test.ts b/src/lib/inference/vllm-dual-station.test.ts new file mode 100644 index 0000000000..fa042efe81 --- /dev/null +++ b/src/lib/inference/vllm-dual-station.test.ts @@ -0,0 +1,728 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { EventEmitter } from "node:events"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; + +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +const mocks = vi.hoisted(() => ({ + areContainersRunning: vi.fn(), + cleanup: vi.fn(), + commitLegacyMigration: vi.fn(), + dockerCapture: vi.fn(), + dockerForceRm: vi.fn(), + dockerImageInspectFormat: vi.fn(), + dockerPullWithProgressWatchdog: vi.fn(), + dockerRunDetached: vi.fn(), + dockerSpawn: vi.fn(), + dockerStop: vi.fn(), + ensureApiKey: vi.fn(), + findUnwritableModelCachePath: vi.fn(), + getManagedBaseUrl: vi.fn(), + getGpuIndicesByName: vi.fn(), + loadApiKey: vi.fn(), + measureDirectorySizeBytes: vi.fn(), + preflightGpuRuntime: vi.fn(), + preflightOwnership: vi.fn(), + probeCapability: vi.fn(), + probeDockerStorage: vi.fn(), + probeHostStorage: vi.fn(), + runCapture: vi.fn(), + runCurlProbe: vi.fn(), + rollbackLegacyMigration: vi.fn(), + startManaged: vi.fn(), + stageModelSnapshot: vi.fn(), + withLifecycle: vi.fn(), +})); + +vi.mock("../runner", async (importOriginal) => ({ + ...(await importOriginal()), + runCapture: mocks.runCapture, +})); + +vi.mock("../adapters/http/probe", () => ({ + runCurlProbe: mocks.runCurlProbe, +})); + +vi.mock("../adapters/docker", () => ({ + dockerCapture: mocks.dockerCapture, + dockerForceRm: mocks.dockerForceRm, + dockerImageInspectFormat: mocks.dockerImageInspectFormat, + dockerPullWithProgressWatchdog: mocks.dockerPullWithProgressWatchdog, + dockerRunDetached: mocks.dockerRunDetached, + dockerSpawn: mocks.dockerSpawn, + dockerStop: mocks.dockerStop, +})); + +vi.mock("./nim", () => ({ + getGpuIndicesByName: mocks.getGpuIndicesByName, +})); + +vi.mock("./vllm-storage", async (importOriginal) => ({ + ...(await importOriginal()), + findUnwritableModelCachePath: mocks.findUnwritableModelCachePath, + measureDirectorySizeBytes: mocks.measureDirectorySizeBytes, + probeDockerStorage: mocks.probeDockerStorage, + probeHostStorage: mocks.probeHostStorage, +})); + +vi.mock("./vllm-station-cluster", async (importOriginal) => ({ + ...(await importOriginal()), + probeDualStationVllmCapability: mocks.probeCapability, +})); + +vi.mock("./vllm-station-model-staging", () => ({ + stageDualStationModelSnapshot: mocks.stageModelSnapshot, +})); + +vi.mock("./vllm-station-cluster-lifecycle", () => ({ + areDualStationManagedVllmContainersRunning: mocks.areContainersRunning, + cleanupDualStationManagedVllm: mocks.cleanup, + commitDualStationLegacyMigration: mocks.commitLegacyMigration, + getDualStationManagedVllmBaseUrl: mocks.getManagedBaseUrl, + preflightDualStationGpuRuntime: mocks.preflightGpuRuntime, + preflightDualStationManagedVllm: mocks.preflightOwnership, + rollbackDualStationLegacyMigration: mocks.rollbackLegacyMigration, + startDualStationManagedVllm: mocks.startManaged, + withDualStationManagedVllmLifecycle: mocks.withLifecycle, +})); + +vi.mock("./vllm-api-key", () => ({ + ensureDualStationVllmApiKey: mocks.ensureApiKey, + loadDualStationVllmApiKey: mocks.loadApiKey, +})); + +import { detectVllmProfile, installVllm } from "./vllm"; +import { DUAL_STATION_VLLM_RUNTIME, type DualStationVllmPlan } from "./vllm-station-cluster"; +import { + createDualStationSshBindingFixture, + type DualStationSshBindingFixture, +} from "./vllm-station-ssh-binding.test-support"; + +const API_KEY = "ab".repeat(32); +const HEAD_ID = "a".repeat(64); +const WORKER_ID = "b".repeat(64); +const HEAD_BASE_URL = "http://192.168.100.1:8000"; +const LEGACY_MIGRATION = { + backupContainerName: `nemoclaw-vllm-legacy-${"1".repeat(32)}`, + legacyContainerId: "c".repeat(64), + transactionId: "1".repeat(32), + headContainerId: HEAD_ID, + workerContainerId: WORKER_ID, +}; + +function plan(): DualStationVllmPlan { + return { + peerSshBinding: sshFixture.binding, + runtime: DUAL_STATION_VLLM_RUNTIME, + local: { + hostname: "station-a", + home: "/home/nvidia", + uid: 1000, + gid: 1000, + gpu: { index: 0, name: "NVIDIA GB300", uuid: "GPU-aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa" }, + }, + peer: { + hostname: "station-b", + home: "/home/nvidia", + uid: 1000, + gid: 1000, + gpu: { index: 0, name: "NVIDIA GB300", uuid: "GPU-bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb" }, + }, + rails: [ + { + index: 0, + subnet: "192.168.100.0/30", + local: { + rdmaDevice: "mlx5_0", + uverbsDevice: "/dev/infiniband/uverbs0", + netdev: "enp1s0f0np0", + macAddress: "02:00:00:00:00:01", + pciAddress: "0000:01:00.0", + address: "192.168.100.1", + }, + peer: { + rdmaDevice: "mlx5_0", + uverbsDevice: "/dev/infiniband/uverbs0", + netdev: "enp1s0f0np0", + macAddress: "02:00:00:00:00:02", + pciAddress: "0000:01:00.0", + address: "192.168.100.2", + }, + }, + { + index: 1, + subnet: "192.168.200.0/30", + local: { + rdmaDevice: "mlx5_1", + uverbsDevice: "/dev/infiniband/uverbs1", + netdev: "enp1s0f1np1", + macAddress: "02:00:00:00:01:01", + pciAddress: "0000:01:00.1", + address: "192.168.200.1", + }, + peer: { + rdmaDevice: "mlx5_1", + uverbsDevice: "/dev/infiniband/uverbs1", + netdev: "enp1s0f1np1", + macAddress: "02:00:00:00:01:02", + pciAddress: "0000:01:00.1", + address: "192.168.200.2", + }, + }, + ], + masterAddress: "192.168.100.1", + roceGidIndex: 3, + }; +} + +function successfulSpawn(): EventEmitter & { stdout: EventEmitter; stderr: EventEmitter } { + const child = new EventEmitter() as EventEmitter & { + stdout: EventEmitter; + stderr: EventEmitter; + }; + child.stdout = new EventEmitter(); + child.stderr = new EventEmitter(); + process.nextTick(() => child.emit("exit", 0)); + return child; +} + +const originalEnv = { ...process.env }; +let logSpy: ReturnType; +let errorSpy: ReturnType; +let mkdirSpy: ReturnType; +let stdoutSpy: ReturnType; +let sshFixture: DualStationSshBindingFixture; + +beforeEach(() => { + vi.clearAllMocks(); + sshFixture = createDualStationSshBindingFixture(); + process.env.NEMOCLAW_VLLM_MODEL = "nemotron-3-ultra-550b-a55b"; + process.env.NEMOCLAW_DGX_STATION_PEER = "nvidia@station-b"; + process.env.HF_TOKEN = "hf_test"; + delete process.env.NEMOCLAW_VLLM_EXTRA_ARGS_JSON; + delete process.env.VLLM_API_KEY; + + logSpy = vi.spyOn(console, "log").mockImplementation(() => {}); + errorSpy = vi.spyOn(console, "error").mockImplementation(() => {}); + mkdirSpy = vi.spyOn(fs, "mkdirSync").mockImplementation(() => undefined); + stdoutSpy = vi.spyOn(process.stdout, "write").mockImplementation(() => true); + + mocks.probeCapability.mockReturnValue({ + kind: "ready", + plan: plan(), + peerModelSnapshot: "ready", + }); + mocks.stageModelSnapshot.mockResolvedValue({ ok: true, transferred: true }); + mocks.findUnwritableModelCachePath.mockReturnValue(null); + mocks.preflightOwnership.mockReturnValue({ ok: true }); + mocks.preflightGpuRuntime.mockReturnValue({ ok: true }); + mocks.getManagedBaseUrl.mockReturnValue(null); + mocks.ensureApiKey.mockReturnValue(API_KEY); + mocks.loadApiKey.mockReturnValue(API_KEY); + mocks.startManaged.mockReturnValue({ + ok: true, + baseUrl: HEAD_BASE_URL, + headContainerId: HEAD_ID, + workerContainerId: WORKER_ID, + reusedExisting: false, + }); + mocks.withLifecycle.mockImplementation(async (operation) => await operation()); + mocks.areContainersRunning.mockReturnValue(true); + mocks.cleanup.mockReturnValue({ ok: true, removedContainerIds: [] }); + mocks.commitLegacyMigration.mockResolvedValue({ ok: true, cleanupWarnings: [] }); + mocks.rollbackLegacyMigration.mockResolvedValue({ ok: true }); + mocks.measureDirectorySizeBytes.mockReturnValue(0n); + mocks.probeDockerStorage.mockReturnValue({ + ok: true, + capacity: { availableBytes: 1_000_000_000_000n, path: "/docker", source: "Docker" }, + }); + mocks.probeHostStorage.mockReturnValue({ + ok: true, + capacity: { + availableBytes: 1_000_000_000_000n, + path: path.join(os.homedir(), ".cache", "huggingface"), + source: "Hugging Face cache", + }, + }); + mocks.dockerImageInspectFormat.mockReturnValue(`sha256:${"c".repeat(64)}`); + mocks.dockerPullWithProgressWatchdog.mockResolvedValue({ + status: 0, + signal: null, + output: "", + timedOut: false, + timeoutKind: null, + }); + mocks.dockerSpawn.mockReturnValue(successfulSpawn()); + mocks.runCurlProbe.mockImplementation((args: string[]) => { + const url = args.at(-1) ?? ""; + const authenticated = args.includes("--config"); + return url.endsWith("/health") + ? { ok: true, httpStatus: 200, message: "ok", body: "" } + : authenticated + ? { + ok: true, + httpStatus: 200, + message: "ok", + body: JSON.stringify({ data: [{ id: "nemotron-ultra" }] }), + } + : { ok: false, httpStatus: 401, message: "unauthorized", body: "" }; + }); + mocks.runCapture.mockImplementation((args: readonly string[]) => { + switch (args[0]) { + case "sh": + return "/usr/bin/tool\n"; + case "curl": + return "200"; + default: + return ""; + } + }); +}); + +afterEach(() => { + logSpy.mockRestore(); + errorSpy.mockRestore(); + mkdirSpy.mockRestore(); + stdoutSpy.mockRestore(); + sshFixture.cleanup(); + process.env = { ...originalEnv }; +}); + +describe("dual DGX Station vLLM install orchestration", () => { + it("pulls both immutable images, preflights GPUs, authenticates, and starts worker plus head", async () => { + const clusterPlan = plan(); + let lifecycleActive = false; + let capabilityCalls = 0; + const probeImplementation = mocks.runCurlProbe.getMockImplementation(); + mocks.withLifecycle.mockImplementation(async (operation) => { + lifecycleActive = true; + try { + return await operation(); + } finally { + lifecycleActive = false; + } + }); + mocks.runCurlProbe.mockImplementation((args: string[], options?: unknown) => { + expect(lifecycleActive).toBe(true); + return probeImplementation?.(args, options); + }); + mocks.areContainersRunning.mockImplementation(() => { + expect(lifecycleActive).toBe(true); + return true; + }); + mocks.stageModelSnapshot.mockImplementation(async () => { + expect(lifecycleActive).toBe(true); + return { ok: true, transferred: false }; + }); + mocks.probeCapability.mockImplementation(() => { + capabilityCalls += 1; + expect(lifecycleActive).toBe(capabilityCalls === 2); + return { + kind: "ready", + plan: clusterPlan, + peerModelSnapshot: "ready", + }; + }); + const profile = detectVllmProfile({ platform: "station", type: "nvidia" }); + + await expect( + installVllm(profile!, { hasImage: true, nonInteractive: true, promptFn: vi.fn() }), + ).resolves.toEqual({ ok: true }); + + expect(mocks.preflightOwnership).toHaveBeenCalledWith(clusterPlan); + expect(mocks.dockerPullWithProgressWatchdog).toHaveBeenCalledTimes(2); + const localPullOptions = mocks.dockerPullWithProgressWatchdog.mock.calls[0][1]; + const peerPullOptions = mocks.dockerPullWithProgressWatchdog.mock.calls[1][1]; + expect(peerPullOptions.env.DOCKER_HOST).toBe("ssh://nvidia@192.168.50.20"); + expect(peerPullOptions.env.PATH.split(":")[0]).toBe( + clusterPlan.peerSshBinding.sshWrapperDirectory, + ); + expect(localPullOptions.env.DOCKER_HOST).not.toBe(peerPullOptions.env.DOCKER_HOST); + expect(localPullOptions.env.DOCKER_CONTEXT).toBe("default"); + expect(peerPullOptions.env.DOCKER_CONTEXT).toBeUndefined(); + expect(localPullOptions.env.VLLM_API_KEY).toBeUndefined(); + expect(peerPullOptions.env.VLLM_API_KEY).toBeUndefined(); + expect(mocks.preflightGpuRuntime).toHaveBeenCalledWith(clusterPlan); + expect(mocks.preflightGpuRuntime.mock.invocationCallOrder[0]).toBeGreaterThan( + mocks.dockerPullWithProgressWatchdog.mock.invocationCallOrder[1], + ); + expect(mocks.ensureApiKey.mock.invocationCallOrder[0]).toBeGreaterThan( + mocks.preflightGpuRuntime.mock.invocationCallOrder[0], + ); + expect(mocks.ensureApiKey.mock.invocationCallOrder[0]).toBeGreaterThan( + mocks.dockerSpawn.mock.invocationCallOrder[0], + ); + expect(mocks.startManaged).toHaveBeenCalledWith(clusterPlan, { apiKey: API_KEY }); + expect(mocks.startManaged.mock.invocationCallOrder[0]).toBeGreaterThan( + mocks.dockerSpawn.mock.invocationCallOrder[0], + ); + expect(mocks.areContainersRunning).toHaveBeenCalledWith(clusterPlan); + expect(mocks.withLifecycle).toHaveBeenCalledTimes(2); + expect(lifecycleActive).toBe(false); + expect(mocks.probeCapability).toHaveBeenCalledTimes(2); + expect(mocks.stageModelSnapshot).toHaveBeenCalledWith(clusterPlan); + + const readinessArgs = mocks.runCurlProbe.mock.calls[0][0] as string[]; + expect(readinessArgs).toContain(`${HEAD_BASE_URL}/health`); + expect(readinessArgs.join(" ")).not.toContain("/v1/models"); + expect(mocks.runCurlProbe.mock.calls[0][1]).toMatchObject({ pinnedAddresses: [] }); + const unauthenticatedArgs = mocks.runCurlProbe.mock.calls[1][0] as string[]; + const authenticatedArgs = mocks.runCurlProbe.mock.calls[2][0] as string[]; + expect(unauthenticatedArgs).toContain(`${HEAD_BASE_URL}/v1/models`); + expect(unauthenticatedArgs).not.toContain("--config"); + expect(authenticatedArgs).toContain("--config"); + expect(authenticatedArgs).not.toContain(API_KEY); + expect(mocks.runCurlProbe.mock.calls[1][1]).toMatchObject({ pinnedAddresses: [] }); + expect(mocks.runCurlProbe.mock.calls[2][1]).toMatchObject({ pinnedAddresses: [] }); + expect(mocks.dockerSpawn.mock.calls[0][1].env.VLLM_API_KEY).toBeUndefined(); + expect(mocks.dockerSpawn.mock.calls[0][1].env.DOCKER_CONTEXT).toBe("default"); + expect(mocks.dockerImageInspectFormat.mock.calls[0][2].env.DOCKER_CONTEXT).toBe("default"); + }); + + it("selects pinned Nemotron Ultra without prompting when an explicit peer qualifies", async () => { + delete process.env.NEMOCLAW_VLLM_MODEL; + const promptFn = vi.fn(); + const profile = detectVllmProfile({ platform: "station", type: "nvidia" }); + + await expect( + installVllm(profile!, { hasImage: true, nonInteractive: true, promptFn }), + ).resolves.toEqual({ ok: true }); + + expect(promptFn).not.toHaveBeenCalled(); + const [downloadArgs] = mocks.dockerSpawn.mock.calls[0] as [string[]]; + expect(downloadArgs).toEqual( + expect.arrayContaining([ + "nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-NVFP4", + "--revision", + DUAL_STATION_VLLM_RUNTIME.modelRevision, + ]), + ); + expect(mocks.probeCapability).toHaveBeenCalledTimes(2); + }); + + it("installs the public Nemotron Ultra model without an HF token", async () => { + delete process.env.NEMOCLAW_VLLM_MODEL; + delete process.env.HF_TOKEN; + delete process.env.HUGGING_FACE_HUB_TOKEN; + const beforeInstall = vi.fn(); + const profile = detectVllmProfile({ platform: "station", type: "nvidia" }); + + await expect( + installVllm(profile!, { + hasImage: true, + nonInteractive: true, + promptFn: vi.fn(), + beforeInstall, + }), + ).resolves.toEqual({ ok: true }); + + expect(mocks.probeCapability).toHaveBeenCalledTimes(2); + expect(beforeInstall).toHaveBeenCalledWith("nemotron-ultra"); + expect(mocks.preflightOwnership).toHaveBeenCalled(); + expect(mocks.dockerSpawn).toHaveBeenCalled(); + expect(mocks.stageModelSnapshot).toHaveBeenCalled(); + expect(errorSpy).not.toHaveBeenCalledWith(expect.stringContaining("gated on Hugging Face")); + }); + + it("skips only the model picker for an interactive qualified-peer install", async () => { + delete process.env.NEMOCLAW_VLLM_MODEL; + const promptFn = vi.fn().mockResolvedValue("y"); + const profile = detectVllmProfile({ platform: "station", type: "nvidia" }); + + await expect( + installVllm(profile!, { hasImage: true, nonInteractive: false, promptFn }), + ).resolves.toEqual({ ok: true }); + + expect(promptFn).toHaveBeenCalledTimes(1); + expect(promptFn).toHaveBeenCalledWith(" Continue? [y/N]: "); + expect(promptFn).not.toHaveBeenCalledWith(expect.stringContaining("Choose model")); + }); + + it("rejects a configured peer combined with an explicit non-Ultra model", async () => { + process.env.NEMOCLAW_VLLM_MODEL = "deepseek-r1-distill-70b"; + delete process.env.HF_TOKEN; + delete process.env.HUGGING_FACE_HUB_TOKEN; + const beforeInstall = vi.fn(); + const promptFn = vi.fn(); + const profile = detectVllmProfile({ platform: "station", type: "nvidia" }); + + await expect( + installVllm(profile!, { + hasImage: true, + nonInteractive: true, + promptFn, + beforeInstall, + }), + ).resolves.toEqual({ ok: false }); + + expect(promptFn).not.toHaveBeenCalled(); + expect(beforeInstall).not.toHaveBeenCalled(); + expect(mocks.probeCapability).not.toHaveBeenCalled(); + expect(mocks.stageModelSnapshot).not.toHaveBeenCalled(); + expect(mocks.runCapture).not.toHaveBeenCalled(); + expect(mocks.dockerPullWithProgressWatchdog).not.toHaveBeenCalled(); + expect(mocks.dockerSpawn).not.toHaveBeenCalled(); + expect(errorSpy).toHaveBeenCalledWith( + " vLLM install failed: NEMOCLAW_DGX_STATION_PEER requires the DGX Station dual-serving model. " + + "Unset NEMOCLAW_VLLM_MODEL or select nemotron-3-ultra-550b-a55b; the explicit model override remains authoritative.", + ); + }); + + it("stages a missing peer snapshot after download and requires a ready re-probe", async () => { + const clusterPlan = plan(); + mocks.probeCapability + .mockReturnValueOnce({ + kind: "ready", + plan: clusterPlan, + peerModelSnapshot: "staging-required", + }) + .mockReturnValueOnce({ + kind: "ready", + plan: clusterPlan, + peerModelSnapshot: "ready", + }); + const profile = detectVllmProfile({ platform: "station", type: "nvidia" }); + + await expect( + installVllm(profile!, { hasImage: true, nonInteractive: true, promptFn: vi.fn() }), + ).resolves.toEqual({ ok: true }); + + expect(mocks.stageModelSnapshot).toHaveBeenCalledWith(clusterPlan); + expect(mocks.stageModelSnapshot.mock.invocationCallOrder[0]).toBeGreaterThan( + mocks.dockerSpawn.mock.invocationCallOrder[0], + ); + expect(mocks.probeCapability.mock.invocationCallOrder[1]).toBeGreaterThan( + mocks.stageModelSnapshot.mock.invocationCallOrder[0], + ); + expect(mocks.startManaged.mock.invocationCallOrder[0]).toBeGreaterThan( + mocks.probeCapability.mock.invocationCallOrder[1], + ); + }); + + it("fails closed before re-probe or launch when peer snapshot staging fails", async () => { + mocks.probeCapability.mockReturnValue({ + kind: "ready", + plan: plan(), + peerModelSnapshot: "staging-required", + }); + mocks.stageModelSnapshot.mockResolvedValue({ ok: false, reason: "peer transfer timed out" }); + const profile = detectVllmProfile({ platform: "station", type: "nvidia" }); + + await expect( + installVllm(profile!, { hasImage: true, nonInteractive: true, promptFn: vi.fn() }), + ).resolves.toEqual({ ok: false }); + + expect(mocks.probeCapability).toHaveBeenCalledTimes(1); + expect(mocks.startManaged).not.toHaveBeenCalled(); + expect(errorSpy).toHaveBeenCalledWith(" vLLM install failed: peer transfer timed out"); + }); + + it("fails before callbacks, prompts, or Docker work when a configured peer is incapable", async () => { + const beforeInstall = vi.fn(); + mocks.probeCapability.mockReturnValue({ + kind: "unavailable", + code: "peer-fabric-unavailable", + reason: "peer fabric is incomplete", + }); + const profile = detectVllmProfile({ platform: "station", type: "nvidia" }); + + await expect( + installVllm(profile!, { + hasImage: true, + nonInteractive: true, + promptFn: vi.fn(), + beforeInstall, + }), + ).resolves.toEqual({ ok: false }); + + expect(beforeInstall).not.toHaveBeenCalled(); + expect(mocks.runCapture).not.toHaveBeenCalled(); + expect(mocks.preflightOwnership).not.toHaveBeenCalled(); + expect(mocks.dockerPullWithProgressWatchdog).not.toHaveBeenCalled(); + expect(mocks.ensureApiKey).not.toHaveBeenCalled(); + expect(errorSpy).toHaveBeenCalledWith( + " Dual DGX Station setup unavailable: peer fabric is incomplete", + ); + }); + + it("stops before key creation and model download when either GPU runtime smoke fails", async () => { + mocks.preflightGpuRuntime.mockReturnValue({ + ok: false, + reason: "worker GPU smoke did not expose exactly the discovered GPU", + }); + const profile = detectVllmProfile({ platform: "station", type: "nvidia" }); + + await expect( + installVllm(profile!, { hasImage: true, nonInteractive: true, promptFn: vi.fn() }), + ).resolves.toEqual({ ok: false }); + + expect(mocks.dockerPullWithProgressWatchdog).toHaveBeenCalledTimes(2); + expect(mocks.ensureApiKey).not.toHaveBeenCalled(); + expect(mocks.dockerSpawn).not.toHaveBeenCalled(); + expect(mocks.startManaged).not.toHaveBeenCalled(); + expect(errorSpy).toHaveBeenCalledWith( + " vLLM install failed: worker GPU smoke did not expose exactly the discovered GPU", + ); + }); + + it("refuses teardown when the qualified topology changes during the model download", async () => { + const originalPlan = plan(); + const changedPlan = plan(); + changedPlan.peer.hostname = "station-b-replaced"; + mocks.probeCapability + .mockReturnValueOnce({ + kind: "ready", + plan: originalPlan, + peerModelSnapshot: "ready", + }) + .mockReturnValueOnce({ + kind: "ready", + plan: changedPlan, + peerModelSnapshot: "ready", + }); + const profile = detectVllmProfile({ platform: "station", type: "nvidia" }); + + await expect( + installVllm(profile!, { hasImage: true, nonInteractive: true, promptFn: vi.fn() }), + ).resolves.toEqual({ ok: false }); + + expect(mocks.dockerSpawn).toHaveBeenCalledTimes(1); + expect(mocks.ensureApiKey).not.toHaveBeenCalled(); + expect(mocks.startManaged).not.toHaveBeenCalled(); + expect(errorSpy).toHaveBeenCalledWith( + " vLLM install failed: dual-Station topology changed during download; rerun setup against a stable pair.", + ); + }); + + it("commits legacy retirement only after readiness, auth, and final pair validation", async () => { + mocks.startManaged.mockReturnValue({ + ok: true, + baseUrl: HEAD_BASE_URL, + headContainerId: HEAD_ID, + workerContainerId: WORKER_ID, + reusedExisting: false, + legacyMigration: LEGACY_MIGRATION, + }); + const profile = detectVllmProfile({ platform: "station", type: "nvidia" }); + + await expect( + installVllm(profile!, { hasImage: true, nonInteractive: true, promptFn: vi.fn() }), + ).resolves.toEqual({ ok: true }); + + expect(mocks.commitLegacyMigration).toHaveBeenCalledWith(plan(), LEGACY_MIGRATION); + expect(mocks.commitLegacyMigration.mock.invocationCallOrder[0]).toBeGreaterThan( + mocks.areContainersRunning.mock.invocationCallOrder[0], + ); + expect(mocks.commitLegacyMigration.mock.invocationCallOrder[0]).toBeGreaterThan( + mocks.runCurlProbe.mock.invocationCallOrder.at(-1) ?? 0, + ); + expect(mocks.rollbackLegacyMigration).not.toHaveBeenCalled(); + expect(mocks.cleanup).not.toHaveBeenCalled(); + }); + + it.each([ + { + failure: "readiness", + configureFailure: () => { + mocks.runCurlProbe.mockReturnValue({ ok: false, httpStatus: 503, message: "loading" }); + mocks.dockerCapture.mockReturnValue(""); + }, + expectedCommitCalls: 0, + }, + { + failure: "authentication", + configureFailure: () => { + mocks.runCurlProbe.mockImplementation((args: string[]) => ({ + ok: true, + httpStatus: 200, + message: "ok", + body: args.at(-1)?.endsWith("/v1/models") + ? JSON.stringify({ data: [{ id: "nemotron-ultra" }] }) + : "", + })); + }, + expectedCommitCalls: 0, + }, + { + failure: "final running check", + configureFailure: () => mocks.areContainersRunning.mockReturnValue(false), + expectedCommitCalls: 0, + }, + { + failure: "commit validation", + configureFailure: () => + mocks.commitLegacyMigration.mockResolvedValue({ + ok: false, + reason: "new dual-Station transaction changed before commit", + }), + expectedCommitCalls: 1, + }, + ])("restores legacy state when external $failure fails", async ({ + configureFailure, + expectedCommitCalls, + }) => { + mocks.startManaged.mockReturnValue({ + ok: true, + baseUrl: HEAD_BASE_URL, + headContainerId: HEAD_ID, + workerContainerId: WORKER_ID, + reusedExisting: false, + legacyMigration: LEGACY_MIGRATION, + }); + configureFailure(); + const profile = detectVllmProfile({ platform: "station", type: "nvidia" }); + + await expect( + installVllm(profile!, { hasImage: true, nonInteractive: true, promptFn: vi.fn() }), + ).resolves.toEqual({ ok: false }); + + expect(mocks.rollbackLegacyMigration).toHaveBeenCalledWith(plan(), LEGACY_MIGRATION); + expect(mocks.cleanup).not.toHaveBeenCalled(); + expect(mocks.commitLegacyMigration).toHaveBeenCalledTimes(expectedCommitCalls); + }); + + it("rolls back a new pair when unauthenticated model inventory is exposed", async () => { + mocks.runCurlProbe.mockImplementation((args: string[]) => ({ + ok: true, + httpStatus: 200, + message: "ok", + body: args.at(-1)?.endsWith("/v1/models") + ? JSON.stringify({ data: [{ id: "nemotron-ultra" }] }) + : "", + })); + const profile = detectVllmProfile({ platform: "station", type: "nvidia" }); + + await expect( + installVllm(profile!, { hasImage: true, nonInteractive: true, promptFn: vi.fn() }), + ).resolves.toEqual({ ok: false }); + + expect(mocks.cleanup).toHaveBeenCalledWith( + expect.objectContaining({ masterAddress: "192.168.100.1" }), + ); + expect(errorSpy).toHaveBeenCalledWith( + " vLLM install failed: unauthenticated model inventory returned HTTP 200; expected vLLM to reject it with HTTP 401", + ); + }); + + it("stops before storage or image work when container ownership is not exact", async () => { + mocks.preflightOwnership.mockReturnValue({ + ok: false, + reason: "worker container ownership is foreign; refusing mutation", + }); + const profile = detectVllmProfile({ platform: "station", type: "nvidia" }); + + await expect( + installVllm(profile!, { hasImage: true, nonInteractive: true, promptFn: vi.fn() }), + ).resolves.toEqual({ ok: false }); + + expect(mocks.probeHostStorage).not.toHaveBeenCalled(); + expect(mocks.dockerImageInspectFormat).not.toHaveBeenCalled(); + expect(mocks.dockerPullWithProgressWatchdog).not.toHaveBeenCalled(); + expect(mocks.dockerSpawn).not.toHaveBeenCalled(); + expect(mocks.startManaged).not.toHaveBeenCalled(); + }); +}); diff --git a/src/lib/inference/vllm-install-storage.test.ts b/src/lib/inference/vllm-install-storage.test.ts index 54e9fa0a62..2499371607 100644 --- a/src/lib/inference/vllm-install-storage.test.ts +++ b/src/lib/inference/vllm-install-storage.test.ts @@ -152,7 +152,7 @@ describe("managed vLLM install storage", () => { stdoutWrite = vi.spyOn(process.stdout, "write").mockImplementation(() => true); delete process.env.NEMOCLAW_VLLM_MODEL; delete process.env.NEMOCLAW_VLLM_EXTRA_ARGS_JSON; - delete process.env.HF_TOKEN; + process.env.HF_TOKEN = "hf_test"; delete process.env.HUGGING_FACE_HUB_TOKEN; }); diff --git a/src/lib/inference/vllm-models.test.ts b/src/lib/inference/vllm-models.test.ts index 4836f1b9d1..bb928a9859 100644 --- a/src/lib/inference/vllm-models.test.ts +++ b/src/lib/inference/vllm-models.test.ts @@ -5,6 +5,7 @@ import { describe, expect, it } from "vitest"; import { assertGatedModelAccess, + buildNemotronUltraDistributedServeCommand, buildVllmServeCommand, DEFAULT_VLLM_MODEL, modelsForPlatform, @@ -121,6 +122,78 @@ describe("vllm model registry", () => { ); }); + it("builds the pinned two-Station Nemotron Ultra vLLM v0.25.1 Ray head command", () => { + const cmd = buildNemotronUltraDistributedServeCommand({ + nodeRank: 0, + masterAddr: "192.168.240.1", + masterPort: 6379, + }); + + expect(cmd).toContain('python3 -m pip install --user --no-cache-dir "ray==2.56.0"'); + expect(cmd).toContain( + "ray start --head --node-ip-address=192.168.240.1 --port=6379 --num-gpus=1", + ); + expect(cmd).toContain("vllm serve nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-NVFP4"); + expect(cmd).toContain("--tensor-parallel-size 1"); + expect(cmd).toContain("--pipeline-parallel-size 2"); + expect(cmd).toContain("--distributed-executor-backend ray"); + expect(cmd).toContain("--kv-cache-dtype fp8"); + expect(cmd).toContain("--max-model-len 262144"); + expect(cmd).toContain("--distributed-timeout-seconds 7200"); + expect(cmd).toContain("--served-model-name nemotron-ultra"); + expect(cmd).toContain("--host 192.168.240.1"); + expect(cmd).toContain("--max-num-seqs 256"); + expect(cmd).toContain("--gpu-memory-utilization 0.9"); + expect(cmd).not.toContain("--kernel_config"); + expect(cmd).not.toContain("--speculative-config"); + expect(cmd).not.toContain("--cpu-offload"); + }); + + it("builds a Ray worker command without starting a second API server", () => { + const head = buildNemotronUltraDistributedServeCommand({ + nodeRank: 0, + masterAddr: "192.168.240.1", + masterPort: 6379, + }); + const worker = buildNemotronUltraDistributedServeCommand({ + nodeRank: 1, + masterAddr: "192.168.240.1", + masterPort: 6379, + nodeAddr: "192.168.240.2", + }); + + expect(worker).toContain( + "ray start --address=192.168.240.1:6379 --node-ip-address=192.168.240.2 --num-gpus=1 --block", + ); + expect(worker).not.toContain("vllm serve"); + expect(head).toContain("vllm serve"); + }); + + it("rejects invalid two-Station Nemotron Ultra rendezvous options", () => { + expect(() => + buildNemotronUltraDistributedServeCommand({ + nodeRank: 0, + masterAddr: "station a", + masterPort: 6379, + }), + ).toThrow(/masterAddr/); + expect(() => + buildNemotronUltraDistributedServeCommand({ + nodeRank: 1, + masterAddr: "192.168.240.1", + masterPort: 70000, + }), + ).toThrow(/masterPort/); + expect(() => + buildNemotronUltraDistributedServeCommand({ + nodeRank: 1, + masterAddr: "192.168.240.1", + masterPort: 6379, + nodeAddr: "worker.example.com", + }), + ).toThrow(/nodeAddr/); + }); + it("rejects an unknown NEMOCLAW_VLLM_MODEL with a helpful message", () => { expect(() => selectVllmModelFromEnv({ NEMOCLAW_VLLM_MODEL: "made-up-model" } as NodeJS.ProcessEnv), @@ -152,6 +225,12 @@ describe("vllm model registry", () => { ); }); + it("keeps the public Nemotron Ultra recipe usable without a Hugging Face token", () => { + const ultra = VLLM_MODELS.find((m) => m.envValue === "nemotron-3-ultra-550b-a55b"); + expect(ultra?.gated).toBe(false); + expect(() => assertGatedModelAccess(ultra!, {} as NodeJS.ProcessEnv)).not.toThrow(); + }); + it("never rejects a non-gated model regardless of token state", () => { const qwen = VLLM_MODELS.find((m) => m.envValue === "qwen3.6-27b"); expect(() => assertGatedModelAccess(qwen!, {} as NodeJS.ProcessEnv)).not.toThrow(); diff --git a/src/lib/inference/vllm-models.ts b/src/lib/inference/vllm-models.ts index 5e181fe51e..6591e80569 100644 --- a/src/lib/inference/vllm-models.ts +++ b/src/lib/inference/vllm-models.ts @@ -27,6 +27,8 @@ * envelope, and tool-call behaviour validated. */ +import net from "node:net"; + export type VllmPlatform = "spark" | "station" | "linux"; export interface VllmRuntimeOverride { @@ -50,6 +52,15 @@ export const NEMOTRON_ULTRA_STATION_IMAGE = { }, } as const; +/** Runtime pinned from the published dual-DGX-Station playbook. */ +export const NEMOTRON_ULTRA_DUAL_STATION_IMAGE = { + tag: "vllm/vllm-openai:v0.25.1-aarch64", + arm64: { + ref: "vllm/vllm-openai@sha256:2cc49b81319f7a66a33dd8bd63a7bfddae079122b33ce51989b6828a1f038c37", + downloadSizeBytes: 10_238_912_364, + }, +} as const; + export interface VllmModelDef { /** Hugging Face model id (also passed to `vllm serve`). */ id: string; @@ -263,8 +274,9 @@ export const VLLM_MODELS: readonly VllmModelDef[] = [ imageDownloadSizeBytes: NEMOTRON_ULTRA_STATION_IMAGE.arm64.downloadSizeBytes, modelDownloadSizeBytes: 352_381_245_521, loadTimeoutSec: 3600, - // Keep NemoClaw's bridge-networked local-inference boundary instead of - // importing the playbook's host-network setting. + // The single-host runtime keeps NemoClaw's bridge-networked local- + // inference boundary. The qualified dual-Station lifecycle intentionally + // builds a separate host-networked launch contract for NCCL/RDMA. dockerRunArgs: ["--shm-size", "16g", "--ulimit", "memlock=-1", "--ulimit", "stack=67108864"], }, // The digest-pinned vLLM image already contains the serving package, and @@ -480,6 +492,168 @@ function shellQuote(value: string): string { return `'${value.replace(/'/g, `'\"'\"'`)}'`; } +function rewriteVllmArgs( + args: readonly string[], + overrides: Readonly>, + omittedFlags: ReadonlySet = new Set(), +): string[] { + const result: string[] = []; + const remainingOverrides = new Set(Object.keys(overrides)); + for (let index = 0; index < args.length; index += 1) { + const arg = args[index]; + if (omittedFlags.has(arg)) { + if (index === args.length - 1) throw new Error(`Missing value for vLLM argument '${arg}'.`); + index += 1; + continue; + } + if (Object.prototype.hasOwnProperty.call(overrides, arg)) { + if (index === args.length - 1) throw new Error(`Missing value for vLLM argument '${arg}'.`); + result.push(arg, overrides[arg]); + remainingOverrides.delete(arg); + index += 1; + continue; + } + result.push(arg); + } + if (remainingOverrides.size > 0) { + throw new Error(`Cannot override missing vLLM argument '${[...remainingOverrides][0]}'.`); + } + return result; +} + +export interface NemotronUltraDistributedServeOptions { + /** Ray role: rank 0 owns the API and rank 1 is the worker. */ + nodeRank: 0 | 1; + /** Routable head address used by both nodes for the Ray control plane. */ + masterAddr: string; + /** Routable Ray head port. */ + masterPort: number; + /** Routable address of the node running this command. */ + nodeAddr?: string; +} + +/** + * Build one side of the published two-Station Nemotron Ultra vLLM v0.25.1 + * Ray pipeline-parallel launch. Existing callers keep the single-node + * registry command unless they opt into this role/address/port API. + */ +export function buildNemotronUltraDistributedServeCommand( + options: NemotronUltraDistributedServeOptions, +): string { + if (options.nodeRank !== 0 && options.nodeRank !== 1) { + throw new Error("Nemotron Ultra distributed nodeRank must be 0 or 1."); + } + const masterAddr = options.masterAddr.trim(); + if (net.isIP(masterAddr) !== 4) { + throw new Error("Nemotron Ultra distributed masterAddr must be a canonical IPv4 address."); + } + if ( + !Number.isInteger(options.masterPort) || + options.masterPort < 1 || + options.masterPort > 65535 + ) { + throw new Error("Nemotron Ultra distributed masterPort must be an integer from 1 to 65535."); + } + const nodeAddr = (options.nodeAddr ?? masterAddr).trim(); + if (net.isIP(nodeAddr) !== 4) { + throw new Error("Nemotron Ultra distributed nodeAddr must be a canonical IPv4 address."); + } + if (options.nodeRank === 0 && nodeAddr !== masterAddr) { + throw new Error("Nemotron Ultra Ray head nodeAddr must match masterAddr."); + } + + const model = VLLM_MODELS.find( + (candidate) => candidate.envValue === "nemotron-3-ultra-550b-a55b", + ); + if (!model?.revision || !model.servedModelId) { + throw new Error( + "Nemotron Ultra distributed serving requires a pinned revision and served model id.", + ); + } + + const sharedArgs = rewriteVllmArgs(SHARED_VLLM_ARGS, { + "--tensor-parallel-size": "1", + "--pipeline-parallel-size": "2", + }); + const modelArgs = rewriteVllmArgs( + model.modelArgs, + { + // Rank 0 binds only to the selected direct-attach RoCE address. This + // keeps the API off the management network while still giving the + // OpenShell route a host-reachable endpoint. The lifecycle also enables + // vLLM bearer authentication. The Ray worker exposes no API. + "--host": masterAddr, + "--max-num-seqs": "256", + "--gpu-memory-utilization": "0.9", + }, + new Set([ + "--cpu-offload-gb", + "--cpu-offload-params", + "--kernel_config", + "--speculative-config", + "--default-chat-template-kwargs", + ]), + ); + const args = [ + ...sharedArgs, + "--distributed-executor-backend", + "ray", + "--kv-cache-dtype", + "fp8", + "--max-model-len", + "262144", + "--distributed-timeout-seconds", + "7200", + "--enable-prefix-caching", + "--revision", + model.revision, + "--served-model-name", + "nemotron-ultra", + ...modelArgs, + ]; + const bootstrap = [ + "set -euo pipefail", + 'export PATH="$HOME/.local/bin:$PATH"', + 'python3 -m pip install --user --no-cache-dir "ray==2.56.0"', + ]; + if (options.nodeRank === 1) { + return [ + ...bootstrap, + "python3 - <<'PY'", + "import socket", + "import time", + `address = (${JSON.stringify(masterAddr)}, ${String(options.masterPort)})`, + "deadline = time.time() + 3600", + "while True:", + " try:", + " with socket.create_connection(address, timeout=5):", + " break", + " except OSError:", + " if time.time() >= deadline:", + ' raise TimeoutError("Ray head did not become reachable within 3600 seconds")', + " time.sleep(5)", + "PY", + `exec ray start --address=${shellQuote(`${masterAddr}:${String(options.masterPort)}`)} --node-ip-address=${shellQuote(nodeAddr)} --num-gpus=1 --block`, + ].join("\n"); + } + return [ + ...bootstrap, + `ray start --head --node-ip-address=${shellQuote(masterAddr)} --port=${String(options.masterPort)} --num-gpus=1`, + "python3 - <<'PY'", + "import time", + "import ray", + 'ray.init(address="auto")', + "deadline = time.time() + 3600", + 'while ray.cluster_resources().get("GPU", 0) < 2:', + " if time.time() >= deadline:", + ' raise TimeoutError("peer DGX Station GPU did not join Ray within 3600 seconds")', + " time.sleep(5)", + "print(ray.cluster_resources())", + "PY", + `exec vllm serve ${model.id} ${args.join(" ")}`, + ].join("\n"); +} + /** * Build the `vllm serve` command line for the supplied model: the shared * serving flags merged with the model-specific args from the registry. diff --git a/src/lib/inference/vllm-ownership.test.ts b/src/lib/inference/vllm-ownership.test.ts new file mode 100644 index 0000000000..1a18e912c8 --- /dev/null +++ b/src/lib/inference/vllm-ownership.test.ts @@ -0,0 +1,129 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +const mocks = vi.hoisted(() => ({ + dockerCapture: vi.fn(), +})); + +vi.mock("../adapters/docker", async (importOriginal) => ({ + ...(await importOriginal()), + dockerCapture: mocks.dockerCapture, +})); + +import { + isNemoClawManagedVllmRunning, + NEMOCLAW_VLLM_CONTAINER_NAME, + NEMOCLAW_VLLM_MANAGED_LABEL, +} from "./vllm"; +import { + DUAL_STATION_VLLM_CLUSTER_LABEL, + DUAL_STATION_VLLM_ENDPOINT_LABEL, + DUAL_STATION_VLLM_ROLE_LABEL, +} from "./vllm-station-cluster-lifecycle"; + +const MANAGED_CONTAINER_ID = "a".repeat(64); + +function vllmContainerRow( + containerName: string, + { + id = MANAGED_CONTAINER_ID, + label = "true", + state = "exited", + dualRole = "", + dualEndpoint = "", + dualCluster = "", + } = {}, +): string { + return [id, containerName, state, label, dualRole, dualEndpoint, dualCluster].join("|"); +} + +beforeEach(() => vi.clearAllMocks()); +afterEach(() => vi.unstubAllEnvs()); + +describe("managed vLLM ownership", () => { + it("recognizes only the exact running container with the managed label", () => { + mocks.dockerCapture.mockReturnValue( + vllmContainerRow(NEMOCLAW_VLLM_CONTAINER_NAME, { state: "running" }), + ); + + expect(isNemoClawManagedVllmRunning()).toBe(true); + expect(mocks.dockerCapture).toHaveBeenCalledWith( + [ + "container", + "ls", + "--all", + "--no-trunc", + "--filter", + `name=^/${NEMOCLAW_VLLM_CONTAINER_NAME}$`, + "--format", + [ + "{{.ID}}", + "{{.Names}}", + "{{.State}}", + `{{.Label "${NEMOCLAW_VLLM_MANAGED_LABEL}"}}`, + `{{.Label "${DUAL_STATION_VLLM_ROLE_LABEL}"}}`, + `{{.Label "${DUAL_STATION_VLLM_ENDPOINT_LABEL}"}}`, + `{{.Label "${DUAL_STATION_VLLM_CLUSTER_LABEL}"}}`, + ].join("|"), + ], + expect.objectContaining({ timeout: 10_000 }), + ); + }); + + it("recognizes a running dual-Station head without treating it as legacy ownership", () => { + mocks.dockerCapture.mockReturnValue( + vllmContainerRow(NEMOCLAW_VLLM_CONTAINER_NAME, { + state: "running", + dualRole: "head", + dualEndpoint: "http://192.168.100.1:8000", + dualCluster: "f".repeat(64), + }), + ); + + expect(isNemoClawManagedVllmRunning()).toBe(true); + }); + + it("checks the canonical local daemon before an ambient remote Docker host", () => { + vi.stubEnv("DOCKER_HOST", "ssh://builder.example.test"); + vi.stubEnv("DOCKER_CONTEXT", undefined); + mocks.dockerCapture.mockImplementation( + (_args: readonly string[], options?: { env?: NodeJS.ProcessEnv }) => + options?.env?.DOCKER_CONTEXT === "default" + ? vllmContainerRow(NEMOCLAW_VLLM_CONTAINER_NAME, { + state: "running", + dualRole: "head", + dualEndpoint: "http://192.168.100.1:8000", + dualCluster: "f".repeat(64), + }) + : "", + ); + + expect(isNemoClawManagedVllmRunning()).toBe(true); + expect(mocks.dockerCapture).toHaveBeenCalledTimes(1); + expect(mocks.dockerCapture.mock.calls[0]?.[1]?.env).toMatchObject({ + DOCKER_CONTEXT: "default", + }); + expect(mocks.dockerCapture.mock.calls[0]?.[1]?.env).not.toHaveProperty("DOCKER_HOST"); + }); + + it.each([ + vllmContainerRow(NEMOCLAW_VLLM_CONTAINER_NAME), + vllmContainerRow(NEMOCLAW_VLLM_CONTAINER_NAME, { label: "" }), + vllmContainerRow(NEMOCLAW_VLLM_CONTAINER_NAME, { label: "false", state: "running" }), + "", + "malformed", + `${vllmContainerRow(NEMOCLAW_VLLM_CONTAINER_NAME)}\n${vllmContainerRow(NEMOCLAW_VLLM_CONTAINER_NAME)}`, + ])("fails closed for inspect output %j", (output) => { + mocks.dockerCapture.mockReturnValue(output); + expect(isNemoClawManagedVllmRunning()).toBe(false); + }); + + it("fails closed when Docker inspection throws", () => { + mocks.dockerCapture.mockImplementation(() => { + throw new Error("docker unavailable"); + }); + expect(isNemoClawManagedVllmRunning()).toBe(false); + }); +}); diff --git a/src/lib/inference/vllm-station-cluster-lifecycle.test.ts b/src/lib/inference/vllm-station-cluster-lifecycle.test.ts new file mode 100644 index 0000000000..97d7ea907a --- /dev/null +++ b/src/lib/inference/vllm-station-cluster-lifecycle.test.ts @@ -0,0 +1,1147 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { + createDualStationLifecycleHarness, + dualStationDockerValues as dockerValues, + type LifecycleFakeContainer as FakeContainer, + type LifecycleHarnessOptions as HarnessOptions, + requireLegacyMigration, +} from "../../../test/support/vllm-station-cluster-lifecycle-test-support"; +import { DUAL_STATION_VLLM_RUNTIME, type DualStationVllmPlan } from "./vllm-station-cluster"; +import { + areDualStationManagedVllmContainersRunning, + buildDualStationGpuSmokeRunArgs, + buildDualStationVllmRunArgs, + cleanupDualStationManagedVllm, + commitDualStationLegacyMigration, + DUAL_STATION_VLLM_API_KEY_FINGERPRINT_LABEL, + DUAL_STATION_VLLM_CLUSTER_LABEL, + DUAL_STATION_VLLM_ENDPOINT_LABEL, + DUAL_STATION_VLLM_GPU_LABEL, + DUAL_STATION_VLLM_GPU_SMOKE_LABEL, + DUAL_STATION_VLLM_HEAD_CONTAINER_NAME, + DUAL_STATION_VLLM_LAUNCH_CONTRACT_LABEL, + DUAL_STATION_VLLM_LAUNCH_SCHEMA_LABEL, + DUAL_STATION_VLLM_MANAGED_LABEL, + DUAL_STATION_VLLM_ROLE_LABEL, + DUAL_STATION_VLLM_TRANSACTION_LABEL, + DUAL_STATION_VLLM_WORKER_CONTAINER_NAME, + type DualStationVllmLifecycleDeps, + dualStationVllmApiKeyFingerprint, + dualStationVllmClusterId, + dualStationVllmLaunchContract, + getDualStationManagedVllmBaseUrl, + preflightDualStationGpuRuntime, + preflightDualStationManagedVllm, + rollbackDualStationLegacyMigration, + startDualStationManagedVllm, + withDualStationManagedVllmLifecycle, +} from "./vllm-station-cluster-lifecycle"; +import { withDualStationVllmLifecycleLock } from "./vllm-station-lifecycle-lock"; +import { + createDualStationSshBindingFixture, + type DualStationSshBindingFixture, +} from "./vllm-station-ssh-binding.test-support"; + +const WORKER_ID = "a".repeat(64); +const HEAD_ID = "b".repeat(64); +const LEGACY_HEAD_ID = "9".repeat(64); +const WORKER_SMOKE_ID = "c".repeat(64); +const HEAD_SMOKE_ID = "d".repeat(64); +const API_KEY = "e".repeat(64); +const START_CONFIG = { apiKey: API_KEY }; +const API_KEY_FINGERPRINT = dualStationVllmApiKeyFingerprint(API_KEY); +const TRANSACTION_ID = "1".repeat(32); +let sshFixture: DualStationSshBindingFixture; + +beforeEach(() => { + sshFixture = createDualStationSshBindingFixture(); +}); + +afterEach(() => { + sshFixture.cleanup(); +}); + +function fixturePlan(): DualStationVllmPlan { + return { + peerSshBinding: sshFixture.binding, + runtime: DUAL_STATION_VLLM_RUNTIME, + local: { + hostname: "station-a", + home: "/home/local", + uid: 1000, + gid: 1000, + gpu: { + index: 0, + name: "NVIDIA GB300", + uuid: "GPU-aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee", + }, + }, + peer: { + hostname: "station-b", + home: "/home/nvidia", + uid: 1001, + gid: 1001, + gpu: { + index: 1, + name: "NVIDIA GB300 Grace Blackwell Superchip", + uuid: "GPU-99999999-8888-7777-6666-555555555555", + }, + }, + rails: [ + { + index: 0, + subnet: "192.168.240.0/30", + local: { + rdmaDevice: "mlx5_0", + netdev: "cx8a0", + macAddress: "02:00:00:00:00:01", + uverbsDevice: "/dev/infiniband/uverbs0", + pciAddress: "0001:03:00.0", + address: "192.168.240.1", + }, + peer: { + rdmaDevice: "mlx5_0", + netdev: "cx8b0", + macAddress: "02:00:00:00:00:02", + uverbsDevice: "/dev/infiniband/uverbs0", + pciAddress: "0002:03:00.0", + address: "192.168.240.2", + }, + }, + { + index: 1, + subnet: "192.168.240.4/30", + local: { + rdmaDevice: "mlx5_1", + netdev: "cx8a1", + macAddress: "02:00:00:00:00:05", + uverbsDevice: "/dev/infiniband/uverbs1", + pciAddress: "0001:03:00.1", + address: "192.168.240.5", + }, + peer: { + rdmaDevice: "mlx5_1", + netdev: "cx8b1", + macAddress: "02:00:00:00:00:06", + uverbsDevice: "/dev/infiniband/uverbs1", + pciAddress: "0002:03:00.1", + address: "192.168.240.6", + }, + }, + ], + masterAddress: "192.168.240.1", + roceGidIndex: 3, + }; +} + +function fakeContainer( + role: "head" | "worker", + overrides: Partial = {}, +): FakeContainer { + return { + id: role === "head" ? HEAD_ID : WORKER_ID, + name: + role === "head" + ? DUAL_STATION_VLLM_HEAD_CONTAINER_NAME + : DUAL_STATION_VLLM_WORKER_CONTAINER_NAME, + state: "running", + image: DUAL_STATION_VLLM_RUNTIME.image, + labels: { + [DUAL_STATION_VLLM_MANAGED_LABEL]: "true", + [DUAL_STATION_VLLM_ROLE_LABEL]: role, + [DUAL_STATION_VLLM_ENDPOINT_LABEL]: + role === "head" ? "http://192.168.240.1:8000" : "headless", + [DUAL_STATION_VLLM_CLUSTER_LABEL]: dualStationVllmClusterId(fixturePlan()), + [DUAL_STATION_VLLM_GPU_LABEL]: + role === "head" ? fixturePlan().local.gpu.uuid : fixturePlan().peer.gpu.uuid, + [DUAL_STATION_VLLM_LAUNCH_SCHEMA_LABEL]: "2", + [DUAL_STATION_VLLM_LAUNCH_CONTRACT_LABEL]: dualStationVllmLaunchContract(fixturePlan(), role), + [DUAL_STATION_VLLM_API_KEY_FINGERPRINT_LABEL]: API_KEY_FINGERPRINT, + [DUAL_STATION_VLLM_TRANSACTION_LABEL]: TRANSACTION_ID, + }, + ...overrides, + }; +} + +function harness(options: HarnessOptions = {}) { + return createDualStationLifecycleHarness( + { + apiKey: API_KEY, + fakeContainer, + headSmokeId: HEAD_SMOKE_ID, + legacyHeadId: LEGACY_HEAD_ID, + plan: fixturePlan, + workerSmokeId: WORKER_SMOKE_ID, + }, + options, + ); +} + +type LifecycleHarness = ReturnType; + +function seedLegacyHead(fake: LifecycleHarness): void { + fake.seed( + "local", + fakeContainer("head", { + id: LEGACY_HEAD_ID, + image: + "vllm/vllm-openai@sha256:0fec7ec5f3e6bc168e54899935fb0557da908a4832a1dbc88e2debcf2f889416", + labels: { [DUAL_STATION_VLLM_MANAGED_LABEL]: "true" }, + }), + ); +} + +function expectRestoredLegacyHead(fake: LifecycleHarness): void { + expect(fake.containers.get(`local:${DUAL_STATION_VLLM_HEAD_CONTAINER_NAME}`)).toEqual([ + expect.objectContaining({ + id: LEGACY_HEAD_ID, + name: DUAL_STATION_VLLM_HEAD_CONTAINER_NAME, + state: "running", + }), + ]); + expect(fake.containers.get(`peer:${DUAL_STATION_VLLM_WORKER_CONTAINER_NAME}`) ?? []).toEqual([]); +} + +describe("dual-Station managed vLLM run argv", () => { + it("derives stable, distinct service-key fingerprints", () => { + expect(API_KEY_FINGERPRINT).toMatch(/^[a-f0-9]{64}$/u); + expect(dualStationVllmApiKeyFingerprint(API_KEY)).toBe(API_KEY_FINGERPRINT); + expect(dualStationVllmApiKeyFingerprint("f".repeat(64))).not.toBe(API_KEY_FINGERPRINT); + }); + + it("rejects a long slash-heavy runtime image", () => { + const plan = { + ...fixturePlan(), + runtime: { + ...DUAL_STATION_VLLM_RUNTIME, + image: `${"!/".repeat(10_000)}image@sha256:${"a".repeat(64)}`, + }, + } as unknown as DualStationVllmPlan; + + expect(() => + buildDualStationVllmRunArgs(plan, "head", TRANSACTION_ID, API_KEY_FINGERPRINT), + ).toThrow("exact pinned runtime contract"); + }); + + it.each(["head", "worker"] as const)("builds the exact %s launch contract", (role) => { + const plan = fixturePlan(); + const args = buildDualStationVllmRunArgs(plan, role, TRANSACTION_ID, API_KEY_FINGERPRINT); + const env = dockerValues(args, "--env"); + const expectedNode = role === "head" ? plan.local : plan.peer; + const expectedNetdev = role === "head" ? "cx8a0" : "cx8b0"; + + expect(args).toEqual( + expect.arrayContaining(["--network", "host", "--shm-size", "16g", "--read-only"]), + ); + expect( + args.some( + (arg) => + arg.startsWith("-p") || + arg.startsWith("-P") || + arg === "--publish" || + arg === "--publish-all" || + arg.startsWith("--publish=") || + arg.startsWith("--publish-all="), + ), + ).toBe(false); + expect(dockerValues(args, "--workdir")).toEqual(["/home/vllm"]); + expect(dockerValues(args, "--tmpfs")).toEqual([ + "/tmp:rw,nosuid,nodev,size=17179869184", + `/usr/local/lib/python3.12/dist-packages/flashinfer_cubin/cubins/flashinfer:rw,nosuid,nodev,noexec,uid=${String(expectedNode.uid)},gid=${String(expectedNode.gid)},mode=0700,size=16777216`, + `/home/vllm:rw,nosuid,nodev,exec,uid=${String(expectedNode.uid)},gid=${String(expectedNode.gid)},mode=0700,size=68719476736`, + ]); + expect(dockerValues(args, "--user")).toEqual([ + `${String(expectedNode.uid)}:${String(expectedNode.gid)}`, + ]); + expect(dockerValues(args, "--security-opt")).toEqual(["no-new-privileges:true"]); + expect(dockerValues(args, "--cap-drop")).toEqual(["ALL"]); + expect(dockerValues(args, "--cap-add")).toEqual([]); + expect(args).not.toContain("DAC_READ_SEARCH"); + expect(args).not.toContain("IPC_LOCK"); + expect(dockerValues(args, "--ulimit")).toEqual([ + "memlock=-1", + "stack=67108864", + "nofile=1048576:1048576", + ]); + expect(dockerValues(args, "--gpus")).toEqual([`device=${expectedNode.gpu.uuid}`]); + expect(args.filter((arg) => arg.startsWith("--device="))).toEqual([ + "--device=/dev/infiniband/uverbs0", + "--device=/dev/infiniband/uverbs1", + ]); + expect(dockerValues(args, "--volume")).toEqual([ + `${expectedNode.home}/.cache/huggingface/hub:/model-cache:ro`, + ]); + expect(dockerValues(args, "--volume").join("\n")).not.toContain("/huggingface/token"); + expect(env).toEqual( + expect.arrayContaining([ + "NCCL_IB_HCA=mlx5_0,mlx5_1", + "NCCL_IB_DISABLE=0", + "NCCL_IB_GID_INDEX=3", + `NCCL_SOCKET_IFNAME=${expectedNetdev}`, + `GLOO_SOCKET_IFNAME=${expectedNetdev}`, + `TP_SOCKET_IFNAME=${expectedNetdev}`, + `OMPI_MCA_btl_tcp_if_include=${expectedNetdev}`, + `MN_IF_NAME=${expectedNetdev}`, + "NCCL_IGNORE_CPU_AFFINITY=1", + "UCX_NET_DEVICES=mlx5_0:1,mlx5_1:1", + "UCX_IB_GID_INDEX=3", + "HF_HUB_OFFLINE=1", + "TRANSFORMERS_OFFLINE=1", + "HF_HOME=/home/vllm/.cache/huggingface", + "HF_HUB_CACHE=/model-cache", + "HUGGINGFACE_HUB_CACHE=/model-cache", + "HOME=/home/vllm", + "USER=vllm", + "LOGNAME=vllm", + ]), + ); + expect(args).toContain(plan.runtime.image); + expect(dockerValues(args, "--label")).toEqual( + expect.arrayContaining([ + `${DUAL_STATION_VLLM_LAUNCH_SCHEMA_LABEL}=2`, + `${DUAL_STATION_VLLM_LAUNCH_CONTRACT_LABEL}=${dualStationVllmLaunchContract(plan, role)}`, + `${DUAL_STATION_VLLM_API_KEY_FINGERPRINT_LABEL}=${API_KEY_FINGERPRINT}`, + `${DUAL_STATION_VLLM_TRANSACTION_LABEL}=${TRANSACTION_ID}`, + ]), + ); + expect(args).not.toContain("--privileged"); + expect(args.join("\n")).not.toContain("--device=/dev/infiniband:"); + expect(dockerValues(args, "--env").filter((name) => name === "VLLM_API_KEY")).toEqual( + role === "head" ? ["VLLM_API_KEY"] : [], + ); + expect(args).not.toContain(API_KEY); + expect(args.join("\n")).not.toMatch(/HF_TOKEN|HUGGING_FACE_HUB_TOKEN|docker run/u); + const command = args.at(-1) ?? ""; + const imageIndex = args.indexOf(plan.runtime.image); + expect(args.slice(imageIndex - 2, imageIndex + 2)).toEqual([ + "--entrypoint", + "/bin/bash", + plan.runtime.image, + "-lc", + ]); + expect(args.filter((arg) => arg === "-lc")).toHaveLength(1); + expect(imageIndex).toBe(args.length - 3); + expect(command).toContain('python3 -m pip install --user --no-cache-dir "ray==2.56.0"'); + expect(command).toContain( + role === "head" + ? "ray start --head --node-ip-address=192.168.240.1 --port=6379 --num-gpus=1" + : "ray start --address=192.168.240.1:6379 --node-ip-address=192.168.240.2 --num-gpus=1 --block", + ); + expect(command.includes("vllm serve")).toBe(role === "head"); + }); + + it.each(["head", "worker"] as const)("builds a bounded, no-pull %s GPU smoke command", (role) => { + const nonce = "1".repeat(32); + const { args, containerName } = buildDualStationGpuSmokeRunArgs(fixturePlan(), role, nonce); + + expect(containerName).toBe(`nemoclaw-vllm-gpu-smoke-${role}-${nonce}`); + expect(dockerValues(args, "--gpus")).toEqual([ + `device=${role === "head" ? fixturePlan().local.gpu.uuid : fixturePlan().peer.gpu.uuid}`, + ]); + expect(dockerValues(args, "--network")).toEqual(["none"]); + const expectedNode = role === "head" ? fixturePlan().local : fixturePlan().peer; + const expectedDevices = fixturePlan().rails.map((rail) => + role === "head" ? rail.local.uverbsDevice : rail.peer.uverbsDevice, + ); + expect(args).toContain("--read-only"); + expect(dockerValues(args, "--workdir")).toEqual(["/home/vllm"]); + expect(dockerValues(args, "--user")).toEqual([ + `${String(expectedNode.uid)}:${String(expectedNode.gid)}`, + ]); + expect(dockerValues(args, "--security-opt")).toEqual(["no-new-privileges:true"]); + expect(dockerValues(args, "--cap-drop")).toEqual(["ALL"]); + expect(dockerValues(args, "--cap-add")).toEqual([]); + expect(dockerValues(args, "--ulimit")).toEqual(["memlock=-1"]); + expect(dockerValues(args, "--device")).toEqual(expectedDevices); + expect(dockerValues(args, "--tmpfs")).toEqual([ + "/tmp:rw,nosuid,nodev,size=17179869184", + `/home/vllm:rw,nosuid,nodev,noexec,uid=${String(expectedNode.uid)},gid=${String(expectedNode.gid)},mode=0700,size=68719476736`, + ]); + expect(dockerValues(args, "--env")).toEqual( + expect.arrayContaining([ + "HF_HOME=/home/vllm/.cache/huggingface", + "HOME=/home/vllm", + "USER=vllm", + "LOGNAME=vllm", + ]), + ); + expect(dockerValues(args, "--label")).toEqual( + expect.arrayContaining([ + `${DUAL_STATION_VLLM_GPU_SMOKE_LABEL}=${nonce}`, + `${DUAL_STATION_VLLM_ROLE_LABEL}=${role}`, + ]), + ); + expect(args).toContain("--pull=never"); + expect(args).toContain(DUAL_STATION_VLLM_RUNTIME.image); + expect(args).toEqual(expect.arrayContaining(["--entrypoint", "/bin/bash"])); + expect(args.slice(-3, -1)).toEqual([DUAL_STATION_VLLM_RUNTIME.image, "-c"]); + const command = args.at(-1) ?? ""; + expect(command).toContain("NoNewPrivs"); + expect(command).toContain("Cap(Inh|Prm|Eff|Bnd|Amb)"); + expect(command).toContain('test "$(ulimit -l)" = "unlimited"'); + expect(command).toContain( + `for device in ${expectedDevices.join(" ")}; do test -c "$device"; test -r "$device"; test -w "$device"; exec 3<>"$device"; exec 3>&-; done`, + ); + expect(command).toContain("$HOME/.cache/torch/.nemoclaw-write-probe"); + expect(command).toContain("$HF_HOME/.nemoclaw-write-probe"); + expect(command).toContain("exec nvidia-smi --query-gpu=uuid --format=csv,noheader"); + expect(dockerValues(args, "--volume")).toEqual([]); + expect(args).not.toContain("--rm"); + expect(args.join("\n")).not.toContain("VLLM_API_KEY"); + }); + + it("binds the managed launch contract to both runtime owner IDs", () => { + const baseline = fixturePlan(); + const changedUid = fixturePlan(); + const changedGid = fixturePlan(); + changedUid.local.uid += 1; + changedGid.local.gid += 1; + + expect(dualStationVllmLaunchContract(changedUid, "head")).not.toBe( + dualStationVllmLaunchContract(baseline, "head"), + ); + expect(dualStationVllmLaunchContract(changedGid, "head")).not.toBe( + dualStationVllmLaunchContract(baseline, "head"), + ); + }); + + it.each([ + ["root uid", (plan: DualStationVllmPlan) => (plan.local.uid = 0)], + ["root gid", (plan: DualStationVllmPlan) => (plan.peer.gid = 0)], + ])("rejects an unsafe %s runtime identity", (_label, mutate) => { + const plan = fixturePlan(); + mutate(plan); + + expect(() => + buildDualStationVllmRunArgs(plan, "head", TRANSACTION_ID, API_KEY_FINGERPRINT), + ).toThrow("runtime identity must use non-root UID and GID values"); + }); + + it.each([ + "/dev/infiniband/rdma_cm", + "/dev/infiniband/uverbs0", + ])("rejects an unsafe or duplicate verbs character device: %s", (uverbsDevice) => { + const plan = fixturePlan(); + plan.rails[1].local.uverbsDevice = uverbsDevice; + + expect(() => + buildDualStationVllmRunArgs(plan, "head", TRANSACTION_ID, API_KEY_FINGERPRINT), + ).toThrow(/rails must use two distinct devices|rail endpoint is invalid/u); + }); +}); + +describe("dual-Station managed vLLM lifecycle", () => { + it("rejects an effective account that differs from the prepared controller", () => { + const fake = harness(); + + expect( + preflightDualStationManagedVllm(fixturePlan(), { + ...fake.deps, + effectiveControllerUid: () => fixturePlan().local.uid + 1, + }), + ).toEqual({ + ok: false, + reason: + "Dual-Station lifecycle effective UID 1001 does not match prepared controller UID 1000", + }); + expect(fake.operations).toEqual([]); + }); + + it("rejects a prepared controller account that does not own the probed local Station plan", () => { + const fake = harness(); + + expect( + preflightDualStationManagedVllm(fixturePlan(), { + ...fake.deps, + effectiveControllerUid: () => fixturePlan().local.uid + 1, + readControllerUid: () => fixturePlan().local.uid + 1, + }), + ).toEqual({ + ok: false, + reason: "Dual-Station lifecycle controller UID must match probed local UID 1000", + }); + expect(fake.operations).toEqual([]); + }); + + it("anchors the default lock under the effective account home instead of mutable HOME", async () => { + const root = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-station-lock-home-")); + const accountHome = path.join(root, "account-home"); + const ambientHome = path.join(root, "ambient-home"); + fs.mkdirSync(accountHome, { mode: 0o700 }); + const userInfo = os.userInfo(); + const userInfoSpy = vi.spyOn(os, "userInfo").mockReturnValue({ + ...userInfo, + homedir: accountHome, + }); + vi.stubEnv("HOME", ambientHome); + try { + await withDualStationVllmLifecycleLock( + () => { + expect( + fs.existsSync(path.join(accountHome, ".nemoclaw", "state", "mcp-lifecycle-locks")), + ).toBe(true); + expect(fs.existsSync(ambientHome)).toBe(false); + }, + { pollIntervalMs: 5, timeoutMs: 250, corruptLockGraceMs: 5 }, + { + readControllerUid: () => userInfo.uid, + effectiveControllerUid: () => userInfo.uid, + }, + ); + } finally { + vi.unstubAllEnvs(); + userInfoSpy.mockRestore(); + fs.rmSync(root, { recursive: true, force: true }); + } + }); + + it("provides a read-only ownership preflight before download work", () => { + const fake = harness(); + + expect(preflightDualStationManagedVllm(fixturePlan(), fake.deps)).toEqual({ ok: true }); + expect(fake.operations.filter((operation) => operation.kind === "capture")).toHaveLength(2); + expect(fake.operations.some((operation) => operation.kind === "run")).toBe(false); + expect(fake.operations.some((operation) => operation.kind === "rm")).toBe(false); + }); + + it("rejects a changed host-key pin before reading or mutating either Docker daemon", () => { + const fake = harness(); + fs.appendFileSync(sshFixture.binding.knownHostsFile, "changed\n"); + + expect(preflightDualStationManagedVllm(fixturePlan(), fake.deps)).toMatchObject({ + ok: false, + reason: expect.stringContaining("known-hosts binding changed"), + }); + expect(fake.operations).toEqual([]); + expect(fake.buildRemoteDockerEnv).not.toHaveBeenCalled(); + }); + + it("proves exact-image GPU execution on both daemons and removes both exact probe IDs", async () => { + const fake = harness(); + + expect(await preflightDualStationGpuRuntime(fixturePlan(), fake.deps)).toEqual({ ok: true }); + expect(fake.operations.filter((operation) => operation.kind === "run")).toHaveLength(2); + expect(fake.operations.filter((operation) => operation.kind === "rm")).toEqual([ + { kind: "rm", target: "peer", value: WORKER_SMOKE_ID }, + { kind: "rm", target: "local", value: HEAD_SMOKE_ID }, + ]); + expect( + fake.operations.filter( + (operation) => operation.kind === "capture" && operation.value.startsWith("image:"), + ), + ).toHaveLength(2); + for (const call of fake.runCalls) { + expect(call.args).toContain("--pull=never"); + expect(call.args).toContain(DUAL_STATION_VLLM_RUNTIME.image); + expect(call.options?.env?.VLLM_API_KEY).toBeUndefined(); + } + for (const options of [...fake.captureOptions, ...fake.rmOptions]) { + expect(options?.env?.VLLM_API_KEY).toBeUndefined(); + } + }); + + it("does not mutate either daemon unless both exact pinned images are present", async () => { + const fake = harness({ missingImageTarget: "local" }); + + expect(await preflightDualStationGpuRuntime(fixturePlan(), fake.deps)).toEqual({ + ok: false, + reason: "head pinned vLLM image is not present or could not be inspected", + }); + expect(fake.operations.some((operation) => operation.kind === "run")).toBe(false); + expect(fake.operations.some((operation) => operation.kind === "rm")).toBe(false); + }); + + it("removes an exact failed GPU probe and never starts the managed containers", async () => { + const fake = harness({ smokeGpuOutput: { peer: "GPU-not-the-discovered-device" } }); + + expect(await startDualStationManagedVllm(fixturePlan(), START_CONFIG, fake.deps)).toEqual({ + ok: false, + reason: "worker GPU smoke did not expose exactly the discovered GPU", + rollbackErrors: [], + }); + expect(fake.operations.filter((operation) => operation.kind === "run")).toHaveLength(1); + expect(fake.operations.filter((operation) => operation.kind === "rm")).toEqual([ + { kind: "rm", target: "peer", value: WORKER_SMOKE_ID }, + ]); + }); + + it.each([ + "short", + "A".repeat(64), + ])("rejects an unsafe API key before probing: %s", async (apiKey) => { + const fake = harness(); + + expect(await startDualStationManagedVllm(fixturePlan(), { apiKey }, fake.deps)).toMatchObject({ + ok: false, + reason: expect.stringContaining("64 lowercase hexadecimal"), + }); + expect(fake.operations).toEqual([]); + }); + + it("starts the worker before the head and returns validated exact IDs", async () => { + const fake = harness(); + + expect(await startDualStationManagedVllm(fixturePlan(), START_CONFIG, fake.deps)).toEqual({ + ok: true, + baseUrl: "http://192.168.240.1:8000", + headContainerId: HEAD_ID, + workerContainerId: WORKER_ID, + reusedExisting: false, + }); + expect(fake.operations.filter((operation) => operation.kind === "run")).toEqual([ + { + kind: "run", + target: "peer", + value: `nemoclaw-vllm-gpu-smoke-worker-${"1".padStart(32, "0")}`, + }, + { + kind: "run", + target: "local", + value: `nemoclaw-vllm-gpu-smoke-head-${"2".padStart(32, "0")}`, + }, + { kind: "run", target: "peer", value: DUAL_STATION_VLLM_WORKER_CONTAINER_NAME }, + { kind: "run", target: "local", value: DUAL_STATION_VLLM_HEAD_CONTAINER_NAME }, + ]); + const headRun = fake.runCalls.find( + ({ args }) => dockerValues(args, "--name")[0] === DUAL_STATION_VLLM_HEAD_CONTAINER_NAME, + ); + expect(headRun?.options?.env?.VLLM_API_KEY).toBe(API_KEY); + expect(headRun?.args).toContain("VLLM_API_KEY"); + expect(headRun?.args).not.toContain(API_KEY); + for (const call of fake.runCalls.filter((call) => call !== headRun)) { + expect(call.options?.env?.VLLM_API_KEY).toBeUndefined(); + expect(call.args).not.toContain(API_KEY); + } + for (const options of [...fake.captureOptions, ...fake.rmOptions]) { + expect(options?.env?.VLLM_API_KEY).toBeUndefined(); + } + expect(fake.buildRemoteDockerEnv).toHaveBeenCalledWith(sshFixture.binding); + }); + + it("reuses an already-running exact pair without tearing down the working service", async () => { + const fake = harness(); + fake.seed("local", fakeContainer("head")); + fake.seed("peer", fakeContainer("worker")); + + expect(await startDualStationManagedVllm(fixturePlan(), START_CONFIG, fake.deps)).toEqual({ + ok: true, + baseUrl: "http://192.168.240.1:8000", + headContainerId: HEAD_ID, + workerContainerId: WORKER_ID, + reusedExisting: true, + }); + expect(fake.operations.filter((operation) => operation.kind === "rm")).toEqual([ + { kind: "rm", target: "peer", value: WORKER_SMOKE_ID }, + { kind: "rm", target: "local", value: HEAD_SMOKE_ID }, + ]); + }); + + it.each([ + [DUAL_STATION_VLLM_LAUNCH_CONTRACT_LABEL, "f".repeat(64)], + [DUAL_STATION_VLLM_API_KEY_FINGERPRINT_LABEL, "d".repeat(64)], + ])("recreates an owned pair whose %s no longer matches", async (label, value) => { + const fake = harness(); + const head = fakeContainer("head"); + const worker = fakeContainer("worker"); + head.labels[label] = value; + worker.labels[label] = value; + fake.seed("local", head); + fake.seed("peer", worker); + + expect(await startDualStationManagedVllm(fixturePlan(), START_CONFIG, fake.deps)).toMatchObject( + { + ok: true, + reusedExisting: false, + }, + ); + expect(fake.operations.filter((operation) => operation.kind === "rm")).toEqual( + expect.arrayContaining([ + { kind: "rm", target: "local", value: HEAD_ID }, + { kind: "rm", target: "peer", value: WORKER_ID }, + ]), + ); + }); + + it("recreates a mixed pair from different lifecycle transactions", async () => { + const fake = harness(); + const worker = fakeContainer("worker"); + worker.labels[DUAL_STATION_VLLM_TRANSACTION_LABEL] = "f".repeat(32); + fake.seed("local", fakeContainer("head")); + fake.seed("peer", worker); + + expect(await startDualStationManagedVllm(fixturePlan(), START_CONFIG, fake.deps)).toMatchObject( + { + ok: true, + reusedExisting: false, + }, + ); + expect(fake.operations.filter((operation) => operation.kind === "rm")).toEqual( + expect.arrayContaining([ + { kind: "rm", target: "local", value: HEAD_ID }, + { kind: "rm", target: "peer", value: WORKER_ID }, + ]), + ); + }); + + it("serializes concurrent same-plan starts so only one worker is created", async () => { + const fake = harness(); + + const results = await Promise.all([ + startDualStationManagedVllm(fixturePlan(), START_CONFIG, fake.deps), + startDualStationManagedVllm(fixturePlan(), START_CONFIG, fake.deps), + ]); + + expect(results).toEqual([ + expect.objectContaining({ ok: true, reusedExisting: false }), + expect.objectContaining({ ok: true, reusedExisting: true }), + ]); + expect( + fake.operations.filter( + (operation) => + operation.kind === "run" && operation.value === DUAL_STATION_VLLM_WORKER_CONTAINER_NAME, + ), + ).toHaveLength(1); + expect(fake.getMaxLifecycleLockActive()).toBe(1); + }); + + it("holds the lifecycle lease after start until validation commits", async () => { + const fake = harness(); + let releaseValidation: () => void = () => undefined; + let reportStarted: () => void = () => undefined; + const validationGate = new Promise((resolve) => { + releaseValidation = resolve; + }); + const started = new Promise((resolve) => { + reportStarted = resolve; + }); + const first = withDualStationManagedVllmLifecycle(async () => { + const result = await startDualStationManagedVllm(fixturePlan(), START_CONFIG, fake.deps); + reportStarted(); + await validationGate; + return result; + }, fake.deps); + await started; + + const second = startDualStationManagedVllm(fixturePlan(), START_CONFIG, fake.deps); + await Promise.resolve(); + expect( + fake.operations.filter( + (operation) => + operation.kind === "run" && operation.value === DUAL_STATION_VLLM_WORKER_CONTAINER_NAME, + ), + ).toHaveLength(1); + + releaseValidation(); + expect(await Promise.all([first, second])).toEqual([ + expect.objectContaining({ ok: true, reusedExisting: false }), + expect.objectContaining({ ok: true, reusedExisting: true }), + ]); + expect(fake.getMaxLifecycleLockActive()).toBe(1); + }); + + it("reconciles a late worker create and rolls back only its transaction", async () => { + const fake = harness({ lateCreateRole: "worker" }); + + expect(await startDualStationManagedVllm(fixturePlan(), START_CONFIG, fake.deps)).toEqual({ + ok: false, + reason: "worker container failed to start", + rollbackErrors: [], + }); + expect(fake.operations.filter((operation) => operation.kind === "rm")).toContainEqual({ + kind: "rm", + target: "peer", + value: WORKER_ID, + }); + }); + + it("never rolls back an ambiguous worker labeled by another transaction", async () => { + const fake = harness({ failedRoleForeignTransaction: "worker" }); + + expect(await startDualStationManagedVllm(fixturePlan(), START_CONFIG, fake.deps)).toEqual({ + ok: false, + reason: "worker container failed to start", + rollbackErrors: [], + }); + expect(fake.operations.filter((operation) => operation.kind === "rm")).not.toContainEqual({ + kind: "rm", + target: "peer", + value: WORKER_ID, + }); + expect(fake.containers.get(`peer:${DUAL_STATION_VLLM_WORKER_CONTAINER_NAME}`)).toHaveLength(1); + }); + + it("migrates only the exact running legacy single-Station head from the rollback window", async () => { + const fake = harness(); + seedLegacyHead(fake); + const started = await startDualStationManagedVllm(fixturePlan(), START_CONFIG, fake.deps); + expect(started).toMatchObject({ + ok: true, + reusedExisting: false, + legacyMigration: expect.objectContaining({ legacyContainerId: LEGACY_HEAD_ID }), + }); + const legacyMigration = requireLegacyMigration(started); + expect( + fake.operations.some(({ kind, value }) => kind === "rm" && value === LEGACY_HEAD_ID), + ).toBe(false); + await expect( + commitDualStationLegacyMigration(fixturePlan(), legacyMigration, fake.deps), + ).resolves.toEqual({ ok: true, cleanupWarnings: [] }); + const cutoverOrder = fake.operations + .filter( + ({ kind, value }) => + (kind === "run" && + (value === DUAL_STATION_VLLM_WORKER_CONTAINER_NAME || + value === DUAL_STATION_VLLM_HEAD_CONTAINER_NAME)) || + kind === "rename" || + ((kind === "stop" || kind === "rm") && value === LEGACY_HEAD_ID), + ) + .map(({ kind, value }) => `${kind}:${value}`); + expect(cutoverOrder).toEqual([ + `run:${DUAL_STATION_VLLM_WORKER_CONTAINER_NAME}`, + expect.stringMatching(`^rename:${LEGACY_HEAD_ID}:nemoclaw-vllm-legacy-`), + `stop:${LEGACY_HEAD_ID}`, + `run:${DUAL_STATION_VLLM_HEAD_CONTAINER_NAME}`, + `rm:${LEGACY_HEAD_ID}`, + ]); + }); + + it("restores the preserved legacy head when external validation rolls back", async () => { + const fake = harness(); + seedLegacyHead(fake); + const started = await startDualStationManagedVllm(fixturePlan(), START_CONFIG, fake.deps); + const legacyMigration = requireLegacyMigration(started); + + await expect( + rollbackDualStationLegacyMigration(fixturePlan(), legacyMigration, fake.deps), + ).resolves.toEqual({ ok: true }); + expectRestoredLegacyHead(fake); + }); + + it("keeps the running legacy head untouched when the peer worker cannot start", async () => { + const fake = harness({ failRole: "worker" }); + seedLegacyHead(fake); + expect(await startDualStationManagedVllm(fixturePlan(), START_CONFIG, fake.deps)).toEqual({ + ok: false, + reason: "worker container failed to start", + rollbackErrors: [], + }); + expectRestoredLegacyHead(fake); + expect(fake.operations.some(({ kind }) => ["rename", "start", "stop"].includes(kind))).toBe( + false, + ); + expect( + fake.operations.some(({ kind, value }) => kind === "rm" && value === LEGACY_HEAD_ID), + ).toBe(false); + }); + + it.each([ + ["new head launch", { failRole: "head" }, "head container failed to start", false], + [ + "final pair verification", + { failFinalInspectionRole: "head" }, + "dual-Station containers did not remain running", + true, + ], + ] as const)("restores the exact legacy head after %s failure", async (_case, options, reason, ranHead) => { + const fake = harness(options); + seedLegacyHead(fake); + expect(await startDualStationManagedVllm(fixturePlan(), START_CONFIG, fake.deps)).toEqual({ + ok: false, + reason, + rollbackErrors: [], + }); + expectRestoredLegacyHead(fake); + expect(fake.operations).toEqual( + expect.arrayContaining([ + { kind: "rm", target: "peer", value: WORKER_ID }, + { kind: "start", target: "local", value: LEGACY_HEAD_ID }, + ]), + ); + expect(fake.operations.some(({ kind, value }) => kind === "rm" && value === HEAD_ID)).toBe( + ranHead, + ); + expect( + fake.operations.some(({ kind, value }) => kind === "rm" && value === LEGACY_HEAD_ID), + ).toBe(false); + }); + + it("keeps the validated new pair when legacy backup removal is ambiguous", async () => { + const fake = harness({ failLegacyBackupRemoval: true }); + seedLegacyHead(fake); + const started = await startDualStationManagedVllm(fixturePlan(), START_CONFIG, fake.deps); + const legacyMigration = requireLegacyMigration(started); + await expect( + commitDualStationLegacyMigration(fixturePlan(), legacyMigration, fake.deps), + ).resolves.toMatchObject({ + ok: true, + cleanupWarnings: [expect.stringContaining("legacy backup")], + }); + expect(fake.containers.get(`local:${DUAL_STATION_VLLM_HEAD_CONTAINER_NAME}`)).toEqual([ + expect.objectContaining({ id: HEAD_ID, state: "running" }), + ]); + expect(fake.containers.get(`peer:${DUAL_STATION_VLLM_WORKER_CONTAINER_NAME}`)).toEqual([ + expect.objectContaining({ id: WORKER_ID, state: "running" }), + ]); + const preservedBackup = [...fake.containers.entries()].flatMap(([containerKey, entries]) => + containerKey.startsWith(`local:${DUAL_STATION_VLLM_HEAD_CONTAINER_NAME}-legacy-`) + ? entries + : [], + ); + expect(preservedBackup).toEqual([ + expect.objectContaining({ id: LEGACY_HEAD_ID, state: "exited" }), + ]); + }); + + it.each([ + ["stopped", { state: "exited" }], + [ + "outside the frozen image window", + { + image: + "vllm/vllm-openai@sha256:2222222222222222222222222222222222222222222222222222222222222222", + }, + ], + ])("refuses a schema-less managed head that is %s", async (_case, override) => { + const fake = harness(); + const plan = fixturePlan(); + fake.seed( + "local", + fakeContainer("head", { + ...override, + labels: { [DUAL_STATION_VLLM_MANAGED_LABEL]: "true" }, + }), + ); + + expect(await startDualStationManagedVllm(plan, START_CONFIG, fake.deps)).toMatchObject({ + ok: false, + reason: expect.stringContaining("foreign"), + }); + expect(fake.operations.some((operation) => operation.kind === "rm")).toBe(false); + expect(fake.operations.some((operation) => operation.kind === "run")).toBe(false); + }); + + it("refuses a same-image worker that belongs to another physical cluster plan", () => { + const fake = harness(); + const unrelatedWorker = fakeContainer("worker"); + unrelatedWorker.labels[DUAL_STATION_VLLM_CLUSTER_LABEL] = "f".repeat(64); + fake.seed("peer", unrelatedWorker); + + expect(preflightDualStationManagedVllm(fixturePlan(), fake.deps)).toMatchObject({ + ok: false, + reason: expect.stringContaining("foreign"), + }); + expect(fake.operations.some((operation) => operation.kind === "rm")).toBe(false); + }); + + it("fails before mutation when an exact name has foreign ownership", async () => { + const fake = harness(); + fake.seed( + "local", + fakeContainer("head", { + labels: { + [DUAL_STATION_VLLM_MANAGED_LABEL]: "false", + [DUAL_STATION_VLLM_ROLE_LABEL]: "head", + [DUAL_STATION_VLLM_ENDPOINT_LABEL]: "http://192.168.240.1:8000", + }, + }), + ); + + expect(await startDualStationManagedVllm(fixturePlan(), START_CONFIG, fake.deps)).toMatchObject( + { + ok: false, + reason: expect.stringContaining("foreign"), + }, + ); + expect(fake.operations.some((operation) => operation.kind === "run")).toBe(false); + expect(fake.operations.some((operation) => operation.kind === "rm")).toBe(false); + }); + + it("rolls back the exact worker ID when the head fails", async () => { + const fake = harness({ failRole: "head" }); + + expect(await startDualStationManagedVllm(fixturePlan(), START_CONFIG, fake.deps)).toEqual({ + ok: false, + reason: "head container failed to start", + rollbackErrors: [], + }); + expect(fake.operations.filter((operation) => operation.kind === "rm")).toEqual([ + { kind: "rm", target: "peer", value: WORKER_SMOKE_ID }, + { kind: "rm", target: "local", value: HEAD_SMOKE_ID }, + { kind: "rm", target: "peer", value: WORKER_ID }, + ]); + }); + + it("uses the real reentrant lease across an outer lifecycle and start rollback", async () => { + const fake = harness({ failRole: "head" }); + const stateDir = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-station-real-lock-")); + const withRealLock: DualStationVllmLifecycleDeps["withLifecycleLock"] = (operation) => + withDualStationVllmLifecycleLock( + operation, + { + stateDir, + pollIntervalMs: 5, + timeoutMs: 250, + corruptLockGraceMs: 5, + }, + { + readControllerUid: () => fixturePlan().local.uid, + effectiveControllerUid: () => fixturePlan().local.uid, + }, + ); + const deps = { ...fake.deps, withLifecycleLock: withRealLock }; + try { + expect( + await withDualStationManagedVllmLifecycle( + () => startDualStationManagedVllm(fixturePlan(), START_CONFIG, deps), + deps, + ), + ).toEqual({ + ok: false, + reason: "head container failed to start", + rollbackErrors: [], + }); + expect(fake.operations.filter((operation) => operation.kind === "rm")).toContainEqual({ + kind: "rm", + target: "peer", + value: WORKER_ID, + }); + } finally { + fs.rmSync(stateDir, { recursive: true, force: true }); + } + }, 2_000); + + it("rejects an invalid docker-run ID and recovers only its exact owned ID", async () => { + const fake = harness({ invalidIdRole: "worker" }); + + expect(await startDualStationManagedVllm(fixturePlan(), START_CONFIG, fake.deps)).toEqual({ + ok: false, + reason: "worker container failed to start", + rollbackErrors: [], + }); + expect(fake.operations.filter((operation) => operation.kind === "rm")).toEqual([ + { kind: "rm", target: "peer", value: WORKER_SMOKE_ID }, + { kind: "rm", target: "local", value: HEAD_SMOKE_ID }, + { kind: "rm", target: "peer", value: WORKER_ID }, + ]); + }); + + it("cleans up exact owned IDs head-first and reports both-running state", async () => { + const fake = harness(); + fake.seed("local", fakeContainer("head")); + fake.seed("peer", fakeContainer("worker")); + + expect(areDualStationManagedVllmContainersRunning(fixturePlan(), fake.deps)).toBe(true); + expect(await cleanupDualStationManagedVllm(fixturePlan(), fake.deps)).toEqual({ + ok: true, + removedContainerIds: [HEAD_ID, WORKER_ID], + }); + expect(fake.operations.filter((operation) => operation.kind === "rm")).toEqual([ + { kind: "rm", target: "local", value: HEAD_ID }, + { kind: "rm", target: "peer", value: WORKER_ID }, + ]); + expect(areDualStationManagedVllmContainersRunning(fixturePlan(), fake.deps)).toBe(false); + }); + + it("refuses all cleanup when either exact name is ambiguous", async () => { + const fake = harness(); + fake.seed("local", fakeContainer("head")); + fake.seed("local", fakeContainer("head", { id: "c".repeat(64) })); + fake.seed("peer", fakeContainer("worker")); + + expect(await cleanupDualStationManagedVllm(fixturePlan(), fake.deps)).toMatchObject({ + ok: false, + reason: expect.stringContaining("ambiguous"), + }); + expect(fake.operations.some((operation) => operation.kind === "rm")).toBe(false); + }); +}); + +describe("managed dual-Station base URL recovery", () => { + it("returns only a running, owned RFC1918 head endpoint with a bounded inspect", () => { + const fake = harness(); + fake.seed("local", fakeContainer("head")); + + expect(getDualStationManagedVllmBaseUrl(fake.deps)).toBe("http://192.168.240.1:8000"); + expect( + fake.operations.some( + (operation) => + operation.kind === "capture" && + operation.target === "local" && + operation.value === DUAL_STATION_VLLM_HEAD_CONTAINER_NAME, + ), + ).toBe(true); + expect(fake.captureOptions.at(-1)).toMatchObject({ timeout: 10_000 }); + expect(fake.captureOptions.at(-1)?.env?.VLLM_API_KEY).toBeUndefined(); + }); + + it("reports a structurally managed running head before API-key fingerprint validation", () => { + const fake = harness(); + const head = fakeContainer("head"); + head.labels[DUAL_STATION_VLLM_API_KEY_FINGERPRINT_LABEL] = "invalid"; + fake.seed("local", head); + const onManagedHeadObserved = vi.fn(); + const loadApiKey = vi.fn(() => API_KEY); + + expect( + getDualStationManagedVllmBaseUrl({ + ...fake.deps, + onManagedHeadObserved, + loadApiKey, + }), + ).toBeNull(); + expect(onManagedHeadObserved).toHaveBeenCalledOnce(); + expect(loadApiKey).not.toHaveBeenCalled(); + }); + + it.each([ + ["missing persisted key", { loadApiKey: () => null }], + ["mismatched persisted key", { loadApiKey: () => "a".repeat(64) }], + ["endpoint absent from local interfaces", { localInterfaceAddresses: () => [] }], + ])("rejects day-2 recovery with %s", (_case, overrides) => { + const fake = harness(); + fake.seed("local", fakeContainer("head")); + + expect(getDualStationManagedVllmBaseUrl({ ...fake.deps, ...overrides })).toBeNull(); + }); + + it.each([ + "http://8.8.8.8:8000", + "http://0.0.0.0:8000", + "http://192.168.240.1:8000/", + ])("rejects unsafe or non-canonical endpoint label %s", (endpoint) => { + const fake = harness(); + fake.seed( + "local", + fakeContainer("head", { + labels: { + [DUAL_STATION_VLLM_MANAGED_LABEL]: "true", + [DUAL_STATION_VLLM_ROLE_LABEL]: "head", + [DUAL_STATION_VLLM_ENDPOINT_LABEL]: endpoint, + }, + }), + ); + + expect(getDualStationManagedVllmBaseUrl(fake.deps)).toBeNull(); + }); + + it("rejects an owned-looking head that does not use the pinned runtime image", () => { + const fake = harness(); + const onManagedHeadObserved = vi.fn(); + fake.seed( + "local", + fakeContainer("head", { + image: + "nvcr.io/nvidia/vllm:forged@sha256:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + }), + ); + + expect(getDualStationManagedVllmBaseUrl({ ...fake.deps, onManagedHeadObserved })).toBeNull(); + expect(onManagedHeadObserved).not.toHaveBeenCalled(); + }); +}); diff --git a/src/lib/inference/vllm-station-cluster-lifecycle.ts b/src/lib/inference/vllm-station-cluster-lifecycle.ts new file mode 100644 index 0000000000..8c0f9b3f6d --- /dev/null +++ b/src/lib/inference/vllm-station-cluster-lifecycle.ts @@ -0,0 +1,1790 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { createHash, createHmac, randomBytes } from "node:crypto"; +import net from "node:net"; +import os from "node:os"; +import path from "node:path"; + +import { dockerCapture, dockerForceRm, dockerRunDetached } from "../adapters/docker"; +import { DUAL_STATION_VLLM_API_KEY_PATTERN, loadDualStationVllmApiKey } from "./vllm-api-key"; +import { buildLocalDualStationDockerEnv, buildRemoteVllmDockerEnv } from "./vllm-docker-env"; +import { buildNemotronUltraDistributedServeCommand } from "./vllm-models"; +import { DUAL_STATION_VLLM_RUNTIME, type DualStationVllmPlan } from "./vllm-station-cluster"; +import { + assertDualStationControllerAccount, + readDualStationControllerUid, + withDualStationVllmLifecycleLock, +} from "./vllm-station-lifecycle-lock"; +import { + assertDualStationSshBindingFiles, + type DualStationSshBinding, +} from "./vllm-station-ssh-binding"; + +export const DUAL_STATION_VLLM_HEAD_CONTAINER_NAME = "nemoclaw-vllm"; +export const DUAL_STATION_VLLM_WORKER_CONTAINER_NAME = "nemoclaw-vllm-worker"; +export const DUAL_STATION_VLLM_MANAGED_LABEL = "com.nvidia.nemoclaw.managed-vllm"; +export const DUAL_STATION_VLLM_ROLE_LABEL = "com.nvidia.nemoclaw.vllm-role"; +export const DUAL_STATION_VLLM_ENDPOINT_LABEL = "com.nvidia.nemoclaw.vllm-endpoint"; +export const DUAL_STATION_VLLM_CLUSTER_LABEL = "com.nvidia.nemoclaw.vllm-cluster"; +export const DUAL_STATION_VLLM_GPU_LABEL = "com.nvidia.nemoclaw.vllm-gpu"; +export const DUAL_STATION_VLLM_LAUNCH_SCHEMA_LABEL = "com.nvidia.nemoclaw.vllm-launch-schema"; +export const DUAL_STATION_VLLM_LAUNCH_CONTRACT_LABEL = "com.nvidia.nemoclaw.vllm-launch-contract"; +export const DUAL_STATION_VLLM_API_KEY_FINGERPRINT_LABEL = + "com.nvidia.nemoclaw.vllm-api-key-fingerprint"; +export const DUAL_STATION_VLLM_TRANSACTION_LABEL = "com.nvidia.nemoclaw.vllm-transaction"; +export const DUAL_STATION_VLLM_GPU_SMOKE_LABEL = "com.nvidia.nemoclaw.gpu-smoke"; +export const DUAL_STATION_VLLM_MASTER_PORT = 6379; + +const HEAD_API_PORT = 8000; +// The pinned Station vLLM images ship a non-root-ready /home/vllm. Each Station +// mounts an owner-only tmpfs there and runs as the probed model-cache owner. +const VLLM_RUNTIME_HOME = "/home/vllm"; +const HF_CACHE_CONTAINER_DIR = `${VLLM_RUNTIME_HOME}/.cache/huggingface`; +const HF_HUB_CACHE_CONTAINER_DIR = "/model-cache"; +const FLASHINFER_CUBIN_RUNTIME_DIR = + "/usr/local/lib/python3.12/dist-packages/flashinfer_cubin/cubins/flashinfer"; +const DOCKER_INSPECT_TIMEOUT_MS = 10_000; +const DOCKER_MUTATION_TIMEOUT_MS = 60_000; +const DOCKER_GPU_SMOKE_TIMEOUT_MS = 30_000; +const DOCKER_LATE_CREATE_RECONCILE_ATTEMPTS = 5; +const DOCKER_LATE_CREATE_RECONCILE_INTERVAL_MS = 250; +const DOCKER_CONTAINER_ID_PATTERN = /^[a-f0-9]{64}$/; +const CLUSTER_ID_PATTERN = /^[a-f0-9]{64}$/; +const SHA256_HEX_PATTERN = /^[a-f0-9]{64}$/; +const IMMUTABLE_IMAGE_PATTERN = /^(?:[^\s/@]+\/)+[^\s/@]+@sha256:[a-f0-9]{64}$/; +const IMAGE_ID_PATTERN = /^sha256:[a-f0-9]{64}$/; +const SAFE_DEVICE_NAME_PATTERN = /^[A-Za-z0-9][A-Za-z0-9_.:-]{0,63}$/; +const SAFE_UVERBS_DEVICE_PATTERN = /^\/dev\/infiniband\/uverbs[0-9]+$/; +const SAFE_GPU_UUID_PATTERN = /^GPU-[A-Za-z0-9-]{8,123}$/; +const GPU_SMOKE_NONCE_PATTERN = /^[a-f0-9]{32}$/; +const TRANSACTION_ID_PATTERN = /^[a-f0-9]{32}$/; +const GPU_SMOKE_CONTAINER_PREFIX = "nemoclaw-vllm-gpu-smoke"; +const DUAL_STATION_VLLM_LAUNCH_SCHEMA = "2"; +const VLLM_FINGERPRINT_CONTEXT = "nemoclaw-dual-station-vllm-api-key\0"; +// Compatibility bridge for schema-less single-Station Ultra containers from +// the v0.0.86 rollback window before dual launch schema 1. +// Do not add future image digests; retire this branch when v0.0.86 support ends. +const LEGACY_SINGLE_STATION_MIGRATION_IMAGE = + "vllm/vllm-openai@sha256:0fec7ec5f3e6bc168e54899935fb0557da908a4832a1dbc88e2debcf2f889416"; + +export type DualStationVllmRole = "head" | "worker"; + +export interface DualStationDockerOptions { + env?: NodeJS.ProcessEnv; + ignoreError?: boolean; + suppressOutput?: boolean; + timeout?: number; +} + +export interface DualStationDockerResult { + status: number | null; + stdout?: string | Buffer | null; + stderr?: string | Buffer | null; + error?: Error; + signal?: NodeJS.Signals | null; +} + +export interface DualStationVllmLifecycleDeps { + dockerCapture(args: readonly string[], options?: DualStationDockerOptions): string; + dockerForceRm(containerId: string, options?: DualStationDockerOptions): DualStationDockerResult; + dockerRunDetached( + args: readonly string[], + options?: DualStationDockerOptions, + ): DualStationDockerResult; + buildLocalDockerEnv(): Record; + buildRemoteDockerEnv(binding: DualStationSshBinding): Record; + createProbeNonce(): string; + createTransactionId(): string; + effectiveControllerUid(): number | null; + readControllerUid(): number; + onManagedHeadObserved?(): void; + waitBeforeReconcile(ms: number): Promise; + withLifecycleLock(operation: () => Promise | T): Promise; + loadApiKey(): string | null; + localInterfaceAddresses(): readonly string[]; +} + +export interface DualStationVllmStartConfig { + /** A caller-owned 256-bit API key. It is never persisted by the lifecycle. */ + apiKey: string; +} + +export interface DualStationLegacyMigration { + backupContainerName: string; + legacyContainerId: string; + transactionId: string; + headContainerId: string; + workerContainerId: string; +} + +export type StartDualStationVllmResult = + | { + ok: true; + baseUrl: string; + headContainerId: string; + workerContainerId: string; + /** True when an already-running exact owned pair was left untouched. */ + reusedExisting: boolean; + /** Present until the caller commits or rolls back external validation. */ + legacyMigration?: DualStationLegacyMigration; + } + | { ok: false; reason: string; rollbackErrors: string[] }; + +export type CleanupDualStationVllmResult = + | { ok: true; removedContainerIds: string[] } + | { ok: false; reason: string }; + +export type PreflightDualStationVllmResult = { ok: true } | { ok: false; reason: string }; + +export type CommitDualStationLegacyMigrationResult = + | { ok: true; cleanupWarnings: string[] } + | { ok: false; reason: string }; + +export type RollbackDualStationLegacyMigrationResult = + | { ok: true } + | { ok: false; rollbackErrors: string[] }; + +type ManagedContainerSpec = { + role: DualStationVllmRole; + name: string; + endpoint: string; + clusterId: string; + gpuUuid: string; + image: string; + launchContract: string; + apiKeyFingerprint: string | null; + env: Record; +}; + +type ManagedContainerInspection = + | { kind: "absent" } + | { + kind: "managed"; + containerId: string; + running: boolean; + transactionId: string; + reusable: boolean; + } + | { kind: "legacy-managed"; containerId: string; running: boolean } + | { kind: "foreign" | "ambiguous" | "unknown" }; + +type LegacyHeadCutover = { + originalSpec: ManagedContainerSpec; + backupSpec: ManagedContainerSpec; + containerId: string; +}; + +type GpuSmokeSpec = { + role: DualStationVllmRole; + containerName: string; + nonce: string; + image: string; + expectedGpuUuid: string; + env: Record; +}; + +type GpuSmokeInspection = + | { kind: "absent" } + | { kind: "owned"; containerId: string } + | { kind: "foreign" | "ambiguous" | "unknown" }; + +const DEFAULT_DEPS: DualStationVllmLifecycleDeps = { + dockerCapture, + dockerForceRm, + dockerRunDetached, + buildLocalDockerEnv: buildLocalDualStationDockerEnv, + buildRemoteDockerEnv: buildRemoteVllmDockerEnv, + createProbeNonce: () => randomBytes(16).toString("hex"), + createTransactionId: () => randomBytes(16).toString("hex"), + effectiveControllerUid: () => process.getuid?.() ?? null, + readControllerUid: readDualStationControllerUid, + onManagedHeadObserved: () => undefined, + waitBeforeReconcile: (ms) => new Promise((resolve) => setTimeout(resolve, ms)), + withLifecycleLock: withDualStationVllmLifecycleLock, + loadApiKey: loadDualStationVllmApiKey, + localInterfaceAddresses: () => + Object.values(os.networkInterfaces()).flatMap((addresses) => + (addresses ?? []).flatMap((address) => + address.family === "IPv4" && !address.internal ? [address.address] : [], + ), + ), +}; + +function depsWith( + overrides: Partial = {}, +): DualStationVllmLifecycleDeps { + return { ...DEFAULT_DEPS, ...overrides }; +} + +function isRfc1918Ipv4(address: string): boolean { + if (net.isIP(address) !== 4) return false; + const octets = address.split(".").map(Number); + return ( + octets[0] === 10 || + (octets[0] === 172 && octets[1] >= 16 && octets[1] <= 31) || + (octets[0] === 192 && octets[1] === 168) + ); +} + +function assertSafePlan(plan: DualStationVllmPlan): void { + assertDualStationSshBindingFiles(plan.peerSshBinding); + if ( + !IMMUTABLE_IMAGE_PATTERN.test(plan.runtime.image) || + plan.runtime.image !== DUAL_STATION_VLLM_RUNTIME.image || + plan.runtime.modelId !== DUAL_STATION_VLLM_RUNTIME.modelId || + plan.runtime.modelRevision !== DUAL_STATION_VLLM_RUNTIME.modelRevision || + plan.runtime.servedModelId !== DUAL_STATION_VLLM_RUNTIME.servedModelId || + plan.runtime.tensorParallelSize !== DUAL_STATION_VLLM_RUNTIME.tensorParallelSize || + plan.runtime.pipelineParallelSize !== DUAL_STATION_VLLM_RUNTIME.pipelineParallelSize || + plan.runtime.nodeCount !== DUAL_STATION_VLLM_RUNTIME.nodeCount + ) { + throw new Error("Dual-Station vLLM requires the exact pinned runtime contract."); + } + if (plan.rails.length !== 2 || plan.masterAddress !== plan.rails[0]?.local.address) { + throw new Error("Dual-Station vLLM requires exactly two ordered rails and a rail-0 master."); + } + if (!Number.isInteger(plan.roceGidIndex) || plan.roceGidIndex < 0 || plan.roceGidIndex > 4095) { + throw new Error("Dual-Station vLLM requires a valid shared RoCE GID index."); + } + for (const [role, node] of [ + ["local", plan.local], + ["peer", plan.peer], + ] as const) { + if ( + !path.posix.isAbsolute(node.home) || + path.posix.normalize(node.home) !== node.home || + node.home.includes(":") + ) { + throw new Error(`Dual-Station ${role} home must be a normalized absolute POSIX path.`); + } + if ( + !Number.isInteger(node.uid) || + node.uid <= 0 || + node.uid > 2_147_483_647 || + !Number.isInteger(node.gid) || + node.gid <= 0 || + node.gid > 2_147_483_647 + ) { + throw new Error( + `Dual-Station ${role} runtime identity must use non-root UID and GID values.`, + ); + } + if (!SAFE_GPU_UUID_PATTERN.test(node.gpu.uuid)) { + throw new Error(`Dual-Station ${role} GB300 UUID is invalid.`); + } + } + for (const side of ["local", "peer"] as const) { + const endpoints = plan.rails.map((rail) => rail[side]); + if ( + new Set(endpoints.map((endpoint) => endpoint.rdmaDevice)).size !== 2 || + new Set(endpoints.map((endpoint) => endpoint.netdev)).size !== 2 || + new Set(endpoints.map((endpoint) => endpoint.uverbsDevice)).size !== 2 + ) { + throw new Error(`Dual-Station ${side} rails must use two distinct devices.`); + } + for (const endpoint of endpoints) { + if ( + !SAFE_DEVICE_NAME_PATTERN.test(endpoint.rdmaDevice) || + !SAFE_DEVICE_NAME_PATTERN.test(endpoint.netdev) || + !SAFE_UVERBS_DEVICE_PATTERN.test(endpoint.uverbsDevice) || + !isRfc1918Ipv4(endpoint.address) + ) { + throw new Error(`Dual-Station ${side} rail endpoint is invalid.`); + } + } + } +} + +function clusterIdForPlan(plan: DualStationVllmPlan): string { + const identity = { + runtime: plan.runtime, + local: { + hostname: plan.local.hostname, + gpuUuid: plan.local.gpu.uuid, + }, + peer: { + hostname: plan.peer.hostname, + gpuUuid: plan.peer.gpu.uuid, + }, + rails: plan.rails.map((rail) => ({ + index: rail.index, + subnet: rail.subnet, + local: rail.local, + peer: rail.peer, + })), + masterAddress: plan.masterAddress, + roceGidIndex: plan.roceGidIndex, + }; + return createHash("sha256").update(JSON.stringify(identity)).digest("hex"); +} + +/** Stable identity binding both managed containers to one exact physical plan. */ +export function dualStationVllmClusterId(plan: DualStationVllmPlan): string { + assertSafePlan(plan); + return clusterIdForPlan(plan); +} + +function assertSafeStartConfig(config: DualStationVllmStartConfig): void { + if (!DUAL_STATION_VLLM_API_KEY_PATTERN.test(config.apiKey)) { + throw new Error( + "Dual-Station vLLM API key must be exactly 64 lowercase hexadecimal characters.", + ); + } +} + +/** Domain-separated, non-secret binding for the host-persisted high-entropy service key. */ +export function dualStationVllmApiKeyFingerprint(apiKey: string): string { + if (!DUAL_STATION_VLLM_API_KEY_PATTERN.test(apiKey)) { + throw new Error( + "Dual-Station vLLM API key must be exactly 64 lowercase hexadecimal characters.", + ); + } + return createHmac("sha256", Buffer.from(apiKey, "hex")) + .update(VLLM_FINGERPRINT_CONTEXT) + .digest("hex"); +} + +function withoutVllmApiKey(env: Record): Record { + const sanitized = { ...env }; + delete sanitized.VLLM_API_KEY; + return sanitized; +} + +function endpointFor(plan: DualStationVllmPlan, role: DualStationVllmRole): string { + return role === "head" ? `http://${plan.masterAddress}:${String(HEAD_API_PORT)}` : "headless"; +} + +function nameFor(role: DualStationVllmRole): string { + return role === "head" + ? DUAL_STATION_VLLM_HEAD_CONTAINER_NAME + : DUAL_STATION_VLLM_WORKER_CONTAINER_NAME; +} + +function appendEnv(args: string[], name: string, value: string): void { + args.push("--env", `${name}=${value}`); +} + +function runtimeUser(node: DualStationVllmPlan["local"]): string { + return `${String(node.uid)}:${String(node.gid)}`; +} + +function runtimeHomeTmpfs( + node: DualStationVllmPlan["local"], + execution: "exec" | "noexec", +): string { + return `${VLLM_RUNTIME_HOME}:rw,nosuid,nodev,${execution},uid=${String(node.uid)},gid=${String(node.gid)},mode=0700,size=68719476736`; +} + +function flashinferCubinTmpfs(node: DualStationVllmPlan["local"]): string { + return `${FLASHINFER_CUBIN_RUNTIME_DIR}:rw,nosuid,nodev,noexec,uid=${String(node.uid)},gid=${String(node.gid)},mode=0700,size=16777216`; +} + +function appendRuntimeHomeEnv(args: string[]): void { + appendEnv(args, "HOME", VLLM_RUNTIME_HOME); + appendEnv(args, "USER", "vllm"); + appendEnv(args, "LOGNAME", "vllm"); +} + +/** Build the deterministic shell-free launch argv before per-operation labels. */ +function buildDualStationVllmBaseRunArgs( + plan: DualStationVllmPlan, + role: DualStationVllmRole, +): string[] { + assertSafePlan(plan); + const node = role === "head" ? plan.local : plan.peer; + const endpoints = plan.rails.map((rail) => (role === "head" ? rail.local : rail.peer)); + const endpoint = endpointFor(plan, role); + const clusterId = clusterIdForPlan(plan); + const args = [ + "--pull=never", + "--init", + "--restart", + "unless-stopped", + "--network", + "host", + "--shm-size", + "16g", + "--read-only", + "--workdir", + VLLM_RUNTIME_HOME, + "--user", + runtimeUser(node), + "--security-opt", + "no-new-privileges:true", + "--tmpfs", + "/tmp:rw,nosuid,nodev,size=17179869184", + "--tmpfs", + // FlashInfer links generated-kernel headers beneath its installed cubin + // package. Keep only that empty subtree writable; packaged cubins and the + // rest of the image remain read-only. + flashinferCubinTmpfs(node), + "--tmpfs", + // Ray is installed under this home and loads its native _raylet module. + // Docker tmpfs mounts default to noexec, so managed serving opts in. + runtimeHomeTmpfs(node, "exec"), + "--cap-drop", + "ALL", + // Non-root GPU/RDMA memory registration is bounded by this explicit rlimit; + // no Linux capabilities remain in the effective or permitted set. + "--ulimit", + "memlock=-1", + "--ulimit", + "stack=67108864", + "--ulimit", + "nofile=1048576:1048576", + "--gpus", + `device=${node.gpu.uuid}`, + ...endpoints.map((endpoint) => `--device=${endpoint.uverbsDevice}`), + "--volume", + `${node.home}/.cache/huggingface/hub:${HF_HUB_CACHE_CONTAINER_DIR}:ro`, + "--label", + `${DUAL_STATION_VLLM_MANAGED_LABEL}=true`, + "--label", + `${DUAL_STATION_VLLM_ROLE_LABEL}=${role}`, + "--label", + `${DUAL_STATION_VLLM_ENDPOINT_LABEL}=${endpoint}`, + "--label", + `${DUAL_STATION_VLLM_CLUSTER_LABEL}=${clusterId}`, + "--label", + `${DUAL_STATION_VLLM_GPU_LABEL}=${node.gpu.uuid}`, + "--name", + nameFor(role), + ]; + + appendEnv(args, "HF_HOME", HF_CACHE_CONTAINER_DIR); + appendEnv(args, "HF_HUB_CACHE", HF_HUB_CACHE_CONTAINER_DIR); + appendEnv(args, "HUGGINGFACE_HUB_CACHE", HF_HUB_CACHE_CONTAINER_DIR); + appendRuntimeHomeEnv(args); + appendEnv(args, "HF_HUB_OFFLINE", "1"); + appendEnv(args, "TRANSFORMERS_OFFLINE", "1"); + appendEnv(args, "VLLM_HOST_IP", endpoints[0].address); + appendEnv(args, "VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS", role === "head" ? "7200" : "3600"); + appendEnv(args, "VLLM_ALLOW_LONG_MAX_MODEL_LEN", "1"); + appendEnv(args, "NCCL_IB_HCA", endpoints.map((item) => item.rdmaDevice).join(",")); + appendEnv(args, "NCCL_IB_DISABLE", "0"); + appendEnv(args, "NCCL_IB_ADDR_FAMILY", "AF_INET"); + appendEnv(args, "NCCL_IB_ROCE_VERSION_NUM", "2"); + appendEnv(args, "NCCL_IB_GID_INDEX", String(plan.roceGidIndex)); + appendEnv(args, "NCCL_IB_TC", "106"); + appendEnv(args, "NCCL_IB_QPS_PER_CONNECTION", "4"); + appendEnv(args, "NCCL_NET_GDR_LEVEL", "PHB"); + appendEnv(args, "NCCL_IB_PCI_RELAXED_ORDERING", "1"); + appendEnv(args, "NCCL_SOCKET_IFNAME", endpoints[0].netdev); + appendEnv(args, "GLOO_SOCKET_IFNAME", endpoints[0].netdev); + appendEnv(args, "TP_SOCKET_IFNAME", endpoints[0].netdev); + appendEnv(args, "OMPI_MCA_btl_tcp_if_include", endpoints[0].netdev); + appendEnv(args, "MN_IF_NAME", endpoints[0].netdev); + appendEnv(args, "NCCL_IGNORE_CPU_AFFINITY", "1"); + appendEnv(args, "UCX_NET_DEVICES", endpoints.map((item) => `${item.rdmaDevice}:1`).join(",")); + appendEnv(args, "UCX_TLS", "rc_x,cuda_copy,cuda_ipc,gdr_copy"); + appendEnv(args, "UCX_IB_GID_INDEX", String(plan.roceGidIndex)); + appendEnv(args, "UCX_RNDV_THRESH", "8192"); + + if (role === "head") { + // Docker resolves this bare name from the head docker-run subprocess. The + // secret value must never enter argv, labels, or the worker environment. + args.push("--env", "VLLM_API_KEY"); + } + + args.push( + "--entrypoint", + "/bin/bash", + plan.runtime.image, + "-lc", + buildNemotronUltraDistributedServeCommand({ + nodeRank: role === "head" ? 0 : 1, + masterAddr: plan.masterAddress, + masterPort: DUAL_STATION_VLLM_MASTER_PORT, + nodeAddr: endpoints[0].address, + }), + ); + return args; +} + +/** Stable digest of the exact role-local launch argv and its schema. */ +export function dualStationVllmLaunchContract( + plan: DualStationVllmPlan, + role: DualStationVllmRole, +): string { + const contract = { + schema: DUAL_STATION_VLLM_LAUNCH_SCHEMA, + role, + args: buildDualStationVllmBaseRunArgs(plan, role), + }; + return createHash("sha256").update(JSON.stringify(contract)).digest("hex"); +} + +/** Build the complete launch argv with non-secret transaction/config bindings. */ +export function buildDualStationVllmRunArgs( + plan: DualStationVllmPlan, + role: DualStationVllmRole, + transactionId: string, + apiKeyFingerprint: string, +): string[] { + if (!TRANSACTION_ID_PATTERN.test(transactionId)) { + throw new Error("Dual-Station vLLM transaction ID is invalid."); + } + if (!SHA256_HEX_PATTERN.test(apiKeyFingerprint)) { + throw new Error("Dual-Station vLLM API key fingerprint is invalid."); + } + const args = buildDualStationVllmBaseRunArgs(plan, role); + const nameIndex = args.indexOf("--name"); + args.splice( + nameIndex, + 0, + "--label", + `${DUAL_STATION_VLLM_LAUNCH_SCHEMA_LABEL}=${DUAL_STATION_VLLM_LAUNCH_SCHEMA}`, + "--label", + `${DUAL_STATION_VLLM_LAUNCH_CONTRACT_LABEL}=${dualStationVllmLaunchContract(plan, role)}`, + "--label", + `${DUAL_STATION_VLLM_API_KEY_FINGERPRINT_LABEL}=${apiKeyFingerprint}`, + "--label", + `${DUAL_STATION_VLLM_TRANSACTION_LABEL}=${transactionId}`, + ); + return args; +} + +/** Build a no-network, no-pull GPU/RDMA runtime probe that finishes with nvidia-smi. */ +export function buildDualStationGpuSmokeRunArgs( + plan: DualStationVllmPlan, + role: DualStationVllmRole, + nonce: string, +): { containerName: string; args: string[] } { + assertSafePlan(plan); + if (!GPU_SMOKE_NONCE_PATTERN.test(nonce)) { + throw new Error("Dual-Station GPU smoke nonce is invalid."); + } + const node = role === "head" ? plan.local : plan.peer; + const endpoints = plan.rails.map((rail) => (role === "head" ? rail.local : rail.peer)); + const containerName = `${GPU_SMOKE_CONTAINER_PREFIX}-${role}-${nonce}`; + const command = [ + "set -euo pipefail", + `test \"$(id -u)\" = \"${String(node.uid)}\"`, + `test \"$(id -g)\" = \"${String(node.gid)}\"`, + "grep -Eq '^NoNewPrivs:[[:space:]]+1$' /proc/self/status", + "! grep -Eq '^Cap(Inh|Prm|Eff|Bnd|Amb):[[:space:]]+[0-9a-fA-F]*[1-9a-fA-F][0-9a-fA-F]*$' /proc/self/status", + 'test "$(ulimit -l)" = "unlimited"', + `for device in ${endpoints.map((endpoint) => endpoint.uverbsDevice).join(" ")}; do test -c "$device"; test -r "$device"; test -w "$device"; exec 3<>"$device"; exec 3>&-; done`, + 'mkdir -p "$HOME/.cache/torch" "$HF_HOME"', + 'probe="$HOME/.cache/torch/.nemoclaw-write-probe"; : > "$probe"; rm -f "$probe"', + 'probe="$HF_HOME/.nemoclaw-write-probe"; : > "$probe"; rm -f "$probe"', + "exec nvidia-smi --query-gpu=uuid --format=csv,noheader", + ].join("; "); + const args = [ + "--pull=never", + "--network", + "none", + "--read-only", + "--workdir", + VLLM_RUNTIME_HOME, + "--user", + runtimeUser(node), + "--security-opt", + "no-new-privileges:true", + "--tmpfs", + "/tmp:rw,nosuid,nodev,size=17179869184", + "--tmpfs", + runtimeHomeTmpfs(node, "noexec"), + "--cap-drop", + "ALL", + "--ulimit", + "memlock=-1", + "--gpus", + `device=${node.gpu.uuid}`, + ...endpoints.flatMap((endpoint) => ["--device", endpoint.uverbsDevice]), + "--label", + `${DUAL_STATION_VLLM_GPU_SMOKE_LABEL}=${nonce}`, + "--label", + `${DUAL_STATION_VLLM_ROLE_LABEL}=${role}`, + "--name", + containerName, + ]; + appendEnv(args, "HF_HOME", HF_CACHE_CONTAINER_DIR); + appendRuntimeHomeEnv(args); + args.push("--entrypoint", "/bin/bash", plan.runtime.image, "-c", command); + return { + containerName, + args, + }; +} + +function specsForPlan( + plan: DualStationVllmPlan, + deps: DualStationVllmLifecycleDeps, + config?: DualStationVllmStartConfig, +): { head: ManagedContainerSpec; worker: ManagedContainerSpec } { + assertSafePlan(plan); + const controllerUid = assertDualStationControllerAccount( + deps.readControllerUid, + deps.effectiveControllerUid, + ); + if (controllerUid !== plan.local.uid) { + throw new Error( + `Dual-Station lifecycle controller UID must match probed local UID ${String(plan.local.uid)}`, + ); + } + const clusterId = clusterIdForPlan(plan); + const apiKeyFingerprint = config ? dualStationVllmApiKeyFingerprint(config.apiKey) : null; + return { + head: { + role: "head", + name: DUAL_STATION_VLLM_HEAD_CONTAINER_NAME, + endpoint: endpointFor(plan, "head"), + clusterId, + gpuUuid: plan.local.gpu.uuid, + image: plan.runtime.image, + launchContract: dualStationVllmLaunchContract(plan, "head"), + apiKeyFingerprint, + env: withoutVllmApiKey(deps.buildLocalDockerEnv()), + }, + worker: { + role: "worker", + name: DUAL_STATION_VLLM_WORKER_CONTAINER_NAME, + endpoint: endpointFor(plan, "worker"), + clusterId, + gpuUuid: plan.peer.gpu.uuid, + image: plan.runtime.image, + launchContract: dualStationVllmLaunchContract(plan, "worker"), + apiKeyFingerprint, + env: withoutVllmApiKey(deps.buildRemoteDockerEnv(plan.peerSshBinding)), + }, + }; +} + +const INSPECTION_FORMAT = [ + "{{.ID}}", + "{{.Names}}", + "{{.State}}", + "{{.Image}}", + `{{.Label \"${DUAL_STATION_VLLM_MANAGED_LABEL}\"}}`, + `{{.Label \"${DUAL_STATION_VLLM_ROLE_LABEL}\"}}`, + `{{.Label \"${DUAL_STATION_VLLM_ENDPOINT_LABEL}\"}}`, + `{{.Label \"${DUAL_STATION_VLLM_CLUSTER_LABEL}\"}}`, + `{{.Label \"${DUAL_STATION_VLLM_GPU_LABEL}\"}}`, + `{{.Label \"${DUAL_STATION_VLLM_LAUNCH_SCHEMA_LABEL}\"}}`, + `{{.Label \"${DUAL_STATION_VLLM_LAUNCH_CONTRACT_LABEL}\"}}`, + `{{.Label \"${DUAL_STATION_VLLM_API_KEY_FINGERPRINT_LABEL}\"}}`, + `{{.Label \"${DUAL_STATION_VLLM_TRANSACTION_LABEL}\"}}`, +].join("\t"); + +function inspectRows( + containerName: string, + env: Record, + deps: DualStationVllmLifecycleDeps, +): string[][] | null { + try { + const output = deps.dockerCapture( + [ + "container", + "ls", + "--all", + "--no-trunc", + "--filter", + `name=^/${containerName}$`, + "--format", + INSPECTION_FORMAT, + ], + { env, timeout: DOCKER_INSPECT_TIMEOUT_MS }, + ); + if (!output.trim()) return []; + return output + .split(/\r?\n/) + .filter(Boolean) + .map((line) => line.split("\t")); + } catch { + return null; + } +} + +function inspectManagedContainer( + spec: ManagedContainerSpec, + deps: DualStationVllmLifecycleDeps, + options: { allowStoppedLegacy?: boolean } = {}, +): ManagedContainerInspection { + const rows = inspectRows(spec.name, spec.env, deps); + if (rows === null) return { kind: "unknown" }; + if (rows.length === 0) return { kind: "absent" }; + if (rows.length !== 1) return { kind: "ambiguous" }; + const [ + containerId, + name, + state, + image, + managed, + role, + endpoint, + clusterId, + gpuUuid, + launchSchema, + launchContract, + apiKeyFingerprint, + transactionId, + ] = rows[0]; + if (!containerId || rows[0].length !== 13 || !DOCKER_CONTAINER_ID_PATTERN.test(containerId)) { + return { kind: "unknown" }; + } + if ( + spec.role === "head" && + name === spec.name && + image === LEGACY_SINGLE_STATION_MIGRATION_IMAGE && + (state === "running" || (options.allowStoppedLegacy && state === "exited")) && + managed === "true" && + !role && + !endpoint && + !clusterId && + !gpuUuid && + !launchSchema && + !launchContract && + !apiKeyFingerprint && + !transactionId + ) { + return { kind: "legacy-managed", containerId, running: state === "running" }; + } + if ( + name !== spec.name || + image !== spec.image || + managed !== "true" || + role !== spec.role || + endpoint !== spec.endpoint || + clusterId !== spec.clusterId || + gpuUuid !== spec.gpuUuid || + launchSchema !== DUAL_STATION_VLLM_LAUNCH_SCHEMA || + !SHA256_HEX_PATTERN.test(launchContract) || + !SHA256_HEX_PATTERN.test(apiKeyFingerprint) || + !TRANSACTION_ID_PATTERN.test(transactionId) + ) { + return { kind: "foreign" }; + } + const reusable = + launchContract === spec.launchContract && + (spec.apiKeyFingerprint === null || apiKeyFingerprint === spec.apiKeyFingerprint); + return { kind: "managed", containerId, running: state === "running", transactionId, reusable }; +} + +const GPU_SMOKE_INSPECTION_FORMAT = [ + "{{.ID}}", + "{{.Names}}", + "{{.Image}}", + `{{.Label \"${DUAL_STATION_VLLM_GPU_SMOKE_LABEL}\"}}`, + `{{.Label \"${DUAL_STATION_VLLM_ROLE_LABEL}\"}}`, +].join("\t"); + +function inspectGpuSmokeContainer( + spec: GpuSmokeSpec, + deps: DualStationVllmLifecycleDeps, +): GpuSmokeInspection { + let output: string; + try { + output = deps.dockerCapture( + [ + "container", + "ls", + "--all", + "--no-trunc", + "--filter", + `name=^/${spec.containerName}$`, + "--format", + GPU_SMOKE_INSPECTION_FORMAT, + ], + { env: spec.env, timeout: DOCKER_INSPECT_TIMEOUT_MS }, + ); + } catch { + return { kind: "unknown" }; + } + if (!output.trim()) return { kind: "absent" }; + const rows = output + .trim() + .split(/\r?\n/) + .filter(Boolean) + .map((line) => line.split("\t")); + if (rows.length !== 1) return { kind: "ambiguous" }; + const [containerId, name, image, nonce, role] = rows[0]; + if (!containerId || rows[0].length !== 5 || !DOCKER_CONTAINER_ID_PATTERN.test(containerId)) { + return { kind: "unknown" }; + } + if ( + name !== spec.containerName || + image !== spec.image || + nonce !== spec.nonce || + role !== spec.role + ) { + return { kind: "foreign" }; + } + return { kind: "owned", containerId }; +} + +function pinnedImageIsPresent( + spec: ManagedContainerSpec, + deps: DualStationVllmLifecycleDeps, +): boolean { + try { + const output = deps.dockerCapture(["image", "inspect", "--format", "{{.Id}}", spec.image], { + env: spec.env, + timeout: DOCKER_INSPECT_TIMEOUT_MS, + }); + const ids = output.trim().split(/\r?\n/).filter(Boolean); + return ids.length === 1 && IMAGE_ID_PATTERN.test(ids[0]); + } catch { + return false; + } +} + +function mutationSucceeded(result: DualStationDockerResult): boolean { + return result.status === 0 && !result.error && !result.signal; +} + +function resultContainerId(result: DualStationDockerResult): string | null { + const value = String(result.stdout ?? "").trim(); + return DOCKER_CONTAINER_ID_PATTERN.test(value) ? value : null; +} + +function removeGpuSmokeExact( + spec: GpuSmokeSpec, + containerId: string, + deps: DualStationVllmLifecycleDeps, +): boolean { + try { + return mutationSucceeded( + deps.dockerForceRm(containerId, { + env: spec.env, + ignoreError: true, + suppressOutput: true, + timeout: DOCKER_MUTATION_TIMEOUT_MS, + }), + ); + } catch { + return false; + } +} + +function runGpuSmoke( + plan: DualStationVllmPlan, + managedSpec: ManagedContainerSpec, + expectedGpuUuid: string, + deps: DualStationVllmLifecycleDeps, +): PreflightDualStationVllmResult { + let nonce: string; + try { + nonce = deps.createProbeNonce(); + } catch { + return { ok: false, reason: `${managedSpec.role} GPU smoke nonce generation failed` }; + } + let built: ReturnType; + try { + built = buildDualStationGpuSmokeRunArgs(plan, managedSpec.role, nonce); + } catch (error) { + return { ok: false, reason: (error as Error).message }; + } + const spec: GpuSmokeSpec = { + role: managedSpec.role, + containerName: built.containerName, + nonce, + image: managedSpec.image, + expectedGpuUuid, + env: managedSpec.env, + }; + const before = inspectGpuSmokeContainer(spec, deps); + if (before.kind !== "absent") { + return { + ok: false, + reason: `${spec.role} GPU smoke name ownership is ${before.kind}; refusing mutation`, + }; + } + + let failureReason: string | null = null; + let runResult: DualStationDockerResult | null = null; + try { + runResult = deps.dockerRunDetached(built.args, { + env: spec.env, + ignoreError: true, + suppressOutput: true, + timeout: DOCKER_GPU_SMOKE_TIMEOUT_MS, + }); + } catch { + failureReason = `${spec.role} GPU smoke container failed to launch`; + } + + const capturedId = runResult ? resultContainerId(runResult) : null; + let cleanupId = capturedId; + const observed = inspectGpuSmokeContainer(spec, deps); + if (!cleanupId && observed.kind === "owned") cleanupId = observed.containerId; + + if (!failureReason && (!runResult || !mutationSucceeded(runResult))) { + failureReason = `${spec.role} GPU smoke container failed to launch`; + } + if (!failureReason && !capturedId) { + failureReason = `${spec.role} GPU smoke returned an invalid container ID`; + } + if (!failureReason && observed.kind !== "owned") { + failureReason = `${spec.role} GPU smoke ownership is ${observed.kind}`; + } + if ( + !failureReason && + observed.kind === "owned" && + capturedId && + observed.containerId !== capturedId + ) { + failureReason = `${spec.role} GPU smoke container ID did not match exact-name inspection`; + } + + if (!failureReason && capturedId) { + try { + const exitCode = deps.dockerCapture(["wait", capturedId], { + env: spec.env, + timeout: DOCKER_GPU_SMOKE_TIMEOUT_MS, + }); + if (exitCode.trim() !== "0") { + failureReason = `${spec.role} GPU smoke exited unsuccessfully`; + } + } catch { + failureReason = `${spec.role} GPU smoke did not finish within the bounded wait`; + } + } + + if (!failureReason && capturedId) { + try { + const visibleGpuUuids = deps + .dockerCapture(["logs", capturedId], { + env: spec.env, + timeout: DOCKER_INSPECT_TIMEOUT_MS, + }) + .trim() + .split(/\r?\n/) + .map((line) => line.trim()) + .filter(Boolean); + if (visibleGpuUuids.length !== 1 || visibleGpuUuids[0] !== spec.expectedGpuUuid) { + failureReason = `${spec.role} GPU smoke did not expose exactly the discovered GPU`; + } + } catch { + failureReason = `${spec.role} GPU smoke output could not be validated`; + } + } + + if (cleanupId && !removeGpuSmokeExact(spec, cleanupId, deps)) { + const cleanupReason = `failed to remove exact ${spec.role} GPU smoke container ${cleanupId}`; + failureReason = failureReason ? `${failureReason}; ${cleanupReason}` : cleanupReason; + } + if (!cleanupId && observed.kind !== "absent") { + const cleanupReason = `${spec.role} GPU smoke could not identify an exact owned ID for cleanup`; + failureReason = failureReason ? `${failureReason}; ${cleanupReason}` : cleanupReason; + } + return failureReason ? { ok: false, reason: failureReason } : { ok: true }; +} + +function removeExact( + spec: ManagedContainerSpec, + containerId: string, + deps: DualStationVllmLifecycleDeps, +): boolean { + try { + return mutationSucceeded( + deps.dockerForceRm(containerId, { + env: spec.env, + ignoreError: true, + suppressOutput: true, + timeout: DOCKER_MUTATION_TIMEOUT_MS, + }), + ); + } catch { + return false; + } +} + +function runExactContainerMutation( + spec: ManagedContainerSpec, + args: readonly string[], + deps: DualStationVllmLifecycleDeps, +): void { + deps.dockerCapture(["container", ...args], { + env: spec.env, + timeout: DOCKER_MUTATION_TIMEOUT_MS, + }); +} + +function exactLegacyGeneration( + spec: ManagedContainerSpec, + containerId: string, + deps: DualStationVllmLifecycleDeps, +): Extract | null { + const inspection = inspectManagedContainer(spec, deps, { allowStoppedLegacy: true }); + return inspection.kind === "legacy-managed" && inspection.containerId === containerId + ? inspection + : null; +} + +function restoreLegacyHead( + cutover: LegacyHeadCutover, + deps: DualStationVllmLifecycleDeps, +): string | null { + let original = exactLegacyGeneration(cutover.originalSpec, cutover.containerId, deps); + let backup = exactLegacyGeneration(cutover.backupSpec, cutover.containerId, deps); + + if (original && backup) { + return "legacy head restoration found the exact container ID under two names"; + } + if (!original) { + if (!backup) return "legacy head restoration could not find the exact owned container ID"; + if (inspectManagedContainer(cutover.originalSpec, deps).kind !== "absent") { + return "legacy head restoration found the original name occupied"; + } + try { + runExactContainerMutation( + cutover.backupSpec, + ["rename", cutover.containerId, cutover.originalSpec.name], + deps, + ); + } catch { + // Docker may commit a rename before its client reports failure. Reconcile + // the exact generation below instead of trusting the command response. + } + original = exactLegacyGeneration(cutover.originalSpec, cutover.containerId, deps); + backup = exactLegacyGeneration(cutover.backupSpec, cutover.containerId, deps); + if (!original || backup) { + return "legacy head restoration could not restore the exact original container name"; + } + } + + if (!original.running) { + try { + runExactContainerMutation(cutover.originalSpec, ["start", cutover.containerId], deps); + } catch { + // As with rename, validate the exact postcondition rather than the Docker + // client's response so a committed start is not reported as lost. + } + original = exactLegacyGeneration(cutover.originalSpec, cutover.containerId, deps); + } + return original?.running + ? null + : "legacy head restoration could not restart the exact owned container ID"; +} + +function prepareLegacyHeadCutover( + spec: ManagedContainerSpec, + containerId: string, + transactionId: string, + deps: DualStationVllmLifecycleDeps, +): + | { ok: true; cutover: LegacyHeadCutover } + | { ok: false; reason: string; rollbackErrors: string[] } { + const cutover: LegacyHeadCutover = { + originalSpec: spec, + backupSpec: { ...spec, name: `${spec.name}-legacy-${transactionId}` }, + containerId, + }; + const fail = (reason: string) => { + const restoreError = restoreLegacyHead(cutover, deps); + return { + ok: false as const, + reason, + rollbackErrors: restoreError ? [restoreError] : [], + }; + }; + + const original = exactLegacyGeneration(spec, containerId, deps); + if (!original?.running) { + return { + ok: false, + reason: "legacy head changed before transactional cutover", + rollbackErrors: [], + }; + } + if (inspectManagedContainer(cutover.backupSpec, deps).kind !== "absent") { + return { + ok: false, + reason: "legacy head backup name is not absent", + rollbackErrors: [], + }; + } + + try { + runExactContainerMutation(spec, ["rename", containerId, cutover.backupSpec.name], deps); + } catch { + // Reconcile below: rename can succeed even when the client times out. + } + const renamed = exactLegacyGeneration(cutover.backupSpec, containerId, deps); + if (!renamed?.running || inspectManagedContainer(cutover.originalSpec, deps).kind !== "absent") { + return fail("failed to preserve the exact legacy head under its transaction backup name"); + } + + try { + runExactContainerMutation(cutover.backupSpec, ["stop", "--time", "30", containerId], deps); + } catch { + // Reconcile and restore below if the exact generation is not stopped. + } + const stopped = exactLegacyGeneration(cutover.backupSpec, containerId, deps); + if (!stopped || stopped.running) { + return fail("failed to stop the preserved exact legacy head for cutover"); + } + return { ok: true, cutover }; +} + +function rollbackNewPairAndRestoreLegacy( + entries: readonly { spec: ManagedContainerSpec; containerId: string }[], + transactionId: string, + cutover: LegacyHeadCutover, + deps: DualStationVllmLifecycleDeps, +): string[] { + const errors = rollbackExact(entries, transactionId, deps); + const restoreError = restoreLegacyHead(cutover, deps); + if (restoreError) errors.push(restoreError); + return errors; +} + +function finalizeLegacyHeadCutover( + cutover: LegacyHeadCutover, + deps: DualStationVllmLifecycleDeps, +): string | null { + const initial = inspectManagedContainer(cutover.backupSpec, deps, { + allowStoppedLegacy: true, + }); + if (initial.kind === "absent") return null; + const backup = exactLegacyGeneration(cutover.backupSpec, cutover.containerId, deps); + if (!backup || backup.running) { + return "could not verify the stopped exact legacy backup for post-commit cleanup"; + } + const removed = removeExact(cutover.backupSpec, cutover.containerId, deps); + const after = inspectManagedContainer(cutover.backupSpec, deps); + // A Docker client error can race a committed removal. Accept either the + // successful exact-ID mutation or an exact-name absence on reconciliation. + return removed || after.kind === "absent" + ? null + : "could not remove the stopped exact legacy backup after the new pair committed"; +} + +function legacyCutoverFromMigration( + specs: ReturnType, + migration: DualStationLegacyMigration, +): LegacyHeadCutover | null { + if ( + !TRANSACTION_ID_PATTERN.test(migration.transactionId) || + !DOCKER_CONTAINER_ID_PATTERN.test(migration.legacyContainerId) || + !DOCKER_CONTAINER_ID_PATTERN.test(migration.headContainerId) || + !DOCKER_CONTAINER_ID_PATTERN.test(migration.workerContainerId) || + new Set([migration.legacyContainerId, migration.headContainerId, migration.workerContainerId]) + .size !== 3 || + migration.backupContainerName !== `${specs.head.name}-legacy-${migration.transactionId}` + ) { + return null; + } + return { + originalSpec: specs.head, + backupSpec: { ...specs.head, name: migration.backupContainerName }, + containerId: migration.legacyContainerId, + }; +} + +function removeTransactionExact( + spec: ManagedContainerSpec, + containerId: string, + transactionId: string, + deps: DualStationVllmLifecycleDeps, +): boolean { + const inspection = inspectManagedContainer(spec, deps); + if (inspection.kind === "absent") return true; + if ( + inspection.kind !== "managed" || + !inspection.reusable || + inspection.containerId !== containerId || + inspection.transactionId !== transactionId + ) { + return false; + } + return removeExact(spec, containerId, deps); +} + +function rollbackExact( + entries: readonly { spec: ManagedContainerSpec; containerId: string }[], + transactionId: string, + deps: DualStationVllmLifecycleDeps, +): string[] { + const errors: string[] = []; + const seen = new Set(); + for (const { spec, containerId } of entries) { + const key = `${spec.role}:${containerId}`; + if (seen.has(key)) continue; + seen.add(key); + if (!removeTransactionExact(spec, containerId, transactionId, deps)) { + errors.push(`failed to remove ${spec.role} container ${containerId}`); + } + } + return errors; +} + +async function recoverManagedId( + spec: ManagedContainerSpec, + transactionId: string, + deps: DualStationVllmLifecycleDeps, +): Promise { + for (let attempt = 0; attempt < DOCKER_LATE_CREATE_RECONCILE_ATTEMPTS; attempt += 1) { + const inspection = inspectManagedContainer(spec, deps); + if ( + inspection.kind === "managed" && + inspection.reusable && + inspection.transactionId === transactionId + ) { + return inspection; + } + if (inspection.kind !== "absent" && inspection.kind !== "unknown") return null; + if (attempt < DOCKER_LATE_CREATE_RECONCILE_ATTEMPTS - 1) { + await deps.waitBeforeReconcile(DOCKER_LATE_CREATE_RECONCILE_INTERVAL_MS); + } + } + return null; +} + +async function startOne( + plan: DualStationVllmPlan, + spec: ManagedContainerSpec, + config: DualStationVllmStartConfig, + transactionId: string, + deps: DualStationVllmLifecycleDeps, +): Promise<{ ok: true; containerId: string } | { ok: false; containerId: string | null }> { + let result: DualStationDockerResult | null = null; + try { + const env = spec.role === "head" ? { ...spec.env, VLLM_API_KEY: config.apiKey } : spec.env; + result = deps.dockerRunDetached( + buildDualStationVllmRunArgs(plan, spec.role, transactionId, spec.apiKeyFingerprint ?? ""), + { + env, + ignoreError: true, + suppressOutput: true, + timeout: DOCKER_MUTATION_TIMEOUT_MS, + }, + ); + } catch { + // A timed-out Docker client may still have committed the create. Reconcile + // only this unguessable transaction before deciding what may be rolled back. + } + const capturedId = result ? resultContainerId(result) : null; + const observed = await recoverManagedId(spec, transactionId, deps); + if ( + result && + mutationSucceeded(result) && + capturedId && + observed?.kind === "managed" && + observed.containerId === capturedId && + observed.running + ) { + return { ok: true, containerId: capturedId }; + } + return { + ok: false, + containerId: observed?.kind === "managed" ? observed.containerId : null, + }; +} + +function unsafeInspectionReason( + role: DualStationVllmRole, + inspection: ManagedContainerInspection, +): string | null { + if ( + inspection.kind === "managed" || + inspection.kind === "legacy-managed" || + inspection.kind === "absent" + ) { + return null; + } + return `${role} container ownership is ${inspection.kind}; refusing mutation`; +} + +function ownershipTopologyReason( + head: ManagedContainerInspection, + worker: ManagedContainerInspection, +): string | null { + const headReason = unsafeInspectionReason("head", head); + if (headReason) return headReason; + const workerReason = unsafeInspectionReason("worker", worker); + if (workerReason) return workerReason; + if (head.kind === "legacy-managed" && worker.kind !== "absent") { + return "legacy single-Station head can only migrate when the peer worker name is absent"; + } + return null; +} + +function ownershipPreflightForSpecs( + specs: ReturnType, + deps: DualStationVllmLifecycleDeps, +): PreflightDualStationVllmResult { + const head = inspectManagedContainer(specs.head, deps); + const worker = inspectManagedContainer(specs.worker, deps); + const reason = ownershipTopologyReason(head, worker); + return reason ? { ok: false, reason } : { ok: true }; +} + +function gpuRuntimePreflightForSpecs( + plan: DualStationVllmPlan, + specs: ReturnType, + deps: DualStationVllmLifecycleDeps, +): PreflightDualStationVllmResult { + const ownershipBefore = ownershipPreflightForSpecs(specs, deps); + if (!ownershipBefore.ok) return ownershipBefore; + + // Validate both exact digest references before either daemon is mutated. + for (const spec of [specs.worker, specs.head]) { + if (!pinnedImageIsPresent(spec, deps)) { + return { + ok: false, + reason: `${spec.role} pinned vLLM image is not present or could not be inspected`, + }; + } + } + + for (const [spec, gpuUuid] of [ + [specs.worker, plan.peer.gpu.uuid], + [specs.head, plan.local.gpu.uuid], + ] as const) { + const smoke = runGpuSmoke(plan, spec, gpuUuid, deps); + if (!smoke.ok) return smoke; + } + + // Exact names may have changed while the bounded probes ran. + return ownershipPreflightForSpecs(specs, deps); +} + +/** Read-only ownership preflight used before downloads and repeated by start. */ +export function preflightDualStationManagedVllm( + plan: DualStationVllmPlan, + overrides: Partial = {}, +): PreflightDualStationVllmResult { + const deps = depsWith(overrides); + let specs: ReturnType; + try { + specs = specsForPlan(plan, deps); + } catch (error) { + return { ok: false, reason: (error as Error).message }; + } + return ownershipPreflightForSpecs(specs, deps); +} + +/** + * After both pinned images are installed, prove GPU-container execution on + * each daemon without pulling an image, starting vLLM, or receiving an API key. + */ +export async function preflightDualStationGpuRuntime( + plan: DualStationVllmPlan, + overrides: Partial = {}, +): Promise { + const deps = depsWith(overrides); + try { + const specs = specsForPlan(plan, deps); + return await deps.withLifecycleLock(() => gpuRuntimePreflightForSpecs(plan, specs, deps)); + } catch (error) { + return { ok: false, reason: `dual-Station lifecycle lock failed: ${(error as Error).message}` }; + } +} + +/** + * Hold the host-global lease across start, readiness/auth validation, and any + * rollback. The callback's successful return is the lifecycle commit point. + */ +export function withDualStationManagedVllmLifecycle( + operation: () => Promise | T, + overrides: Partial = {}, +): Promise { + return depsWith(overrides).withLifecycleLock(operation); +} + +/** Start rank 1 first, then rank 0, rolling back only exact newly created IDs. */ +export async function startDualStationManagedVllm( + plan: DualStationVllmPlan, + config: DualStationVllmStartConfig, + overrides: Partial = {}, +): Promise { + const deps = depsWith(overrides); + let specs: ReturnType; + try { + assertSafeStartConfig(config); + specs = specsForPlan(plan, deps, config); + } catch (error) { + return { ok: false, reason: (error as Error).message, rollbackErrors: [] }; + } + + try { + return await deps.withLifecycleLock(async () => { + // Deliberately repeat the post-pull smoke here: model download can be long, + // and this closes that TOCTOU window immediately before owned replacement. + const gpuRuntimePreflight = gpuRuntimePreflightForSpecs(plan, specs, deps); + if (!gpuRuntimePreflight.ok) { + return { ...gpuRuntimePreflight, rollbackErrors: [] }; + } + + const existingHead = inspectManagedContainer(specs.head, deps); + const existingWorker = inspectManagedContainer(specs.worker, deps); + const topologyReason = ownershipTopologyReason(existingHead, existingWorker); + if (topologyReason) return { ok: false, reason: topologyReason, rollbackErrors: [] }; + + if ( + existingHead.kind === "managed" && + existingHead.running && + existingHead.reusable && + existingWorker.kind === "managed" && + existingWorker.running && + existingWorker.reusable && + existingHead.transactionId === existingWorker.transactionId + ) { + return { + ok: true, + baseUrl: specs.head.endpoint, + headContainerId: existingHead.containerId, + workerContainerId: existingWorker.containerId, + reusedExisting: true, + }; + } + + let transactionId: string; + try { + transactionId = deps.createTransactionId(); + } catch { + return { + ok: false, + reason: "dual-Station lifecycle transaction generation failed", + rollbackErrors: [], + }; + } + if (!TRANSACTION_ID_PATTERN.test(transactionId)) { + return { + ok: false, + reason: "dual-Station lifecycle transaction ID is invalid", + rollbackErrors: [], + }; + } + + for (const [spec, inspection] of [ + [specs.head, existingHead], + [specs.worker, existingWorker], + ] as const) { + // A legacy head remains live until the newly created peer worker is + // proven running. Its cutover is handled transactionally below. + if (inspection.kind === "managed" && !removeExact(spec, inspection.containerId, deps)) { + return { + ok: false, + reason: `failed to remove existing owned ${spec.role} container`, + rollbackErrors: [], + }; + } + } + + const worker = await startOne(plan, specs.worker, config, transactionId, deps); + if (!worker.ok) { + const rollbackErrors = worker.containerId + ? rollbackExact( + [{ spec: specs.worker, containerId: worker.containerId }], + transactionId, + deps, + ) + : []; + return { ok: false, reason: "worker container failed to start", rollbackErrors }; + } + + let legacyCutover: LegacyHeadCutover | null = null; + if (existingHead.kind === "legacy-managed") { + const prepared = prepareLegacyHeadCutover( + specs.head, + existingHead.containerId, + transactionId, + deps, + ); + if (!prepared.ok) { + return { + ok: false, + reason: prepared.reason, + rollbackErrors: [ + ...rollbackExact( + [{ spec: specs.worker, containerId: worker.containerId }], + transactionId, + deps, + ), + ...prepared.rollbackErrors, + ], + }; + } + legacyCutover = prepared.cutover; + } + + const head = await startOne(plan, specs.head, config, transactionId, deps); + if (!head.ok) { + const rollback = [ + ...(head.containerId ? [{ spec: specs.head, containerId: head.containerId }] : []), + { spec: specs.worker, containerId: worker.containerId }, + ]; + return { + ok: false, + reason: "head container failed to start", + rollbackErrors: legacyCutover + ? rollbackNewPairAndRestoreLegacy(rollback, transactionId, legacyCutover, deps) + : rollbackExact(rollback, transactionId, deps), + }; + } + + const finalHead = inspectManagedContainer(specs.head, deps); + const finalWorker = inspectManagedContainer(specs.worker, deps); + if ( + finalHead.kind !== "managed" || + !finalHead.running || + !finalHead.reusable || + finalHead.transactionId !== transactionId || + finalHead.containerId !== head.containerId || + finalWorker.kind !== "managed" || + !finalWorker.running || + !finalWorker.reusable || + finalWorker.transactionId !== transactionId || + finalWorker.containerId !== worker.containerId + ) { + const rollback = [ + { spec: specs.head, containerId: head.containerId }, + { spec: specs.worker, containerId: worker.containerId }, + ]; + return { + ok: false, + reason: "dual-Station containers did not remain running", + rollbackErrors: legacyCutover + ? rollbackNewPairAndRestoreLegacy(rollback, transactionId, legacyCutover, deps) + : rollbackExact(rollback, transactionId, deps), + }; + } + + return { + ok: true, + baseUrl: specs.head.endpoint, + headContainerId: head.containerId, + workerContainerId: worker.containerId, + reusedExisting: false, + ...(legacyCutover + ? { + legacyMigration: { + backupContainerName: legacyCutover.backupSpec.name, + legacyContainerId: legacyCutover.containerId, + transactionId, + headContainerId: head.containerId, + workerContainerId: worker.containerId, + }, + } + : {}), + }; + }); + } catch (error) { + return { + ok: false, + reason: `dual-Station lifecycle lock failed: ${(error as Error).message}`, + rollbackErrors: [], + }; + } +} + +/** Retire the preserved legacy generation only after external readiness/auth commit. */ +export async function commitDualStationLegacyMigration( + plan: DualStationVllmPlan, + migration: DualStationLegacyMigration, + overrides: Partial = {}, +): Promise { + const deps = depsWith(overrides); + try { + const specs = specsForPlan(plan, deps); + const cutover = legacyCutoverFromMigration(specs, migration); + if (!cutover) return { ok: false, reason: "legacy migration handle is invalid" }; + return await deps.withLifecycleLock(() => { + const head = inspectManagedContainer(specs.head, deps); + const worker = inspectManagedContainer(specs.worker, deps); + if ( + head.kind !== "managed" || + !head.running || + !head.reusable || + head.containerId !== migration.headContainerId || + head.transactionId !== migration.transactionId || + worker.kind !== "managed" || + !worker.running || + !worker.reusable || + worker.containerId !== migration.workerContainerId || + worker.transactionId !== migration.transactionId + ) { + return { ok: false, reason: "new dual-Station transaction changed before commit" }; + } + const warning = finalizeLegacyHeadCutover(cutover, deps); + return { ok: true, cleanupWarnings: warning ? [warning] : [] }; + }); + } catch (error) { + return { ok: false, reason: `legacy migration commit failed: ${(error as Error).message}` }; + } +} + +/** Remove only the new transaction and restore its exact preserved legacy generation. */ +export async function rollbackDualStationLegacyMigration( + plan: DualStationVllmPlan, + migration: DualStationLegacyMigration, + overrides: Partial = {}, +): Promise { + const deps = depsWith(overrides); + try { + const specs = specsForPlan(plan, deps); + const cutover = legacyCutoverFromMigration(specs, migration); + if (!cutover) return { ok: false, rollbackErrors: ["legacy migration handle is invalid"] }; + return await deps.withLifecycleLock(() => { + const rollbackErrors = rollbackNewPairAndRestoreLegacy( + [ + { spec: specs.head, containerId: migration.headContainerId }, + { spec: specs.worker, containerId: migration.workerContainerId }, + ], + migration.transactionId, + cutover, + deps, + ); + return rollbackErrors.length === 0 ? { ok: true } : { ok: false, rollbackErrors }; + }); + } catch (error) { + return { + ok: false, + rollbackErrors: [`legacy migration rollback failed: ${(error as Error).message}`], + }; + } +} + +/** Remove only containers whose complete dual-Station ownership tuple matches. */ +function cleanupDualStationManagedVllmUnlocked( + plan: DualStationVllmPlan, + deps: DualStationVllmLifecycleDeps, +): CleanupDualStationVllmResult { + let specs: ReturnType; + try { + specs = specsForPlan(plan, deps); + } catch (error) { + return { ok: false, reason: (error as Error).message }; + } + const head = inspectManagedContainer(specs.head, deps); + const worker = inspectManagedContainer(specs.worker, deps); + if (head.kind === "legacy-managed" || worker.kind === "legacy-managed") { + return { ok: false, reason: "refusing dual-Station cleanup of a legacy single-host container" }; + } + for (const [role, inspection] of [ + ["head", head], + ["worker", worker], + ] as const) { + const reason = unsafeInspectionReason(role, inspection); + if (reason) return { ok: false, reason }; + } + + const removedContainerIds: string[] = []; + for (const [spec, inspection] of [ + [specs.head, head], + [specs.worker, worker], + ] as const) { + if (inspection.kind !== "managed") continue; + if (!removeExact(spec, inspection.containerId, deps)) { + return { ok: false, reason: `failed to remove owned ${spec.role} container` }; + } + removedContainerIds.push(inspection.containerId); + } + return { ok: true, removedContainerIds }; +} + +/** Serialize cleanup with start so fixed names cannot transfer between owners. */ +export async function cleanupDualStationManagedVllm( + plan: DualStationVllmPlan, + overrides: Partial = {}, +): Promise { + const deps = depsWith(overrides); + try { + return await deps.withLifecycleLock(() => cleanupDualStationManagedVllmUnlocked(plan, deps)); + } catch (error) { + return { ok: false, reason: `dual-Station lifecycle lock failed: ${(error as Error).message}` }; + } +} + +/** Read-only exact-ownership/running check across both Docker daemons. */ +export function areDualStationManagedVllmContainersRunning( + plan: DualStationVllmPlan, + overrides: Partial = {}, +): boolean { + const deps = depsWith(overrides); + try { + const specs = specsForPlan(plan, deps); + const head = inspectManagedContainer(specs.head, deps); + const worker = inspectManagedContainer(specs.worker, deps); + return ( + head.kind === "managed" && + head.running && + head.reusable && + worker.kind === "managed" && + worker.running && + worker.reusable && + head.transactionId === worker.transactionId + ); + } catch { + return false; + } +} + +function validatedManagedBaseUrl(value: string): string | null { + let parsed: URL; + try { + parsed = new URL(value); + } catch { + return null; + } + if ( + parsed.protocol !== "http:" || + parsed.username || + parsed.password || + parsed.pathname !== "/" || + parsed.search || + parsed.hash || + parsed.port !== String(HEAD_API_PORT) || + !isRfc1918Ipv4(parsed.hostname) + ) { + return null; + } + const canonical = `http://${parsed.hostname}:${String(HEAD_API_PORT)}`; + return value === canonical ? canonical : null; +} + +/** Recover the local managed head endpoint without trusting persisted user input. */ +export function getDualStationManagedVllmBaseUrl( + overrides: Partial = {}, +): string | null { + const deps = depsWith(overrides); + const env = withoutVllmApiKey(deps.buildLocalDockerEnv()); + const rows = inspectRows(DUAL_STATION_VLLM_HEAD_CONTAINER_NAME, env, deps); + if (!rows || rows.length !== 1 || rows[0].length !== 13) return null; + const [ + containerId, + name, + state, + image, + managed, + role, + endpoint, + clusterId, + gpuUuid, + launchSchema, + launchContract, + apiKeyFingerprint, + transactionId, + ] = rows[0]; + if ( + !DOCKER_CONTAINER_ID_PATTERN.test(containerId) || + name !== DUAL_STATION_VLLM_HEAD_CONTAINER_NAME || + state !== "running" || + image !== DUAL_STATION_VLLM_RUNTIME.image || + managed !== "true" || + role !== "head" || + !CLUSTER_ID_PATTERN.test(clusterId) || + !SAFE_GPU_UUID_PATTERN.test(gpuUuid) || + launchSchema !== DUAL_STATION_VLLM_LAUNCH_SCHEMA || + !SHA256_HEX_PATTERN.test(launchContract) || + !TRANSACTION_ID_PATTERN.test(transactionId) + ) { + return null; + } + deps.onManagedHeadObserved?.(); + if (!SHA256_HEX_PATTERN.test(apiKeyFingerprint)) return null; + let apiKey: string | null; + try { + apiKey = deps.loadApiKey(); + } catch { + return null; + } + if (!apiKey || dualStationVllmApiKeyFingerprint(apiKey) !== apiKeyFingerprint) return null; + const baseUrl = validatedManagedBaseUrl(endpoint); + if (!baseUrl || !deps.localInterfaceAddresses().includes(new URL(baseUrl).hostname)) return null; + return baseUrl; +} diff --git a/src/lib/inference/vllm-station-cluster.test.ts b/src/lib/inference/vllm-station-cluster.test.ts new file mode 100644 index 0000000000..8b0f32bbbc --- /dev/null +++ b/src/lib/inference/vllm-station-cluster.test.ts @@ -0,0 +1,1262 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { type SpawnSyncOptionsWithStringEncoding, spawnSync } from "node:child_process"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; + +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +import { buildRemoteVllmDockerEnv } from "./vllm-docker-env"; +import { + createStationClusterProbeDeps, + DUAL_STATION_VLLM_RUNTIME, + NEMOCLAW_DGX_STATION_PEER_ENV, + parseStationHostProbe, + probeDualStationVllmCapability, + type StationClusterProbeDeps, + type StationHostProbe, + type StationProbeCommandResult, + type StationRailConnectivityRequest, + validatePeerTarget, +} from "./vllm-station-cluster"; +import { + type DualStationSshBinding, + loadDualStationSshBindingHandoff, + NEMOCLAW_DGX_STATION_SSH_BINDING_ENV, +} from "./vllm-station-ssh-binding"; +import { + createDualStationSshBindingFixture, + retargetDualStationSshBindingFixture, + type DualStationSshBindingFixture, +} from "./vllm-station-ssh-binding.test-support"; +import { resolveStationFixturePython } from "./vllm-station-fixture.test-support"; + +const LOCAL_HOME = "/home/local"; +const PEER_HOME = "/home/nvidia"; +const STATION_ACCEPTANCE_IMAGE = + "docker.io/library/ubuntu@sha256:7f622ca8766bccb22f04242ecb6f19f770b2f08827dc4b8c707de5e78a6da7ab"; + +function strictDockerSshConfig(binding: DualStationSshBinding): string { + return [ + `hostname ${binding.resolvedHost}`, + `user ${binding.sshUser}`, + `port ${String(binding.port)}`, + `hostkeyalias ${binding.lookupHost}`, + `userknownhostsfile ${binding.knownHostsFile}`, + "globalknownhostsfile /dev/null", + "batchmode yes", + "stricthostkeychecking true", + "permitlocalcommand no", + "forwardagent no", + "forwardx11 no", + "forwardx11trusted no", + "tunnel false", + "updatehostkeys no", + "controlmaster false", + "controlpersist no", + "sendenv LANG", + "sendenv LC_*", + ].join("\n"); +} + +let sshFixture: DualStationSshBindingFixture; + +beforeEach(() => { + sshFixture = createDualStationSshBindingFixture(); +}); + +afterEach(() => { + vi.unstubAllEnvs(); + sshFixture.cleanup(); +}); + +function snapshotPath(home: string): string { + return [ + home, + ".cache/huggingface/hub", + `models--${DUAL_STATION_VLLM_RUNTIME.modelId.replace("/", "--")}`, + "snapshots", + DUAL_STATION_VLLM_RUNTIME.modelRevision, + ].join("/"); +} + +function rail( + rdmaDevice: string, + netdev: string, + pciAddress: string, + address: string, + gidIndexes: number[] = [3, 5], +) { + const hostOctet = netdev.includes("a") ? "aa" : "bb"; + const railOctet = netdev.endsWith("0") ? "00" : "01"; + return { + rdmaDevice, + port: 1, + netdev, + macAddress: `02:00:00:${hostOctet}:00:${railOctet}`, + uverbsDevice: `/dev/infiniband/uverbs${rdmaDevice.endsWith("0") ? "0" : "1"}`, + pciAddress, + pciName: `${pciAddress} Ethernet controller: NVIDIA ConnectX-8 SuperNIC`, + state: "4: ACTIVE", + linkLayer: "Ethernet", + speedMbps: 400_000, + mtu: 9000, + ipv4Addresses: [{ address, prefixLength: 30 }], + roceV2Ipv4Gids: gidIndexes.map((index) => ({ index, address })), + }; +} + +function setRailAddress( + item: ReturnType, + address: string, + prefixLength: number, +): void { + item.ipv4Addresses = [{ address, prefixLength }]; + item.roceV2Ipv4Gids = item.roceV2Ipv4Gids.map((gid) => ({ ...gid, address })); +} + +function hostFixture(side: "local" | "peer"): StationHostProbe { + const isLocal = side === "local"; + const home = isLocal ? LOCAL_HOME : PEER_HOME; + return { + schemaVersion: 1, + hostname: isLocal ? "station-a" : "station-b", + productName: "NVIDIA DGX Station GB300", + architecture: "aarch64", + home, + uid: isLocal ? 1000 : 1001, + gid: isLocal ? 1000 : 1001, + gpus: isLocal + ? [{ index: 0, name: "NVIDIA GB300", uuid: "GPU-aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee" }] + : [ + { + index: 0, + name: "NVIDIA RTX PRO 6000", + uuid: "GPU-11111111-2222-3333-4444-555555555555", + }, + { + index: 1, + name: "NVIDIA GB300 Grace Blackwell Superchip", + uuid: "GPU-99999999-8888-7777-6666-555555555555", + }, + ], + docker: { reachable: true, nvidiaRuntime: true }, + rsyncAvailable: true, + rails: isLocal + ? [ + rail("mlx5_0", "cx8a0", "0001:03:00.0", "192.168.240.1"), + rail("mlx5_1", "cx8a1", "0001:03:00.1", "192.168.240.5"), + ] + : [ + // Deliberately reverse inventory order; matching is by subnet. + rail("mlx5_1", "cx8b1", "0002:03:00.1", "192.168.240.6"), + rail("mlx5_0", "cx8b0", "0002:03:00.0", "192.168.240.2"), + ], + modelSnapshot: { + modelId: DUAL_STATION_VLLM_RUNTIME.modelId, + revision: DUAL_STATION_VLLM_RUNTIME.modelRevision, + path: snapshotPath(home), + directoryExists: !isLocal, + complete: !isLocal, + shardCount: isLocal ? 0 : 113, + reason: isLocal ? "not staged yet" : "", + }, + }; +} + +function command(stdout: unknown, status = 0): StationProbeCommandResult { + return { status, stdout: typeof stdout === "string" ? stdout : JSON.stringify(stdout) }; +} + +function connectivityResponse( + requests: readonly StationRailConnectivityRequest[], + alter: (check: Record, index: number) => void = () => undefined, +): StationProbeCommandResult { + const checks = requests.map((request, index) => { + const check: Record = { + ...request, + routeDevice: request.netdev, + routeSource: request.sourceAddress, + routeGateway: null, + routeScope: "link", + peerMac: request.expectedPeerMac, + peerNeighborState: "REACHABLE", + jumboPing: true, + }; + alter(check, index); + return check; + }); + return command({ schemaVersion: 1, checks }); +} + +type FixtureDeps = StationClusterProbeDeps & { + calls: { + sshConfig: ReturnType; + localHost: ReturnType; + peerHost: ReturnType; + localConnectivity: ReturnType; + peerConnectivity: ReturnType; + }; +}; + +function fixtureDeps( + local = hostFixture("local"), + peer = hostFixture("peer"), + options: { + localConnectivityAlter?: (check: Record, index: number) => void; + peerConnectivityAlter?: (check: Record, index: number) => void; + } = {}, +): FixtureDeps { + const localHost = vi.fn(() => command(local)); + const peerHost = vi.fn(() => command(peer)); + const sshConfig = vi.fn((binding: DualStationSshBinding) => + command(strictDockerSshConfig(binding)), + ); + const localConnectivity = vi.fn((requests: readonly StationRailConnectivityRequest[]) => + connectivityResponse(requests, options.localConnectivityAlter), + ); + const peerConnectivity = vi.fn( + (_binding: DualStationSshBinding, requests: readonly StationRailConnectivityRequest[]) => + connectivityResponse(requests, options.peerConnectivityAlter), + ); + return { + loadPeerSshBinding: loadDualStationSshBindingHandoff, + probePeerSshConfig: sshConfig, + probeLocalHost: localHost, + probePeerHost: peerHost, + probeLocalConnectivity: localConnectivity, + probePeerConnectivity: peerConnectivity, + calls: { sshConfig, localHost, peerHost, localConnectivity, peerConnectivity }, + }; +} + +function runWith(deps: StationClusterProbeDeps, target = "nvidia@station-b") { + sshFixture = retargetDualStationSshBindingFixture( + sshFixture, + target, + validatePeerTarget(target).ok, + ); + return probeDualStationVllmCapability({ + env: { + [NEMOCLAW_DGX_STATION_PEER_ENV]: target, + [NEMOCLAW_DGX_STATION_SSH_BINDING_ENV]: sshFixture.token, + }, + deps, + }); +} + +describe("probeDualStationVllmCapability", () => { + it.each([ + undefined, + "", + " ", + ])("does no work when the explicit peer is absent or blank (%s)", (value) => { + const deps = fixtureDeps(); + const env = value === undefined ? {} : { [NEMOCLAW_DGX_STATION_PEER_ENV]: value }; + + expect(probeDualStationVllmCapability({ env, deps })).toEqual({ kind: "not-configured" }); + expect(deps.calls.sshConfig).not.toHaveBeenCalled(); + expect(deps.calls.localHost).not.toHaveBeenCalled(); + expect(deps.calls.peerHost).not.toHaveBeenCalled(); + expect(deps.calls.localConnectivity).not.toHaveBeenCalled(); + expect(deps.calls.peerConnectivity).not.toHaveBeenCalled(); + }); + + it.each([ + "ssh://station-b", + "-oProxyCommand=bad", + "station-a,station-b", + "station-b:2222", + "user name@station-b", + "station-b;id", + "station-b$(id)", + "user@station-b@other", + " station-b", + "station-b\nother", + "Station-B", + "station_b", + "station..b", + "station-b.", + "1user@station-b", + ])("rejects a non-single-host peer value without executing: %s", (target) => { + const deps = fixtureDeps(); + + expect(runWith(deps, target)).toMatchObject({ kind: "unavailable", code: "invalid-peer" }); + expect(deps.calls.sshConfig).not.toHaveBeenCalled(); + expect(deps.calls.localHost).not.toHaveBeenCalled(); + expect(deps.calls.peerHost).not.toHaveBeenCalled(); + }); + + it("requires the installer-qualified SSH binding before any peer probe", () => { + const deps = fixtureDeps(); + + expect( + probeDualStationVllmCapability({ + env: { [NEMOCLAW_DGX_STATION_PEER_ENV]: "nvidia@station-b" }, + deps, + }), + ).toMatchObject({ kind: "unavailable", code: "peer-ssh-config-unsafe" }); + expect(deps.calls.sshConfig).not.toHaveBeenCalled(); + expect(deps.calls.localHost).not.toHaveBeenCalled(); + expect(deps.calls.peerHost).not.toHaveBeenCalled(); + }); + + it("rejects a changed qualified host-key pin before any peer probe", () => { + const deps = fixtureDeps(); + fs.appendFileSync(sshFixture.binding.knownHostsFile, "changed\n"); + + expect(runWith(deps)).toMatchObject({ + kind: "unavailable", + code: "peer-ssh-config-unsafe", + }); + expect(deps.calls.sshConfig).not.toHaveBeenCalled(); + expect(deps.calls.localHost).not.toHaveBeenCalled(); + expect(deps.calls.peerHost).not.toHaveBeenCalled(); + }); + + it("rejects Docker-over-SSH when the effective operator config weakens peer trust", () => { + const deps = fixtureDeps(); + deps.probePeerSshConfig = (binding) => + command( + strictDockerSshConfig(binding).replace( + "stricthostkeychecking true", + "stricthostkeychecking false", + ), + ); + + expect(runWith(deps)).toMatchObject({ + kind: "unavailable", + code: "peer-ssh-config-unsafe", + }); + expect(deps.calls.localHost).not.toHaveBeenCalled(); + expect(deps.calls.peerHost).not.toHaveBeenCalled(); + }); + + it("rejects an SSH config that can SendEnv arbitrary NemoClaw secrets", () => { + const deps = fixtureDeps(); + deps.probePeerSshConfig = (binding) => command(`${strictDockerSshConfig(binding)}\nsendenv *`); + + expect(runWith(deps)).toMatchObject({ + kind: "unavailable", + code: "peer-ssh-config-unsafe", + }); + expect(deps.calls.localHost).not.toHaveBeenCalled(); + }); + + it.each([ + ["named Docker context", { DOCKER_CONTEXT: "remote-builder" }], + ["remote Docker host", { DOCKER_HOST: "ssh://builder.example" }], + ])("rejects an ambient %s before mixing it with local hardware", (_label, dockerEnv) => { + const deps = fixtureDeps(); + + expect( + probeDualStationVllmCapability({ + env: { + [NEMOCLAW_DGX_STATION_PEER_ENV]: "nvidia@station-b", + [NEMOCLAW_DGX_STATION_SSH_BINDING_ENV]: sshFixture.token, + ...dockerEnv, + }, + deps, + }), + ).toMatchObject({ kind: "unavailable", code: "local-docker-unavailable" }); + expect(deps.calls.sshConfig).not.toHaveBeenCalled(); + expect(deps.calls.localHost).not.toHaveBeenCalled(); + }); + + it.each([ + ["context selected by Station host preparation", { DOCKER_CONTEXT: "default" }], + ["socket detected by the NemoClaw runner", { DOCKER_HOST: "unix:///run/docker.sock" }], + [ + "socket alias detected by the NemoClaw runner", + { DOCKER_HOST: "unix:///var/run/docker.sock" }, + ], + ])("accepts the local default Docker %s", (_label, dockerEnv) => { + const deps = fixtureDeps(); + + expect( + probeDualStationVllmCapability({ + env: { + [NEMOCLAW_DGX_STATION_PEER_ENV]: "nvidia@station-b", + [NEMOCLAW_DGX_STATION_SSH_BINDING_ENV]: sshFixture.token, + ...dockerEnv, + }, + deps, + }), + ).toMatchObject({ kind: "ready" }); + expect(deps.calls.sshConfig).toHaveBeenCalledOnce(); + expect(deps.calls.localHost).toHaveBeenCalledOnce(); + }); + + it.each([ + "station-b", + "nvidia@station-b", + "_svc@192.168.50.20", + ])("returns the qualified binding for %s", (target) => { + const result = runWith(fixtureDeps(), target); + expect(result).toMatchObject({ + kind: "ready", + peerModelSnapshot: "ready", + plan: { peerSshBinding: { peerTarget: target } }, + }); + expect(result.kind).toBe("ready"); + const ready = result as Extract; + expect(buildRemoteVllmDockerEnv(ready.plan.peerSshBinding, {}).DOCKER_HOST).toBe( + `ssh://${ready.plan.peerSshBinding.sshUser}@${ready.plan.peerSshBinding.resolvedHost}`, + ); + }); + + it("returns a deterministic two-rail Ray PP2 plan and permits one auxiliary non-GB300 GPU", () => { + const deps = fixtureDeps(); + + const result = runWith(deps); + + expect(result).toMatchObject({ + kind: "ready", + plan: { + peerSshBinding: { + peerTarget: "nvidia@station-b", + resolvedHost: "192.168.50.20", + sshUser: "nvidia", + }, + runtime: { + image: + "vllm/vllm-openai@sha256:2cc49b81319f7a66a33dd8bd63a7bfddae079122b33ce51989b6828a1f038c37", + modelRevision: "183968f87ae4cedce3039313cac1fd43d112c578", + servedModelId: "nemotron-ultra", + tensorParallelSize: 1, + pipelineParallelSize: 2, + nodeCount: 2, + }, + local: { + home: LOCAL_HOME, + uid: 1000, + gid: 1000, + gpu: { index: 0, name: "NVIDIA GB300" }, + }, + peer: { + home: PEER_HOME, + uid: 1001, + gid: 1001, + gpu: { index: 1, name: "NVIDIA GB300 Grace Blackwell Superchip" }, + }, + masterAddress: "192.168.240.1", + roceGidIndex: 3, + rails: [ + { + subnet: "192.168.240.0/30", + local: { + rdmaDevice: "mlx5_0", + netdev: "cx8a0", + uverbsDevice: "/dev/infiniband/uverbs0", + address: "192.168.240.1", + }, + peer: { + rdmaDevice: "mlx5_0", + netdev: "cx8b0", + uverbsDevice: "/dev/infiniband/uverbs0", + address: "192.168.240.2", + }, + }, + { + subnet: "192.168.240.4/30", + local: { + rdmaDevice: "mlx5_1", + netdev: "cx8a1", + uverbsDevice: "/dev/infiniband/uverbs1", + address: "192.168.240.5", + }, + peer: { + rdmaDevice: "mlx5_1", + netdev: "cx8b1", + uverbsDevice: "/dev/infiniband/uverbs1", + address: "192.168.240.6", + }, + }, + ], + }, + }); + expect(deps.calls.peerHost).toHaveBeenCalledWith(sshFixture.binding); + expect(deps.calls.localConnectivity).toHaveBeenCalledWith([ + { + netdev: "cx8a0", + sourceAddress: "192.168.240.1", + peerAddress: "192.168.240.2", + expectedPeerMac: "02:00:00:bb:00:00", + }, + { + netdev: "cx8a1", + sourceAddress: "192.168.240.5", + peerAddress: "192.168.240.6", + expectedPeerMac: "02:00:00:bb:00:01", + }, + ]); + expect(deps.calls.peerConnectivity).toHaveBeenCalledWith(sshFixture.binding, [ + { + netdev: "cx8b0", + sourceAddress: "192.168.240.2", + peerAddress: "192.168.240.1", + expectedPeerMac: "02:00:00:aa:00:00", + }, + { + netdev: "cx8b1", + sourceAddress: "192.168.240.6", + peerAddress: "192.168.240.5", + expectedPeerMac: "02:00:00:aa:00:01", + }, + ]); + }); + + it("uses the lowest common RoCEv2 IPv4 GID when preferred index 3 is unavailable", () => { + const local = hostFixture("local"); + const peer = hostFixture("peer"); + for (const item of [...local.rails, ...peer.rails]) { + item.roceV2Ipv4Gids = item.roceV2Ipv4Gids.map((gid) => ({ ...gid, index: 5 })); + } + + expect(runWith(fixtureDeps(local, peer))).toMatchObject({ + kind: "ready", + plan: { roceGidIndex: 5 }, + }); + }); + + it.each([ + "DGX-Station", + "P3830", + "NVIDIA Station GB300", + ])("accepts an existing Station firmware product identifier: %s", (productName) => { + const peer = hostFixture("peer"); + peer.productName = productName; + + expect(runWith(fixtureDeps(hostFixture("local"), peer))).toMatchObject({ + kind: "ready", + peerModelSnapshot: "ready", + }); + }); + + it.each([ + { + name: "non-Station peer", + code: "peer-not-station", + mutate: (host: StationHostProbe) => { + host.productName = "Generic Linux Workstation"; + }, + }, + { + name: "more than one peer GB300", + code: "peer-gpu-unavailable", + mutate: (host: StationHostProbe) => { + host.gpus.push({ index: 2, name: "NVIDIA GB300", uuid: "GPU-aaaa-bbbb-cccc-dddd" }); + }, + }, + { + name: "missing peer NVIDIA runtime", + code: "peer-docker-unavailable", + mutate: (host: StationHostProbe) => { + host.docker.nvidiaRuntime = false; + }, + }, + { + name: "slow peer rail", + code: "peer-fabric-unavailable", + mutate: (host: StationHostProbe) => { + host.rails[0].speedMbps = 200_000; + }, + }, + { + name: "non-jumbo peer rail", + code: "peer-fabric-unavailable", + mutate: (host: StationHostProbe) => { + host.rails[1].mtu = 1500; + }, + }, + { + name: "unsupported peer RDMA port", + code: "peer-fabric-unavailable", + mutate: (host: StationHostProbe) => { + host.rails[0].port = 2; + }, + }, + { + name: "missing peer uverbs character device", + code: "peer-fabric-unavailable", + mutate: (host: StationHostProbe) => { + host.rails[0].uverbsDevice = ""; + }, + }, + { + name: "duplicate peer uverbs mapping", + code: "peer-fabric-unavailable", + mutate: (host: StationHostProbe) => { + host.rails[1].uverbsDevice = host.rails[0].uverbsDevice; + }, + }, + { + name: "incomplete peer snapshot", + code: "peer-model-cache-unavailable", + mutate: (host: StationHostProbe) => { + host.modelSnapshot.complete = false; + }, + }, + { + name: "wrong peer snapshot revision", + code: "peer-model-cache-unavailable", + mutate: (host: StationHostProbe) => { + host.modelSnapshot.revision = "f".repeat(40); + }, + }, + { + name: "truncated peer snapshot manifest", + code: "peer-model-cache-unavailable", + mutate: (host: StationHostProbe) => { + host.modelSnapshot.shardCount = 1; + }, + }, + ])("fails closed for $name", ({ code, mutate }) => { + const peer = hostFixture("peer"); + mutate(peer); + + expect(runWith(fixtureDeps(hostFixture("local"), peer))).toMatchObject({ + kind: "unavailable", + code, + }); + }); + + it("qualifies a missing peer snapshot when exact staging prerequisites exist", () => { + const peer = hostFixture("peer"); + peer.modelSnapshot.directoryExists = false; + peer.modelSnapshot.complete = false; + peer.modelSnapshot.shardCount = 0; + peer.modelSnapshot.reason = "snapshot directory is missing"; + + expect(runWith(fixtureDeps(hostFixture("local"), peer))).toMatchObject({ + kind: "ready", + peerModelSnapshot: "staging-required", + }); + }); + + it.each([ + { side: "local", code: "local-model-staging-unavailable" }, + { side: "peer", code: "peer-model-staging-unavailable" }, + ])("requires rsync on the $side host before qualifying a missing snapshot", ({ side, code }) => { + const local = hostFixture("local"); + const peer = hostFixture("peer"); + peer.modelSnapshot.directoryExists = false; + peer.modelSnapshot.complete = false; + peer.modelSnapshot.shardCount = 0; + peer.modelSnapshot.reason = "snapshot directory is missing"; + (side === "local" ? local : peer).rsyncAvailable = false; + + expect(runWith(fixtureDeps(local, peer))).toMatchObject({ kind: "unavailable", code }); + }); + + it("rejects rails that do not form two distinct shared direct subnets", () => { + const peer = hostFixture("peer"); + peer.rails[0].ipv4Addresses = [{ address: "10.20.30.2", prefixLength: 30 }]; + peer.rails[0].roceV2Ipv4Gids = [{ index: 3, address: "10.20.30.2" }]; + + expect(runWith(fixtureDeps(hostFixture("local"), peer))).toMatchObject({ + kind: "unavailable", + code: "fabric-mismatch", + }); + }); + + it("rejects two otherwise matching switched /24 rail networks", () => { + const local = hostFixture("local"); + const peer = hostFixture("peer"); + setRailAddress(local.rails[0], "192.168.100.1", 24); + setRailAddress(local.rails[1], "192.168.101.1", 24); + setRailAddress(peer.rails.find((item) => item.rdmaDevice === "mlx5_0")!, "192.168.100.2", 24); + setRailAddress(peer.rails.find((item) => item.rdmaDevice === "mlx5_1")!, "192.168.101.2", 24); + + expect(runWith(fixtureDeps(local, peer))).toMatchObject({ + kind: "unavailable", + code: "fabric-mismatch", + }); + }); + + it("rejects public addresses even when they form matching /30 rail networks", () => { + const local = hostFixture("local"); + const peer = hostFixture("peer"); + setRailAddress(local.rails[0], "203.0.113.1", 30); + setRailAddress(local.rails[1], "198.51.100.5", 30); + setRailAddress(peer.rails.find((item) => item.rdmaDevice === "mlx5_0")!, "203.0.113.2", 30); + setRailAddress(peer.rails.find((item) => item.rdmaDevice === "mlx5_1")!, "198.51.100.6", 30); + + expect(runWith(fixtureDeps(local, peer))).toMatchObject({ + kind: "unavailable", + code: "fabric-mismatch", + }); + }); + + it("rejects asymmetric RoCEv2 IPv4 GID indexes", () => { + const peer = hostFixture("peer"); + peer.rails[0].roceV2Ipv4Gids = peer.rails[0].roceV2Ipv4Gids.map((gid) => ({ + ...gid, + index: 7, + })); + + expect(runWith(fixtureDeps(hostFixture("local"), peer))).toMatchObject({ + kind: "unavailable", + code: "gid-mismatch", + }); + }); + + it("rejects an SSH target that resolves back to the local Station identity", () => { + const local = hostFixture("local"); + const peer = hostFixture("peer"); + peer.gpus[1].uuid = local.gpus[0].uuid; + + expect(runWith(fixtureDeps(local, peer))).toMatchObject({ + kind: "unavailable", + code: "fabric-mismatch", + }); + }); + + it("allows distinct factory-imaged Stations that report the same hostname", () => { + const local = hostFixture("local"); + const peer = hostFixture("peer"); + peer.hostname = local.hostname; + + expect(runWith(fixtureDeps(local, peer))).toMatchObject({ kind: "ready" }); + }); + + it("rejects a local route that traverses a gateway", () => { + const deps = fixtureDeps(hostFixture("local"), hostFixture("peer"), { + localConnectivityAlter: (check, index) => { + check.routeGateway = index === 0 ? "192.168.240.254" : check.routeGateway; + }, + }); + + expect(runWith(deps)).toMatchObject({ + kind: "unavailable", + code: "local-connectivity-failed", + }); + expect(deps.calls.peerConnectivity).not.toHaveBeenCalled(); + }); + + it("rejects a peer rail that cannot pass an MTU-9000 ping", () => { + const deps = fixtureDeps(hostFixture("local"), hostFixture("peer"), { + peerConnectivityAlter: (check, index) => { + check.jumboPing = index === 1 ? false : check.jumboPing; + }, + }); + + expect(runWith(deps)).toMatchObject({ + kind: "unavailable", + code: "peer-connectivity-failed", + }); + }); + + it("rejects a route without a scope-link connected-prefix proof", () => { + const deps = fixtureDeps(hostFixture("local"), hostFixture("peer"), { + localConnectivityAlter: (check, index) => { + check.routeScope = index === 0 ? "global" : check.routeScope; + }, + }); + + expect(runWith(deps)).toMatchObject({ + kind: "unavailable", + code: "local-connectivity-failed", + }); + }); + + it("rejects a neighbor MAC that does not identify the matched peer rail", () => { + const deps = fixtureDeps(hostFixture("local"), hostFixture("peer"), { + peerConnectivityAlter: (check, index) => { + check.peerMac = index === 0 ? "02:00:00:cc:00:00" : check.peerMac; + }, + }); + + expect(runWith(deps)).toMatchObject({ + kind: "unavailable", + code: "peer-connectivity-failed", + }); + }); + + it("fails closed on command failure or malformed host JSON", () => { + const localFailure = fixtureDeps(); + localFailure.probeLocalHost = () => command("", 1); + expect(runWith(localFailure)).toMatchObject({ + kind: "unavailable", + code: "local-probe-failed", + }); + + const peerMalformed = fixtureDeps(); + peerMalformed.probePeerHost = () => command("not-json"); + expect(runWith(peerMalformed)).toMatchObject({ + kind: "unavailable", + code: "peer-probe-failed", + }); + }); +}); + +describe("probe command boundary", () => { + it("audits the exact effective SSH config used later by Docker transport", () => { + const spawn = vi.fn( + ( + _file: string, + _args: readonly string[], + _options: SpawnSyncOptionsWithStringEncoding, + ): StationProbeCommandResult => command(strictDockerSshConfig(sshFixture.binding)), + ); + const deps = createStationClusterProbeDeps(spawn); + + deps.probePeerSshConfig(sshFixture.binding); + + const [file, args, options] = spawn.mock.calls[0]; + expect(file).toBe("ssh"); + expect(args).toEqual( + expect.arrayContaining([ + "-G", + "BatchMode=yes", + `UserKnownHostsFile=${sshFixture.binding.knownHostsFile}`, + `HostKeyAlias=${sshFixture.binding.lookupHost}`, + `Hostname=${sshFixture.binding.resolvedHost}`, + "User=nvidia", + "Port=22", + "--", + "nvidia@station-b", + ]), + ); + expect(options.input).toBe(""); + expect(options.timeout).toBe(20_000); + }); + + it("uses a fixed stdin script and strict pretrusted SSH without discovery or prompting", () => { + const spawn = vi.fn( + ( + _file: string, + _args: readonly string[], + _options: SpawnSyncOptionsWithStringEncoding, + ): StationProbeCommandResult => command({}), + ); + const deps = createStationClusterProbeDeps(spawn); + + deps.probePeerHost(sshFixture.binding); + + const [file, args, options] = spawn.mock.calls[0]; + expect(file).toBe("ssh"); + expect(args).toEqual( + expect.arrayContaining([ + "BatchMode=yes", + "StrictHostKeyChecking=yes", + "NumberOfPasswordPrompts=0", + "ConnectTimeout=5", + "ClearAllForwardings=yes", + "--", + "nvidia@station-b", + "python3 -", + ]), + ); + expect(args.join(" ")).not.toMatch(/keyscan|accept-new|StrictHostKeyChecking=no/); + expect(options.input).toEqual(expect.stringContaining('docker", "info')); + expect(options.input).toEqual(expect.stringContaining('"docker", "run", "--rm"')); + expect(options.input).toEqual(expect.stringContaining('"--pull=never"')); + expect(options.input).toEqual(expect.stringContaining('"--gpus", "all"')); + expect(options.input).toEqual(expect.stringContaining(STATION_ACCEPTANCE_IMAGE)); + expect(options.input).toEqual(expect.stringContaining("/sys/firmware/devicetree/base/model")); + expect(options.timeout).toBe(20_000); + expect(options.maxBuffer).toBe(1024 * 1024); + }); + + it("executes the host probe and reports malformed weights plus missing staging tools", () => { + const root = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-station-probe-fixture-")); + const home = path.join(root, "home"); + const bin = path.join(root, "bin"); + const snapshot = snapshotPath(home); + fs.mkdirSync(snapshot, { mode: 0o700, recursive: true }); + fs.mkdirSync(bin, { mode: 0o700 }); + fs.writeFileSync(path.join(snapshot, "config.json"), "{}"); + fs.writeFileSync(path.join(snapshot, "tokenizer.json"), "{}"); + const shards = Array.from( + { length: 113 }, + (_, index) => `model-${String(index + 1).padStart(5, "0")}-of-00113.safetensors`, + ); + fs.writeFileSync( + path.join(snapshot, "model.safetensors.index.json"), + JSON.stringify({ + metadata: { total_size: 1 }, + weight_map: Object.fromEntries( + shards.map((shard, index) => [`model.layers.${String(index)}.weight`, shard]), + ), + }), + ); + fs.writeFileSync(path.join(snapshot, shards[0]), "malformed"); + + let probeScript = ""; + const recordingSpawn = vi.fn( + ( + _file: string, + _args: readonly string[], + options: SpawnSyncOptionsWithStringEncoding, + ): StationProbeCommandResult => { + probeScript = typeof options.input === "string" ? options.input : ""; + return command({}); + }, + ); + createStationClusterProbeDeps(recordingSpawn).probeLocalHost(); + const python = resolveStationFixturePython(); + + try { + const executed = spawnSync(python, ["-"], { + encoding: "utf8", + env: { ...process.env, HOME: home, PATH: bin }, + input: probeScript, + timeout: 20_000, + }); + expect(executed.status, executed.stderr).toBe(0); + const observed = JSON.parse(executed.stdout) as StationHostProbe; + expect(observed.modelSnapshot).toMatchObject({ + complete: false, + shardCount: 113, + reason: expect.stringContaining("weight shards are unreadable or malformed"), + }); + expect(observed.rsyncAvailable).toBe(false); + } finally { + fs.rmSync(root, { recursive: true, force: true }); + } + }); + + it("executes the host probe against the digest-pinned no-pull GPU runtime contract", () => { + const root = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-station-runtime-fixture-")); + const home = path.join(root, "home"); + const bin = path.join(root, "bin"); + fs.mkdirSync(home, { mode: 0o700 }); + fs.mkdirSync(bin, { mode: 0o700 }); + + let probeScript = ""; + const recordingSpawn = vi.fn( + ( + _file: string, + _args: readonly string[], + options: SpawnSyncOptionsWithStringEncoding, + ): StationProbeCommandResult => { + probeScript = typeof options.input === "string" ? options.input : ""; + return command({}); + }, + ); + createStationClusterProbeDeps(recordingSpawn).probeLocalHost(); + const python = resolveStationFixturePython(); + const fixturePrelude = String.raw` +import subprocess + +class FixtureResult: + def __init__(self, returncode, stdout=""): + self.returncode = returncode + self.stdout = stdout + self.stderr = "" + +def fixture_run(argv, **_kwargs): + if argv and argv[0] == "nvidia-smi": + return FixtureResult(0, "0, NVIDIA GB300, GPU-11111111-2222-3333-4444-555555555555") + if argv[:2] == ["docker", "info"]: + return FixtureResult(0, "29.2.1") + if argv[:2] == ["docker", "run"]: + return FixtureResult(0, "GPU-11111111-2222-3333-4444-555555555555") + return FixtureResult(127) + +subprocess.run = fixture_run +`; + + try { + const executed = spawnSync(python, ["-"], { + encoding: "utf8", + env: { ...process.env, HOME: home, PATH: bin }, + input: `${fixturePrelude}\n${probeScript}`, + timeout: 20_000, + }); + expect(executed.status, executed.stderr).toBe(0); + const observed = JSON.parse(executed.stdout) as StationHostProbe; + expect(observed.docker).toEqual({ reachable: true, nvidiaRuntime: true }); + } finally { + fs.rmSync(root, { recursive: true, force: true }); + } + }); + + it("executes the host probe and refuses a non-character uverbs device", () => { + const root = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-station-verbs-fixture-")); + const home = path.join(root, "home"); + const bin = path.join(root, "bin"); + fs.mkdirSync(home, { mode: 0o700 }); + fs.mkdirSync(bin, { mode: 0o700 }); + + let probeScript = ""; + const recordingSpawn = vi.fn( + ( + _file: string, + _args: readonly string[], + options: SpawnSyncOptionsWithStringEncoding, + ): StationProbeCommandResult => { + probeScript = typeof options.input === "string" ? options.input : ""; + return command({}); + }, + ); + createStationClusterProbeDeps(recordingSpawn).probeLocalHost(); + const python = resolveStationFixturePython(); + const fixturePrelude = String.raw` +import pathlib +import stat as fixture_stat +import subprocess + +class FixtureResult: + def __init__(self, returncode, stdout=""): + self.returncode = returncode + self.stdout = stdout + self.stderr = "" + +original_iterdir = pathlib.Path.iterdir +original_stat = pathlib.Path.stat + +def fixture_iterdir(candidate): + if str(candidate) == "/sys/class/infiniband/mlx5_0/device/infiniband_verbs": + return iter([candidate / "uverbs0"]) + return original_iterdir(candidate) + +def fixture_path_stat(candidate, *args, **kwargs): + if str(candidate) == "/dev/infiniband/uverbs0": + return type("FixtureStat", (), {"st_mode": fixture_stat.S_IFREG | 0o600})() + return original_stat(candidate, *args, **kwargs) + +def fixture_run(argv, **_kwargs): + if argv and argv[0] == "ibdev2netdev": + return FixtureResult(0, "mlx5_0 port 1 ==> cx8p0 (Up)") + return FixtureResult(127) + +pathlib.Path.iterdir = fixture_iterdir +pathlib.Path.stat = fixture_path_stat +subprocess.run = fixture_run +`; + + try { + const executed = spawnSync(python, ["-"], { + encoding: "utf8", + env: { ...process.env, HOME: home, PATH: bin }, + input: `${fixturePrelude}\n${probeScript}`, + timeout: 20_000, + }); + expect(executed.status, executed.stderr).toBe(0); + const observed = JSON.parse(executed.stdout) as StationHostProbe; + expect(observed.rails).toHaveLength(1); + expect(observed.rails[0]).toMatchObject({ + rdmaDevice: "mlx5_0", + netdev: "cx8p0", + uverbsDevice: "", + }); + } finally { + fs.rmSync(root, { recursive: true, force: true }); + } + }); + + it("removes OPENSHELL secrets from the direct SSH probe environment", () => { + vi.stubEnv("OPENSHELL_GATEWAY_AUTH_TOKEN", "must-not-cross-ssh"); + const spawn = vi.fn( + ( + _file: string, + _args: readonly string[], + _options: SpawnSyncOptionsWithStringEncoding, + ): StationProbeCommandResult => command({}), + ); + const deps = createStationClusterProbeDeps(spawn); + + deps.probePeerHost(sshFixture.binding); + + const [, , options] = spawn.mock.calls[0]; + expect(options.env?.OPENSHELL_GATEWAY_AUTH_TOKEN).toBeUndefined(); + expect(options.env?.PATH).toBeTruthy(); + }); + + it("pins the local hardware probe to Docker's physical default context", () => { + const spawn = vi.fn( + ( + _file: string, + _args: readonly string[], + _options: SpawnSyncOptionsWithStringEncoding, + ): StationProbeCommandResult => command({}), + ); + const deps = createStationClusterProbeDeps(spawn); + + deps.probeLocalHost(); + + const [file, , options] = spawn.mock.calls[0]; + expect(file).toBe("python3"); + expect(options.env?.DOCKER_CONTEXT).toBe("default"); + expect(options.env?.DOCKER_HOST).toBeUndefined(); + expect(options.env?.DOCKER_CONFIG).toBeUndefined(); + expect(options.input).toEqual(expect.stringContaining('"gid": os.getgid()')); + }); + + it("passes only validated discovered rail values to the fixed peer connectivity script", () => { + const spawn = vi.fn( + ( + _file: string, + _args: readonly string[], + _options: SpawnSyncOptionsWithStringEncoding, + ): StationProbeCommandResult => command({}), + ); + const deps = createStationClusterProbeDeps(spawn); + const requests = [ + { + netdev: "cx8b0", + sourceAddress: "192.168.240.2", + peerAddress: "192.168.240.1", + expectedPeerMac: "02:00:00:aa:00:00", + }, + { + netdev: "cx8b1", + sourceAddress: "192.168.240.6", + peerAddress: "192.168.240.5", + expectedPeerMac: "02:00:00:aa:00:01", + }, + ]; + + deps.probePeerConnectivity(sshFixture.binding, requests); + + const [, args, options] = spawn.mock.calls[0]; + expect(args.at(-1)).toBe( + "python3 - cx8b0 192.168.240.2 192.168.240.1 cx8b1 192.168.240.6 192.168.240.5", + ); + expect(options.input).toEqual(expect.stringContaining('"ping", "-4", "-M", "do"')); + expect(options.input).toEqual( + expect.stringContaining('"route", "get", peer, "from", source, "oif", netdev'), + ); + expect(options.input).toEqual( + expect.stringContaining('"route", "show", "exact", network, "dev", netdev'), + ); + expect(options.input).toEqual( + expect.stringContaining('"neighbor", "show", "to", peer, "dev", netdev'), + ); + expect(options.input).toEqual(expect.stringContaining('"-I", source, peer')); + }); + + it("executes the connectivity probe against the route JSON emitted on Station", () => { + const requests = [ + { + netdev: "cx8r0", + sourceAddress: "192.168.240.1", + peerAddress: "192.168.240.2", + expectedPeerMac: "ac:3a:e2:de:3a:23", + }, + { + netdev: "cx8r1", + sourceAddress: "192.168.240.5", + peerAddress: "192.168.240.6", + expectedPeerMac: "ac:3a:e2:de:3a:24", + }, + ]; + let probeScript = ""; + let probeArgs: readonly string[] = []; + const recordingSpawn = vi.fn( + ( + _file: string, + args: readonly string[], + options: SpawnSyncOptionsWithStringEncoding, + ): StationProbeCommandResult => { + probeArgs = args; + probeScript = typeof options.input === "string" ? options.input : ""; + return command({}); + }, + ); + createStationClusterProbeDeps(recordingSpawn).probeLocalConnectivity(requests); + const fixturePrelude = String.raw` +import json +import subprocess + +class FixtureResult: + def __init__(self, returncode, stdout=""): + self.returncode = returncode + self.stdout = stdout + self.stderr = "" + +def fixture_run(argv, **_kwargs): + if argv[:4] == ["ip", "-j", "route", "get"]: + peer, source, netdev = argv[4], argv[6], argv[8] + return FixtureResult(0, json.dumps([{ + "dst": peer, + "from": source, + "dev": netdev, + "flags": [], + "uid": 1000, + "cache": [], + }])) + if argv[:5] == ["ip", "-j", "route", "show", "exact"]: + network = argv[5] + return FixtureResult(0, json.dumps([{ + "dst": network, + "protocol": "kernel", + "scope": "link", + "prefsrc": "192.168.240.1", + "metric": 100, + "flags": [], + }])) + if argv and argv[0] == "ping": + return FixtureResult(0) + if argv[:3] == ["ip", "-j", "neighbor"]: + peer, netdev = argv[5], argv[7] + mac = { + "192.168.240.2": "ac:3a:e2:de:3a:23", + "192.168.240.6": "ac:3a:e2:de:3a:24", + }[peer] + return FixtureResult(0, json.dumps([{ + "dst": peer, + "lladdr": mac, + "state": ["REACHABLE"], + }])) + return FixtureResult(127) + +subprocess.run = fixture_run +`; + const executed = spawnSync(resolveStationFixturePython(), probeArgs, { + encoding: "utf8", + input: `${fixturePrelude}\n${probeScript}`, + timeout: 20_000, + }); + + expect(executed.status, executed.stderr).toBe(0); + const observed = JSON.parse(executed.stdout) as { + checks: Array>; + }; + expect(observed.checks).toEqual([ + expect.objectContaining({ + netdev: "cx8r0", + routeDevice: "cx8r0", + routeSource: "192.168.240.1", + routeScope: "link", + jumboPing: true, + }), + expect.objectContaining({ + netdev: "cx8r1", + routeDevice: "cx8r1", + routeSource: "192.168.240.5", + routeScope: "link", + jumboPing: true, + }), + ]); + }); +}); + +describe("parseStationHostProbe", () => { + it("rejects unsupported schemas and unsafe device names", () => { + const unsupported = { ...hostFixture("local"), schemaVersion: 2 }; + expect(() => parseStationHostProbe(JSON.stringify(unsupported))).toThrow(/schema version/); + + const unsafe = hostFixture("local"); + unsafe.rails[0].netdev = "cx8;touch /tmp/pwned"; + expect(() => parseStationHostProbe(JSON.stringify(unsafe))).toThrow(/unsafe device name/); + + const unsafeUverbs = hostFixture("local"); + unsafeUverbs.rails[0].uverbsDevice = "/dev/infiniband/../mem"; + expect(() => parseStationHostProbe(JSON.stringify(unsafeUverbs))).toThrow(/uverbs/); + }); + + it("rejects root or invalid runtime cache-owner identities", () => { + const root = hostFixture("local"); + root.uid = 0; + expect(() => parseStationHostProbe(JSON.stringify(root))).toThrow(/host probe\.uid/); + + const rootGroup = hostFixture("peer"); + rootGroup.gid = 0; + expect(() => parseStationHostProbe(JSON.stringify(rootGroup))).toThrow(/host probe\.gid/); + }); +}); diff --git a/src/lib/inference/vllm-station-cluster.ts b/src/lib/inference/vllm-station-cluster.ts new file mode 100644 index 0000000000..3ce57481f2 --- /dev/null +++ b/src/lib/inference/vllm-station-cluster.ts @@ -0,0 +1,1633 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { type SpawnSyncOptionsWithStringEncoding, spawnSync } from "node:child_process"; +import net from "node:net"; +import path from "node:path"; + +import { buildSubprocessEnv } from "../subprocess-env"; +import { isDgxStationGb300Product } from "./dgx-station-identity"; +import { buildVllmSshTransportEnv } from "./vllm-docker-env"; +import { NEMOTRON_ULTRA_DUAL_STATION_IMAGE, VLLM_MODELS } from "./vllm-models"; +import { + type DualStationSshBinding, + dualStationPinnedSshArgs, + loadDualStationSshBindingHandoff, + NEMOCLAW_DGX_STATION_SSH_BINDING_ENV, +} from "./vllm-station-ssh-binding"; + +export const NEMOCLAW_DGX_STATION_PEER_ENV = "NEMOCLAW_DGX_STATION_PEER"; + +const HOST_PROBE_SCHEMA_VERSION = 1; +const CONNECTIVITY_PROBE_SCHEMA_VERSION = 1; +const COMMAND_TIMEOUT_MS = 20_000; +const MAX_PROBE_OUTPUT_BYTES = 1024 * 1024; +const PREFERRED_ROCE_GID_INDEX = 3; +const EXPECTED_ULTRA_WEIGHT_SHARDS = 113; +const DIRECT_RAIL_PREFIX_LENGTH = 30; +const STATION_RUNTIME_PROBE_IMAGE = + "docker.io/library/ubuntu@sha256:7f622ca8766bccb22f04242ecb6f19f770b2f08827dc4b8c707de5e78a6da7ab"; +const DUAL_STATION_LOCAL_DOCKER_OVERRIDE_ENV_NAMES = [ + "DOCKER_API_VERSION", + "DOCKER_CERT_PATH", + "DOCKER_CONFIG", + "DOCKER_CONTEXT", + "DOCKER_HOST", + "DOCKER_TLS", + "DOCKER_TLS_VERIFY", +] as const; +const CANONICAL_SSH_HOST_PATTERN = + /^[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?(?:\.[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?)*$/; +const CANONICAL_SSH_USERNAME_PATTERN = /^[A-Za-z_][A-Za-z0-9._-]*$/; + +const ultraModel = VLLM_MODELS.find((model) => model.envValue === "nemotron-3-ultra-550b-a55b"); +if (!ultraModel?.revision) { + throw new Error("Nemotron Ultra must have an immutable Hugging Face revision"); +} + +export const DUAL_STATION_VLLM_RUNTIME = Object.freeze({ + image: NEMOTRON_ULTRA_DUAL_STATION_IMAGE.arm64.ref, + modelId: ultraModel.id, + modelRevision: ultraModel.revision, + servedModelId: "nemotron-ultra", + tensorParallelSize: 1 as const, + pipelineParallelSize: 2 as const, + nodeCount: 2 as const, +}); + +export interface StationGpuProbe { + index: number; + name: string; + uuid: string; +} + +export interface StationIpv4AddressProbe { + address: string; + prefixLength: number; +} + +export interface StationRoceGidProbe { + index: number; + address: string; +} + +export interface StationRailProbe { + rdmaDevice: string; + port: number; + netdev: string; + macAddress: string; + uverbsDevice: string; + pciAddress: string; + pciName: string; + state: string; + linkLayer: string; + speedMbps: number; + mtu: number; + ipv4Addresses: StationIpv4AddressProbe[]; + roceV2Ipv4Gids: StationRoceGidProbe[]; +} + +export interface StationModelSnapshotProbe { + modelId: string; + revision: string; + path: string; + directoryExists: boolean; + complete: boolean; + shardCount: number; + reason: string; +} + +export interface StationHostProbe { + schemaVersion: 1; + hostname: string; + productName: string; + architecture: string; + home: string; + uid: number; + gid: number; + gpus: StationGpuProbe[]; + docker: { + reachable: boolean; + nvidiaRuntime: boolean; + }; + rsyncAvailable: boolean; + rails: StationRailProbe[]; + modelSnapshot: StationModelSnapshotProbe; +} + +export interface StationRailConnectivityRequest { + netdev: string; + sourceAddress: string; + peerAddress: string; + expectedPeerMac: string; +} + +export interface StationRailConnectivityProbe { + netdev: string; + sourceAddress: string; + peerAddress: string; + routeDevice: string; + routeSource: string; + routeGateway: string | null; + routeScope: string; + peerMac: string; + peerNeighborState: string; + jumboPing: boolean; +} + +export interface StationProbeCommandResult { + status: number | null; + stdout: string; + stderr?: string; + error?: string; +} + +export interface StationClusterProbeDeps { + loadPeerSshBinding(token: string, expectedPeerTarget: string): DualStationSshBinding; + probePeerSshConfig(binding: DualStationSshBinding): StationProbeCommandResult; + probeLocalHost(): StationProbeCommandResult; + probePeerHost(binding: DualStationSshBinding): StationProbeCommandResult; + probeLocalConnectivity( + requests: readonly StationRailConnectivityRequest[], + ): StationProbeCommandResult; + probePeerConnectivity( + binding: DualStationSshBinding, + requests: readonly StationRailConnectivityRequest[], + ): StationProbeCommandResult; +} + +export type StationClusterFailureCode = + | "invalid-peer" + | "local-probe-failed" + | "peer-probe-failed" + | "local-not-station" + | "peer-not-station" + | "local-gpu-unavailable" + | "peer-gpu-unavailable" + | "local-docker-unavailable" + | "peer-docker-unavailable" + | "peer-ssh-config-unsafe" + | "local-model-staging-unavailable" + | "peer-model-staging-unavailable" + | "local-fabric-unavailable" + | "peer-fabric-unavailable" + | "fabric-mismatch" + | "gid-mismatch" + | "peer-model-cache-unavailable" + | "local-connectivity-failed" + | "peer-connectivity-failed" + | "probe-error"; + +export interface DualStationPlanNode { + hostname: string; + home: string; + uid: number; + gid: number; + gpu: StationGpuProbe; +} + +export interface DualStationPlanRailEndpoint { + rdmaDevice: string; + netdev: string; + macAddress: string; + uverbsDevice: string; + pciAddress: string; + address: string; +} + +export interface DualStationPlanRail { + index: number; + subnet: string; + local: DualStationPlanRailEndpoint; + peer: DualStationPlanRailEndpoint; +} + +export interface DualStationVllmPlan { + peerSshBinding: DualStationSshBinding; + runtime: typeof DUAL_STATION_VLLM_RUNTIME; + local: DualStationPlanNode; + peer: DualStationPlanNode; + rails: DualStationPlanRail[]; + masterAddress: string; + roceGidIndex: number; +} + +export type StationClusterCapability = + | { kind: "not-configured" } + | { kind: "unavailable"; code: StationClusterFailureCode; reason: string } + | { + kind: "ready"; + plan: DualStationVllmPlan; + peerModelSnapshot: "ready" | "staging-required"; + }; + +type PlanFailure = Extract; + +type MatchedRail = { + localRail: StationRailProbe; + peerRail: StationRailProbe; + localAddress: StationIpv4AddressProbe; + peerAddress: StationIpv4AddressProbe; + subnet: string; +}; + +type StaticPlan = { + plan: DualStationVllmPlan; + peerModelSnapshot: "ready" | "staging-required"; + localConnectivity: StationRailConnectivityRequest[]; + peerConnectivity: StationRailConnectivityRequest[]; +}; + +type StationProbeSpawn = ( + file: string, + args: readonly string[], + options: SpawnSyncOptionsWithStringEncoding, +) => StationProbeCommandResult; + +const HOST_PROBE_SCRIPT = String.raw` +import csv +import ipaddress +import json +import os +from pathlib import Path +import platform +import re +import shutil +import socket +import stat +import subprocess + +MODEL_ID = ${JSON.stringify(DUAL_STATION_VLLM_RUNTIME.modelId)} +MODEL_REVISION = ${JSON.stringify(DUAL_STATION_VLLM_RUNTIME.modelRevision)} +MODEL_CACHE_NAME = "models--" + MODEL_ID.replace("/", "--") + +def read_text(path): + try: + return Path(path).read_text(encoding="utf-8").rstrip("\x00").strip() + except (OSError, UnicodeError): + return "" + +def run(argv, timeout=5): + try: + result = subprocess.run( + argv, + stdin=subprocess.DEVNULL, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + timeout=timeout, + check=False, + ) + return result.returncode, result.stdout.strip() + except (FileNotFoundError, OSError, subprocess.TimeoutExpired): + return 127, "" + +def product_name(): + for candidate in ( + "/sys/class/dmi/id/product_name", + "/sys/devices/virtual/dmi/id/product_name", + "/sys/firmware/devicetree/base/model", + ): + value = read_text(candidate) + if value: + return value + return "" + +def gpu_inventory(): + rc, output = run([ + "nvidia-smi", + "--query-gpu=index,name,uuid", + "--format=csv,noheader,nounits", + ]) + if rc != 0: + return [] + result = [] + for row in csv.reader(output.splitlines()): + if len(row) != 3: + continue + try: + index = int(row[0].strip()) + except ValueError: + continue + result.append({ + "index": index, + "name": row[1].strip(), + "uuid": row[2].strip(), + }) + return result + +def docker_state(): + rc, _ = run(["docker", "info", "--format", "{{.ServerVersion}}"]) + if rc != 0: + return {"reachable": False, "nvidiaRuntime": False} + runtime_rc, runtime_output = run([ + "docker", "run", "--rm", "--pull=never", "--network", "none", + "--read-only", "--cap-drop", "ALL", "--security-opt", "no-new-privileges", + "--pids-limit", "64", "--memory", "512m", "--cpus", "1", + "--gpus", "all", "${STATION_RUNTIME_PROBE_IMAGE}", + "nvidia-smi", "--query-gpu=uuid", "--format=csv,noheader", + ], timeout=15) + return { + "reachable": True, + "nvidiaRuntime": runtime_rc == 0 and any( + line.strip().startswith("GPU-") + for line in runtime_output.splitlines() + ), + } + +def ipv4_addresses(netdev): + rc, output = run(["ip", "-j", "-4", "address", "show", "dev", netdev]) + if rc != 0: + return [] + try: + links = json.loads(output) + except json.JSONDecodeError: + return [] + result = [] + for link in links if isinstance(links, list) else []: + for address in link.get("addr_info", []): + if address.get("family") != "inet" or address.get("scope") == "host": + continue + local = address.get("local") + prefix = address.get("prefixlen") + if isinstance(local, str) and isinstance(prefix, int): + result.append({"address": local, "prefixLength": prefix}) + return result + +def roce_v2_ipv4_gids(rdma_device, port, netdev): + base = Path("/sys/class/infiniband") / rdma_device / "ports" / str(port) + types_dir = base / "gid_attrs" / "types" + result = [] + try: + indexes = sorted(types_dir.iterdir(), key=lambda item: int(item.name)) + except (OSError, ValueError): + return result + for entry in indexes: + try: + index = int(entry.name) + except ValueError: + continue + if read_text(entry).lower() != "roce v2": + continue + observed_netdev = read_text(base / "gid_attrs" / "ndevs" / str(index)) + if observed_netdev and observed_netdev != netdev: + continue + raw_gid = read_text(base / "gids" / str(index)) + try: + mapped = ipaddress.IPv6Address(raw_gid).ipv4_mapped + except ipaddress.AddressValueError: + mapped = None + if mapped is not None: + result.append({"index": index, "address": str(mapped)}) + return result + +def uverbs_device(rdma_device): + verbs_dir = Path("/sys/class/infiniband") / rdma_device / "device" / "infiniband_verbs" + try: + names = sorted({ + entry.name + for entry in verbs_dir.iterdir() + if re.fullmatch(r"uverbs[0-9]+", entry.name) + }) + except OSError: + return "" + if len(names) != 1: + return "" + device = Path("/dev/infiniband") / names[0] + try: + if not stat.S_ISCHR(device.stat().st_mode): + return "" + except OSError: + return "" + return str(device) + +def rail_inventory(): + rc, output = run(["ibdev2netdev"]) + if rc != 0: + return [] + rails = [] + pattern = re.compile(r"^(\S+)\s+port\s+(\d+)\s+==>\s+(\S+)\s+\(([^)]*)\)") + for line in output.splitlines(): + match = pattern.match(line.strip()) + if not match: + continue + rdma_device, raw_port, netdev, _reported_state = match.groups() + port = int(raw_port) + device_path = Path("/sys/class/net") / netdev / "device" + try: + pci_address = device_path.resolve(strict=True).name + except OSError: + pci_address = "" + rc_lspci, pci_name = run(["lspci", "-D", "-s", pci_address]) if pci_address else (127, "") + if rc_lspci != 0: + pci_name = "" + try: + speed_mbps = int(read_text(Path("/sys/class/net") / netdev / "speed")) + except ValueError: + speed_mbps = -1 + try: + mtu = int(read_text(Path("/sys/class/net") / netdev / "mtu")) + except ValueError: + mtu = -1 + ib_port = Path("/sys/class/infiniband") / rdma_device / "ports" / str(port) + rails.append({ + "rdmaDevice": rdma_device, + "port": port, + "netdev": netdev, + "macAddress": read_text(Path("/sys/class/net") / netdev / "address").lower(), + "uverbsDevice": uverbs_device(rdma_device), + "pciAddress": pci_address, + "pciName": pci_name, + "state": read_text(ib_port / "state"), + "linkLayer": read_text(ib_port / "link_layer"), + "speedMbps": speed_mbps, + "mtu": mtu, + "ipv4Addresses": ipv4_addresses(netdev), + "roceV2Ipv4Gids": roce_v2_ipv4_gids(rdma_device, port, netdev), + }) + return rails + +def snapshot_state(): + home = Path.home() + snapshot = home / ".cache" / "huggingface" / "hub" / MODEL_CACHE_NAME / "snapshots" / MODEL_REVISION + reasons = [] + shard_count = 0 + index_path = snapshot / "model.safetensors.index.json" + if not snapshot.is_dir(): + reasons.append("snapshot directory is missing") + if not (snapshot / "config.json").is_file(): + reasons.append("config.json is missing") + tokenizer_present = any( + (snapshot / name).is_file() + for name in ("tokenizer.json", "tokenizer.model", "vocab.json") + ) + if not tokenizer_present: + reasons.append("tokenizer assets are missing") + try: + index = json.loads(index_path.read_text(encoding="utf-8")) + weight_map = index.get("weight_map", {}) + metadata = index.get("metadata", {}) + shards = sorted(set(weight_map.values())) if isinstance(weight_map, dict) else [] + if not shards or not all(isinstance(item, str) for item in shards): + reasons.append("weight index is empty or malformed") + shards = [] + if len(shards) != ${String(113)}: + reasons.append("weight index does not list the pinned shard count") + expected_total_size = metadata.get("total_size") if isinstance(metadata, dict) else None + if not isinstance(expected_total_size, int) or expected_total_size <= 0: + reasons.append("weight index total_size is missing or malformed") + shard_count = len(shards) + observed_tensor_size = 0 + for shard in shards: + shard_path = Path(shard) + if ( + shard_path.is_absolute() + or shard_path.name != shard + or shard in (".", "..") + or shard_path.suffix != ".safetensors" + ): + reasons.append("weight index contains an unsafe shard path") + break + candidate = snapshot / shard_path + try: + if not candidate.is_file() or candidate.stat().st_size <= 0: + reasons.append("one or more weight shards are missing") + break + with candidate.open("rb") as handle: + raw_header_size = handle.read(8) + if len(raw_header_size) != 8: + raise ValueError("short safetensors header") + header_size = int.from_bytes(raw_header_size, "little") + if header_size <= 0 or header_size > 128 * 1024 * 1024: + raise ValueError("invalid safetensors header size") + header = json.loads(handle.read(header_size)) + tensor_ranges = [] + for tensor_name, tensor in header.items(): + if tensor_name == "__metadata__": + continue + offsets = tensor.get("data_offsets") if isinstance(tensor, dict) else None + if ( + not isinstance(offsets, list) + or len(offsets) != 2 + or not all(isinstance(offset, int) for offset in offsets) + or offsets[0] < 0 + or offsets[1] < offsets[0] + ): + raise ValueError("invalid safetensors data offsets") + tensor_ranges.append(offsets) + if not tensor_ranges: + raise ValueError("empty safetensors shard") + payload_size = max(offsets[1] for offsets in tensor_ranges) + if candidate.stat().st_size != 8 + header_size + payload_size: + raise ValueError("truncated safetensors shard") + observed_tensor_size += sum(offsets[1] - offsets[0] for offsets in tensor_ranges) + except (OSError, UnicodeError, ValueError, json.JSONDecodeError): + reasons.append("one or more weight shards are unreadable or malformed") + break + if isinstance(expected_total_size, int) and observed_tensor_size != expected_total_size: + reasons.append("weight shard sizes do not match the pinned index") + except (OSError, UnicodeError, json.JSONDecodeError): + reasons.append("model.safetensors.index.json is missing or malformed") + return { + "modelId": MODEL_ID, + "revision": MODEL_REVISION, + "path": str(snapshot), + "directoryExists": snapshot.is_dir(), + "complete": not reasons, + "shardCount": shard_count, + "reason": "; ".join(dict.fromkeys(reasons)), + } + +payload = { + "schemaVersion": ${String(HOST_PROBE_SCHEMA_VERSION)}, + "hostname": socket.gethostname(), + "productName": product_name(), + "architecture": platform.machine(), + "home": str(Path.home()), + "uid": os.getuid(), + "gid": os.getgid(), + "gpus": gpu_inventory(), + "docker": docker_state(), + "rsyncAvailable": shutil.which("rsync") is not None, + "rails": rail_inventory(), + "modelSnapshot": snapshot_state(), +} +print(json.dumps(payload, separators=(",", ":"))) +`; + +const CONNECTIVITY_PROBE_SCRIPT = String.raw` +import ipaddress +import json +import subprocess +import sys + +def run(argv, timeout=5): + try: + result = subprocess.run( + argv, + stdin=subprocess.DEVNULL, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + timeout=timeout, + check=False, + ) + return result.returncode, result.stdout.strip() + except (FileNotFoundError, OSError, subprocess.TimeoutExpired): + return 127, "" + +if len(sys.argv[1:]) != 6: + raise SystemExit("expected two netdev/source/peer triples") + +checks = [] +for offset in (0, 3): + netdev, source, peer = sys.argv[1 + offset:4 + offset] + route_device = "" + route_source = "" + route_gateway = None + route_scope = "" + peer_mac = "" + peer_neighbor_state = "" + rc, output = run(["ip", "-j", "route", "get", peer, "from", source, "oif", netdev]) + if rc == 0: + try: + routes = json.loads(output) + route = routes[0] if isinstance(routes, list) and routes else {} + route_device = route.get("dev", "") if isinstance(route, dict) else "" + route_source = ( + route.get("prefsrc") or route.get("src") or route.get("from", "") + if isinstance(route, dict) + else "" + ) + route_gateway = route.get("gateway") if isinstance(route, dict) else None + except json.JSONDecodeError: + pass + network = str(ipaddress.ip_network(source + "/${String(DIRECT_RAIL_PREFIX_LENGTH)}", strict=False)) + link_rc, link_output = run(["ip", "-j", "route", "show", "exact", network, "dev", netdev]) + if link_rc == 0: + try: + link_routes = json.loads(link_output) + link_route = link_routes[0] if isinstance(link_routes, list) and link_routes else {} + if ( + isinstance(link_route, dict) + and link_route.get("dst") == network + and link_route.get("dev", netdev) == netdev + and link_route.get("gateway") is None + ): + route_scope = link_route.get("scope", "") + except json.JSONDecodeError: + pass + ping_rc, _ = run([ + "ping", "-4", "-M", "do", "-s", "8972", "-c", "1", "-W", "2", + "-I", source, peer, + ]) + neighbor_rc, neighbor_output = run(["ip", "-j", "neighbor", "show", "to", peer, "dev", netdev]) + if neighbor_rc == 0: + try: + neighbors = json.loads(neighbor_output) + neighbor = neighbors[0] if isinstance(neighbors, list) and neighbors else {} + if ( + isinstance(neighbor, dict) + and neighbor.get("dst") == peer + and neighbor.get("dev", netdev) == netdev + ): + peer_mac = str(neighbor.get("lladdr", "")).lower() + raw_state = neighbor.get("state", "") + peer_neighbor_state = ",".join(raw_state) if isinstance(raw_state, list) else str(raw_state) + except json.JSONDecodeError: + pass + checks.append({ + "netdev": netdev, + "sourceAddress": source, + "peerAddress": peer, + "routeDevice": route_device, + "routeSource": route_source, + "routeGateway": route_gateway, + "routeScope": route_scope, + "peerMac": peer_mac, + "peerNeighborState": peer_neighbor_state, + "jumboPing": ping_rc == 0, + }) + +print(json.dumps({ + "schemaVersion": ${String(CONNECTIVITY_PROBE_SCHEMA_VERSION)}, + "checks": checks, +}, separators=(",", ":"))) +`; + +function defaultSpawn( + file: string, + args: readonly string[], + options: SpawnSyncOptionsWithStringEncoding, +): StationProbeCommandResult { + const result = spawnSync(file, [...args], options); + return { + status: result.status, + stdout: result.stdout ?? "", + stderr: result.stderr ?? "", + error: result.error?.message, + }; +} + +function commandOptions(input: string): SpawnSyncOptionsWithStringEncoding { + return { + encoding: "utf8", + input, + timeout: COMMAND_TIMEOUT_MS, + maxBuffer: MAX_PROBE_OUTPUT_BYTES, + killSignal: "SIGKILL", + windowsHide: true, + env: buildSubprocessEnv({ LC_ALL: "C" }), + }; +} + +function localProbeCommandOptions(input: string): SpawnSyncOptionsWithStringEncoding { + const options = commandOptions(input); + for (const name of DUAL_STATION_LOCAL_DOCKER_OVERRIDE_ENV_NAMES) delete options.env?.[name]; + if (options.env) options.env.DOCKER_CONTEXT = "default"; + return options; +} + +function sshCommandOptions(input: string): SpawnSyncOptionsWithStringEncoding { + return { + ...commandOptions(input), + env: buildVllmSshTransportEnv({ LC_ALL: "C" }), + }; +} + +function strictSshArgs(binding: DualStationSshBinding, remoteCommand: string): string[] { + return [...dualStationPinnedSshArgs(binding), "--", binding.peerTarget, remoteCommand]; +} + +function connectivityArgv(requests: readonly StationRailConnectivityRequest[]): string[] { + if (requests.length !== 2) { + throw new Error("dual-Station connectivity requires exactly two rail requests"); + } + const args: string[] = []; + for (const request of requests) { + if (!isSafeDeviceName(request.netdev)) throw new Error("unsafe connectivity netdev"); + if (!isIpv4(request.sourceAddress) || !isIpv4(request.peerAddress)) { + throw new Error("invalid connectivity address"); + } + normalizeMacAddress(request.expectedPeerMac, "connectivity peer MAC"); + args.push(request.netdev, request.sourceAddress, request.peerAddress); + } + return args; +} + +/** + * Construct the real read-only probe boundary. The optional spawn injection is + * intentionally lower-level than StationClusterProbeDeps so tests can assert + * the exact SSH trust flags and fixed-stdin behavior without making a network + * connection. + */ +export function createStationClusterProbeDeps( + spawn: StationProbeSpawn = defaultSpawn, +): StationClusterProbeDeps { + return { + loadPeerSshBinding: loadDualStationSshBindingHandoff, + probePeerSshConfig: (binding) => { + return spawn( + "ssh", + ["-G", ...dualStationPinnedSshArgs(binding), "--", binding.peerTarget], + sshCommandOptions(""), + ); + }, + probeLocalHost: () => spawn("python3", ["-"], localProbeCommandOptions(HOST_PROBE_SCRIPT)), + probePeerHost: (binding) => { + return spawn( + "ssh", + strictSshArgs(binding, "python3 -"), + sshCommandOptions(HOST_PROBE_SCRIPT), + ); + }, + probeLocalConnectivity: (requests) => { + const args = connectivityArgv(requests); + return spawn("python3", ["-", ...args], localProbeCommandOptions(CONNECTIVITY_PROBE_SCRIPT)); + }, + probePeerConnectivity: (binding, requests) => { + const args = connectivityArgv(requests); + const remoteCommand = ["python3", "-", ...args].join(" "); + return spawn( + "ssh", + strictSshArgs(binding, remoteCommand), + sshCommandOptions(CONNECTIVITY_PROBE_SCRIPT), + ); + }, + }; +} + +const defaultStationClusterProbeDeps = createStationClusterProbeDeps(); + +type PeerValidation = { ok: true; target: string } | { ok: false; reason: string }; + +export function validatePeerTarget(raw: string): PeerValidation { + if (raw.length === 0 || raw !== raw.trim()) { + return { ok: false, reason: `${NEMOCLAW_DGX_STATION_PEER_ENV} must not contain whitespace` }; + } + if (raw.length > 286) { + return { ok: false, reason: `${NEMOCLAW_DGX_STATION_PEER_ENV} is too long` }; + } + if (/[/,:;`'"\\$(){}[\]<>|&!?*\s\u0000-\u001f\u007f]/.test(raw)) { + return { + ok: false, + reason: `${NEMOCLAW_DGX_STATION_PEER_ENV} must name one SSH host or user@host`, + }; + } + const parts = raw.split("@"); + if (parts.length > 2) { + return { + ok: false, + reason: `${NEMOCLAW_DGX_STATION_PEER_ENV} must name one SSH host or user@host`, + }; + } + const username = parts.length === 2 ? parts[0] : ""; + const hostname = parts.at(-1) ?? ""; + const validHostname = net.isIP(hostname) === 4 || CANONICAL_SSH_HOST_PATTERN.test(hostname); + if ( + !validHostname || + (username.length > 0 && !CANONICAL_SSH_USERNAME_PATTERN.test(username)) || + (parts.length === 2 && username.length === 0) + ) { + return { + ok: false, + reason: `${NEMOCLAW_DGX_STATION_PEER_ENV} must name one canonical SSH host or user@host`, + }; + } + return { ok: true, target: raw }; +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +function requireRecord(value: unknown, label: string): Record { + if (!isRecord(value)) throw new Error(`${label} must be an object`); + return value; +} + +function requireString(value: unknown, label: string, maxLength = 1024): string { + if ( + typeof value !== "string" || + value.length === 0 || + value.length > maxLength || + /[\u0000-\u001f\u007f]/.test(value) + ) { + throw new Error(`${label} must be a non-empty printable string`); + } + return value; +} + +function requireBoolean(value: unknown, label: string): boolean { + if (typeof value !== "boolean") throw new Error(`${label} must be a boolean`); + return value; +} + +function requireInteger(value: unknown, label: string, min: number, max: number): number { + if (!Number.isInteger(value) || (value as number) < min || (value as number) > max) { + throw new Error(`${label} must be an integer between ${String(min)} and ${String(max)}`); + } + return value as number; +} + +function requireArray(value: unknown, label: string, maxLength: number): unknown[] { + if (!Array.isArray(value) || value.length > maxLength) { + throw new Error(`${label} must be an array with at most ${String(maxLength)} entries`); + } + return value; +} + +function isSafeDeviceName(value: string): boolean { + return /^[A-Za-z0-9][A-Za-z0-9_.:-]{0,63}$/.test(value); +} + +function isIpv4(value: string): boolean { + return net.isIP(value) === 4; +} + +function normalizeMacAddress(value: unknown, label: string, allowEmpty = false): string { + if (allowEmpty && value === "") return ""; + const mac = requireString(value, label, 17).toLowerCase(); + if (!/^(?:[0-9a-f]{2}:){5}[0-9a-f]{2}$/.test(mac)) { + throw new Error(`${label} must be a canonical MAC address`); + } + return mac; +} + +function parseUverbsDevice(value: unknown, label: string): string { + if (value === "") return ""; + const device = requireString(value, label, 64); + if (!/^\/dev\/infiniband\/uverbs[0-9]+$/.test(device)) { + throw new Error(`${label} must be a safe /dev/infiniband/uverbs* character-device path`); + } + return device; +} + +function parseGpu(value: unknown, label: string): StationGpuProbe { + const record = requireRecord(value, label); + const name = requireString(record.name, `${label}.name`, 256); + const uuid = requireString(record.uuid, `${label}.uuid`, 128); + if (!/^GPU-[A-Za-z0-9-]+$/.test(uuid)) throw new Error(`${label}.uuid is invalid`); + return { + index: requireInteger(record.index, `${label}.index`, 0, 1024), + name, + uuid, + }; +} + +function parseIpv4Address(value: unknown, label: string): StationIpv4AddressProbe { + const record = requireRecord(value, label); + const address = requireString(record.address, `${label}.address`, 15); + if (!isIpv4(address)) throw new Error(`${label}.address must be IPv4`); + return { + address, + prefixLength: requireInteger(record.prefixLength, `${label}.prefixLength`, 1, 32), + }; +} + +function parseGid(value: unknown, label: string): StationRoceGidProbe { + const record = requireRecord(value, label); + const address = requireString(record.address, `${label}.address`, 15); + if (!isIpv4(address)) throw new Error(`${label}.address must be IPv4`); + return { + index: requireInteger(record.index, `${label}.index`, 0, 4095), + address, + }; +} + +function parseRail(value: unknown, label: string): StationRailProbe { + const record = requireRecord(value, label); + const rdmaDevice = requireString(record.rdmaDevice, `${label}.rdmaDevice`, 64); + const netdev = requireString(record.netdev, `${label}.netdev`, 64); + if (!isSafeDeviceName(rdmaDevice) || !isSafeDeviceName(netdev)) { + throw new Error(`${label} contains an unsafe device name`); + } + const pciAddress = requireString(record.pciAddress, `${label}.pciAddress`, 32); + if (!/^[0-9A-Fa-f]{4}:[0-9A-Fa-f]{2}:[0-9A-Fa-f]{2}\.[0-7]$/.test(pciAddress)) { + throw new Error(`${label}.pciAddress is invalid`); + } + return { + rdmaDevice, + port: requireInteger(record.port, `${label}.port`, 1, 255), + netdev, + macAddress: normalizeMacAddress(record.macAddress, `${label}.macAddress`), + uverbsDevice: parseUverbsDevice(record.uverbsDevice, `${label}.uverbsDevice`), + pciAddress, + pciName: requireString(record.pciName, `${label}.pciName`, 512), + state: requireString(record.state, `${label}.state`, 128), + linkLayer: requireString(record.linkLayer, `${label}.linkLayer`, 128), + speedMbps: requireInteger(record.speedMbps, `${label}.speedMbps`, -1, 1_000_000), + mtu: requireInteger(record.mtu, `${label}.mtu`, -1, 1_000_000), + ipv4Addresses: requireArray(record.ipv4Addresses, `${label}.ipv4Addresses`, 16).map( + (entry, index) => parseIpv4Address(entry, `${label}.ipv4Addresses[${String(index)}]`), + ), + roceV2Ipv4Gids: requireArray(record.roceV2Ipv4Gids, `${label}.roceV2Ipv4Gids`, 128).map( + (entry, index) => parseGid(entry, `${label}.roceV2Ipv4Gids[${String(index)}]`), + ), + }; +} + +function parseSnapshot(value: unknown, label: string): StationModelSnapshotProbe { + const record = requireRecord(value, label); + const snapshotPath = requireString(record.path, `${label}.path`, 4096); + if ( + !path.posix.isAbsolute(snapshotPath) || + path.posix.normalize(snapshotPath) !== snapshotPath || + snapshotPath.includes(":") + ) { + throw new Error(`${label}.path must be a normalized absolute POSIX path`); + } + return { + modelId: requireString(record.modelId, `${label}.modelId`, 512), + revision: requireString(record.revision, `${label}.revision`, 128), + path: snapshotPath, + directoryExists: requireBoolean(record.directoryExists, `${label}.directoryExists`), + complete: requireBoolean(record.complete, `${label}.complete`), + shardCount: requireInteger(record.shardCount, `${label}.shardCount`, 0, 100_000), + reason: typeof record.reason === "string" ? record.reason.slice(0, 1024) : "", + }; +} + +export function parseStationHostProbe(stdout: string): StationHostProbe { + if (Buffer.byteLength(stdout, "utf8") > MAX_PROBE_OUTPUT_BYTES) { + throw new Error("host probe output is too large"); + } + let parsed: unknown; + try { + parsed = JSON.parse(stdout); + } catch { + throw new Error("host probe did not return valid JSON"); + } + const record = requireRecord(parsed, "host probe"); + if (record.schemaVersion !== HOST_PROBE_SCHEMA_VERSION) { + throw new Error("host probe schema version is unsupported"); + } + const home = requireString(record.home, "host probe.home", 4096); + if (!path.posix.isAbsolute(home) || path.posix.normalize(home) !== home || home.includes(":")) { + throw new Error("host probe.home must be a normalized absolute POSIX path"); + } + const docker = requireRecord(record.docker, "host probe.docker"); + return { + schemaVersion: 1, + hostname: requireString(record.hostname, "host probe.hostname", 256), + productName: requireString(record.productName, "host probe.productName", 512), + architecture: requireString(record.architecture, "host probe.architecture", 64), + home, + uid: requireInteger(record.uid, "host probe.uid", 1, 2_147_483_647), + gid: requireInteger(record.gid, "host probe.gid", 1, 2_147_483_647), + gpus: requireArray(record.gpus, "host probe.gpus", 64).map((entry, index) => + parseGpu(entry, `host probe.gpus[${String(index)}]`), + ), + docker: { + reachable: requireBoolean(docker.reachable, "host probe.docker.reachable"), + nvidiaRuntime: requireBoolean(docker.nvidiaRuntime, "host probe.docker.nvidiaRuntime"), + }, + rsyncAvailable: requireBoolean(record.rsyncAvailable, "host probe.rsyncAvailable"), + rails: requireArray(record.rails, "host probe.rails", 32).map((entry, index) => + parseRail(entry, `host probe.rails[${String(index)}]`), + ), + modelSnapshot: parseSnapshot(record.modelSnapshot, "host probe.modelSnapshot"), + }; +} + +function parseConnectivityProbe(stdout: string): StationRailConnectivityProbe[] { + if (Buffer.byteLength(stdout, "utf8") > MAX_PROBE_OUTPUT_BYTES) { + throw new Error("connectivity probe output is too large"); + } + let parsed: unknown; + try { + parsed = JSON.parse(stdout); + } catch { + throw new Error("connectivity probe did not return valid JSON"); + } + const record = requireRecord(parsed, "connectivity probe"); + if (record.schemaVersion !== CONNECTIVITY_PROBE_SCHEMA_VERSION) { + throw new Error("connectivity probe schema version is unsupported"); + } + return requireArray(record.checks, "connectivity probe.checks", 2).map((value, index) => { + const label = `connectivity probe.checks[${String(index)}]`; + const check = requireRecord(value, label); + const netdev = requireString(check.netdev, `${label}.netdev`, 64); + const sourceAddress = requireString(check.sourceAddress, `${label}.sourceAddress`, 15); + const peerAddress = requireString(check.peerAddress, `${label}.peerAddress`, 15); + const routeDevice = + check.routeDevice === "" ? "" : requireString(check.routeDevice, `${label}.routeDevice`, 64); + const routeSource = + check.routeSource === "" ? "" : requireString(check.routeSource, `${label}.routeSource`, 15); + const routeScope = + check.routeScope === "" ? "" : requireString(check.routeScope, `${label}.routeScope`, 32); + const peerNeighborState = + check.peerNeighborState === "" + ? "" + : requireString(check.peerNeighborState, `${label}.peerNeighborState`, 128); + if ( + !isSafeDeviceName(netdev) || + !isIpv4(sourceAddress) || + !isIpv4(peerAddress) || + (routeDevice !== "" && !isSafeDeviceName(routeDevice)) || + (routeSource !== "" && !isIpv4(routeSource)) + ) { + throw new Error(`${label} contains invalid route data`); + } + let routeGateway: string | null = null; + if (check.routeGateway !== null && check.routeGateway !== undefined) { + routeGateway = requireString(check.routeGateway, `${label}.routeGateway`, 15); + if (!isIpv4(routeGateway)) throw new Error(`${label}.routeGateway must be IPv4`); + } + return { + netdev, + sourceAddress, + peerAddress, + routeDevice, + routeSource, + routeGateway, + routeScope, + peerMac: normalizeMacAddress(check.peerMac, `${label}.peerMac`, true), + peerNeighborState, + jumboPing: requireBoolean(check.jumboPing, `${label}.jumboPing`), + }; + }); +} + +function commandSucceeded(result: StationProbeCommandResult): boolean { + return result.status === 0 && !result.error && result.stdout.trim().length > 0; +} + +function unavailable(code: StationClusterFailureCode, reason: string): PlanFailure { + return { kind: "unavailable", code, reason }; +} + +function dockerSshConfigIsStrict( + result: StationProbeCommandResult, + binding: DualStationSshBinding, +): boolean { + if (!commandSucceeded(result)) return false; + const values = new Map(); + for (const rawLine of result.stdout.split(/\r?\n/)) { + const line = rawLine.trim(); + if (!line) continue; + const separator = line.search(/\s/); + if (separator <= 0) return false; + const key = line.slice(0, separator).toLowerCase(); + const value = line.slice(separator).trim(); + values.set(key, [...(values.get(key) ?? []), value]); + } + const exactly = (key: string, allowed: readonly string[]): boolean => { + const observed = (values.get(key) ?? []).map((value) => value.toLowerCase()); + return observed.length === 1 && allowed.includes(observed[0]); + }; + const exactlyValue = (key: string, expected: string): boolean => { + const observed = values.get(key) ?? []; + return observed.length === 1 && observed[0] === expected; + }; + const absentOrNone = (key: string): boolean => { + const observed = (values.get(key) ?? []).map((value) => value.toLowerCase()); + return observed.length === 0 || (observed.length === 1 && observed[0] === "none"); + }; + const sendEnv = (values.get("sendenv") ?? []).map((value) => value.toLowerCase()); + return ( + exactlyValue("hostname", binding.resolvedHost) && + exactlyValue("user", binding.sshUser) && + exactlyValue("port", String(binding.port)) && + exactlyValue("hostkeyalias", binding.lookupHost) && + exactlyValue("userknownhostsfile", binding.knownHostsFile) && + exactlyValue("globalknownhostsfile", "/dev/null") && + exactly("batchmode", ["yes"]) && + exactly("stricthostkeychecking", ["yes", "true"]) && + exactly("permitlocalcommand", ["no"]) && + exactly("forwardagent", ["no"]) && + exactly("forwardx11", ["no"]) && + exactly("forwardx11trusted", ["no"]) && + exactly("tunnel", ["false", "no"]) && + exactly("updatehostkeys", ["false", "no"]) && + exactly("controlmaster", ["false", "no"]) && + exactly("controlpersist", ["no", "0"]) && + absentOrNone("controlpath") && + absentOrNone("remotecommand") && + absentOrNone("proxycommand") && + absentOrNone("proxyjump") && + absentOrNone("localcommand") && + !values.has("localforward") && + !values.has("remoteforward") && + !values.has("dynamicforward") && + absentOrNone("knownhostscommand") && + !values.has("setenv") && + sendEnv.every((value) => value === "lang" || value === "lc_*") + ); +} + +function selectGb300(host: StationHostProbe): StationGpuProbe | null { + const matches = host.gpus.filter((gpu) => /\bGB300\b/i.test(gpu.name)); + return matches.length === 1 ? matches[0] : null; +} + +function qualifiedRails(host: StationHostProbe): StationRailProbe[] | null { + const cx8Rails = host.rails.filter((rail) => /ConnectX[- ]?8|\bCX-?8\b/i.test(rail.pciName)); + if (cx8Rails.length !== 2) return null; + const uniqueRdma = new Set(cx8Rails.map((rail) => rail.rdmaDevice)); + const uniqueNetdev = new Set(cx8Rails.map((rail) => rail.netdev)); + const uniquePci = new Set(cx8Rails.map((rail) => rail.pciAddress)); + const uniqueMac = new Set(cx8Rails.map((rail) => rail.macAddress)); + const uniqueUverbs = new Set(cx8Rails.map((rail) => rail.uverbsDevice)); + if ( + uniqueRdma.size !== 2 || + uniqueNetdev.size !== 2 || + uniquePci.size !== 2 || + uniqueMac.size !== 2 || + uniqueUverbs.size !== 2 || + uniqueUverbs.has("") + ) { + return null; + } + for (const rail of cx8Rails) { + const firstMacOctet = Number.parseInt(rail.macAddress.slice(0, 2), 16); + if ( + rail.macAddress === "00:00:00:00:00:00" || + (firstMacOctet & 1) !== 0 || + rail.port !== 1 || + !/\bACTIVE\b/i.test(rail.state) || + rail.linkLayer.toLowerCase() !== "ethernet" || + rail.speedMbps !== 400_000 || + rail.mtu !== 9000 || + rail.ipv4Addresses.length === 0 || + rail.roceV2Ipv4Gids.length === 0 + ) { + return null; + } + } + return cx8Rails; +} + +function ipv4ToNumber(address: string): number { + return address + .split(".") + .map(Number) + .reduce((value, octet) => value * 256 + octet, 0); +} + +function numberToIpv4(value: number): string { + return [24, 16, 8, 0].map((shift) => Math.floor(value / 2 ** shift) % 256).join("."); +} + +function subnetKey(address: StationIpv4AddressProbe): string { + const hostBits = 32 - address.prefixLength; + const divisor = 2 ** hostBits; + const network = Math.floor(ipv4ToNumber(address.address) / divisor) * divisor; + return `${numberToIpv4(network)}/${String(address.prefixLength)}`; +} + +function isPrivateFabricIpv4(address: string): boolean { + const value = ipv4ToNumber(address); + return ( + (value >= ipv4ToNumber("10.0.0.0") && value <= ipv4ToNumber("10.255.255.255")) || + (value >= ipv4ToNumber("172.16.0.0") && value <= ipv4ToNumber("172.31.255.255")) || + (value >= ipv4ToNumber("192.168.0.0") && value <= ipv4ToNumber("192.168.255.255")) + ); +} + +function sharedAddressPairs( + localRail: StationRailProbe, + peerRail: StationRailProbe, +): Array<{ + localAddress: StationIpv4AddressProbe; + peerAddress: StationIpv4AddressProbe; + subnet: string; +}> { + const matches: Array<{ + localAddress: StationIpv4AddressProbe; + peerAddress: StationIpv4AddressProbe; + subnet: string; + }> = []; + for (const localAddress of localRail.ipv4Addresses) { + for (const peerAddress of peerRail.ipv4Addresses) { + const localSubnet = subnetKey(localAddress); + if ( + localAddress.address !== peerAddress.address && + localAddress.prefixLength === DIRECT_RAIL_PREFIX_LENGTH && + peerAddress.prefixLength === DIRECT_RAIL_PREFIX_LENGTH && + isPrivateFabricIpv4(localAddress.address) && + isPrivateFabricIpv4(peerAddress.address) && + localSubnet === subnetKey(peerAddress) + ) { + matches.push({ localAddress, peerAddress, subnet: localSubnet }); + } + } + } + return matches; +} + +function matchRails( + localRails: readonly StationRailProbe[], + peerRails: readonly StationRailProbe[], +): MatchedRail[] | null { + const permutations = [ + [0, 1], + [1, 0], + ] as const; + const candidates: MatchedRail[][] = []; + for (const permutation of permutations) { + const candidate: MatchedRail[] = []; + let valid = true; + for (let index = 0; index < 2; index += 1) { + const localRail = localRails[index]; + const peerRail = peerRails[permutation[index]]; + const addressMatches = sharedAddressPairs(localRail, peerRail); + if (addressMatches.length !== 1) { + valid = false; + break; + } + candidate.push({ localRail, peerRail, ...addressMatches[0] }); + } + if (valid && new Set(candidate.map((match) => match.subnet)).size === 2) { + candidates.push(candidate); + } + } + if (candidates.length !== 1) return null; + return candidates[0].sort((left, right) => + left.localRail.rdmaDevice.localeCompare(right.localRail.rdmaDevice, undefined, { + numeric: true, + }), + ); +} + +function gidsForMatchedAddress(rail: StationRailProbe, address: string): Set { + return new Set( + rail.roceV2Ipv4Gids.filter((gid) => gid.address === address).map((gid) => gid.index), + ); +} + +function commonGidIndex(matches: readonly MatchedRail[]): number | null { + const sets = matches.flatMap((match) => [ + gidsForMatchedAddress(match.localRail, match.localAddress.address), + gidsForMatchedAddress(match.peerRail, match.peerAddress.address), + ]); + if (sets.some((set) => set.size === 0)) return null; + const common = [...sets[0]].filter((index) => sets.slice(1).every((set) => set.has(index))); + if (common.includes(PREFERRED_ROCE_GID_INDEX)) return PREFERRED_ROCE_GID_INDEX; + return common.sort((left, right) => left - right)[0] ?? null; +} + +function expectedPeerSnapshotPath(peer: StationHostProbe): string { + return path.posix.join( + peer.home, + ".cache", + "huggingface", + "hub", + `models--${DUAL_STATION_VLLM_RUNTIME.modelId.replace("/", "--")}`, + "snapshots", + DUAL_STATION_VLLM_RUNTIME.modelRevision, + ); +} + +function buildStaticPlan( + peerSshBinding: DualStationSshBinding, + local: StationHostProbe, + peer: StationHostProbe, +): StaticPlan | PlanFailure { + if ( + !isDgxStationGb300Product(local.productName) || + !/^(?:aarch64|arm64)$/i.test(local.architecture) + ) { + return unavailable("local-not-station", "local host is not a verified arm64 DGX Station"); + } + if ( + !isDgxStationGb300Product(peer.productName) || + !/^(?:aarch64|arm64)$/i.test(peer.architecture) + ) { + return unavailable("peer-not-station", "configured peer is not a verified arm64 DGX Station"); + } + + const localGpu = selectGb300(local); + if (!localGpu) { + return unavailable("local-gpu-unavailable", "local host must expose exactly one GB300 GPU"); + } + const peerGpu = selectGb300(peer); + if (!peerGpu) { + return unavailable("peer-gpu-unavailable", "configured peer must expose exactly one GB300 GPU"); + } + if (localGpu.uuid === peerGpu.uuid) { + return unavailable( + "fabric-mismatch", + "configured peer resolved to the local Station instead of a distinct host", + ); + } + if (!local.docker.reachable || !local.docker.nvidiaRuntime) { + return unavailable( + "local-docker-unavailable", + "local Docker daemon and NVIDIA runtime could not both be verified", + ); + } + if (!peer.docker.reachable || !peer.docker.nvidiaRuntime) { + return unavailable( + "peer-docker-unavailable", + "peer Docker daemon and NVIDIA runtime could not both be verified", + ); + } + const localRails = qualifiedRails(local); + if (!localRails) { + return unavailable( + "local-fabric-unavailable", + "local host does not have two active Ethernet 400G MTU-9000 CX8 RDMA rails", + ); + } + const peerRails = qualifiedRails(peer); + if (!peerRails) { + return unavailable( + "peer-fabric-unavailable", + "peer does not have two active Ethernet 400G MTU-9000 CX8 RDMA rails", + ); + } + + const snapshot = peer.modelSnapshot; + if ( + snapshot.modelId !== DUAL_STATION_VLLM_RUNTIME.modelId || + snapshot.revision !== DUAL_STATION_VLLM_RUNTIME.modelRevision || + snapshot.path !== expectedPeerSnapshotPath(peer) + ) { + return unavailable( + "peer-model-cache-unavailable", + "peer reported an unexpected Nemotron Ultra snapshot identity or path", + ); + } + const peerSnapshotReady = + snapshot.directoryExists && + snapshot.complete && + snapshot.shardCount === EXPECTED_ULTRA_WEIGHT_SHARDS; + if (!peerSnapshotReady && snapshot.directoryExists) { + return unavailable( + "peer-model-cache-unavailable", + "peer has an incomplete pinned Nemotron Ultra snapshot; refusing to overwrite it", + ); + } + if (!peerSnapshotReady && !local.rsyncAvailable) { + return unavailable( + "local-model-staging-unavailable", + "local rsync is required to stage the pinned Nemotron Ultra snapshot", + ); + } + if (!peerSnapshotReady && !peer.rsyncAvailable) { + return unavailable( + "peer-model-staging-unavailable", + "peer rsync is required to stage the pinned Nemotron Ultra snapshot", + ); + } + + const matches = matchRails(localRails, peerRails); + if (!matches) { + return unavailable( + "fabric-mismatch", + "the two hosts do not expose one unambiguous pair of distinct private /30 CX8 subnets", + ); + } + const gidIndex = commonGidIndex(matches); + if (gidIndex === null) { + return unavailable( + "gid-mismatch", + "the four matched rail endpoints do not share a RoCEv2 IPv4 GID index", + ); + } + + const rails = matches.map( + (match, index): DualStationPlanRail => ({ + index, + subnet: match.subnet, + local: { + rdmaDevice: match.localRail.rdmaDevice, + netdev: match.localRail.netdev, + macAddress: match.localRail.macAddress, + uverbsDevice: match.localRail.uverbsDevice, + pciAddress: match.localRail.pciAddress, + address: match.localAddress.address, + }, + peer: { + rdmaDevice: match.peerRail.rdmaDevice, + netdev: match.peerRail.netdev, + macAddress: match.peerRail.macAddress, + uverbsDevice: match.peerRail.uverbsDevice, + pciAddress: match.peerRail.pciAddress, + address: match.peerAddress.address, + }, + }), + ); + + return { + plan: { + peerSshBinding, + runtime: DUAL_STATION_VLLM_RUNTIME, + local: { + hostname: local.hostname, + home: local.home, + uid: local.uid, + gid: local.gid, + gpu: localGpu, + }, + peer: { + hostname: peer.hostname, + home: peer.home, + uid: peer.uid, + gid: peer.gid, + gpu: peerGpu, + }, + rails, + masterAddress: rails[0].local.address, + roceGidIndex: gidIndex, + }, + peerModelSnapshot: peerSnapshotReady ? "ready" : "staging-required", + localConnectivity: rails.map((rail) => ({ + netdev: rail.local.netdev, + sourceAddress: rail.local.address, + peerAddress: rail.peer.address, + expectedPeerMac: rail.peer.macAddress, + })), + peerConnectivity: rails.map((rail) => ({ + netdev: rail.peer.netdev, + sourceAddress: rail.peer.address, + peerAddress: rail.local.address, + expectedPeerMac: rail.local.macAddress, + })), + }; +} + +function connectivityMatches( + requests: readonly StationRailConnectivityRequest[], + observed: readonly StationRailConnectivityProbe[], +): boolean { + if (observed.length !== requests.length) return false; + const byKey = new Map( + observed.map((check) => [`${check.netdev}|${check.sourceAddress}|${check.peerAddress}`, check]), + ); + if (byKey.size !== observed.length) return false; + return requests.every((request) => { + const check = byKey.get(`${request.netdev}|${request.sourceAddress}|${request.peerAddress}`); + return Boolean( + check && + check.routeDevice === request.netdev && + check.routeSource === request.sourceAddress && + check.routeGateway === null && + check.routeScope.toLowerCase() === "link" && + check.peerMac === request.expectedPeerMac && + /^(?:REACHABLE|STALE|DELAY|PROBE|PERMANENT|NOARP)(?:,(?:REACHABLE|STALE|DELAY|PROBE|PERMANENT|NOARP))*$/i.test( + check.peerNeighborState, + ) && + check.jumboPing, + ); + }); +} + +function parseHostCommand( + result: StationProbeCommandResult, + code: "local-probe-failed" | "peer-probe-failed", +): StationHostProbe | PlanFailure { + if (!commandSucceeded(result)) { + return unavailable( + code, + code === "local-probe-failed" ? "local host probe failed" : "peer host probe failed", + ); + } + try { + return parseStationHostProbe(result.stdout.trim()); + } catch { + return unavailable( + code, + code === "local-probe-failed" + ? "local host probe returned invalid data" + : "peer host probe returned invalid data", + ); + } +} + +export interface ProbeDualStationVllmOptions { + env?: NodeJS.ProcessEnv; + deps?: StationClusterProbeDeps; +} + +/** + * Read-only, fail-closed capability probe for the explicit two-Station path. + * + * An unset/blank NEMOCLAW_DGX_STATION_PEER returns before touching deps. A + * configured peer is contacted only through the installer-qualified binding; + * the probe never discovers hosts, changes known_hosts, prompts, or mutates + * either machine. + */ +export function probeDualStationVllmCapability( + options: ProbeDualStationVllmOptions = {}, +): StationClusterCapability { + const env = options.env ?? process.env; + const rawPeer = env[NEMOCLAW_DGX_STATION_PEER_ENV]; + if (rawPeer === undefined || rawPeer === "" || rawPeer.trim() === "") { + return { kind: "not-configured" }; + } + const peerValidation = validatePeerTarget(rawPeer); + if (!peerValidation.ok) return unavailable("invalid-peer", peerValidation.reason); + const rawBinding = env[NEMOCLAW_DGX_STATION_SSH_BINDING_ENV]; + if (rawBinding === undefined || rawBinding === "" || rawBinding.trim() === "") { + return unavailable( + "peer-ssh-config-unsafe", + `${NEMOCLAW_DGX_STATION_SSH_BINDING_ENV} must identify the installer-qualified peer`, + ); + } + const localDockerOverride = DUAL_STATION_LOCAL_DOCKER_OVERRIDE_ENV_NAMES.find((name) => { + const value = String(env[name] ?? "").trim(); + const isLocalDefaultSelection = + (name === "DOCKER_CONTEXT" && value === "default") || + (name === "DOCKER_HOST" && + (value === "unix:///run/docker.sock" || value === "unix:///var/run/docker.sock")); + return value !== "" && !isLocalDefaultSelection; + }); + if (localDockerOverride) { + return unavailable( + "local-docker-unavailable", + `${localDockerOverride} must be unset so dual-Station setup can bind the physical local Docker daemon`, + ); + } + + const deps = options.deps ?? defaultStationClusterProbeDeps; + let peerSshBinding: DualStationSshBinding; + try { + peerSshBinding = deps.loadPeerSshBinding(rawBinding, peerValidation.target); + } catch { + return unavailable( + "peer-ssh-config-unsafe", + "installer-qualified Station SSH binding is invalid or changed", + ); + } + if (peerSshBinding.peerTarget !== peerValidation.target) { + return unavailable( + "peer-ssh-config-unsafe", + "qualified Station SSH binding does not match the configured peer", + ); + } + try { + const sshConfig = deps.probePeerSshConfig(peerSshBinding); + if (!dockerSshConfigIsStrict(sshConfig, peerSshBinding)) { + return unavailable( + "peer-ssh-config-unsafe", + "configured peer SSH options are not safe for Docker transport; require BatchMode=yes, StrictHostKeyChecking=yes, no forwarding/proxy/local commands, and no connection sharing", + ); + } + const localResult = parseHostCommand(deps.probeLocalHost(), "local-probe-failed"); + if ("kind" in localResult) return localResult; + const peerResult = parseHostCommand(deps.probePeerHost(peerSshBinding), "peer-probe-failed"); + if ("kind" in peerResult) return peerResult; + + const staticPlan = buildStaticPlan(peerSshBinding, localResult, peerResult); + if ("kind" in staticPlan) return staticPlan; + + const localConnectivityResult = deps.probeLocalConnectivity(staticPlan.localConnectivity); + if (!commandSucceeded(localConnectivityResult)) { + return unavailable( + "local-connectivity-failed", + "local dual-rail route and jumbo-frame probe failed", + ); + } + let localConnectivity: StationRailConnectivityProbe[]; + try { + localConnectivity = parseConnectivityProbe(localConnectivityResult.stdout.trim()); + } catch { + return unavailable( + "local-connectivity-failed", + "local dual-rail route probe returned invalid data", + ); + } + if (!connectivityMatches(staticPlan.localConnectivity, localConnectivity)) { + return unavailable( + "local-connectivity-failed", + "local routes are not direct on both matched rails or jumbo ping failed", + ); + } + + const peerConnectivityResult = deps.probePeerConnectivity( + peerSshBinding, + staticPlan.peerConnectivity, + ); + if (!commandSucceeded(peerConnectivityResult)) { + return unavailable( + "peer-connectivity-failed", + "peer dual-rail route and jumbo-frame probe failed", + ); + } + let peerConnectivity: StationRailConnectivityProbe[]; + try { + peerConnectivity = parseConnectivityProbe(peerConnectivityResult.stdout.trim()); + } catch { + return unavailable( + "peer-connectivity-failed", + "peer dual-rail route probe returned invalid data", + ); + } + if (!connectivityMatches(staticPlan.peerConnectivity, peerConnectivity)) { + return unavailable( + "peer-connectivity-failed", + "peer routes are not direct on both matched rails or jumbo ping failed", + ); + } + + return { + kind: "ready", + plan: staticPlan.plan, + peerModelSnapshot: staticPlan.peerModelSnapshot, + }; + } catch { + return unavailable("probe-error", "dual-Station capability probe failed closed"); + } +} diff --git a/src/lib/inference/vllm-station-fixture.test-support.ts b/src/lib/inference/vllm-station-fixture.test-support.ts new file mode 100644 index 0000000000..d3d335efa1 --- /dev/null +++ b/src/lib/inference/vllm-station-fixture.test-support.ts @@ -0,0 +1,21 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import fs from "node:fs"; +import path from "node:path"; + +export function resolveStationFixturePython(): string { + for (const directory of (process.env.PATH ?? "").split(path.delimiter)) { + if (!directory || !path.isAbsolute(directory) || path.normalize(directory) !== directory) { + continue; + } + try { + const candidate = fs.realpathSync(path.join(directory, "python3")); + fs.accessSync(candidate, fs.constants.X_OK); + if (fs.statSync(candidate).isFile()) return candidate; + } catch { + // Keep searching PATH for an executable fixture interpreter. + } + } + throw new Error("python3 is required for the Station fixtures"); +} diff --git a/src/lib/inference/vllm-station-lifecycle-lock.test.ts b/src/lib/inference/vllm-station-lifecycle-lock.test.ts new file mode 100644 index 0000000000..f654e12cef --- /dev/null +++ b/src/lib/inference/vllm-station-lifecycle-lock.test.ts @@ -0,0 +1,116 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; +import { describe, expect, it, vi } from "vitest"; +import { + DUAL_STATION_CONTROLLER_UID_FILE, + type DualStationControllerUidFileStat, + readDualStationControllerUid, + withDualStationVllmLifecycleLock, +} from "./vllm-station-lifecycle-lock"; + +function controllerUidStat( + kind: "directory" | "file", + overrides: Partial = {}, +): DualStationControllerUidFileStat { + return { + uid: 0, + gid: 0, + mode: kind === "directory" ? 0o40755 : 0o100644, + size: kind === "directory" ? 0 : 5, + isDirectory: () => kind === "directory", + isFile: () => kind === "file", + ...overrides, + }; +} + +describe("dual-Station controller UID binding", () => { + it("reads one root-owned UID through an O_NOFOLLOW descriptor and fstat", () => { + const close = vi.fn(); + const open = vi.fn((_pathname: string, _flags: number) => 17); + const fstat = vi.fn(() => controllerUidStat("file")); + + expect( + readDualStationControllerUid({ + lstat: (pathname) => { + expect(pathname).toBe(path.dirname(DUAL_STATION_CONTROLLER_UID_FILE)); + return controllerUidStat("directory"); + }, + open, + fstat, + read: (fd) => { + expect(fd).toBe(17); + return "1001\n"; + }, + close, + }), + ).toBe(1001); + expect(open).toHaveBeenCalledWith(DUAL_STATION_CONTROLLER_UID_FILE, expect.any(Number)); + const flags = open.mock.calls[0]?.[1] ?? 0; + expect(flags & fs.constants.O_NOFOLLOW).toBe(fs.constants.O_NOFOLLOW); + expect(fstat).toHaveBeenCalledWith(17); + expect(close).toHaveBeenCalledWith(17); + }); + + it.each([ + [ + "group-writable parent", + controllerUidStat("directory", { mode: 0o40775 }), + null, + "1001\n", + [], + ], + [ + "non-traversable parent", + controllerUidStat("directory", { mode: 0o40700 }), + null, + "1001\n", + [], + ], + ["non-root-owned file", controllerUidStat("directory"), { uid: 1001 }, "1001\n", [[19]]], + ["wrong file mode", controllerUidStat("directory"), { mode: 0o100664 }, "1001\n", [[19]]], + ["root UID content", controllerUidStat("directory"), null, "0\n", [[19]]], + ["multiple UID lines", controllerUidStat("directory"), { size: 10 }, "1001\n1002\n", [[19]]], + ])("rejects an unsafe controller binding: %s", (_case, directory, fileOverride, contents, expectedCloseCalls) => { + const close = vi.fn(); + expect(() => + readDualStationControllerUid({ + lstat: () => directory, + open: () => 19, + fstat: () => controllerUidStat("file", fileOverride ?? {}), + read: () => contents, + close, + }), + ).toThrow(/Dual-Station controller/u); + expect(close.mock.calls).toEqual(expectedCloseCalls); + }); + + it("refuses a direct lock call from an account other than the prepared controller", () => { + const root = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-station-refused-lock-")); + const stateDir = path.join(root, "state"); + const effectiveUid = process.getuid?.() ?? 0; + const operation = vi.fn(); + + try { + expect(() => + withDualStationVllmLifecycleLock( + operation, + { stateDir, pollIntervalMs: 5, timeoutMs: 250, corruptLockGraceMs: 5 }, + { + readControllerUid: () => (effectiveUid > 0 ? effectiveUid + 1 : 1), + effectiveControllerUid: () => effectiveUid, + }, + ), + ).toThrow( + /requires a non-root effective controller UID|does not match prepared controller UID/u, + ); + expect(operation).not.toHaveBeenCalled(); + expect(fs.existsSync(stateDir)).toBe(false); + } finally { + fs.rmSync(root, { recursive: true, force: true }); + } + }); +}); diff --git a/src/lib/inference/vllm-station-lifecycle-lock.ts b/src/lib/inference/vllm-station-lifecycle-lock.ts new file mode 100644 index 0000000000..1db4076f21 --- /dev/null +++ b/src/lib/inference/vllm-station-lifecycle-lock.ts @@ -0,0 +1,154 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; + +import { type McpLifecycleLockOptions, withMcpLifecycleLock } from "../state/mcp-lifecycle-lock"; +import { STATE_DIR_NAME } from "../state/state-root"; + +const DUAL_STATION_VLLM_LIFECYCLE_LOCK = "dual-station-vllm:host-global"; +const DUAL_STATION_CONTROLLER_CONFIG_DIR = "/etc/nemoclaw"; +export const DUAL_STATION_CONTROLLER_UID_FILE = path.join( + DUAL_STATION_CONTROLLER_CONFIG_DIR, + "dual-station-controller-uid", +); +const DUAL_STATION_CONTROLLER_UID_FILE_MODE = 0o644; +const MAX_POSIX_UID = 0xffff_ffff; + +/** @internal Filesystem seam for security-focused unit tests. */ +export interface DualStationControllerUidFileStat { + uid: number; + gid: number; + mode: number; + size: number; + isDirectory(): boolean; + isFile(): boolean; +} + +/** @internal Filesystem seam for security-focused unit tests. */ +export interface DualStationControllerUidReaderDeps { + lstat(pathname: string): DualStationControllerUidFileStat; + open(pathname: string, flags: number): number; + fstat(fd: number): DualStationControllerUidFileStat; + read(fd: number): string; + close(fd: number): void; +} + +/** @internal Identity seam for tests that cannot use the prepared host account. */ +export interface DualStationControllerIdentityDeps { + readControllerUid(): number; + effectiveControllerUid(): number | null; +} + +const DEFAULT_CONTROLLER_UID_READER_DEPS: DualStationControllerUidReaderDeps = { + lstat: (pathname) => fs.lstatSync(pathname), + open: (pathname, flags) => fs.openSync(pathname, flags), + fstat: (fd) => fs.fstatSync(fd), + read: (fd) => fs.readFileSync(fd, "utf8"), + close: (fd) => fs.closeSync(fd), +}; + +function isNonRootUid(value: number | null): value is number { + return Number.isSafeInteger(value) && value !== null && value > 0 && value <= MAX_POSIX_UID; +} + +/** + * Read the host-preparation controller binding without following the file's + * final path component. Metadata and contents are consumed from the same open + * descriptor so a replacement cannot change what is authorized. + */ +export function readDualStationControllerUid( + deps: DualStationControllerUidReaderDeps = DEFAULT_CONTROLLER_UID_READER_DEPS, +): number { + const directory = deps.lstat(DUAL_STATION_CONTROLLER_CONFIG_DIR); + if ( + !directory.isDirectory() || + directory.uid !== 0 || + directory.gid !== 0 || + (directory.mode & 0o7777) !== 0o755 + ) { + throw new Error( + `Dual-Station controller directory must be root-owned with mode 0755: ${DUAL_STATION_CONTROLLER_CONFIG_DIR}`, + ); + } + + const fd = deps.open( + DUAL_STATION_CONTROLLER_UID_FILE, + fs.constants.O_RDONLY | fs.constants.O_NOFOLLOW, + ); + try { + const binding = deps.fstat(fd); + if ( + !binding.isFile() || + binding.uid !== 0 || + binding.gid !== 0 || + (binding.mode & 0o7777) !== DUAL_STATION_CONTROLLER_UID_FILE_MODE || + binding.size < 2 || + binding.size > 11 + ) { + throw new Error( + `Dual-Station controller UID binding must be a root-owned regular file with mode 0644: ${DUAL_STATION_CONTROLLER_UID_FILE}`, + ); + } + const match = /^([1-9][0-9]*)\n$/.exec(deps.read(fd)); + const controllerUid = match ? Number(match[1]) : Number.NaN; + if (!isNonRootUid(controllerUid)) { + throw new Error( + `Dual-Station controller UID binding must contain exactly one non-root UID: ${DUAL_STATION_CONTROLLER_UID_FILE}`, + ); + } + return controllerUid; + } finally { + deps.close(fd); + } +} + +export function assertDualStationControllerAccount( + readControllerUid: () => number = readDualStationControllerUid, + effectiveControllerUid: () => number | null = () => process.getuid?.() ?? null, +): number { + const controllerUid = readControllerUid(); + if (!isNonRootUid(controllerUid)) { + throw new Error("Dual-Station host preparation returned an invalid controller UID"); + } + const effectiveUid = effectiveControllerUid(); + if (!isNonRootUid(effectiveUid)) { + throw new Error("Dual-Station lifecycle requires a non-root effective controller UID"); + } + if (effectiveUid !== controllerUid) { + throw new Error( + `Dual-Station lifecycle effective UID ${String(effectiveUid)} does not match prepared controller UID ${String(controllerUid)}`, + ); + } + return controllerUid; +} + +/** + * Serialize the host-managed dual-Station service across gateway instances. + * + * Dual-Station lifecycle supports one effective controller account per host. + * This anchors every supported caller at that account's passwd home instead of + * mutable HOME or a gateway-specific root. Host preparation binds that account + * in root-owned state before the lease can be acquired. + */ +export function withDualStationVllmLifecycleLock( + operation: () => Promise | T, + options: McpLifecycleLockOptions = {}, + /** @internal Explicit identity injection keeps test storage overrides authorized. */ + identityDeps: DualStationControllerIdentityDeps = { + readControllerUid: readDualStationControllerUid, + effectiveControllerUid: () => process.getuid?.() ?? null, + }, +): Promise { + assertDualStationControllerAccount( + identityDeps.readControllerUid, + identityDeps.effectiveControllerUid, + ); + const stateDir = options.stateDir ?? path.join(os.userInfo().homedir, STATE_DIR_NAME, "state"); + return withMcpLifecycleLock(DUAL_STATION_VLLM_LIFECYCLE_LOCK, operation, { + ...options, + stateDir, + }); +} diff --git a/src/lib/inference/vllm-station-model-staging.test.ts b/src/lib/inference/vllm-station-model-staging.test.ts new file mode 100644 index 0000000000..5b80725a02 --- /dev/null +++ b/src/lib/inference/vllm-station-model-staging.test.ts @@ -0,0 +1,725 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { createHash } from "node:crypto"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; + +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { + createAtomicIdentityReplacementRunner, + createBetweenAuditMutationRunner, + createManifestPeerPythonRunner, + createPeerIntegrityRunner, + createPostAuditMutationRunner, + createPythonOnlyRunner, +} from "../../../test/support/vllm-station-model-staging-test-support"; +import { DUAL_STATION_VLLM_RUNTIME, type DualStationVllmPlan } from "./vllm-station-cluster"; +import { + type ModelStagingCommandResult, + stageDualStationModelSnapshot, +} from "./vllm-station-model-staging"; +import { + createDualStationSshBindingFixture, + type DualStationSshBindingFixture, +} from "./vllm-station-ssh-binding.test-support"; + +let sshFixture: DualStationSshBindingFixture; +let mockedLocalRoot: string; +let mockedLocalHome: string; +let mockedLocalModelRoot: string; + +function modelRootForHome(home: string): string { + return path.join( + home, + ".cache", + "huggingface", + "hub", + `models--${DUAL_STATION_VLLM_RUNTIME.modelId.replace("/", "--")}`, + ); +} + +function snapshotForHome(home: string): string { + return path.join(modelRootForHome(home), "snapshots", DUAL_STATION_VLLM_RUNTIME.modelRevision); +} + +beforeEach(() => { + sshFixture = createDualStationSshBindingFixture(); + mockedLocalRoot = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-model-staging-test-")); + mockedLocalHome = path.join(mockedLocalRoot, "home"); + mockedLocalModelRoot = modelRootForHome(mockedLocalHome); + fs.mkdirSync(mockedLocalModelRoot, { mode: 0o700, recursive: true }); +}); + +function plan(): DualStationVllmPlan { + return { + peerSshBinding: sshFixture.binding, + runtime: DUAL_STATION_VLLM_RUNTIME, + local: { + hostname: "station-a", + home: mockedLocalHome, + uid: 1000, + gid: 1000, + gpu: { index: 0, name: "NVIDIA GB300", uuid: "GPU-a" }, + }, + peer: { + hostname: "station-b", + home: "/home/nvidia", + uid: 1000, + gid: 1000, + gpu: { index: 0, name: "NVIDIA GB300", uuid: "GPU-b" }, + }, + rails: [ + { + index: 0, + subnet: "192.168.100.0/30", + local: { + rdmaDevice: "mlx5_0", + netdev: "cx8a0", + macAddress: "02:00:00:00:00:01", + uverbsDevice: "/dev/infiniband/uverbs0", + pciAddress: "0000:01:00.0", + address: "192.168.100.1", + }, + peer: { + rdmaDevice: "mlx5_0", + netdev: "cx8b0", + macAddress: "02:00:00:00:00:02", + uverbsDevice: "/dev/infiniband/uverbs0", + pciAddress: "0000:01:00.0", + address: "192.168.100.2", + }, + }, + { + index: 1, + subnet: "192.168.200.0/30", + local: { + rdmaDevice: "mlx5_1", + netdev: "cx8a1", + macAddress: "02:00:00:00:01:01", + uverbsDevice: "/dev/infiniband/uverbs1", + pciAddress: "0000:01:00.1", + address: "192.168.200.1", + }, + peer: { + rdmaDevice: "mlx5_1", + netdev: "cx8b1", + macAddress: "02:00:00:00:01:02", + uverbsDevice: "/dev/infiniband/uverbs1", + pciAddress: "0000:01:00.1", + address: "192.168.200.2", + }, + }, + ], + masterAddress: "192.168.100.1", + roceGidIndex: 3, + }; +} + +function result(stdout = "", status = 0): ModelStagingCommandResult { + return { status, stdout, stderr: "" }; +} + +function manifest(): string { + return JSON.stringify({ + schemaVersion: 1, + files: [ + { + path: "config.json", + size: 2, + sha256: createHash("sha256").update("{}").digest("hex"), + }, + ], + directories: [], + totalBytes: 2, + }); +} + +function peerStagingForPlan(value: DualStationVllmPlan): string { + const identity = [ + "nemoclaw-dual-station-model-staging-v1", + DUAL_STATION_VLLM_RUNTIME.image, + DUAL_STATION_VLLM_RUNTIME.modelId, + DUAL_STATION_VLLM_RUNTIME.modelRevision, + DUAL_STATION_VLLM_RUNTIME.servedModelId, + String(DUAL_STATION_VLLM_RUNTIME.tensorParallelSize), + String(DUAL_STATION_VLLM_RUNTIME.pipelineParallelSize), + String(DUAL_STATION_VLLM_RUNTIME.nodeCount), + value.local.gpu.uuid, + value.peer.gpu.uuid, + ].join("\0"); + const transaction = createHash("sha256").update(identity, "utf8").digest("hex").slice(0, 32); + return path.join( + path.dirname(snapshotForHome(value.peer.home)), + `.nemoclaw-staging-${transaction}`, + ); +} + +function sufficientStatfs() { + return vi.fn().mockResolvedValue({ bavail: 1024n * 1024n * 1024n, bsize: 4096n }); +} + +function createLocalSnapshotFixture(): { root: string; home: string; snapshot: string } { + const root = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-model-snapshot-")); + const home = path.join(root, "home"); + const snapshot = snapshotForHome(home); + fs.mkdirSync(snapshot, { mode: 0o700, recursive: true }); + const shards = Array.from( + { length: 113 }, + (_, index) => `model-${String(index + 1).padStart(5, "0")}-of-00113.safetensors`, + ); + for (const [index, shard] of shards.entries()) { + fs.writeFileSync(path.join(snapshot, shard), `shard-${String(index)}`); + } + fs.writeFileSync(path.join(snapshot, "config.json"), "{}"); + fs.writeFileSync(path.join(snapshot, "tokenizer.json"), "{}"); + fs.writeFileSync( + path.join(snapshot, "model.safetensors.index.json"), + JSON.stringify({ + weight_map: Object.fromEntries( + shards.map((shard, index) => [`model.layers.${String(index)}.weight`, shard]), + ), + }), + ); + return { root, home, snapshot }; +} + +function successfulTransferRunner() { + return vi + .fn() + .mockResolvedValueOnce(result(manifest())) + .mockResolvedValueOnce(result('{"state":"transfer"}')) + .mockResolvedValueOnce(result(manifest())) + .mockResolvedValueOnce(result('{"state":"transfer"}')) + .mockResolvedValueOnce(result()) + .mockResolvedValueOnce(result('{"state":"ready"}')); +} + +function stagingSuffix(runCommand: ReturnType): string { + const destination = String(runCommand.mock.calls[4][1].at(-1)); + return destination.match(/\.nemoclaw-staging-[a-f0-9]{32}/)?.[0] ?? ""; +} + +afterEach(() => { + const stagingLeftovers = fs + .readdirSync(mockedLocalModelRoot) + .filter((entry) => entry.startsWith(".nemoclaw-vllm-model-staging-")); + vi.unstubAllEnvs(); + sshFixture.cleanup(); + fs.rmSync(mockedLocalRoot, { force: true, recursive: true }); + expect(stagingLeftovers).toEqual([]); +}); + +describe("dual-Station pinned model staging", () => { + it("copies only the audited snapshot through strict SSH and verifies it before install", async () => { + vi.stubEnv("OPENSHELL_GATEWAY_AUTH_TOKEN", "must-not-cross-ssh"); + vi.stubEnv("HF_TOKEN", "must-not-cross-ssh"); + const runCommand = successfulTransferRunner(); + const statfs = sufficientStatfs(); + + await expect(stageDualStationModelSnapshot(plan(), { runCommand, statfs })).resolves.toEqual({ + ok: true, + transferred: true, + }); + + expect(runCommand).toHaveBeenCalledTimes(6); + expect(runCommand.mock.calls.map((call) => call[0])).toEqual([ + "python3", + "ssh", + "python3", + "ssh", + "rsync", + "ssh", + ]); + const auditArgs = runCommand.mock.calls[0][1] as string[]; + expect(auditArgs).toEqual(["-", snapshotForHome(mockedLocalHome), mockedLocalModelRoot]); + const materializeArgs = runCommand.mock.calls[2][1] as string[]; + expect(materializeArgs.slice(0, 3)).toEqual(auditArgs); + expect(materializeArgs[3]).toMatch(/nemoclaw-vllm-model-staging-[^/]+\/snapshot$/); + expect(path.dirname(path.dirname(materializeArgs[3]))).toBe(mockedLocalModelRoot); + expect(statfs).toHaveBeenCalledWith(mockedLocalModelRoot); + + const sshArgs = runCommand.mock.calls[1][1] as string[]; + expect(sshArgs).toEqual( + expect.arrayContaining([ + "BatchMode=yes", + "StrictHostKeyChecking=yes", + "ClearAllForwardings=yes", + "ControlMaster=no", + `UserKnownHostsFile=${sshFixture.binding.knownHostsFile}`, + "GlobalKnownHostsFile=/dev/null", + `HostKeyAlias=${sshFixture.binding.lookupHost}`, + `Hostname=${sshFixture.binding.resolvedHost}`, + "User=nvidia", + "Port=22", + "--", + "nvidia@station-b", + "python3 -", + ]), + ); + const rsyncArgs = runCommand.mock.calls[4][1] as string[]; + expect(rsyncArgs).toEqual( + expect.arrayContaining(["--checksum", "--partial", "--protect-args", "--"]), + ); + expect(rsyncArgs).not.toContain("--copy-links"); + expect(rsyncArgs).not.toContain("--delete"); + expect(rsyncArgs.at(-2)).toBe(`${materializeArgs[3]}/`); + expect(rsyncArgs.at(-1)).toMatch( + new RegExp( + `^nvidia@station-b:/home/nvidia/\\.cache/huggingface/hub/models--${DUAL_STATION_VLLM_RUNTIME.modelId.replace("/", "--")}/snapshots/\\.nemoclaw-staging-[a-f0-9]{32}/$`, + ), + ); + expect(rsyncArgs.join(" ")).toContain("StrictHostKeyChecking=yes"); + expect(runCommand.mock.calls[4][2]).toMatchObject({ + idleTimeoutMs: 30 * 60 * 1000, + streamOutput: true, + }); + expect(runCommand.mock.calls[4][2].env.OPENSHELL_GATEWAY_AUTH_TOKEN).toBeUndefined(); + expect(runCommand.mock.calls[4][2].env.HF_TOKEN).toBeUndefined(); + }); + + it("transfers the materialized bytes if the source changes after the second audit", async () => { + const fixture = createLocalSnapshotFixture(); + const fixturePlan = plan(); + fixturePlan.local.home = fixture.home; + const { runCommand, state } = createPostAuditMutationRunner(fixture.snapshot); + const statfs = sufficientStatfs(); + + try { + await expect( + stageDualStationModelSnapshot(fixturePlan, { runCommand, statfs }), + ).resolves.toEqual({ ok: true, transferred: true }); + expect(state.transferredConfig).toBe("{}"); + expect(state.snapshotMode).toBe(0o500); + expect(state.configMode).toBe(0o400); + expect(state.transferSource).not.toBe(fixture.snapshot); + expect(path.dirname(path.dirname(state.transferSource))).toBe(modelRootForHome(fixture.home)); + expect(fs.existsSync(state.transferSource)).toBe(false); + } finally { + fs.rmSync(fixture.root, { force: true, recursive: true }); + } + }); + + it("rejects source mutation between audits and cleans both staging roots", async () => { + const fixture = createLocalSnapshotFixture(); + const fixturePlan = plan(); + fixturePlan.local.home = fixture.home; + const { runCommand, state } = createBetweenAuditMutationRunner(fixture.snapshot); + + try { + await expect( + stageDualStationModelSnapshot(fixturePlan, { + runCommand, + statfs: sufficientStatfs(), + }), + ).resolves.toEqual({ + ok: false, + reason: "local pinned snapshot changed between audit and materialization", + }); + expect(runCommand.mock.calls.map((call) => call[0])).toEqual([ + "python3", + "ssh", + "python3", + "ssh", + ]); + expect(fs.existsSync(state.materializedSnapshot)).toBe(false); + expect(runCommand).not.toHaveBeenCalledWith("rsync", expect.anything(), expect.anything()); + } finally { + fs.rmSync(fixture.root, { force: true, recursive: true }); + } + }); + + it("rejects a snapshot symlink escape through the public staging boundary", async () => { + const fixture = createLocalSnapshotFixture(); + const fixturePlan = plan(); + fixturePlan.local.home = fixture.home; + const outside = path.join(fixture.root, "outside-tokenizer.json"); + const tokenizer = path.join(fixture.snapshot, "tokenizer.json"); + fs.writeFileSync(outside, "{}"); + fs.unlinkSync(tokenizer); + fs.symlinkSync(outside, tokenizer); + const runCommand = createPythonOnlyRunner(); + + try { + await expect( + stageDualStationModelSnapshot(fixturePlan, { runCommand }), + ).resolves.toMatchObject({ + ok: false, + reason: expect.stringContaining("snapshot symlink escapes the pinned model cache"), + }); + expect(runCommand).toHaveBeenCalledTimes(1); + expect(runCommand.mock.calls[0][1] as string[]).toHaveLength(3); + expect( + fs + .readdirSync(modelRootForHome(fixture.home)) + .some((entry) => entry.startsWith(".nemoclaw-vllm-model-staging-")), + ).toBe(false); + } finally { + fs.rmSync(fixture.root, { force: true, recursive: true }); + } + }); + + it("does no transfer when the peer already has the exact byte manifest", async () => { + const runCommand = vi + .fn() + .mockResolvedValueOnce(result(manifest())) + .mockResolvedValueOnce(result('{"state":"ready"}')); + const statfs = sufficientStatfs(); + + await expect(stageDualStationModelSnapshot(plan(), { runCommand, statfs })).resolves.toEqual({ + ok: true, + transferred: false, + }); + expect(runCommand).toHaveBeenCalledTimes(2); + expect(statfs).not.toHaveBeenCalled(); + expect(runCommand.mock.calls[0][1] as string[]).toHaveLength(3); + expect(runCommand).not.toHaveBeenCalledWith("rsync", expect.anything(), expect.anything()); + expect( + fs + .readdirSync(mockedLocalModelRoot) + .some((entry) => entry.startsWith(".nemoclaw-vllm-model-staging-")), + ).toBe(false); + }); + + it("removes deterministic peer staging when the final snapshot is already exact", async () => { + const remoteRoot = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-model-peer-ready-")); + const peerHome = path.join(remoteRoot, "peer-home"); + fs.mkdirSync(peerHome, { mode: 0o700 }); + const fixturePlan = plan(); + fixturePlan.peer.home = peerHome; + const finalSnapshot = snapshotForHome(peerHome); + const peerStaging = peerStagingForPlan(fixturePlan); + fs.mkdirSync(finalSnapshot, { mode: 0o700, recursive: true }); + fs.writeFileSync(path.join(finalSnapshot, "config.json"), "{}"); + fs.mkdirSync(peerStaging, { mode: 0o700 }); + fs.writeFileSync(path.join(peerStaging, "config.json"), "{"); + const statfs = sufficientStatfs(); + const runCommand = createManifestPeerPythonRunner({ + localManifest: manifest(), + peerHome, + }); + + try { + await expect( + stageDualStationModelSnapshot(fixturePlan, { runCommand, statfs }), + ).resolves.toEqual({ ok: true, transferred: false }); + expect(runCommand.mock.calls.map((call) => call[0])).toEqual(["python3", "ssh"]); + expect(statfs).not.toHaveBeenCalled(); + expect(fs.existsSync(finalSnapshot)).toBe(true); + expect(fs.existsSync(peerStaging)).toBe(false); + expect(runCommand).not.toHaveBeenCalledWith("rsync", expect.anything(), expect.anything()); + } finally { + fs.rmSync(remoteRoot, { force: true, recursive: true }); + } + }); + + it("cleans peer staging without materializing when local capacity is insufficient", async () => { + const runCommand = vi + .fn() + .mockResolvedValueOnce(result(manifest())) + .mockResolvedValueOnce(result('{"state":"transfer"}')) + .mockResolvedValueOnce(result('{"state":"cleaned"}')); + const statfs = vi.fn().mockResolvedValue({ bavail: 1n, bsize: 4096n }); + + await expect(stageDualStationModelSnapshot(plan(), { runCommand, statfs })).resolves.toEqual({ + ok: false, + reason: "local model cache does not have enough free space for the audited snapshot copy", + }); + expect(statfs).toHaveBeenCalledWith(mockedLocalModelRoot); + expect(runCommand.mock.calls.map((call) => call[0])).toEqual(["python3", "ssh", "ssh"]); + expect(runCommand).not.toHaveBeenCalledWith("rsync", expect.anything(), expect.anything()); + expect( + fs + .readdirSync(mockedLocalModelRoot) + .some((entry) => entry.startsWith(".nemoclaw-vllm-model-staging-")), + ).toBe(false); + }); + + it("fails remote capacity preflight before creating either staging copy", async () => { + const remoteRoot = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-model-peer-capacity-")); + const peerHome = path.join(remoteRoot, "peer-home"); + fs.mkdirSync(peerHome, { mode: 0o700 }); + const fixturePlan = plan(); + fixturePlan.peer.home = peerHome; + const statfs = sufficientStatfs(); + const noCapacity = `import shutil +class _NemoClawDiskUsage: + free = 0 +shutil.disk_usage = lambda _path: _NemoClawDiskUsage() +`; + const runCommand = createManifestPeerPythonRunner({ + localManifest: manifest(), + peerHome, + peerInputPrefix: noCapacity, + }); + + try { + await expect( + stageDualStationModelSnapshot(fixturePlan, { runCommand, statfs }), + ).resolves.toEqual({ + ok: false, + reason: + "peer snapshot preflight failed: peer model cache does not have enough free space for the pinned snapshot", + }); + expect(runCommand.mock.calls.map((call) => call[0])).toEqual(["python3", "ssh"]); + expect(statfs).not.toHaveBeenCalled(); + expect( + fs + .readdirSync(path.dirname(snapshotForHome(peerHome))) + .some((entry) => entry.startsWith(".nemoclaw-staging-")), + ).toBe(false); + } finally { + fs.rmSync(remoteRoot, { force: true, recursive: true }); + } + }); + + it("does not credit a full-sized corrupt partial file toward remote capacity", async () => { + const remoteRoot = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-model-peer-corrupt-")); + const peerHome = path.join(remoteRoot, "peer-home"); + fs.mkdirSync(peerHome, { mode: 0o700 }); + const fixturePlan = plan(); + fixturePlan.peer.home = peerHome; + const peerStaging = peerStagingForPlan(fixturePlan); + fs.mkdirSync(peerStaging, { mode: 0o700, recursive: true }); + const corruptSnapshotBytes = 64 * 1024 * 1024; + const corruptPartial = path.join(peerStaging, "config.json"); + fs.writeFileSync(corruptPartial, ""); + fs.truncateSync(corruptPartial, corruptSnapshotBytes); + const capacityManifest = JSON.stringify({ + schemaVersion: 1, + files: [ + { + path: "config.json", + size: corruptSnapshotBytes, + sha256: createHash("sha256").update("expected snapshot contents").digest("hex"), + }, + ], + directories: [], + totalBytes: corruptSnapshotBytes, + }); + const headroomOnly = `import shutil +class _NemoClawDiskUsage: + free = 5 * 1024 * 1024 * 1024 +shutil.disk_usage = lambda _path: _NemoClawDiskUsage() +`; + const runCommand = createManifestPeerPythonRunner({ + localManifest: capacityManifest, + peerHome, + peerInputPrefix: headroomOnly, + }); + + try { + await expect( + stageDualStationModelSnapshot(fixturePlan, { + runCommand, + statfs: sufficientStatfs(), + }), + ).resolves.toEqual({ + ok: false, + reason: + "peer snapshot preflight failed: peer model cache does not have enough free space for the pinned snapshot", + }); + expect(runCommand.mock.calls.map((call) => call[0])).toEqual(["python3", "ssh"]); + expect(fs.statSync(corruptPartial).size).toBe(corruptSnapshotBytes); + } finally { + fs.rmSync(remoteRoot, { force: true, recursive: true }); + } + }); + + it("gives concurrent reversed and different local heads disjoint retry-safe staging paths", async () => { + const original = plan(); + const reversed = plan(); + const reversedLocal = reversed.local; + reversed.local = reversed.peer; + reversed.peer = reversedLocal; + reversed.local.home = mockedLocalHome; + reversed.peer.home = "/home/nvidia"; + const reversedSshFixture = createDualStationSshBindingFixture("nvidia@station-a"); + reversed.peerSshBinding = reversedSshFixture.binding; + const differentHead = plan(); + differentHead.local.gpu.uuid = "GPU-c"; + const originalRunner = successfulTransferRunner(); + const reversedRunner = successfulTransferRunner(); + const differentHeadRunner = successfulTransferRunner(); + + try { + await Promise.all([ + stageDualStationModelSnapshot(original, { + runCommand: originalRunner, + statfs: sufficientStatfs(), + }), + stageDualStationModelSnapshot(reversed, { + runCommand: reversedRunner, + statfs: sufficientStatfs(), + }), + stageDualStationModelSnapshot(differentHead, { + runCommand: differentHeadRunner, + statfs: sufficientStatfs(), + }), + ]); + } finally { + reversedSshFixture.cleanup(); + } + + const suffixes = [ + stagingSuffix(originalRunner), + stagingSuffix(reversedRunner), + stagingSuffix(differentHeadRunner), + ]; + expect(suffixes.every((suffix) => /^\.nemoclaw-staging-[a-f0-9]{32}$/.test(suffix))).toBe(true); + expect(new Set(suffixes).size).toBe(3); + }); + + it("reuses the same deterministic partial path for the same ordered pair", async () => { + const firstRunner = successfulTransferRunner(); + const retryRunner = successfulTransferRunner(); + + await stageDualStationModelSnapshot(plan(), { + runCommand: firstRunner, + statfs: sufficientStatfs(), + }); + await stageDualStationModelSnapshot(plan(), { + runCommand: retryRunner, + statfs: sufficientStatfs(), + }); + + expect(stagingSuffix(firstRunner)).toBe(stagingSuffix(retryRunner)); + }); + + it("cleans the private peer staging tree when rsync fails", async () => { + const runCommand = vi + .fn() + .mockResolvedValueOnce(result(manifest())) + .mockResolvedValueOnce(result('{"state":"transfer"}')) + .mockResolvedValueOnce(result(manifest())) + .mockResolvedValueOnce(result('{"state":"transfer"}')) + .mockResolvedValueOnce({ status: 23, stdout: "", stderr: "partial transfer" }) + .mockResolvedValueOnce(result('{"state":"cleaned"}')); + + await expect( + stageDualStationModelSnapshot(plan(), { runCommand, statfs: sufficientStatfs() }), + ).resolves.toEqual({ ok: false, reason: "peer snapshot transfer failed: partial transfer" }); + expect(runCommand).toHaveBeenCalledTimes(6); + expect(runCommand.mock.calls.map((call) => call[0])).toEqual([ + "python3", + "ssh", + "python3", + "ssh", + "rsync", + "ssh", + ]); + }); + + it("removes peer bytes that fail real manifest verification", async () => { + const remoteRoot = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-model-peer-")); + const peerHome = path.join(remoteRoot, "peer-home"); + fs.mkdirSync(peerHome, { mode: 0o700 }); + const fixturePlan = plan(); + fixturePlan.peer.home = peerHome; + const { runCommand, state } = createPeerIntegrityRunner({ + localManifest: manifest(), + peerHome, + peerInputPrefix: `import shutil +class _NemoClawDiskUsage: + free = 1024 * 1024 * 1024 * 1024 +shutil.disk_usage = lambda _path: _NemoClawDiskUsage() +`, + }); + const statfs = sufficientStatfs(); + + try { + await expect( + stageDualStationModelSnapshot(fixturePlan, { runCommand, statfs }), + ).resolves.toEqual({ + ok: false, + reason: + "peer snapshot verification failed: peer staged snapshot failed byte-integrity verification", + }); + expect(runCommand.mock.calls.map((call) => call[0])).toEqual([ + "python3", + "ssh", + "python3", + "ssh", + "rsync", + "ssh", + "ssh", + ]); + expect(fs.existsSync(state.stagingPath)).toBe(false); + } finally { + fs.rmSync(remoteRoot, { force: true, recursive: true }); + } + }); + + it("fails closed when atomic install resolves to a different directory identity", async () => { + const fixture = createLocalSnapshotFixture(); + const remoteRoot = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-model-peer-identity-")); + const peerHome = path.join(remoteRoot, "peer-home"); + fs.mkdirSync(peerHome, { mode: 0o700 }); + const fixturePlan = plan(); + fixturePlan.local.home = fixture.home; + fixturePlan.peer.home = peerHome; + const { runCommand, state } = createAtomicIdentityReplacementRunner(peerHome); + + try { + await expect( + stageDualStationModelSnapshot(fixturePlan, { + runCommand, + statfs: sufficientStatfs(), + }), + ).resolves.toEqual({ + ok: false, + reason: + "peer snapshot verification failed: peer pinned snapshot identity changed during atomic install", + }); + expect(runCommand.mock.calls.map((call) => call[0])).toEqual([ + "python3", + "ssh", + "python3", + "ssh", + "rsync", + "ssh", + "ssh", + ]); + expect(fs.existsSync(state.materializedSnapshot)).toBe(false); + expect(fs.existsSync(snapshotForHome(peerHome))).toBe(true); + expect(fs.existsSync(`${snapshotForHome(peerHome)}.nemoclaw-test-original`)).toBe(true); + } finally { + fs.rmSync(fixture.root, { force: true, recursive: true }); + fs.rmSync(remoteRoot, { force: true, recursive: true }); + } + }); + + it("rejects a peer home that cannot be represented without remote shell syntax", async () => { + const unsafe = plan(); + unsafe.peer.home = "/home/nvidia;touch-pwned"; + const runCommand = vi.fn(); + + await expect(stageDualStationModelSnapshot(unsafe, { runCommand })).resolves.toEqual({ + ok: false, + reason: "peer home is unsafe for model staging", + }); + expect(runCommand).not.toHaveBeenCalled(); + }); + + it("fails before any command when the qualified host-key pin changes", async () => { + fs.appendFileSync(sshFixture.binding.knownHostsFile, "changed\n"); + const runCommand = vi.fn(); + + await expect(stageDualStationModelSnapshot(plan(), { runCommand })).resolves.toMatchObject({ + ok: false, + reason: expect.stringContaining("known-hosts binding changed"), + }); + expect(runCommand).not.toHaveBeenCalled(); + }); + + it("fails before SSH when the local manifest is malformed", async () => { + const runCommand = vi.fn().mockResolvedValueOnce(result('{"schemaVersion":1}')); + + await expect(stageDualStationModelSnapshot(plan(), { runCommand })).resolves.toEqual({ + ok: false, + reason: "local pinned snapshot audit returned an invalid manifest", + }); + expect(runCommand).toHaveBeenCalledTimes(1); + }); +}); diff --git a/src/lib/inference/vllm-station-model-staging.ts b/src/lib/inference/vllm-station-model-staging.ts new file mode 100644 index 0000000000..1992e95263 --- /dev/null +++ b/src/lib/inference/vllm-station-model-staging.ts @@ -0,0 +1,1052 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { type ChildProcess, spawn } from "node:child_process"; +import { createHash } from "node:crypto"; +import fs from "node:fs"; +import path from "node:path"; + +import { buildVllmSshTransportEnv } from "./vllm-docker-env"; +import { DUAL_STATION_VLLM_RUNTIME, type DualStationVllmPlan } from "./vllm-station-cluster"; +import { dualStationPinnedSshArgs } from "./vllm-station-ssh-binding"; + +const MANIFEST_SCHEMA_VERSION = 1; +const MAX_MANIFEST_BYTES = 2 * 1024 * 1024; +const MAX_COMMAND_OUTPUT_BYTES = 256 * 1024; +const SNAPSHOT_AUDIT_TIMEOUT_MS = 6 * 60 * 60 * 1000; +const RSYNC_TIMEOUT_MS = 12 * 60 * 60 * 1000; +const RSYNC_IDLE_TIMEOUT_MS = 30 * 60 * 1000; +const LOCAL_STAGING_MIN_HEADROOM_BYTES = 5n * 1024n * 1024n * 1024n; +const MAX_SNAPSHOT_BYTES = 1024 * 1024 * 1024 * 1024; +const SAFE_POSIX_PATH_PATTERN = /^\/(?:[A-Za-z0-9._-]+\/)*[A-Za-z0-9._-]+$/; + +export interface ModelStagingCommandOptions { + env: Record; + input?: string; + timeoutMs: number; + idleTimeoutMs?: number; + streamOutput?: boolean; +} + +export interface ModelStagingCommandResult { + status: number | null; + stdout: string; + stderr: string; + error?: string; + timedOut?: boolean; +} + +export interface DualStationModelStagingDeps { + runCommand( + file: string, + args: readonly string[], + options: ModelStagingCommandOptions, + ): Promise; + statfs?( + filePath: string, + ): Promise<{ bavail: bigint; bsize: bigint }> | { bavail: bigint; bsize: bigint }; +} + +export type DualStationModelStagingResult = + | { ok: true; transferred: boolean } + | { ok: false; reason: string }; + +interface SnapshotManifest { + schemaVersion: 1; + files: Array<{ path: string; size: number; sha256: string }>; + directories: string[]; + totalBytes: number; +} + +interface StagingPaths { + localModelRoot: string; + localSnapshot: string; + peerSnapshot: string; + peerStaging: string; +} + +function appendBounded(current: string, chunk: string, limit: number): string { + const next = current + chunk; + return next.length <= limit ? next : next.slice(next.length - limit); +} + +function defaultRunCommand( + file: string, + args: readonly string[], + options: ModelStagingCommandOptions, +): Promise { + return new Promise((resolve) => { + let child: ChildProcess; + try { + child = spawn(file, [...args], { + env: options.env, + shell: false, + stdio: [options.input === undefined ? "ignore" : "pipe", "pipe", "pipe"], + windowsHide: true, + }); + } catch (err) { + resolve({ status: null, stdout: "", stderr: "", error: (err as Error).message }); + return; + } + + let stdout = ""; + let stderr = ""; + let settled = false; + let timedOut = false; + let idleTimer: NodeJS.Timeout | undefined; + const absoluteTimer = setTimeout(() => { + timedOut = true; + child.kill("SIGKILL"); + }, options.timeoutMs); + absoluteTimer.unref?.(); + + const resetIdleTimer = (): void => { + if (idleTimer) clearTimeout(idleTimer); + if (options.idleTimeoutMs === undefined) return; + idleTimer = setTimeout(() => { + timedOut = true; + child.kill("SIGKILL"); + }, options.idleTimeoutMs); + idleTimer.unref?.(); + }; + resetIdleTimer(); + + child.stdout?.on("data", (value: Buffer | string) => { + const chunk = String(value); + stdout = appendBounded(stdout, chunk, MAX_MANIFEST_BYTES); + resetIdleTimer(); + if (options.streamOutput) process.stdout.write(chunk); + }); + child.stderr?.on("data", (value: Buffer | string) => { + const chunk = String(value); + stderr = appendBounded(stderr, chunk, MAX_COMMAND_OUTPUT_BYTES); + resetIdleTimer(); + if (options.streamOutput) process.stderr.write(chunk); + }); + child.once("error", (err) => { + if (settled) return; + settled = true; + clearTimeout(absoluteTimer); + if (idleTimer) clearTimeout(idleTimer); + resolve({ status: null, stdout, stderr, error: err.message, timedOut }); + }); + child.once("close", (status) => { + if (settled) return; + settled = true; + clearTimeout(absoluteTimer); + if (idleTimer) clearTimeout(idleTimer); + resolve({ status, stdout, stderr, timedOut }); + }); + child.stdin?.on("error", () => undefined); + child.stdin?.end(options.input); + }); +} + +const defaultDeps: DualStationModelStagingDeps = { + runCommand: defaultRunCommand, + statfs(filePath) { + const stats = fs.statfsSync(filePath, { bigint: true }); + return { bavail: stats.bavail, bsize: stats.bsize }; + }, +}; + +function snapshotPath(home: string): string { + return path.posix.join( + home, + ".cache", + "huggingface", + "hub", + `models--${DUAL_STATION_VLLM_RUNTIME.modelId.replace("/", "--")}`, + "snapshots", + DUAL_STATION_VLLM_RUNTIME.modelRevision, + ); +} + +function stagingTransactionId(plan: DualStationVllmPlan): string { + const identity = [ + "nemoclaw-dual-station-model-staging-v1", + DUAL_STATION_VLLM_RUNTIME.image, + DUAL_STATION_VLLM_RUNTIME.modelId, + DUAL_STATION_VLLM_RUNTIME.modelRevision, + DUAL_STATION_VLLM_RUNTIME.servedModelId, + String(DUAL_STATION_VLLM_RUNTIME.tensorParallelSize), + String(DUAL_STATION_VLLM_RUNTIME.pipelineParallelSize), + String(DUAL_STATION_VLLM_RUNTIME.nodeCount), + plan.local.gpu.uuid, + plan.peer.gpu.uuid, + ].join("\0"); + return createHash("sha256").update(identity, "utf8").digest("hex").slice(0, 32); +} + +function stagingPaths(plan: DualStationVllmPlan): StagingPaths { + for (const [label, home] of [ + ["local", plan.local.home], + ["peer", plan.peer.home], + ] as const) { + if ( + !SAFE_POSIX_PATH_PATTERN.test(home) || + path.posix.normalize(home) !== home || + home.split("/").some((component) => component === "." || component === "..") + ) { + throw new Error(`${label} home is unsafe for model staging`); + } + } + dualStationPinnedSshArgs(plan.peerSshBinding); + if ( + plan.runtime.modelId !== DUAL_STATION_VLLM_RUNTIME.modelId || + plan.runtime.modelRevision !== DUAL_STATION_VLLM_RUNTIME.modelRevision || + plan.runtime.image !== DUAL_STATION_VLLM_RUNTIME.image + ) { + throw new Error("dual-Station plan does not identify the pinned runtime"); + } + + const localSnapshot = snapshotPath(plan.local.home); + const peerSnapshot = snapshotPath(plan.peer.home); + return { + localModelRoot: path.posix.dirname(path.posix.dirname(localSnapshot)), + localSnapshot, + peerSnapshot, + peerStaging: path.posix.join( + path.posix.dirname(peerSnapshot), + `.nemoclaw-staging-${stagingTransactionId(plan)}`, + ), + }; +} + +const LOCAL_MANIFEST_SCRIPT = String.raw` +import hashlib +import json +import os +from pathlib import Path +import re +import stat +import sys + +if len(sys.argv) not in (3, 4): + raise SystemExit("expected snapshot and model-root paths, plus optional materialized-snapshot path") + +snapshot = Path(sys.argv[1]) +model_root = Path(sys.argv[2]) +materialized_snapshot = Path(sys.argv[3]) if len(sys.argv) == 4 else None +safe_component = re.compile(r"^[A-Za-z0-9._-]+$") +weight_index_name = "model.safetensors.index.json" +max_weight_index_bytes = 64 * 1024 * 1024 + +def fail(message): + raise SystemExit(message) + +def relative_name(candidate): + relative = candidate.relative_to(snapshot) + if not relative.parts or any(not safe_component.fullmatch(part) for part in relative.parts): + raise SystemExit("snapshot contains an unsafe relative path") + return relative.as_posix() + +def resolved_regular_file(candidate): + metadata = candidate.lstat() + resolved = candidate.resolve(strict=True) + try: + common = os.path.commonpath((str(canonical_model_root), str(resolved))) + except (OSError, ValueError): + raise SystemExit("snapshot file could not be resolved") + if common != str(canonical_model_root): + raise SystemExit("snapshot symlink escapes the pinned model cache") + if not resolved.is_file(): + raise SystemExit("snapshot contains a non-regular file") + return resolved + +def create_materialized_directory(relative): + if materialized_snapshot is None: + return + destination = materialized_snapshot / relative + try: + destination.mkdir(mode=0o700) + except FileExistsError: + fail("materialized snapshot path changed during audit") + +def audit_regular_file(source, relative): + if not hasattr(os, "O_NOFOLLOW"): + fail("snapshot audit requires O_NOFOLLOW support") + source_flags = os.O_RDONLY | os.O_NOFOLLOW + destination_flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW + digest = hashlib.sha256() + index_bytes = bytearray() if relative == weight_index_name else None + size = 0 + try: + source_handle = os.fdopen(os.open(source, source_flags), "rb") + except OSError: + fail("snapshot file changed while it was being opened") + with source_handle: + if not stat.S_ISREG(os.fstat(source_handle.fileno()).st_mode): + fail("snapshot contains a non-regular file") + destination_handle = None + if materialized_snapshot is not None: + destination = materialized_snapshot / relative + try: + destination_handle = os.fdopen( + os.open(destination, destination_flags, 0o400), + "wb", + ) + except OSError: + fail("materialized snapshot path changed during audit") + try: + for chunk in iter(lambda: source_handle.read(8 * 1024 * 1024), b""): + size += len(chunk) + digest.update(chunk) + if index_bytes is not None: + if len(index_bytes) + len(chunk) > max_weight_index_bytes: + fail("pinned local weight index exceeds the safety bound") + index_bytes.extend(chunk) + if destination_handle is not None: + destination_handle.write(chunk) + if destination_handle is not None: + destination_handle.flush() + os.fsync(destination_handle.fileno()) + os.fchmod(destination_handle.fileno(), 0o400) + finally: + if destination_handle is not None: + destination_handle.close() + return size, digest.hexdigest(), index_bytes + +if snapshot.is_symlink() or not snapshot.is_dir(): + raise SystemExit("pinned local snapshot directory is missing or unsafe") +if model_root.is_symlink() or not model_root.is_dir(): + raise SystemExit("pinned local model cache root is missing or unsafe") +canonical_model_root = model_root.resolve(strict=True) +if materialized_snapshot is not None: + if os.path.lexists(materialized_snapshot): + fail("private materialized snapshot path already exists") + try: + materialized_snapshot.mkdir(mode=0o700) + except OSError: + fail("private materialized snapshot could not be created") + +files = [] +directories = [] +total_bytes = 0 +entry_count = 0 +weight_index_bytes = None +for current, dirnames, filenames in os.walk(snapshot, topdown=True, followlinks=False): + current_path = Path(current) + dirnames.sort() + filenames.sort() + for dirname in dirnames: + directory = current_path / dirname + relative = relative_name(directory) + mode = directory.lstat().st_mode + if stat.S_ISLNK(mode) or not stat.S_ISDIR(mode): + raise SystemExit("snapshot contains an unsafe directory") + create_materialized_directory(relative) + directories.append(relative) + entry_count += 1 + for filename in filenames: + candidate = current_path / filename + relative = relative_name(candidate) + resolved = resolved_regular_file(candidate) + size, digest, audited_index_bytes = audit_regular_file(resolved, relative) + if audited_index_bytes is not None: + weight_index_bytes = audited_index_bytes + files.append({"path": relative, "size": size, "sha256": digest}) + total_bytes += size + entry_count += 1 + if entry_count > 4096 or total_bytes > ${String(MAX_SNAPSHOT_BYTES)}: + raise SystemExit("snapshot exceeds the staging safety bounds") + +file_paths = {item["path"] for item in files} +try: + if weight_index_bytes is None: + fail("pinned local weight index is missing or malformed") + index = json.loads(weight_index_bytes.decode("utf-8")) + weight_map = index.get("weight_map", {}) + shards = sorted(set(weight_map.values())) if isinstance(weight_map, dict) else [] +except (OSError, UnicodeError, json.JSONDecodeError): + fail("pinned local weight index is missing or malformed") +if len(shards) != 113 or any( + not isinstance(item, str) + or Path(item).name != item + or not safe_component.fullmatch(item) + or not item.endswith(".safetensors") + for item in shards +): + fail("pinned local weight index has an unexpected shard set") +if "config.json" not in file_paths: + fail("pinned local config.json is missing") +if not any(name in file_paths for name in ("tokenizer.json", "tokenizer.model", "vocab.json")): + fail("pinned local tokenizer assets are missing") +if not files or any(shard not in file_paths for shard in shards): + raise SystemExit("pinned local snapshot is incomplete") +if materialized_snapshot is not None: + for current, dirnames, _filenames in os.walk(materialized_snapshot, topdown=False): + for dirname in dirnames: + os.chmod(Path(current) / dirname, 0o500, follow_symlinks=False) + os.chmod(materialized_snapshot, 0o500, follow_symlinks=False) +print(json.dumps({ + "schemaVersion": 1, + "files": files, + "directories": directories, + "totalBytes": total_bytes, +}, separators=(",", ":"))) +`; + +function parseManifest(result: ModelStagingCommandResult): SnapshotManifest { + if (result.status !== 0 || result.timedOut || result.error) { + throw new Error(commandFailure("local pinned snapshot audit", result)); + } + if (Buffer.byteLength(result.stdout, "utf8") > MAX_MANIFEST_BYTES) { + throw new Error("local pinned snapshot manifest exceeded the safety bound"); + } + let parsed: unknown; + try { + parsed = JSON.parse(result.stdout.trim()); + } catch { + throw new Error("local pinned snapshot audit returned invalid JSON"); + } + if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) { + throw new Error("local pinned snapshot audit returned an invalid manifest"); + } + const record = parsed as Record; + if ( + record.schemaVersion !== MANIFEST_SCHEMA_VERSION || + !Array.isArray(record.files) || + !Array.isArray(record.directories) || + !Number.isSafeInteger(record.totalBytes) || + Number(record.totalBytes) <= 0 || + Number(record.totalBytes) > MAX_SNAPSHOT_BYTES || + record.files.length === 0 || + record.files.length + record.directories.length > 4096 + ) { + throw new Error("local pinned snapshot audit returned an invalid manifest"); + } + const files = record.files.map((value) => { + if (!value || typeof value !== "object" || Array.isArray(value)) { + throw new Error("local pinned snapshot manifest contains an invalid file"); + } + const file = value as Record; + if ( + typeof file.path !== "string" || + !/^(?:[A-Za-z0-9._-]+\/)*[A-Za-z0-9._-]+$/.test(file.path) || + !Number.isSafeInteger(file.size) || + Number(file.size) < 0 || + typeof file.sha256 !== "string" || + !/^[a-f0-9]{64}$/.test(file.sha256) + ) { + throw new Error("local pinned snapshot manifest contains an invalid file"); + } + return { path: file.path, size: Number(file.size), sha256: file.sha256 }; + }); + const directories = record.directories.map((value) => { + if (typeof value !== "string" || !/^(?:[A-Za-z0-9._-]+\/)*[A-Za-z0-9._-]+$/.test(value)) { + throw new Error("local pinned snapshot manifest contains an invalid directory"); + } + return value; + }); + const filePaths = files.map((file) => file.path); + if ( + new Set(filePaths).size !== files.length || + new Set(directories).size !== directories.length || + directories.some((directory) => filePaths.includes(directory)) || + files.reduce((total, file) => total + file.size, 0) !== Number(record.totalBytes) + ) { + throw new Error("local pinned snapshot manifest is internally inconsistent"); + } + return { + schemaVersion: 1, + files, + directories, + totalBytes: Number(record.totalBytes), + }; +} + +function remoteScript( + plan: DualStationVllmPlan, + paths: StagingPaths, + manifest: SnapshotManifest, + operation: "prepare" | "finalize" | "cleanup", +): string { + const manifestBase64 = Buffer.from(JSON.stringify(manifest), "utf8").toString("base64"); + return String.raw` +import base64 +import hashlib +import json +import os +from pathlib import Path +import re +import shutil +import stat + +EXPECTED_HOME = Path(${JSON.stringify(plan.peer.home)}) +FINAL = Path(${JSON.stringify(paths.peerSnapshot)}) +STAGING = Path(${JSON.stringify(paths.peerStaging)}) +OPERATION = ${JSON.stringify(operation)} +EXPECTED = json.loads(base64.b64decode(${JSON.stringify(manifestBase64)}, validate=True)) +SAFE_COMPONENT = re.compile(r"^[A-Za-z0-9._-]+$") + +def fail(message): + raise SystemExit(message) + +def relative_name(root, candidate): + relative = candidate.relative_to(root) + if not relative.parts or any(not SAFE_COMPONENT.fullmatch(part) for part in relative.parts): + fail("peer snapshot contains an unsafe relative path") + return relative.as_posix() + +def manifest(root): + if root.is_symlink() or not root.is_dir(): + fail("peer snapshot path is missing or unsafe") + files = [] + directories = [] + total_bytes = 0 + entry_count = 0 + for current, dirnames, filenames in os.walk(root, topdown=True, followlinks=False): + current_path = Path(current) + dirnames.sort() + filenames.sort() + for dirname in dirnames: + directory = current_path / dirname + relative = relative_name(root, directory) + mode = directory.lstat().st_mode + if stat.S_ISLNK(mode) or not stat.S_ISDIR(mode): + fail("peer snapshot contains an unsafe directory") + directories.append(relative) + entry_count += 1 + for filename in filenames: + candidate = current_path / filename + relative = relative_name(root, candidate) + mode = candidate.lstat().st_mode + if not stat.S_ISREG(mode): + fail("peer snapshot contains a symlink or non-regular file") + digest = hashlib.sha256() + size = 0 + with candidate.open("rb") as handle: + for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""): + size += len(chunk) + digest.update(chunk) + files.append({"path": relative, "size": size, "sha256": digest.hexdigest()}) + total_bytes += size + entry_count += 1 + if entry_count > 4096 or total_bytes > 1024 * 1024 * 1024 * 1024: + fail("peer snapshot exceeds the staging safety bounds") + return {"schemaVersion": 1, "files": files, "directories": directories, "totalBytes": total_bytes} + +def path_exists(candidate): + return os.path.lexists(candidate) + +def verify_parent_chain(): + observed_home = Path.home() + if observed_home != EXPECTED_HOME or EXPECTED_HOME.is_symlink() or not EXPECTED_HOME.is_dir(): + fail("peer home identity changed during model staging") + current = EXPECTED_HOME + relative_parent = FINAL.parent.relative_to(EXPECTED_HOME) + for component in relative_parent.parts: + if not SAFE_COMPONENT.fullmatch(component): + fail("peer snapshot parent contains an unsafe component") + current = current / component + if path_exists(current): + mode = current.lstat().st_mode + if stat.S_ISLNK(mode) or not stat.S_ISDIR(mode): + fail("peer snapshot parent is a symlink or non-directory") + else: + current.mkdir(mode=0o700) + +def verify_partial(root): + if root.is_symlink() or not root.is_dir(): + fail("peer staging path is unsafe") + expected_files = {item["path"]: item for item in EXPECTED["files"]} + expected_directories = set(EXPECTED["directories"]) + entry_count = 0 + reusable_bytes = 0 + for current, dirnames, filenames in os.walk(root, topdown=True, followlinks=False): + current_path = Path(current) + for name in sorted(dirnames): + candidate = current_path / name + relative = relative_name(root, candidate) + mode = candidate.lstat().st_mode + if stat.S_ISLNK(mode) or not stat.S_ISDIR(mode) or relative not in expected_directories: + fail("peer staging path contains an unsafe entry") + entry_count += 1 + if entry_count > 8192: + fail("peer staging path exceeds the safety bound") + for name in sorted(filenames): + candidate = current_path / name + relative = relative_name(root, candidate) + mode = candidate.lstat().st_mode + size = candidate.stat().st_size + expected = expected_files.get(relative) + if ( + stat.S_ISLNK(mode) + or not stat.S_ISREG(mode) + or expected is None + or size > expected["size"] + ): + fail("peer staging path contains an unsafe entry") + if size == expected["size"]: + digest = hashlib.sha256() + with candidate.open("rb") as handle: + for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""): + digest.update(chunk) + if digest.hexdigest() == expected["sha256"]: + reusable_bytes += size + entry_count += 1 + if entry_count > 8192: + fail("peer staging path exceeds the safety bound") + return reusable_bytes + +verify_parent_chain() +if OPERATION == "prepare": + if path_exists(FINAL): + if manifest(FINAL) != EXPECTED: + fail("peer pinned snapshot already exists with different content") + if path_exists(STAGING): + verify_partial(STAGING) + shutil.rmtree(STAGING) + print(json.dumps({"state": "ready"}, separators=(",", ":"))) + else: + if path_exists(STAGING): + reusable_bytes = verify_partial(STAGING) + else: + reusable_bytes = 0 + remaining_bytes = max(0, EXPECTED["totalBytes"] - reusable_bytes) + headroom_bytes = max(5 * 1024 * 1024 * 1024, EXPECTED["totalBytes"] // 20) + if shutil.disk_usage(STAGING.parent).free < remaining_bytes + headroom_bytes: + fail("peer model cache does not have enough free space for the pinned snapshot") + if not path_exists(STAGING): + STAGING.mkdir(mode=0o700) + print(json.dumps({"state": "transfer"}, separators=(",", ":"))) +elif OPERATION == "finalize": + if path_exists(FINAL): + if manifest(FINAL) != EXPECTED: + fail("peer pinned snapshot appeared with different content") + if path_exists(STAGING): + verify_partial(STAGING) + shutil.rmtree(STAGING) + else: + if manifest(STAGING) != EXPECTED: + fail("peer staged snapshot failed byte-integrity verification") + staged_identity = (STAGING.stat().st_dev, STAGING.stat().st_ino) + os.rename(STAGING, FINAL) + installed_identity = (FINAL.stat().st_dev, FINAL.stat().st_ino) + if installed_identity != staged_identity: + fail("peer pinned snapshot identity changed during atomic install") + print(json.dumps({"state": "ready"}, separators=(",", ":"))) +elif OPERATION == "cleanup": + if path_exists(STAGING): + verify_partial(STAGING) + shutil.rmtree(STAGING) + print(json.dumps({"state": "cleaned"}, separators=(",", ":"))) +else: + fail("unsupported model staging operation") +`; +} + +function commandFailure(label: string, result: ModelStagingCommandResult): string { + if (result.timedOut) return `${label} timed out`; + const detail = ( + result.error || + result.stderr.trim().split("\n").at(-1) || + "command failed" + ).slice(0, 512); + return `${label} failed: ${detail}`; +} + +function parseRemoteState( + label: string, + result: ModelStagingCommandResult, + expected: "cleaned" | "ready" | "transfer", +): void { + if (result.status !== 0 || result.timedOut || result.error) { + throw new Error(commandFailure(label, result)); + } + let parsed: unknown; + try { + parsed = JSON.parse(result.stdout.trim()); + } catch { + throw new Error(`${label} returned invalid JSON`); + } + if ( + !parsed || + typeof parsed !== "object" || + Array.isArray(parsed) || + (parsed as Record).state !== expected + ) { + throw new Error(`${label} returned an unexpected state`); + } +} + +function parseRemotePrepareState(result: ModelStagingCommandResult): "ready" | "transfer" { + if (result.status !== 0 || result.timedOut || result.error) { + throw new Error(commandFailure("peer snapshot preflight", result)); + } + let parsed: unknown; + try { + parsed = JSON.parse(result.stdout.trim()); + } catch { + throw new Error("peer snapshot preflight returned invalid JSON"); + } + if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) { + throw new Error("peer snapshot preflight returned an unexpected state"); + } + const state = (parsed as Record).state; + if (state !== "ready" && state !== "transfer") { + throw new Error("peer snapshot preflight returned an unexpected state"); + } + return state; +} + +function sshArgs(plan: DualStationVllmPlan): string[] { + return [ + ...dualStationPinnedSshArgs(plan.peerSshBinding), + "--", + plan.peerSshBinding.peerTarget, + "python3 -", + ]; +} + +function shellQuote(value: string): string { + return `'${value.replaceAll("'", `'"'"'`)}'`; +} + +function rsyncRsh(plan: DualStationVllmPlan): string { + return ["ssh", ...dualStationPinnedSshArgs(plan.peerSshBinding)].map(shellQuote).join(" "); +} + +function createLocalStagingRoot(modelRoot: string): { root: string; snapshot: string } { + const modelRootMetadata = fs.lstatSync(modelRoot); + if (modelRootMetadata.isSymbolicLink() || !modelRootMetadata.isDirectory()) { + throw new Error("pinned local model cache root is missing or unsafe"); + } + // Keep the potentially large immutable copy on the model-cache filesystem; + // the OS temporary filesystem is commonly too small for the Ultra snapshot. + const root = fs.mkdtempSync(path.join(modelRoot, ".nemoclaw-vllm-model-staging-")); + try { + fs.chmodSync(root, 0o700); + } catch (err) { + try { + fs.rmSync(root, { force: false, recursive: true }); + } catch (cleanupError) { + throw new Error( + `${(err as Error).message}; local staging setup cleanup failed: ${(cleanupError as Error).message}`, + ); + } + throw err; + } + return { root, snapshot: path.join(root, "snapshot") }; +} + +function makeLocalStagingTreeRemovable(candidate: string): void { + const metadata = fs.lstatSync(candidate); + if (metadata.isSymbolicLink()) { + throw new Error("local audited snapshot cleanup encountered a symbolic link"); + } + if (metadata.isDirectory()) { + fs.chmodSync(candidate, 0o700); + for (const entry of fs.readdirSync(candidate)) { + makeLocalStagingTreeRemovable(path.join(candidate, entry)); + } + return; + } + if (!metadata.isFile()) { + throw new Error("local audited snapshot cleanup encountered a non-file entry"); + } +} + +function clearLocalStagingRoot(root: string): void { + makeLocalStagingTreeRemovable(root); + fs.rmSync(root, { force: false, recursive: true }); +} + +async function cleanupPeerStaging( + plan: DualStationVllmPlan, + paths: StagingPaths, + manifest: SnapshotManifest, + deps: DualStationModelStagingDeps, + env: Record, +): Promise { + try { + const cleanup = await deps.runCommand("ssh", sshArgs(plan), { + env, + input: remoteScript(plan, paths, manifest, "cleanup"), + timeoutMs: SNAPSHOT_AUDIT_TIMEOUT_MS, + }); + parseRemoteState("peer snapshot cleanup", cleanup, "cleaned"); + return null; + } catch (err) { + return (err as Error).message; + } +} + +async function failAfterPeerCleanup( + reason: string, + plan: DualStationVllmPlan, + paths: StagingPaths, + manifest: SnapshotManifest, + deps: DualStationModelStagingDeps, + env: Record, +): Promise { + const cleanupFailure = await cleanupPeerStaging(plan, paths, manifest, deps, env); + return { + ok: false, + reason: cleanupFailure ? `${reason}; ${cleanupFailure}` : reason, + }; +} + +async function auditLocalSnapshot( + paths: StagingPaths, + deps: DualStationModelStagingDeps, + env: Record, + localMaterializedSnapshot?: string, +): Promise { + const args = ["-", paths.localSnapshot, paths.localModelRoot]; + if (localMaterializedSnapshot !== undefined) args.push(localMaterializedSnapshot); + const audit = await deps.runCommand("python3", args, { + env, + input: LOCAL_MANIFEST_SCRIPT, + timeoutMs: SNAPSHOT_AUDIT_TIMEOUT_MS, + }); + return parseManifest(audit); +} + +async function preparePeerSnapshot( + plan: DualStationVllmPlan, + paths: StagingPaths, + manifest: SnapshotManifest, + deps: DualStationModelStagingDeps, + env: Record, +): Promise<"ready" | "transfer"> { + const prepare = await deps.runCommand("ssh", sshArgs(plan), { + env, + input: remoteScript(plan, paths, manifest, "prepare"), + timeoutMs: SNAPSHOT_AUDIT_TIMEOUT_MS, + }); + return parseRemotePrepareState(prepare); +} + +function manifestsEqual(left: SnapshotManifest, right: SnapshotManifest): boolean { + return ( + left.schemaVersion === right.schemaVersion && + left.totalBytes === right.totalBytes && + left.directories.length === right.directories.length && + left.directories.every((directory, index) => directory === right.directories[index]) && + left.files.length === right.files.length && + left.files.every((file, index) => { + const other = right.files[index]; + return ( + other !== undefined && + file.path === other.path && + file.size === other.size && + file.sha256 === other.sha256 + ); + }) + ); +} + +async function requireLocalStagingCapacity( + modelRoot: string, + manifest: SnapshotManifest, + deps: DualStationModelStagingDeps, +): Promise { + let stats: { bavail: bigint; bsize: bigint }; + try { + const statfs = deps.statfs ?? defaultDeps.statfs; + if (!statfs) throw new Error("statfs dependency is unavailable"); + stats = await statfs(modelRoot); + } catch (err) { + throw new Error(`local model cache capacity check failed: ${(err as Error).message}`); + } + if ( + typeof stats.bavail !== "bigint" || + typeof stats.bsize !== "bigint" || + stats.bavail < 0n || + stats.bsize <= 0n + ) { + throw new Error("local model cache capacity check returned invalid filesystem data"); + } + const snapshotBytes = BigInt(manifest.totalBytes); + const proportionalHeadroom = (snapshotBytes + 19n) / 20n; + const headroom = + proportionalHeadroom > LOCAL_STAGING_MIN_HEADROOM_BYTES + ? proportionalHeadroom + : LOCAL_STAGING_MIN_HEADROOM_BYTES; + if (stats.bavail * stats.bsize < snapshotBytes + headroom) { + throw new Error( + "local model cache does not have enough free space for the audited snapshot copy", + ); + } +} + +async function stagePreparedSnapshot( + plan: DualStationVllmPlan, + paths: StagingPaths, + auditedManifest: SnapshotManifest, + localMaterializedSnapshot: string, + deps: DualStationModelStagingDeps, + env: Record, +): Promise { + let transferManifest: SnapshotManifest; + try { + transferManifest = await auditLocalSnapshot(paths, deps, env, localMaterializedSnapshot); + } catch (err) { + return failAfterPeerCleanup((err as Error).message, plan, paths, auditedManifest, deps, env); + } + if (!manifestsEqual(auditedManifest, transferManifest)) { + return failAfterPeerCleanup( + "local pinned snapshot changed between audit and materialization", + plan, + paths, + auditedManifest, + deps, + env, + ); + } + + let prepareState: "ready" | "transfer"; + try { + prepareState = await preparePeerSnapshot(plan, paths, transferManifest, deps, env); + } catch (err) { + return failAfterPeerCleanup((err as Error).message, plan, paths, transferManifest, deps, env); + } + if (prepareState === "ready") { + return { ok: true, transferred: false }; + } + + let rsync: ModelStagingCommandResult; + try { + rsync = await deps.runCommand( + "rsync", + [ + "--recursive", + "--times", + "--checksum", + "--partial", + "--protect-args", + "--no-owner", + "--no-group", + "--chmod=Du=rwx,Dgo=,Fu=rw,Fgo=", + "--info=progress2", + `--rsh=${rsyncRsh(plan)}`, + "--", + `${localMaterializedSnapshot}/`, + `${plan.peerSshBinding.peerTarget}:${paths.peerStaging}/`, + ], + { + env, + timeoutMs: RSYNC_TIMEOUT_MS, + idleTimeoutMs: RSYNC_IDLE_TIMEOUT_MS, + streamOutput: true, + }, + ); + } catch (err) { + return failAfterPeerCleanup((err as Error).message, plan, paths, transferManifest, deps, env); + } + if (rsync.status !== 0 || rsync.timedOut || rsync.error) { + return failAfterPeerCleanup( + commandFailure("peer snapshot transfer", rsync), + plan, + paths, + transferManifest, + deps, + env, + ); + } + + let finalize: ModelStagingCommandResult; + try { + finalize = await deps.runCommand("ssh", sshArgs(plan), { + env, + input: remoteScript(plan, paths, transferManifest, "finalize"), + timeoutMs: SNAPSHOT_AUDIT_TIMEOUT_MS, + }); + } catch (err) { + return failAfterPeerCleanup((err as Error).message, plan, paths, transferManifest, deps, env); + } + try { + parseRemoteState("peer snapshot verification", finalize, "ready"); + } catch (err) { + return failAfterPeerCleanup((err as Error).message, plan, paths, transferManifest, deps, env); + } + return { ok: true, transferred: true }; +} + +/** + * Audit before peer preflight so an already-ready peer avoids the large local + * copy. Transfers use only a second, immutable audit that exactly matches the + * first manifest, so no token or unrelated blob cache crosses the SSH boundary. + */ +export async function stageDualStationModelSnapshot( + plan: DualStationVllmPlan, + deps: DualStationModelStagingDeps = defaultDeps, +): Promise { + let paths: StagingPaths; + try { + paths = stagingPaths(plan); + } catch (err) { + return { ok: false, reason: (err as Error).message }; + } + const env = buildVllmSshTransportEnv({ LC_ALL: "C" }); + + let auditedManifest: SnapshotManifest; + try { + auditedManifest = await auditLocalSnapshot(paths, deps, env); + } catch (err) { + return { ok: false, reason: (err as Error).message }; + } + + let prepareState: "ready" | "transfer"; + try { + prepareState = await preparePeerSnapshot(plan, paths, auditedManifest, deps, env); + } catch (err) { + return { ok: false, reason: (err as Error).message }; + } + if (prepareState === "ready") { + return { ok: true, transferred: false }; + } + + try { + await requireLocalStagingCapacity(paths.localModelRoot, auditedManifest, deps); + } catch (err) { + return failAfterPeerCleanup((err as Error).message, plan, paths, auditedManifest, deps, env); + } + + let localStaging: { root: string; snapshot: string }; + try { + localStaging = createLocalStagingRoot(paths.localModelRoot); + } catch (err) { + return failAfterPeerCleanup( + `local audited snapshot setup failed: ${(err as Error).message}`, + plan, + paths, + auditedManifest, + deps, + env, + ); + } + + let result: DualStationModelStagingResult; + try { + result = await stagePreparedSnapshot( + plan, + paths, + auditedManifest, + localStaging.snapshot, + deps, + env, + ); + } catch (err) { + result = await failAfterPeerCleanup( + (err as Error).message, + plan, + paths, + auditedManifest, + deps, + env, + ); + } + try { + clearLocalStagingRoot(localStaging.root); + } catch (err) { + const cleanupFailure = `local audited snapshot cleanup failed: ${(err as Error).message}`; + return { + ok: false, + reason: result.ok ? cleanupFailure : `${result.reason}; ${cleanupFailure}`, + }; + } + return result; +} diff --git a/src/lib/inference/vllm-station-ssh-binding.test-support.ts b/src/lib/inference/vllm-station-ssh-binding.test-support.ts new file mode 100644 index 0000000000..b7cd16ce35 --- /dev/null +++ b/src/lib/inference/vllm-station-ssh-binding.test-support.ts @@ -0,0 +1,65 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; + +import { + type DualStationSshBinding, + encodeDualStationSshBindingHandoff, + type QualifiedStationSshIdentity, + stationKnownHostsDigest, + writeDualStationSshBinding, +} from "./vllm-station-ssh-binding"; + +export interface DualStationSshBindingFixture { + binding: DualStationSshBinding; + dockerCliFile: string; + identity: QualifiedStationSshIdentity; + resumeStatePath: string; + token: string; + cleanup(): void; +} + +export function createDualStationSshBindingFixture( + peerTarget = "nvidia@station-b", +): DualStationSshBindingFixture { + const root = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-station-ssh-test-")); + fs.chmodSync(root, 0o700); + const dockerCliFile = path.join(root, "docker-cli"); + fs.writeFileSync(dockerCliFile, "#!/bin/bash\nexit 0\n", { mode: 0o700 }); + fs.chmodSync(dockerCliFile, 0o700); + const knownHostsLines = ["192.168.50.20 ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIGZpeHR1cmU="]; + const hostKeyDigest = stationKnownHostsDigest(`${knownHostsLines.join("\n")}\n`); + const resumeStatePath = path.join(root, "pair.json"); + const identity: QualifiedStationSshIdentity = { + requestedTarget: peerTarget, + sshTarget: peerTarget, + resolvedHost: "192.168.50.20", + sshUser: peerTarget.includes("@") ? peerTarget.split("@", 1)[0] : "nvidia", + port: 22, + lookupHost: "192.168.50.20", + hostKeyDigest, + knownHostsLines, + }; + const binding = writeDualStationSshBinding(resumeStatePath, identity, { dockerCliFile }); + return { + binding, + dockerCliFile, + identity, + resumeStatePath, + token: encodeDualStationSshBindingHandoff(binding), + cleanup: () => fs.rmSync(root, { recursive: true, force: true }), + }; +} + +export function retargetDualStationSshBindingFixture( + fixture: DualStationSshBindingFixture, + peerTarget: string, + validTarget: boolean, +): DualStationSshBindingFixture { + if (!validTarget || fixture.binding.peerTarget === peerTarget) return fixture; + fixture.cleanup(); + return createDualStationSshBindingFixture(peerTarget); +} diff --git a/src/lib/inference/vllm-station-ssh-binding.test.ts b/src/lib/inference/vllm-station-ssh-binding.test.ts new file mode 100644 index 0000000000..5d548b7da9 --- /dev/null +++ b/src/lib/inference/vllm-station-ssh-binding.test.ts @@ -0,0 +1,430 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { spawnSync } from "node:child_process"; +import { createHash } from "node:crypto"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; + +import { afterEach, beforeEach, describe, expect, it } from "vitest"; + +import { + assertDualStationSshBindingFiles, + clearDualStationSshBinding, + type DualStationSshBinding, + dualStationDockerSshUri, + dualStationPinnedSshArgs, + dualStationSshBindingDirectory, + encodeDualStationSshBindingHandoff, + loadDualStationSshBinding, + loadDualStationSshBindingHandoff, + type QualifiedStationSshIdentity, + stationKnownHostsDigest, + strictStationSshTransportArgs, + writeDualStationSshBinding, +} from "./vllm-station-ssh-binding"; + +const PEER_TARGET = "station@10.10.0.2"; +const PEER_HOST = "10.10.0.2"; +const PEER_PORT = 2222; +const ED25519_KEY = "AAAAC3NzaC1lZDI1NTE5AAAAIFirstQualifiedStationKey"; +const RSA_KEY = "AAAAB3NzaC1yc2EAAAADAQABAAABAQCRevokedStationKey"; +const KNOWN_HOSTS_LINES = [ + `[${PEER_HOST}]:${String(PEER_PORT)} ssh-ed25519 ${ED25519_KEY}`, + `@revoked [${PEER_HOST}]:${String(PEER_PORT)} ssh-rsa ${RSA_KEY}`, +] as const; + +function sha256(value: string): string { + return createHash("sha256").update(value).digest("hex"); +} + +function fileMode(filePath: string): number { + return fs.lstatSync(filePath).mode & 0o777; +} + +describe("qualified dual-Station SSH binding", () => { + let dockerCliFile: string; + let root: string; + let resumeStatePath: string; + + beforeEach(() => { + root = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-station-ssh-binding-")); + fs.chmodSync(root, 0o700); + dockerCliFile = path.join(root, "docker-cli"); + fs.writeFileSync( + dockerCliFile, + `#!/bin/bash +set -Eeuo pipefail +printf '%s\\0' "$@" > "${"${NEMOCLAW_TEST_DOCKER_RECORD:?}"}" +exit "${"${NEMOCLAW_TEST_DOCKER_EXIT:-0}"}" +`, + { mode: 0o700 }, + ); + fs.chmodSync(dockerCliFile, 0o700); + resumeStatePath = path.join(root, "station-dual-pair-resume.json"); + }); + + afterEach(() => { + fs.rmSync(root, { force: true, recursive: true }); + }); + + function identity(overrides: Partial = {}) { + const knownHostsLines = overrides.knownHostsLines ?? KNOWN_HOSTS_LINES; + return { + requestedTarget: PEER_TARGET, + sshTarget: PEER_TARGET, + resolvedHost: PEER_HOST, + sshUser: "station", + port: PEER_PORT, + lookupHost: `[${PEER_HOST}]:${String(PEER_PORT)}`, + hostKeyDigest: stationKnownHostsDigest(knownHostsLines.join("\n")), + knownHostsLines, + ...overrides, + } satisfies QualifiedStationSshIdentity; + } + + function writeBinding( + overrides: Partial = {}, + ): DualStationSshBinding { + return writeDualStationSshBinding(resumeStatePath, identity(overrides), { + dockerCliFile, + }); + } + + it("matches the coordinator digest over sorted unique key identities", () => { + const raw = [ + "# ignored comment", + KNOWN_HOSTS_LINES[1], + KNOWN_HOSTS_LINES[0], + KNOWN_HOSTS_LINES[0], + "", + ].join("\n"); + const expected = sha256( + [`@revoked|ssh-rsa|${RSA_KEY}`, `|ssh-ed25519|${ED25519_KEY}`].sort().join("\n"), + ); + + expect(stationKnownHostsDigest(raw)).toBe(expected); + expect(() => stationKnownHostsDigest(`@revoked ${PEER_HOST} ssh-rsa ${RSA_KEY}`)).toThrow( + "no trusted key", + ); + expect(() => + stationKnownHostsDigest( + [`@cert-authority ${PEER_HOST} ssh-ed25519 ${ED25519_KEY}`, KNOWN_HOSTS_LINES[0]].join( + "\n", + ), + ), + ).toThrow("marker is not allowed"); + }); + + it("persists owner-only evidence and reloads one canonical handoff", () => { + const binding = writeBinding(); + const runtimeDirectory = dualStationSshBindingDirectory(resumeStatePath); + const versionDirectory = path.dirname(binding.bindingFile); + const knownHosts = `${[...KNOWN_HOSTS_LINES].sort().join("\n")}\n`; + + expect(binding).toEqual( + expect.objectContaining({ + schemaVersion: 2, + peerTarget: PEER_TARGET, + resolvedHost: PEER_HOST, + sshUser: "station", + port: PEER_PORT, + lookupHost: `[${PEER_HOST}]:${String(PEER_PORT)}`, + hostKeyDigest: stationKnownHostsDigest(knownHosts), + bindingFile: path.join(versionDirectory, "binding.json"), + dockerCliFile: fs.realpathSync(dockerCliFile), + dockerShimFile: path.join(versionDirectory, "bin", "docker"), + knownHostsFile: path.join(versionDirectory, "known_hosts"), + sshWrapperDirectory: path.join(versionDirectory, "bin"), + sshWrapperFile: path.join(versionDirectory, "bin", "ssh"), + }), + ); + expect(path.dirname(versionDirectory)).toBe(runtimeDirectory); + expect(path.basename(versionDirectory)).toMatch(/^v2-[a-f0-9]{32}$/); + expect(fs.readFileSync(binding.knownHostsFile, "utf8")).toBe(knownHosts); + expect(fileMode(runtimeDirectory)).toBe(0o700); + expect(fileMode(versionDirectory)).toBe(0o700); + expect(fileMode(binding.sshWrapperDirectory)).toBe(0o700); + expect(fileMode(binding.bindingFile)).toBe(0o600); + expect(fileMode(binding.dockerShimFile)).toBe(0o700); + expect(fileMode(binding.knownHostsFile)).toBe(0o600); + expect(fileMode(binding.sshWrapperFile)).toBe(0o700); + + const token = encodeDualStationSshBindingHandoff(binding); + expect(loadDualStationSshBindingHandoff(token, PEER_TARGET)).toEqual(binding); + expect( + loadDualStationSshBinding(binding.bindingFile, PEER_TARGET, binding.hostKeyDigest), + ).toEqual(binding); + }); + + it("pins direct SSH and Docker-over-SSH to the qualified endpoint", () => { + const binding = writeBinding(); + const args = dualStationPinnedSshArgs(binding); + + expect(args.slice(0, strictStationSshTransportArgs().length)).toEqual( + strictStationSshTransportArgs(), + ); + expect(args).toEqual( + expect.arrayContaining([ + `UserKnownHostsFile=${binding.knownHostsFile}`, + "GlobalKnownHostsFile=/dev/null", + `HostKeyAlias=${binding.lookupHost}`, + `Hostname=${PEER_HOST}`, + "User=station", + `Port=${String(PEER_PORT)}`, + ]), + ); + expect(dualStationDockerSshUri(binding)).toBe( + `ssh://station@${PEER_HOST}:${String(PEER_PORT)}`, + ); + + expect(spawnSync("/bin/bash", ["-n", binding.sshWrapperFile]).status).toBe(0); + expect(spawnSync("/bin/bash", ["-n", binding.dockerShimFile]).status).toBe(0); + + const dockerRecord = path.join(root, "docker-args"); + const dockerResult = spawnSync( + binding.dockerShimFile, + ["context", "inspect", "value with spaces"], + { + env: { + ...process.env, + NEMOCLAW_TEST_DOCKER_EXIT: "37", + NEMOCLAW_TEST_DOCKER_RECORD: dockerRecord, + }, + }, + ); + expect(dockerResult.status).toBe(37); + expect(fs.readFileSync(dockerRecord).toString("utf8").split("\0").slice(0, -1)).toEqual([ + "context", + "inspect", + "value with spaces", + ]); + + fs.appendFileSync(binding.knownHostsFile, "# tampered\n"); + const sshResult = spawnSync(binding.sshWrapperFile, ["-V"], { encoding: "utf8" }); + expect(sshResult.status).toBe(255); + expect(sshResult.stderr).toContain("refused a changed dual-Station SSH host-key pin"); + }); + + it("omits only the default SSH port from the Docker URI", () => { + const binding = writeDualStationSshBinding( + resumeStatePath, + identity({ port: 22, lookupHost: PEER_HOST }), + { dockerCliFile }, + ); + + expect(dualStationDockerSshUri(binding)).toBe(`ssh://station@${PEER_HOST}`); + }); + + it("keeps an earlier accepted host-key set in its immutable version", () => { + const first = writeBinding(); + const replacementKey = "BBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB"; + const replacementLines = [`[${PEER_HOST}]:${String(PEER_PORT)} ssh-ed25519 ${replacementKey}`]; + const second = writeDualStationSshBinding( + resumeStatePath, + identity({ + knownHostsLines: replacementLines, + hostKeyDigest: stationKnownHostsDigest(replacementLines.join("\n")), + }), + { dockerCliFile }, + ); + + expect(second.hostKeyDigest).not.toBe(first.hostKeyDigest); + expect(path.dirname(second.bindingFile)).not.toBe(path.dirname(first.bindingFile)); + expect(() => assertDualStationSshBindingFiles(first)).not.toThrow(); + expect(() => assertDualStationSshBindingFiles(second)).not.toThrow(); + expect(fs.readFileSync(first.knownHostsFile, "utf8")).toContain(ED25519_KEY); + expect(fs.readFileSync(second.knownHostsFile, "utf8")).toContain(replacementKey); + }); + + it.each([ + { + name: "mismatched requested target", + override: { requestedTarget: "station@10.10.0.3" }, + }, + { name: "non-canonical resolved host", override: { resolvedHost: "999.999.999.999" } }, + { name: "unsafe SSH user", override: { sshUser: "station root" } }, + { name: "mismatched explicit SSH user", override: { sshUser: "other-user" } }, + { name: "out-of-range port", override: { port: 65_536 } }, + { name: "mismatched lookup host", override: { lookupHost: PEER_HOST } }, + { name: "mismatched digest", override: { hostKeyDigest: "0".repeat(64) } }, + { name: "comment evidence", override: { knownHostsLines: ["# not evidence"] } }, + { + name: "multiline evidence", + override: { knownHostsLines: [`${KNOWN_HOSTS_LINES[0]}\n${KNOWN_HOSTS_LINES[1]}`] }, + }, + ])("rejects $name before persisting it", ({ override }) => { + expect(() => + writeDualStationSshBinding( + resumeStatePath, + identity(override as Partial), + { dockerCliFile }, + ), + ).toThrow(); + expect(fs.existsSync(dualStationSshBindingDirectory(resumeStatePath))).toBe(false); + }); + + it("rejects a state path that cannot be one Docker helper PATH entry", () => { + const incompatiblePath = path.join(root, `pair${path.delimiter}resume.json`); + + expect(() => + writeDualStationSshBinding(incompatiblePath, identity(), { dockerCliFile }), + ).toThrow("normalized absolute path"); + }); + + it.each([ + ["binding file", (binding: DualStationSshBinding) => binding.bindingFile, 0o644], + ["Docker CLI", (binding: DualStationSshBinding) => binding.dockerCliFile, 0o722], + ["Docker shim", (binding: DualStationSshBinding) => binding.dockerShimFile, 0o755], + ["known-hosts file", (binding: DualStationSshBinding) => binding.knownHostsFile, 0o644], + ["wrapper file", (binding: DualStationSshBinding) => binding.sshWrapperFile, 0o755], + ["wrapper directory", (binding: DualStationSshBinding) => binding.sshWrapperDirectory, 0o755], + [ + "binding version", + (binding: DualStationSshBinding) => path.dirname(binding.bindingFile), + 0o755, + ], + [ + "binding root", + (binding: DualStationSshBinding) => path.dirname(path.dirname(binding.bindingFile)), + 0o755, + ], + ["binding parent", (_binding: DualStationSshBinding) => root, 0o755], + ] as const)("rejects unsafe %s permissions", (_name, target, mode) => { + const binding = writeBinding(); + fs.chmodSync(target(binding), mode); + + expect(() => + loadDualStationSshBinding(binding.bindingFile, PEER_TARGET, binding.hostKeyDigest), + ).toThrow(); + }); + + it.each([ + ["binding file", (binding: DualStationSshBinding) => binding.bindingFile], + ["Docker shim", (binding: DualStationSshBinding) => binding.dockerShimFile], + ["known-hosts file", (binding: DualStationSshBinding) => binding.knownHostsFile], + ["wrapper file", (binding: DualStationSshBinding) => binding.sshWrapperFile], + ] as const)("refuses a symbolic-link %s", (_name, selectedPath) => { + const binding = writeBinding(); + const filePath = selectedPath(binding); + const copyPath = path.join(root, `copy-${path.basename(filePath)}`); + fs.copyFileSync(filePath, copyPath); + fs.chmodSync(copyPath, fileMode(filePath)); + fs.unlinkSync(filePath); + fs.symlinkSync(copyPath, filePath); + + expect(() => + loadDualStationSshBinding(binding.bindingFile, PEER_TARGET, binding.hostKeyDigest), + ).toThrow(); + }); + + it("rejects known-hosts and wrapper tampering before returning transport data", () => { + const binding = writeBinding(); + fs.appendFileSync(binding.knownHostsFile, "# changed\n"); + + expect(() => dualStationPinnedSshArgs(binding)).toThrow( + "known-hosts binding changed after qualification", + ); + expect(() => dualStationDockerSshUri(binding)).toThrow( + "known-hosts binding changed after qualification", + ); + + const restored = writeBinding(); + fs.appendFileSync(restored.sshWrapperFile, "# changed\n"); + expect(() => dualStationPinnedSshArgs(restored)).toThrow("wrapper changed after qualification"); + + const dockerTampered = writeBinding(); + fs.appendFileSync(dockerTampered.dockerShimFile, "# changed\n"); + expect(() => dualStationDockerSshUri(dockerTampered)).toThrow( + "Docker shim changed after qualification", + ); + }); + + it.each([ + ["binding file", (binding: DualStationSshBinding) => binding.bindingFile, 16 * 1024 + 1], + ["known-hosts file", (binding: DualStationSshBinding) => binding.knownHostsFile, 64 * 1024 + 1], + ["wrapper file", (binding: DualStationSshBinding) => binding.sshWrapperFile, 16 * 1024 + 1], + ["Docker shim", (binding: DualStationSshBinding) => binding.dockerShimFile, 16 * 1024 + 1], + ] as const)("rejects an oversized %s", (_name, selectedPath, size) => { + const binding = writeBinding(); + fs.writeFileSync(selectedPath(binding), "x".repeat(size)); + + expect(() => + loadDualStationSshBinding(binding.bindingFile, PEER_TARGET, binding.hostKeyDigest), + ).toThrow(); + }); + + it.each([ + ["unexpected field", (value: Record): void => void (value.extra = true)], + [ + "unsupported schema", + (value: Record): void => void (value.schemaVersion = 1), + ], + [ + "string port", + (value: Record): void => void (value.port = String(PEER_PORT)), + ], + [ + "changed peer", + (value: Record): void => void (value.peerTarget = "station@10.10.0.3"), + ], + ] as const)("rejects a binding JSON %s", (_name, mutate) => { + const binding = writeBinding(); + const value = JSON.parse(fs.readFileSync(binding.bindingFile, "utf8")) as Record< + string, + unknown + >; + mutate(value); + fs.writeFileSync(binding.bindingFile, `${JSON.stringify(value)}\n`); + + expect(() => + loadDualStationSshBinding(binding.bindingFile, PEER_TARGET, binding.hostKeyDigest), + ).toThrow(); + }); + + it("rejects forged handoffs and in-memory endpoint fields", () => { + const binding = writeBinding(); + const token = encodeDualStationSshBindingHandoff(binding); + const extraFieldToken = Buffer.from( + JSON.stringify({ + bindingFile: binding.bindingFile, + hostKeyDigest: binding.hostKeyDigest, + peerTarget: PEER_TARGET, + }), + "utf8", + ).toString("base64url"); + + expect(() => loadDualStationSshBindingHandoff(`${token}=`, PEER_TARGET)).toThrow( + "NEMOCLAW_DGX_STATION_SSH_BINDING is invalid", + ); + expect(() => loadDualStationSshBindingHandoff(extraFieldToken, PEER_TARGET)).toThrow( + "unexpected fields", + ); + expect(() => loadDualStationSshBindingHandoff(token, "station@10.10.0.3")).toThrow( + "does not match the qualified peer identity", + ); + expect(() => + dualStationPinnedSshArgs({ + ...binding, + resolvedHost: "10.10.0.2 ProxyCommand=attacker", + }), + ).toThrow("identity is invalid"); + }); + + it("clears only an owner-only regular binding tree", () => { + writeBinding(); + const runtimeDirectory = dualStationSshBindingDirectory(resumeStatePath); + clearDualStationSshBinding(resumeStatePath); + clearDualStationSshBinding(resumeStatePath); + expect(fs.existsSync(runtimeDirectory)).toBe(false); + + const outside = path.join(root, "outside"); + fs.mkdirSync(outside, { mode: 0o700 }); + const marker = path.join(outside, "keep"); + fs.writeFileSync(marker, "keep", { mode: 0o600 }); + fs.symlinkSync(outside, runtimeDirectory); + expect(() => clearDualStationSshBinding(resumeStatePath)).toThrow("unsafe to remove"); + expect(fs.readFileSync(marker, "utf8")).toBe("keep"); + }); +}); diff --git a/src/lib/inference/vllm-station-ssh-binding.ts b/src/lib/inference/vllm-station-ssh-binding.ts new file mode 100644 index 0000000000..ceb413ec27 --- /dev/null +++ b/src/lib/inference/vllm-station-ssh-binding.ts @@ -0,0 +1,781 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { createHash, randomBytes } from "node:crypto"; +import fs from "node:fs"; +import net from "node:net"; +import path from "node:path"; + +export const NEMOCLAW_DGX_STATION_SSH_BINDING_ENV = "NEMOCLAW_DGX_STATION_SSH_BINDING"; + +const BINDING_SCHEMA_VERSION = 2; +const SHA256_HEX_PATTERN = /^[a-f0-9]{64}$/; +const SSH_USERNAME_PATTERN = /^[A-Za-z_][A-Za-z0-9._-]*$/; +const SSH_HOST_PATTERN = + /^[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?(?:\.[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?)*$/; +const SSH_KEY_TYPE_PATTERN = /^(?:ssh-|ecdsa-|sk-)[A-Za-z0-9@._+-]+$/; +const SSH_KEY_DATA_PATTERN = /^[A-Za-z0-9+/]+={0,3}$/; +const MAX_KNOWN_HOSTS_LINE_BYTES = 16 * 1024; +const MAX_BINDING_BYTES = 16 * 1024; +const MAX_KNOWN_HOSTS_BYTES = 64 * 1024; +const MAX_WRAPPER_BYTES = 16 * 1024; +const VERSION_DIRECTORY_PATTERN = /^v2-[a-f0-9]{32}$/; +const BINDING_KEYS = [ + "bindingFile", + "dockerCliFile", + "dockerShimFile", + "dockerShimSha256", + "hostKeyDigest", + "knownHostsFile", + "knownHostsSha256", + "lookupHost", + "peerTarget", + "port", + "resolvedHost", + "schemaVersion", + "sshUser", + "sshWrapperDirectory", + "sshWrapperFile", + "sshWrapperSha256", +] as const; + +export interface QualifiedStationSshIdentity { + requestedTarget: string; + sshTarget: string; + resolvedHost: string; + sshUser: string; + port: number; + lookupHost: string; + hostKeyDigest: string; + knownHostsLines: readonly string[]; +} + +export interface DualStationSshBinding { + schemaVersion: 2; + peerTarget: string; + resolvedHost: string; + sshUser: string; + port: number; + lookupHost: string; + hostKeyDigest: string; + bindingFile: string; + dockerCliFile: string; + dockerShimFile: string; + dockerShimSha256: string; + knownHostsFile: string; + knownHostsSha256: string; + sshWrapperDirectory: string; + sshWrapperFile: string; + sshWrapperSha256: string; +} + +interface BindingHandoff { + bindingFile: string; + hostKeyDigest: string; +} + +export interface WriteDualStationSshBindingOptions { + /** Test seam; production resolves the effective Docker CLI from PATH. */ + dockerCliFile?: string; +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +function requireString(value: unknown, label: string, maxLength: number): string { + if ( + typeof value !== "string" || + value.length === 0 || + value.length > maxLength || + value !== value.trim() || + /[\u0000-\u001f\u007f]/.test(value) + ) { + throw new Error(`${label} is invalid`); + } + return value; +} + +function isCanonicalHost(value: string): boolean { + return net.isIP(value) === 4 || (!/^[0-9.]+$/.test(value) && SSH_HOST_PATTERN.test(value)); +} + +function validatePeerTarget(value: string): string { + requireString(value, "Station SSH peer target", 286); + if (/[/,:;`'"\\$(){}[\]<>|&!?*\s]/.test(value)) { + throw new Error("Station SSH peer target is invalid"); + } + const parts = value.split("@"); + if (parts.length > 2) throw new Error("Station SSH peer target is invalid"); + const user = parts.length === 2 ? parts[0] : ""; + const host = parts.at(-1) ?? ""; + if ((user && !SSH_USERNAME_PATTERN.test(user)) || !isCanonicalHost(host)) { + throw new Error("Station SSH peer target is invalid"); + } + const canonical = user ? `${user}@${host}` : host; + if (canonical !== value) throw new Error("Station SSH peer target is not canonical"); + return canonical; +} + +function explicitPeerUser(peerTarget: string): string | null { + const separator = peerTarget.indexOf("@"); + return separator === -1 ? null : peerTarget.slice(0, separator); +} + +function requireAbsolutePath(value: unknown, label: string): string { + const parsed = requireString(value, label, 4096); + if ( + !path.isAbsolute(parsed) || + path.normalize(parsed) !== parsed || + parsed.includes(path.delimiter) + ) { + throw new Error(`${label} must be a normalized absolute path`); + } + return parsed; +} + +function requireInteger(value: unknown, label: string, minimum: number, maximum: number): number { + if (typeof value !== "number" || !Number.isInteger(value) || value < minimum || value > maximum) { + throw new Error(`${label} is invalid`); + } + return value; +} + +function requireExactKeys( + value: Record, + expected: readonly string[], + label: string, +): void { + const actual = Object.keys(value).sort(); + if (actual.length !== expected.length || actual.some((key, index) => key !== expected[index])) { + throw new Error(`${label} has unexpected fields`); + } +} + +function currentUid(): number { + const uid = process.getuid?.(); + if (uid === undefined) throw new Error("Station SSH binding requires a POSIX user identity"); + return uid; +} + +function assertDirectory(filePath: string, mode: number, label: string): void { + const metadata = fs.lstatSync(filePath); + if ( + metadata.isSymbolicLink() || + !metadata.isDirectory() || + metadata.uid !== currentUid() || + (metadata.mode & 0o777) !== mode + ) { + throw new Error(`${label} must be an owner-only directory`); + } +} + +function readBoundedFile(filePath: string, mode: number, maxBytes: number, label: string): Buffer { + if (typeof fs.constants.O_NOFOLLOW !== "number") { + throw new Error("Station SSH binding requires O_NOFOLLOW support"); + } + const flags = fs.constants.O_RDONLY | fs.constants.O_NOFOLLOW | fs.constants.O_NONBLOCK; + let fd: number; + try { + fd = fs.openSync(filePath, flags); + } catch (error) { + if ((error as NodeJS.ErrnoException).code === "ELOOP") { + throw new Error(`${label} must not be a symbolic link`); + } + throw error; + } + try { + const metadata = fs.fstatSync(fd); + if ( + !metadata.isFile() || + metadata.uid !== currentUid() || + (metadata.mode & 0o777) !== mode || + metadata.size <= 0 || + metadata.size > maxBytes + ) { + throw new Error(`${label} metadata is invalid`); + } + const content = fs.readFileSync(fd); + if (content.length !== metadata.size || content.length > maxBytes) { + throw new Error(`${label} changed while it was being read`); + } + return content; + } finally { + fs.closeSync(fd); + } +} + +function fsyncDirectory(directory: string): void { + const fd = fs.openSync(directory, fs.constants.O_RDONLY); + try { + fs.fsyncSync(fd); + } finally { + fs.closeSync(fd); + } +} + +function sha256(value: Buffer | string): string { + return createHash("sha256").update(value).digest("hex"); +} + +function validateDockerCliFile(filePath: string): string { + const requested = requireAbsolutePath(filePath, "Station Docker CLI"); + let canonical: string; + try { + canonical = fs.realpathSync(requested); + } catch { + throw new Error("Station Docker CLI could not be resolved"); + } + requireAbsolutePath(canonical, "Station Docker CLI"); + const metadata = fs.lstatSync(canonical); + const uid = currentUid(); + if ( + metadata.isSymbolicLink() || + !metadata.isFile() || + (metadata.uid !== 0 && metadata.uid !== uid) || + (metadata.mode & 0o111) === 0 || + (metadata.mode & 0o022) !== 0 || + metadata.size <= 0 + ) { + throw new Error("Station Docker CLI is not a safe executable file"); + } + return canonical; +} + +function resolveDockerCliFile(explicit?: string): string { + if (explicit !== undefined) return validateDockerCliFile(explicit); + for (const directory of (process.env.PATH ?? "").split(path.delimiter)) { + if (!directory || !path.isAbsolute(directory) || path.normalize(directory) !== directory) { + continue; + } + const candidate = path.join(directory, "docker"); + try { + return validateDockerCliFile(candidate); + } catch { + // Continue through the effective PATH until one safe Docker CLI is found. + } + } + throw new Error("Station Docker CLI could not be resolved to a safe absolute executable"); +} + +export function stationKnownHostsDigest(raw: string): string { + const keys = new Set(); + let positiveKeys = 0; + for (const rawLine of raw.split(/\r?\n/)) { + const line = rawLine.trim(); + if (!line || line.startsWith("#") || /[\u0000\r\n]/.test(line)) continue; + const fields = line.split(/\s+/); + const marker = fields[0]?.startsWith("@") ? (fields.shift() ?? "") : ""; + if (marker !== "" && marker !== "@revoked") { + throw new Error("Pinned Station known-hosts marker is not allowed"); + } + if (fields.length < 3) throw new Error("Pinned Station known-hosts data is invalid"); + const [_hosts, keyType, keyData] = fields; + if (!SSH_KEY_TYPE_PATTERN.test(keyType) || !SSH_KEY_DATA_PATTERN.test(keyData)) { + throw new Error("Pinned Station known-hosts key is invalid"); + } + keys.add(`${marker}|${keyType}|${keyData}`); + if (marker === "") positiveKeys += 1; + } + if (keys.size === 0 || positiveKeys === 0) { + throw new Error("Pinned Station known-hosts data has no trusted key"); + } + return sha256([...keys].sort().join("\n")); +} + +function shellQuote(value: string): string { + return `'${value.replaceAll("'", `'"'"'`)}'`; +} + +export function strictStationSshTransportArgs(): string[] { + return [ + "-T", + "-o", + "BatchMode=yes", + "-o", + "StrictHostKeyChecking=yes", + "-o", + "VerifyHostKeyDNS=no", + "-o", + "NoHostAuthenticationForLocalhost=no", + "-o", + "NumberOfPasswordPrompts=0", + "-o", + "PasswordAuthentication=no", + "-o", + "KbdInteractiveAuthentication=no", + "-o", + "PreferredAuthentications=publickey", + "-o", + "ConnectTimeout=5", + "-o", + "ConnectionAttempts=1", + "-o", + "ServerAliveInterval=5", + "-o", + "ServerAliveCountMax=1", + "-o", + "ClearAllForwardings=yes", + "-o", + "ForwardAgent=no", + "-o", + "ForwardX11=no", + "-o", + "ForwardX11Trusted=no", + "-o", + "Tunnel=no", + "-o", + "UpdateHostKeys=no", + "-o", + "ControlMaster=no", + "-o", + "ControlPath=none", + "-o", + "PermitLocalCommand=no", + "-o", + "RemoteCommand=none", + "-o", + "ProxyCommand=none", + "-o", + "ProxyJump=none", + "-o", + "KnownHostsCommand=none", + "-o", + "LogLevel=ERROR", + ]; +} + +type PinnedStationEndpoint = Pick< + DualStationSshBinding, + "knownHostsFile" | "lookupHost" | "port" | "resolvedHost" | "sshUser" +>; + +function pinnedOptionArgs(binding: PinnedStationEndpoint): string[] { + return [ + ...strictStationSshTransportArgs(), + "-o", + `UserKnownHostsFile=${binding.knownHostsFile}`, + "-o", + "GlobalKnownHostsFile=/dev/null", + "-o", + `HostKeyAlias=${binding.lookupHost}`, + "-o", + `Hostname=${binding.resolvedHost}`, + "-o", + `User=${binding.sshUser}`, + "-o", + `Port=${String(binding.port)}`, + ]; +} + +function renderSshWrapper(binding: PinnedStationEndpoint & { knownHostsSha256: string }): string { + const args = pinnedOptionArgs(binding).map(shellQuote).join(" "); + return `#!/bin/bash +set -Eeuo pipefail +readonly known_hosts=${shellQuote(binding.knownHostsFile)} +readonly expected_sha256=${shellQuote(binding.knownHostsSha256)} +if [[ -x /usr/bin/sha256sum ]]; then + actual_sha256="$(/usr/bin/sha256sum < "$known_hosts")" || exit 255 +elif [[ -x /usr/bin/shasum ]]; then + actual_sha256="$(/usr/bin/shasum -a 256 < "$known_hosts")" || exit 255 +else + printf '%s\n' 'NemoClaw could not verify the dual-Station SSH host-key pin.' >&2 + exit 255 +fi +actual_sha256="${"${actual_sha256%% *}"}" +if [[ "$actual_sha256" != "$expected_sha256" ]]; then + printf '%s\n' 'NemoClaw refused a changed dual-Station SSH host-key pin.' >&2 + exit 255 +fi +exec /usr/bin/ssh ${args} "$@" +`; +} + +function renderDockerShim(binding: Pick): string { + return `#!/bin/bash +set -Eeuo pipefail +exec ${shellQuote(binding.dockerCliFile)} "$@" +`; +} + +function writeExclusive(filePath: string, content: string, mode: number): void { + if (typeof fs.constants.O_NOFOLLOW !== "number") { + throw new Error("Station SSH binding requires O_NOFOLLOW support"); + } + const flags = + fs.constants.O_WRONLY | fs.constants.O_CREAT | fs.constants.O_EXCL | fs.constants.O_NOFOLLOW; + const fd = fs.openSync(filePath, flags, mode); + try { + fs.writeFileSync(fd, content, "utf8"); + fs.fsyncSync(fd); + fs.fchmodSync(fd, mode); + } finally { + fs.closeSync(fd); + } +} + +function assertRemovableTree(root: string): void { + const metadata = fs.lstatSync(root); + if ( + metadata.isSymbolicLink() || + !metadata.isDirectory() || + metadata.uid !== currentUid() || + (metadata.mode & 0o777) !== 0o700 + ) { + throw new Error("Station SSH binding directory is unsafe to remove"); + } + for (const entry of fs.readdirSync(root, { withFileTypes: true })) { + const child = path.join(root, entry.name); + const childMetadata = fs.lstatSync(child); + if (childMetadata.isSymbolicLink() || childMetadata.uid !== currentUid()) { + throw new Error("Station SSH binding tree contains an unsafe entry"); + } + if (childMetadata.isDirectory()) assertRemovableTree(child); + else if (!childMetadata.isFile()) { + throw new Error("Station SSH binding tree contains a non-file entry"); + } + } +} + +export function dualStationSshBindingDirectory(resumeStatePath: string): string { + const statePath = requireAbsolutePath(resumeStatePath, "Dual-Station resume-state path"); + return `${statePath}.ssh-binding`; +} + +export function clearDualStationSshBinding(resumeStatePath: string): void { + const runtimeDirectory = dualStationSshBindingDirectory(resumeStatePath); + try { + fs.lstatSync(runtimeDirectory); + } catch (error) { + if ((error as NodeJS.ErrnoException).code === "ENOENT") return; + throw error; + } + assertRemovableTree(runtimeDirectory); + fs.rmSync(runtimeDirectory, { recursive: true, force: false }); + fsyncDirectory(path.dirname(runtimeDirectory)); +} + +function ensureBindingRoot(runtimeDirectory: string, parent: string): void { + try { + fs.mkdirSync(runtimeDirectory, { mode: 0o700 }); + fs.chmodSync(runtimeDirectory, 0o700); + fsyncDirectory(parent); + } catch (error) { + if ((error as NodeJS.ErrnoException).code !== "EEXIST") throw error; + } + assertDirectory(runtimeDirectory, 0o700, "Station SSH binding root"); +} + +export function writeDualStationSshBinding( + resumeStatePath: string, + identity: QualifiedStationSshIdentity, + options: WriteDualStationSshBindingOptions = {}, +): DualStationSshBinding { + const runtimeDirectory = dualStationSshBindingDirectory(resumeStatePath); + const parent = path.dirname(runtimeDirectory); + assertDirectory(parent, 0o700, "Station SSH binding parent"); + const peerTarget = validatePeerTarget(identity.sshTarget); + if (identity.requestedTarget !== peerTarget || !isCanonicalHost(identity.resolvedHost)) { + throw new Error("Qualified Station SSH endpoint is invalid"); + } + if (!SSH_USERNAME_PATTERN.test(identity.sshUser)) { + throw new Error("Qualified Station SSH user is invalid"); + } + const requestedUser = explicitPeerUser(peerTarget); + if (requestedUser !== null && requestedUser !== identity.sshUser) { + throw new Error("Qualified Station SSH target and resolved user do not match"); + } + if (!Number.isInteger(identity.port) || identity.port < 1 || identity.port > 65535) { + throw new Error("Qualified Station SSH port is invalid"); + } + const expectedLookup = + identity.port === 22 + ? identity.resolvedHost + : `[${identity.resolvedHost}]:${String(identity.port)}`; + if (identity.lookupHost !== expectedLookup || !SHA256_HEX_PATTERN.test(identity.hostKeyDigest)) { + throw new Error("Qualified Station SSH host-key identity is invalid"); + } + if ( + !Array.isArray(identity.knownHostsLines) || + identity.knownHostsLines.length === 0 || + identity.knownHostsLines.some( + (line) => + typeof line !== "string" || + line.length === 0 || + Buffer.byteLength(line, "utf8") > MAX_KNOWN_HOSTS_LINE_BYTES || + line !== line.trim() || + line.startsWith("#") || + /[\u0000\r\n]/.test(line), + ) + ) { + throw new Error("Qualified Station known-hosts evidence is invalid"); + } + const knownHosts = `${[...new Set(identity.knownHostsLines)].sort().join("\n")}\n`; + if (Buffer.byteLength(knownHosts, "utf8") > MAX_KNOWN_HOSTS_BYTES) { + throw new Error("Qualified Station known-hosts evidence is too large"); + } + if (stationKnownHostsDigest(knownHosts) !== identity.hostKeyDigest) { + throw new Error("Qualified Station known-hosts evidence does not match its accepted digest"); + } + + const dockerCliFile = resolveDockerCliFile(options.dockerCliFile); + ensureBindingRoot(runtimeDirectory, parent); + const versionDirectory = path.join(runtimeDirectory, `v2-${randomBytes(16).toString("hex")}`); + fs.mkdirSync(versionDirectory, { mode: 0o700 }); + fs.chmodSync(versionDirectory, 0o700); + try { + const wrapperDirectory = path.join(versionDirectory, "bin"); + fs.mkdirSync(wrapperDirectory, { mode: 0o700 }); + const finalKnownHostsFile = path.join(versionDirectory, "known_hosts"); + const finalBindingFile = path.join(versionDirectory, "binding.json"); + const finalWrapperDirectory = path.join(versionDirectory, "bin"); + const finalWrapperFile = path.join(finalWrapperDirectory, "ssh"); + const finalDockerShimFile = path.join(finalWrapperDirectory, "docker"); + const knownHostsSha256 = sha256(knownHosts); + const persisted = { + schemaVersion: 2, + peerTarget, + resolvedHost: identity.resolvedHost, + sshUser: identity.sshUser, + port: identity.port, + lookupHost: identity.lookupHost, + hostKeyDigest: identity.hostKeyDigest, + bindingFile: finalBindingFile, + dockerCliFile, + dockerShimFile: finalDockerShimFile, + knownHostsFile: finalKnownHostsFile, + knownHostsSha256, + sshWrapperDirectory: finalWrapperDirectory, + sshWrapperFile: finalWrapperFile, + } satisfies Omit; + const wrapper = renderSshWrapper(persisted); + const dockerShim = renderDockerShim(persisted); + const binding: DualStationSshBinding = { + ...persisted, + dockerShimSha256: sha256(dockerShim), + sshWrapperSha256: sha256(wrapper), + }; + writeExclusive(path.join(versionDirectory, "known_hosts"), knownHosts, 0o600); + writeExclusive(path.join(wrapperDirectory, "ssh"), wrapper, 0o700); + writeExclusive(path.join(wrapperDirectory, "docker"), dockerShim, 0o700); + writeExclusive( + path.join(versionDirectory, "binding.json"), + `${JSON.stringify(binding)}\n`, + 0o600, + ); + fsyncDirectory(wrapperDirectory); + fsyncDirectory(versionDirectory); + fsyncDirectory(runtimeDirectory); + return loadDualStationSshBinding(binding.bindingFile, peerTarget, identity.hostKeyDigest); + } catch (error) { + try { + fs.rmSync(versionDirectory, { recursive: true, force: true }); + fsyncDirectory(runtimeDirectory); + } catch { + // Preserve the original validation or persistence failure. + } + throw error; + } +} + +function parseBinding(value: unknown): DualStationSshBinding { + if (!isRecord(value) || value.schemaVersion !== BINDING_SCHEMA_VERSION) { + throw new Error("Station SSH binding schema is unsupported"); + } + requireExactKeys(value, BINDING_KEYS, "Station SSH binding"); + const binding: DualStationSshBinding = { + schemaVersion: 2, + peerTarget: validatePeerTarget(requireString(value.peerTarget, "Station SSH peer", 286)), + resolvedHost: requireString(value.resolvedHost, "Station SSH resolved host", 253), + sshUser: requireString(value.sshUser, "Station SSH user", 64), + port: requireInteger(value.port, "Station SSH port", 1, 65535), + lookupHost: requireString(value.lookupHost, "Station SSH lookup host", 300), + hostKeyDigest: requireString(value.hostKeyDigest, "Station SSH host-key digest", 64), + bindingFile: requireAbsolutePath(value.bindingFile, "Station SSH binding file"), + dockerCliFile: requireAbsolutePath(value.dockerCliFile, "Station Docker CLI"), + dockerShimFile: requireAbsolutePath(value.dockerShimFile, "Station Docker shim"), + dockerShimSha256: requireString(value.dockerShimSha256, "Station Docker shim SHA-256", 64), + knownHostsFile: requireAbsolutePath(value.knownHostsFile, "Station SSH known-hosts file"), + knownHostsSha256: requireString(value.knownHostsSha256, "Station SSH known-hosts SHA-256", 64), + sshWrapperDirectory: requireAbsolutePath( + value.sshWrapperDirectory, + "Station SSH wrapper directory", + ), + sshWrapperFile: requireAbsolutePath(value.sshWrapperFile, "Station SSH wrapper file"), + sshWrapperSha256: requireString(value.sshWrapperSha256, "Station SSH wrapper SHA-256", 64), + }; + if ( + !isCanonicalHost(binding.resolvedHost) || + !SSH_USERNAME_PATTERN.test(binding.sshUser) || + (explicitPeerUser(binding.peerTarget) !== null && + explicitPeerUser(binding.peerTarget) !== binding.sshUser) || + !SHA256_HEX_PATTERN.test(binding.hostKeyDigest) || + !SHA256_HEX_PATTERN.test(binding.dockerShimSha256) || + !SHA256_HEX_PATTERN.test(binding.knownHostsSha256) || + !SHA256_HEX_PATTERN.test(binding.sshWrapperSha256) + ) { + throw new Error("Station SSH binding identity is invalid"); + } + const expectedLookup = + binding.port === 22 + ? binding.resolvedHost + : `[${binding.resolvedHost}]:${String(binding.port)}`; + if (binding.lookupHost !== expectedLookup) { + throw new Error("Station SSH binding lookup host is invalid"); + } + return binding; +} + +function canonicalBinding(value: unknown): DualStationSshBinding { + return parseBinding(value); +} + +function validateDualStationSshBindingFiles(binding: DualStationSshBinding): DualStationSshBinding { + const canonical = canonicalBinding(binding); + const versionDirectory = path.dirname(canonical.bindingFile); + const runtimeDirectory = path.dirname(versionDirectory); + if ( + !path.basename(runtimeDirectory).endsWith(".ssh-binding") || + !VERSION_DIRECTORY_PATTERN.test(path.basename(versionDirectory)) || + canonical.bindingFile !== path.join(versionDirectory, "binding.json") || + canonical.knownHostsFile !== path.join(versionDirectory, "known_hosts") || + canonical.sshWrapperDirectory !== path.join(versionDirectory, "bin") || + canonical.sshWrapperFile !== path.join(versionDirectory, "bin", "ssh") || + canonical.dockerShimFile !== path.join(versionDirectory, "bin", "docker") + ) { + throw new Error("Station SSH binding paths are inconsistent"); + } + assertDirectory(path.dirname(runtimeDirectory), 0o700, "Station SSH binding parent"); + assertDirectory(runtimeDirectory, 0o700, "Station SSH binding root"); + assertDirectory(versionDirectory, 0o700, "Station SSH binding version"); + assertDirectory(canonical.sshWrapperDirectory, 0o700, "Station SSH wrapper directory"); + if (validateDockerCliFile(canonical.dockerCliFile) !== canonical.dockerCliFile) { + throw new Error("Station Docker CLI changed after qualification"); + } + const dockerShim = readBoundedFile( + canonical.dockerShimFile, + 0o700, + MAX_WRAPPER_BYTES, + "Station Docker shim", + ); + if ( + sha256(dockerShim) !== canonical.dockerShimSha256 || + dockerShim.toString("utf8") !== renderDockerShim(canonical) + ) { + throw new Error("Station Docker shim changed after qualification"); + } + const knownHosts = readBoundedFile( + canonical.knownHostsFile, + 0o600, + MAX_KNOWN_HOSTS_BYTES, + "Station SSH known-hosts file", + ); + if ( + sha256(knownHosts) !== canonical.knownHostsSha256 || + stationKnownHostsDigest(knownHosts.toString("utf8")) !== canonical.hostKeyDigest + ) { + throw new Error("Station SSH known-hosts binding changed after qualification"); + } + const wrapper = readBoundedFile( + canonical.sshWrapperFile, + 0o700, + MAX_WRAPPER_BYTES, + "Station SSH wrapper", + ); + if ( + sha256(wrapper) !== canonical.sshWrapperSha256 || + wrapper.toString("utf8") !== renderSshWrapper(canonical) + ) { + throw new Error("Station SSH wrapper changed after qualification"); + } + return canonical; +} + +export function assertDualStationSshBindingFiles(binding: DualStationSshBinding): void { + validateDualStationSshBindingFiles(binding); +} + +export function loadDualStationSshBinding( + bindingFile: string, + expectedPeerTarget: string, + expectedHostKeyDigest: string, +): DualStationSshBinding { + const expectedPeer = validatePeerTarget(expectedPeerTarget); + if (!SHA256_HEX_PATTERN.test(expectedHostKeyDigest)) { + throw new Error("Expected Station SSH host-key digest is invalid"); + } + const normalizedBindingFile = requireAbsolutePath(bindingFile, "Station SSH binding file"); + const raw = readBoundedFile( + normalizedBindingFile, + 0o600, + MAX_BINDING_BYTES, + "Station SSH binding file", + ); + let value: unknown; + try { + value = JSON.parse(raw.toString("utf8")); + } catch { + throw new Error("Station SSH binding file contains invalid JSON"); + } + const binding = parseBinding(value); + if ( + binding.bindingFile !== normalizedBindingFile || + binding.peerTarget !== expectedPeer || + binding.hostKeyDigest !== expectedHostKeyDigest + ) { + throw new Error("Station SSH binding does not match the qualified peer identity"); + } + return validateDualStationSshBindingFiles(binding); +} + +export function encodeDualStationSshBindingHandoff(binding: DualStationSshBinding): string { + const canonical = validateDualStationSshBindingFiles(binding); + const handoff: BindingHandoff = { + bindingFile: canonical.bindingFile, + hostKeyDigest: canonical.hostKeyDigest, + }; + return Buffer.from(JSON.stringify(handoff), "utf8").toString("base64url"); +} + +export function loadDualStationSshBindingHandoff( + token: string, + expectedPeerTarget: string, +): DualStationSshBinding { + if ( + typeof token !== "string" || + token.length === 0 || + token.length > 8192 || + !/^[A-Za-z0-9_-]+$/.test(token) + ) { + throw new Error(`${NEMOCLAW_DGX_STATION_SSH_BINDING_ENV} is invalid`); + } + let value: unknown; + try { + const decoded = Buffer.from(token, "base64url"); + if (decoded.toString("base64url") !== token) throw new Error("non-canonical base64url"); + value = JSON.parse(decoded.toString("utf8")); + } catch { + throw new Error(`${NEMOCLAW_DGX_STATION_SSH_BINDING_ENV} is invalid`); + } + if (!isRecord(value)) { + throw new Error(`${NEMOCLAW_DGX_STATION_SSH_BINDING_ENV} is invalid`); + } + requireExactKeys(value, ["bindingFile", "hostKeyDigest"], NEMOCLAW_DGX_STATION_SSH_BINDING_ENV); + const bindingFile = requireAbsolutePath(value.bindingFile, "Station SSH binding handoff file"); + const hostKeyDigest = requireString( + value.hostKeyDigest, + "Station SSH binding handoff digest", + 64, + ); + return loadDualStationSshBinding(bindingFile, expectedPeerTarget, hostKeyDigest); +} + +export function dualStationPinnedSshArgs(binding: DualStationSshBinding): string[] { + return pinnedOptionArgs(validateDualStationSshBindingFiles(binding)); +} + +export function dualStationDockerSshUri(binding: DualStationSshBinding): string { + const canonical = validateDualStationSshBindingFiles(binding); + const port = canonical.port === 22 ? "" : `:${String(canonical.port)}`; + return `ssh://${canonical.sshUser}@${canonical.resolvedHost}${port}`; +} diff --git a/src/lib/inference/vllm.test.ts b/src/lib/inference/vllm.test.ts index 6b365186ca..42d7cc3b5b 100644 --- a/src/lib/inference/vllm.test.ts +++ b/src/lib/inference/vllm.test.ts @@ -150,7 +150,7 @@ function vllmContainerRow( containerName: string, { id = MANAGED_CONTAINER_ID, label = "true", state = "exited" } = {}, ): string { - return `${id}|${containerName}|${state}|${label}`; + return `${id}|${containerName}|${state}|${label}|||`; } function mockSuccessfulVllmInstall( @@ -174,8 +174,19 @@ function mockSuccessfulVllmInstall( mocks.dockerSpawn.mockReturnValue(mockDockerSpawnSuccess()); mocks.dockerRunDetached.mockReturnValue({ status: 0, stdout: "", stderr: "", error: null }); const ownershipQueue = [...ownershipResponses]; + let ownershipCallIndex = 0; + const ownershipHandlers = [ + (): string => "", + (): string => + ( + ownershipQueue.shift() ?? + (() => { + throw new Error("Unexpected extra ambient vLLM ownership inspection"); + }) + )(), + ]; const dockerCaptureByCommand = new Map string>([ - ["container", () => (ownershipQueue.shift() ?? (() => ""))()], + ["container", () => ownershipHandlers[ownershipCallIndex++ % ownershipHandlers.length]()], ["ps", () => `${containerName}\n`], ]); mocks.dockerCapture.mockImplementation((args: readonly string[]) => @@ -579,7 +590,15 @@ describe("managed vLLM ownership", () => { "--filter", `name=^/${NEMOCLAW_VLLM_CONTAINER_NAME}$`, "--format", - `{{.ID}}|{{.Names}}|{{.State}}|{{.Label "${NEMOCLAW_VLLM_MANAGED_LABEL}"}}`, + [ + "{{.ID}}", + "{{.Names}}", + "{{.State}}", + `{{.Label "${NEMOCLAW_VLLM_MANAGED_LABEL}"}}`, + '{{.Label "com.nvidia.nemoclaw.vllm-role"}}', + '{{.Label "com.nvidia.nemoclaw.vllm-endpoint"}}', + '{{.Label "com.nvidia.nemoclaw.vllm-cluster"}}', + ].join("|"), ], expect.objectContaining({ timeout: 10_000 }), ); @@ -739,6 +758,7 @@ describe("installVllm model resolution", () => { it("rejects a Station-only runtime override before side effects on generic Linux", async () => { process.env.NEMOCLAW_VLLM_MODEL = "nemotron-3-ultra-550b-a55b"; + process.env.HF_TOKEN = "hf_test"; const profile = detectVllmProfile({ platform: "linux", type: "nvidia" })!; const beforeInstall = vi.fn(); @@ -822,6 +842,7 @@ describe("installVllm model resolution", () => { it("installs the complete Nemotron Ultra Station recipe without another selection", async () => { process.env.NEMOCLAW_VLLM_MODEL = "nemotron-3-ultra-550b-a55b"; + process.env.HF_TOKEN = "hf_test"; mocks.getGpuIndicesByName.mockReturnValue([0]); const profile = detectVllmProfile({ platform: "station", type: "nvidia" })!; const beforeInstall = vi.fn(); @@ -1007,8 +1028,16 @@ describe("installVllm model resolution", () => { ...mocks.dockerRunDetached.mock.calls.map((call) => call[1]), ...mocks.dockerCapture.mock.calls.map((call) => call[1]), ]; - expect(dockerAdapterOptions).toHaveLength(7); - for (const options of dockerAdapterOptions) { + expect(dockerAdapterOptions).toHaveLength(9); + const canonicalOwnershipOptions = dockerAdapterOptions.filter( + (options) => options.env?.DOCKER_CONTEXT === "default", + ); + expect(canonicalOwnershipOptions).toHaveLength(2); + const ambientDockerOptions = dockerAdapterOptions.filter( + (options) => options.env?.DOCKER_CONTEXT !== "default", + ); + expect(ambientDockerOptions).toHaveLength(7); + for (const options of ambientDockerOptions) { expect(options).toEqual( expect.objectContaining({ env: expect.objectContaining({ DOCKER_CONTEXT: "local-test-context" }), diff --git a/src/lib/inference/vllm.ts b/src/lib/inference/vllm.ts index 20a2bbe5ce..b72835c853 100644 --- a/src/lib/inference/vllm.ts +++ b/src/lib/inference/vllm.ts @@ -9,6 +9,7 @@ import fs from "node:fs"; import os from "node:os"; import path from "node:path"; import { StringDecoder } from "node:string_decoder"; +import { isDeepStrictEqual } from "node:util"; import { dockerCapture, dockerForceRm, @@ -18,7 +19,9 @@ import { dockerSpawn, dockerStop, } from "../adapters/docker"; +import { createBearerAuthConfig } from "../adapters/http/auth-config"; import { buildValidatedCurlCommandArgs } from "../adapters/http/curl-args"; +import { runCurlProbe } from "../adapters/http/probe"; import { CLI_NAME } from "../cli/branding"; import { warnLine } from "../cli/terminal-style"; import { markPhaseActivity } from "../core/phase-activity"; @@ -29,9 +32,15 @@ import { runCapture } from "../runner"; import { redactFull } from "../security/redact"; import { isSafeModelId } from "../validation"; import { getGpuIndicesByName } from "./nim"; -import { buildVllmDockerEnv } from "./vllm-docker-env"; +import { ensureDualStationVllmApiKey, loadDualStationVllmApiKey } from "./vllm-api-key"; +import { + buildLocalDualStationDockerEnv, + buildRemoteVllmDockerEnv, + buildVllmDockerEnv, +} from "./vllm-docker-env"; import { buildVllmServeCommand, + NEMOTRON_ULTRA_DUAL_STATION_IMAGE, NEMOTRON_ULTRA_STATION_IMAGE, parseVllmExtraServeArgs, VLLM_EXTRA_ARGS_ENV, @@ -40,6 +49,26 @@ import { type VllmPlatform, } from "./vllm-models"; import { resolveVllmInstallModel } from "./vllm-prompt"; +import { + type DualStationVllmPlan, + NEMOCLAW_DGX_STATION_PEER_ENV, + probeDualStationVllmCapability, +} from "./vllm-station-cluster"; +import { + areDualStationManagedVllmContainersRunning, + cleanupDualStationManagedVllm, + commitDualStationLegacyMigration, + DUAL_STATION_VLLM_CLUSTER_LABEL, + DUAL_STATION_VLLM_ENDPOINT_LABEL, + DUAL_STATION_VLLM_ROLE_LABEL, + getDualStationManagedVllmBaseUrl, + preflightDualStationGpuRuntime, + preflightDualStationManagedVllm, + rollbackDualStationLegacyMigration, + startDualStationManagedVllm, + withDualStationManagedVllmLifecycle, +} from "./vllm-station-cluster-lifecycle"; +import { stageDualStationModelSnapshot } from "./vllm-station-model-staging"; import { findUnwritableModelCachePath, formatStorageBytes, @@ -447,7 +476,10 @@ function dockerPrereqsOk(): { ok: boolean; reason?: string } { return { ok: true }; } -export async function pullImage(profile: VllmProfile): Promise<{ ok: boolean; reason?: string }> { +export async function pullImage( + profile: VllmProfile, + dockerEnv: Record = buildVllmDockerEnv(), +): Promise<{ ok: boolean; reason?: string }> { try { assertVllmRegistryDigestRef(profile.image); } catch (err) { @@ -458,7 +490,7 @@ export async function pullImage(profile: VllmProfile): Promise<{ ok: boolean; re // profile, so all profiles intentionally share the 15-minute stall default. // The profile-specific maximum still bounds the complete pull operation. const result = await dockerPullWithProgressWatchdog(profile.image, { - env: buildVllmDockerEnv(), + env: dockerEnv, maxTimeoutMs: profile.pullTimeoutSec * 1000, logLine: emit, }); @@ -481,6 +513,7 @@ export async function pullImage(profile: VllmProfile): Promise<{ ok: boolean; re function downloadModel( profile: VllmProfile, model: VllmModelDef, + dockerEnv: Record = buildVllmDockerEnv(), ): Promise<{ ok: boolean; reason?: string }> { emit(`Pre-downloading model with hf: ${model.id}`); return new Promise((resolve) => { @@ -505,7 +538,7 @@ function downloadModel( ...(model.revision ? ["--revision", model.revision] : []), ], { - env: buildVllmDockerEnv(buildHfTokenForwardEnv()), + env: { ...dockerEnv, ...buildHfTokenForwardEnv() }, stdio: ["ignore", "pipe", "pipe"], }, ); @@ -738,12 +771,24 @@ export function assertVllmRegistryDigestRef(image: string): void { type VllmContainerOwnership = | { kind: "absent" } + | { kind: "dual-managed"; containerId: string; running: boolean } | { kind: "foreign" } | { kind: "managed"; containerId: string; running: boolean } | { kind: "unknown" }; -function inspectVllmContainerOwnership(containerName: string): VllmContainerOwnership { - const format = `{{.ID}}|{{.Names}}|{{.State}}|{{.Label "${NEMOCLAW_VLLM_MANAGED_LABEL}"}}`; +function inspectVllmContainerOwnershipInDockerEnv( + containerName: string, + env: Record, +): VllmContainerOwnership { + const format = [ + "{{.ID}}", + "{{.Names}}", + "{{.State}}", + `{{.Label "${NEMOCLAW_VLLM_MANAGED_LABEL}"}}`, + `{{.Label "${DUAL_STATION_VLLM_ROLE_LABEL}"}}`, + `{{.Label "${DUAL_STATION_VLLM_ENDPOINT_LABEL}"}}`, + `{{.Label "${DUAL_STATION_VLLM_CLUSTER_LABEL}"}}`, + ].join("|"); try { const output = dockerCapture( [ @@ -756,25 +801,54 @@ function inspectVllmContainerOwnership(containerName: string): VllmContainerOwne "--format", format, ], - { env: buildVllmDockerEnv(), timeout: 10_000 }, + { env, timeout: 10_000 }, ).trim(); if (!output) return { kind: "absent" }; const rows = output.split(/\r?\n/); if (rows.length !== 1) return { kind: "unknown" }; const fields = rows[0].split("|"); - if (fields.length !== 4) return { kind: "unknown" }; - const [containerId, observedName, state, managedLabel] = fields; + if (fields.length !== 7) return { kind: "unknown" }; + const [containerId, observedName, state, managedLabel, dualRole, dualEndpoint, dualCluster] = + fields; if (observedName !== containerName || !DOCKER_CONTAINER_ID_PATTERN.test(containerId)) { return { kind: "unknown" }; } if (managedLabel !== "true") return { kind: "foreign" }; + const hasAnyDualLabel = Boolean(dualRole || dualEndpoint || dualCluster); + if (hasAnyDualLabel) { + const exactDualHead = + dualRole === "head" && + /^http:\/\/192\.168\.|^http:\/\/10\.|^http:\/\/172\.(?:1[6-9]|2[0-9]|3[01])\./.test( + dualEndpoint, + ) && + /^[a-f0-9]{64}$/.test(dualCluster); + return exactDualHead + ? { kind: "dual-managed", containerId, running: state === "running" } + : { kind: "unknown" }; + } return { kind: "managed", containerId, running: state === "running" }; } catch { return { kind: "unknown" }; } } +function inspectVllmContainerOwnership(containerName: string): VllmContainerOwnership { + // A managed dual-Station head always lives on the physical host's default + // daemon. Inspect it before following ambient single-host Docker routing so + // DOCKER_HOST, DOCKER_CONTEXT, or Docker's persisted currentContext cannot + // hide the pair from running-state detection or replacement guards. + const canonicalOwnership = inspectVllmContainerOwnershipInDockerEnv( + containerName, + buildLocalDualStationDockerEnv(), + ); + if (canonicalOwnership.kind === "dual-managed" || canonicalOwnership.kind === "unknown") { + return canonicalOwnership; + } + + return inspectVllmContainerOwnershipInDockerEnv(containerName, buildVllmDockerEnv()); +} + function vllmContainerReplacementTarget( containerName: string, ): { ok: true; containerId?: string } | { ok: false; reason: string } { @@ -791,6 +865,14 @@ function vllmContainerReplacementTarget( reason: `Could not verify ownership of Docker container "${containerName}". NemoClaw will not remove it. Check Docker access and retry.`, }; } + if (ownership.kind === "dual-managed") { + return { + ok: false, + reason: + `Container "${containerName}" is the head of a managed dual-Station deployment. ` + + `Refusing single-host replacement because it would orphan the peer worker. Restore ${NEMOCLAW_DGX_STATION_PEER_ENV} and select Nemotron Ultra to manage the pair.`, + }; + } return ownership.kind === "managed" ? { ok: true, containerId: ownership.containerId } : { ok: true }; @@ -798,7 +880,7 @@ function vllmContainerReplacementTarget( export function isNemoClawManagedVllmRunning(): boolean { const ownership = inspectVllmContainerOwnership(NEMOCLAW_VLLM_CONTAINER_NAME); - return ownership.kind === "managed" && ownership.running; + return (ownership.kind === "managed" || ownership.kind === "dual-managed") && ownership.running; } function startContainer( @@ -839,11 +921,17 @@ function startContainer( return { ok: true }; } -function vllmModelsEndpoint(): string { - return `http://127.0.0.1:${String(VLLM_PORT)}/v1/models`; -} - -function vllmEndpointReady(): boolean { +function vllmEndpointReady(baseUrl?: string): boolean { + if (baseUrl) { + // The dual-Station /v1 surface is bearer-protected. vLLM deliberately + // leaves /health outside its auth middleware, so readiness can stay + // secret-free while onboarding separately validates model inventory with + // the persisted key. + return runCurlProbe( + ["-sS", "--connect-timeout", "2", "--max-time", "5", `${baseUrl.replace(/\/+$/, "")}/health`], + { pinnedAddresses: [] }, + ).ok; + } const response = runCapture( [ "curl", @@ -853,7 +941,7 @@ function vllmEndpointReady(): boolean { "2", "--max-time", "5", - vllmModelsEndpoint(), + `http://127.0.0.1:${String(VLLM_PORT)}/v1/models`, ]), ], { ignoreError: true }, @@ -867,27 +955,102 @@ function vllmEndpointReady(): boolean { } } -function readContainerLogTail(profile: VllmProfile, lineCount = 80): string[] { +function verifyDualStationVllmAuthBoundary( + baseUrl: string, + apiKey: string, + expectedModelId: string, +): { ok: true } | { ok: false; reason: string } { + const modelsUrl = `${baseUrl.replace(/\/+$/, "")}/v1/models`; + const unauthenticated = runCurlProbe( + ["-sS", "--connect-timeout", "3", "--max-time", "5", modelsUrl], + { pinnedAddresses: [] }, + ); + if (unauthenticated.httpStatus !== 401) { + return { + ok: false, + reason: + `unauthenticated model inventory returned HTTP ${String(unauthenticated.httpStatus)}; ` + + "expected vLLM to reject it with HTTP 401", + }; + } + + let authConfig: ReturnType | undefined; + try { + authConfig = createBearerAuthConfig(apiKey, { prefix: "nemoclaw-vllm-install-auth" }); + const authenticated = runCurlProbe( + ["-sS", "--connect-timeout", "3", "--max-time", "5", ...authConfig.args, modelsUrl], + { + trustedConfigFiles: authConfig.trustedConfigFiles, + pinnedAddresses: [], + }, + ); + if (!authenticated.ok) { + return { + ok: false, + reason: `authenticated model inventory failed: ${authenticated.message}`, + }; + } + let parsed: unknown; + try { + parsed = JSON.parse(authenticated.body); + } catch { + return { ok: false, reason: "authenticated model inventory returned malformed JSON" }; + } + const data = (parsed as { data?: unknown } | null)?.data; + const ids = (Array.isArray(data) ? data : []).flatMap((entry) => { + if (typeof entry !== "object" || entry === null) return []; + const id = (entry as { id?: unknown }).id; + return typeof id === "string" ? [id] : []; + }); + if (ids.length !== 1 || ids[0] !== expectedModelId) { + return { + ok: false, + reason: `authenticated model inventory did not expose exactly '${expectedModelId}'`, + }; + } + return { ok: true }; + } catch (error) { + return { + ok: false, + reason: `authenticated model inventory failed: ${(error as Error).message}`, + }; + } finally { + authConfig?.cleanup(); + } +} + +function readContainerLogTail( + profile: VllmProfile, + lineCount = 80, + dockerEnv: Record = buildVllmDockerEnv(), +): string[] { const output = dockerCapture(["logs", "--tail", String(lineCount), profile.containerName], { - env: buildVllmDockerEnv(), + env: dockerEnv, ignoreError: true, }).trim(); if (!output) return []; return output.split(/\r?\n/).slice(-lineCount); } -function printContainerLogTail(profile: VllmProfile): void { - const tail = readContainerLogTail(profile); +function printContainerLogTail( + profile: VllmProfile, + dockerEnv: Record = buildVllmDockerEnv(), +): void { + const tail = readContainerLogTail(profile, 80, dockerEnv); if (tail.length === 0) return; process.stderr.write(` --- Last ${String(tail.length)} vLLM log lines: ---\n`); for (const line of tail) process.stderr.write(` ${line}\n`); process.stderr.write(" ---\n"); } -// Poll the real OpenAI-compatible models endpoint instead of interpreting -// vLLM startup logs. Logs stay quiet on the happy path and print only on -// failure. -function waitForVllmReady(profile: VllmProfile): Promise<{ ok: boolean; reason?: string }> { +// Poll the real OpenAI-compatible models endpoint for the legacy local path, +// or the secret-free vLLM health endpoint for authenticated dual-Station +// serving. Logs stay quiet on the happy path and print only on failure. +function waitForVllmReady( + profile: VllmProfile, + baseUrl?: string, + dockerEnv: Record = buildVllmDockerEnv(), +): Promise<{ ok: boolean; reason?: string }> { return new Promise((resolve) => { let resolved = false; const start = Date.now(); @@ -907,7 +1070,7 @@ function waitForVllmReady(profile: VllmProfile): Promise<{ ok: boolean; reason?: function poll(): void { if (resolved) return; - if (vllmEndpointReady()) { + if (vllmEndpointReady(baseUrl)) { emit(`vLLM is serving on :${String(VLLM_PORT)}`); done({ ok: true }); return; @@ -920,7 +1083,7 @@ function waitForVllmReady(profile: VllmProfile): Promise<{ ok: boolean; reason?: }); return; } - if (!containerStillRunning(profile)) { + if (!containerStillRunning(profile, dockerEnv)) { done({ ok: false, reason: "vLLM container exited before readiness" }); return; } @@ -935,10 +1098,13 @@ function waitForVllmReady(profile: VllmProfile): Promise<{ ok: boolean; reason?: }); } -function containerStillRunning(profile: VllmProfile): boolean { +function containerStillRunning( + profile: VllmProfile, + dockerEnv: Record = buildVllmDockerEnv(), +): boolean { const out = dockerCapture( ["ps", "--filter", `name=${profile.containerName}`, "--format", "{{.Names}}"], - { env: buildVllmDockerEnv(), ignoreError: true }, + { env: dockerEnv, ignoreError: true }, ).trim(); return out === profile.containerName; } @@ -1144,6 +1310,7 @@ async function managedStorageAccepted( model: VllmModelDef, hasImage: boolean, opts: InstallVllmOptions, + dockerEnv: Record = buildVllmDockerEnv(), ): Promise { const includeImage = !hasImage; const modelDownloadSizeBytes = profile.modelDownloadSizeBytes ?? model.downloadSizeBytes; @@ -1163,7 +1330,18 @@ async function managedStorageAccepted( modelBytes: includeModel ? Number(remainingModelBytes) : modelDownloadSizeBytes, writableAllowanceBytes: VLLM_WRITABLE_ALLOWANCE_BYTES, }); - const dockerProbe = includeImage ? probeDockerStorage() : null; + const dockerProbe = includeImage + ? probeDockerStorage({ + dockerContext: dockerEnv.DOCKER_CONTEXT, + dockerHost: dockerEnv.DOCKER_HOST, + dockerInfo: () => + dockerCapture(["info", "--format", "{{json .}}"], { + env: dockerEnv, + ignoreError: true, + timeout: 10_000, + }), + }) + : null; const modelProbe = includeModel ? probeHostStorage(hostHfCacheDir(), "Hugging Face cache") : null; const requirements = managedStorageRequirements({ dockerProbe, @@ -1267,10 +1445,13 @@ interface InstallVllmOptions { beforeInstall?: (modelId: string) => void; } -function imageIsCached(profile: VllmProfile): boolean { +function imageIsCached( + profile: VllmProfile, + dockerEnv: Record = buildVllmDockerEnv(), +): boolean { return Boolean( dockerImageInspectFormat("{{.Id}}", profile.image, { - env: buildVllmDockerEnv(), + env: dockerEnv, ignoreError: true, timeout: 10_000, }).trim(), @@ -1319,13 +1500,66 @@ async function runVllmInstall( profile: VllmProfile, opts: InstallVllmOptions, ): Promise<{ ok: boolean }> { + let dualStationPlan: DualStationVllmPlan | null = null; + let peerModelSnapshot: "ready" | "staging-required" | null = null; + const explicitModel = String(process.env.NEMOCLAW_VLLM_MODEL ?? "").trim(); + const configuredPeer = String(process.env[NEMOCLAW_DGX_STATION_PEER_ENV] ?? "").trim(); + const ultra = + profile.platform === "station" && configuredPeer + ? VLLM_MODELS.find((candidate) => candidate.envValue === "nemotron-3-ultra-550b-a55b") + : undefined; + + if (profile.platform === "station" && configuredPeer) { + if (!ultra) { + console.error(" vLLM install failed: Nemotron Ultra is missing from the model registry"); + return { ok: false }; + } + const normalizedExplicitModel = explicitModel.toLowerCase(); + if ( + normalizedExplicitModel && + normalizedExplicitModel !== ultra.envValue.toLowerCase() && + normalizedExplicitModel !== ultra.id.toLowerCase() + ) { + console.error( + ` vLLM install failed: ${NEMOCLAW_DGX_STATION_PEER_ENV} requires the DGX Station dual-serving model. ` + + "Unset NEMOCLAW_VLLM_MODEL or select nemotron-3-ultra-550b-a55b; the explicit model override remains authoritative.", + ); + return { ok: false }; + } + } // Model selection lives in `resolveVllmInstallModel` so this entry point // stays focused on the docker side effects. Gated-model access is checked // there before any docker work happens. - const resolved = await resolveVllmInstallModel(profile, { - nonInteractive: opts.nonInteractive, - promptFn: opts.promptFn, - }); + let resolved: Awaited>; + if (profile.platform === "station" && configuredPeer && !explicitModel && ultra) { + const capability = probeDualStationVllmCapability(); + if (capability.kind !== "ready") { + const reason = + capability.kind === "unavailable" + ? capability.reason + : "the explicit peer configuration disappeared"; + console.error(` Dual DGX Station setup unavailable: ${reason}`); + return { ok: false }; + } + resolved = await resolveVllmInstallModel( + { ...profile, defaultModel: ultra }, + { + // A qualified explicit peer is the model-selection signal. The normal + // resolver still owns access validation, but no second model choice is + // presented after hardware qualification. + nonInteractive: true, + promptFn: opts.promptFn, + }, + ); + if (!resolved) return { ok: false }; + dualStationPlan = capability.plan; + peerModelSnapshot = capability.peerModelSnapshot; + } else { + resolved = await resolveVllmInstallModel(profile, { + nonInteractive: opts.nonInteractive, + promptFn: opts.promptFn, + }); + } if (!resolved) return { ok: false }; const { model, source: modelSource } = resolved; // Platform-restricted models are filtered out of the interactive picker, @@ -1355,6 +1589,43 @@ async function runVllmInstall( console.error(` vLLM install failed: ${(err as Error).message}`); return { ok: false }; } + + if (profile.platform === "station" && model.envValue === "nemotron-3-ultra-550b-a55b") { + if (!dualStationPlan) { + const capability = probeDualStationVllmCapability(); + if (capability.kind === "unavailable") { + console.error(` Dual DGX Station setup unavailable: ${capability.reason}`); + return { ok: false }; + } + if (capability.kind === "ready") { + dualStationPlan = capability.plan; + peerModelSnapshot = capability.peerModelSnapshot; + } + } + if (dualStationPlan) { + servedModelId = dualStationPlan.runtime.servedModelId; + runtimeProfile = { + ...runtimeProfile, + image: dualStationPlan.runtime.image, + imageDownloadSizeBytes: NEMOTRON_ULTRA_DUAL_STATION_IMAGE.arm64.downloadSizeBytes, + imageUnpackedSizeBytes: undefined, + loadTimeoutSec: 7200, + }; + if (VLLM_PORT !== 8000) { + console.error( + " Dual DGX Station setup requires the default vLLM port 8000; unset NEMOCLAW_VLLM_PORT and retry.", + ); + return { ok: false }; + } + if (extraServeArgs.length > 0) { + console.error( + ` Dual DGX Station setup does not accept ${VLLM_EXTRA_ARGS_ENV}; the verified distributed launch is fixed.`, + ); + return { ok: false }; + } + } + } + const localDockerEnv = dualStationPlan ? buildLocalDualStationDockerEnv() : buildVllmDockerEnv(); opts.beforeInstall?.(servedModelId); console.log(""); @@ -1368,6 +1639,14 @@ async function runVllmInstall( ` Extra serve args: ${String(extraServeArgs.length)} token(s) from ${VLLM_EXTRA_ARGS_ENV}`, ); } + if (dualStationPlan) { + console.log( + ` Topology: 2× DGX Station (${dualStationPlan.local.hostname} + ${dualStationPlan.peer.hostname})`, + ); + console.log( + ` Fabric: ${dualStationPlan.rails.map((rail) => rail.subnet).join(", ")} (RoCEv2 GID ${String(dualStationPlan.roceGidIndex)})`, + ); + } if (!opts.hasImage) console.log(" Image download on first run, cached after"); console.log(" Model download on first run, cached after"); printHfDownloadAuthentication(opts.nonInteractive); @@ -1387,19 +1666,28 @@ async function runVllmInstall( return { ok: false }; } - // Fail before large downloads when the fixed name belongs to another - // operator. startContainer repeats this check to close the teardown race. - const replacement = vllmContainerReplacementTarget(runtimeProfile.containerName); - if (!replacement.ok) { - console.error(` vLLM install failed: ${replacement.reason}`); - return { ok: false }; + // Fail before large downloads when either daemon has an ambiguous or + // foreign fixed-name container. Each launch path repeats this ownership + // check immediately before teardown to close the name-transfer race. + if (dualStationPlan) { + const preflight = preflightDualStationManagedVllm(dualStationPlan); + if (!preflight.ok) { + console.error(` vLLM install failed: ${preflight.reason}`); + return { ok: false }; + } + } else { + const replacement = vllmContainerReplacementTarget(runtimeProfile.containerName); + if (!replacement.ok) { + console.error(` vLLM install failed: ${replacement.reason}`); + return { ok: false }; + } } // Guard the host filesystem before an image pull or model-download // container can start. The cache path itself is created only after both // storage decisions pass, so Docker never creates it as root. - const hasImage = imageIsCached(runtimeProfile); - if (!(await managedStorageAccepted(runtimeProfile, model, hasImage, opts))) { + const hasImage = imageIsCached(runtimeProfile, localDockerEnv); + if (!(await managedStorageAccepted(runtimeProfile, model, hasImage, opts, localDockerEnv))) { return { ok: false }; } @@ -1409,12 +1697,33 @@ async function runVllmInstall( return { ok: false }; } - const pull = await pullImage(runtimeProfile); + const pull = await pullImage(runtimeProfile, localDockerEnv); if (!pull.ok) { console.error(` vLLM install failed: ${String(pull.reason)}`); return { ok: false }; } + if (dualStationPlan) { + let peerDockerEnv: Record; + try { + peerDockerEnv = buildRemoteVllmDockerEnv(dualStationPlan.peerSshBinding); + } catch (err) { + console.error(` vLLM install failed: ${(err as Error).message}`); + return { ok: false }; + } + emit(`Pulling the pinned vLLM image on peer ${dualStationPlan.peer.hostname}`); + const peerPull = await pullImage(runtimeProfile, peerDockerEnv); + if (!peerPull.ok) { + console.error(` vLLM install failed on peer: ${String(peerPull.reason)}`); + return { ok: false }; + } + const gpuPreflight = await preflightDualStationGpuRuntime(dualStationPlan); + if (!gpuPreflight.ok) { + console.error(` vLLM install failed: ${gpuPreflight.reason}`); + return { ok: false }; + } + } + // A cold image pull can consume the same host filesystem that backs the // Hugging Face cache. Re-probe the model destination after the pull before // `hf download` starts. @@ -1422,12 +1731,170 @@ async function runVllmInstall( return { ok: false }; } - const modelDownload = await downloadModel(runtimeProfile, model); + const modelDownload = await downloadModel(runtimeProfile, model, localDockerEnv); if (!modelDownload.ok) { console.error(` vLLM install failed: ${String(modelDownload.reason)}`); return { ok: false }; } + if (dualStationPlan) { + const stagingPlan = dualStationPlan; + try { + const verification = await withDualStationManagedVllmLifecycle(async () => { + emit( + peerModelSnapshot === "staging-required" + ? `Staging the pinned model snapshot on peer ${stagingPlan.peer.hostname}` + : `Verifying the pinned model snapshot on peer ${stagingPlan.peer.hostname}`, + ); + const staging = await stageDualStationModelSnapshot(stagingPlan); + if (!staging.ok) return { ok: false as const, reason: staging.reason }; + + const refreshedCapability = probeDualStationVllmCapability(); + if (refreshedCapability.kind !== "ready") { + const reason = + refreshedCapability.kind === "unavailable" + ? refreshedCapability.reason + : "the explicit peer configuration disappeared"; + return { + ok: false as const, + reason: `dual-Station capability changed: ${reason}`, + }; + } + if (!isDeepStrictEqual(refreshedCapability.plan, stagingPlan)) { + return { + ok: false as const, + reason: + "dual-Station topology changed during download; rerun setup against a stable pair.", + }; + } + if (refreshedCapability.peerModelSnapshot !== "ready") { + return { + ok: false as const, + reason: "peer pinned model snapshot was not verified after staging.", + }; + } + return { ok: true as const, plan: refreshedCapability.plan }; + }); + if (!verification.ok) { + console.error(` vLLM install failed: ${verification.reason}`); + return { ok: false }; + } + dualStationPlan = verification.plan; + } catch (error) { + console.error( + ` vLLM install failed: dual-Station lifecycle lock failed during model verification: ${(error as Error).message}`, + ); + return { ok: false }; + } + } + + let dualStationApiKey: string | null = null; + if (dualStationPlan) { + try { + const existingManagedBaseUrl = getDualStationManagedVllmBaseUrl(); + const existingApiKey = existingManagedBaseUrl ? loadDualStationVllmApiKey() : null; + // If the key file alone was lost, create a new host-global key. The + // lifecycle fingerprint then forces a coordinated pair replacement + // under its lock instead of reusing containers bound to an unknown key. + dualStationApiKey = existingApiKey ?? ensureDualStationVllmApiKey(); + } catch (err) { + console.error(` vLLM install failed: ${(err as Error).message}`); + return { ok: false }; + } + } + + if (dualStationPlan) { + if (!dualStationApiKey) { + console.error(" vLLM install failed: dual-Station API key was not provisioned"); + return { ok: false }; + } + try { + return await withDualStationManagedVllmLifecycle(async () => { + const start = await startDualStationManagedVllm(dualStationPlan, { + apiKey: dualStationApiKey, + }); + if (!start.ok) { + console.error(` vLLM install failed: ${start.reason}`); + for (const rollbackError of start.rollbackErrors) { + console.error(` vLLM rollback warning: ${rollbackError}`); + } + return { ok: false }; + } + + const rollbackStartedPair = async (): Promise => { + if (start.reusedExisting) return; + if (start.legacyMigration) { + const rollback = await rollbackDualStationLegacyMigration( + dualStationPlan, + start.legacyMigration, + ); + if (!rollback.ok) { + for (const rollbackError of rollback.rollbackErrors) { + console.error(` vLLM rollback warning: ${rollbackError}`); + } + } + return; + } + const cleanup = await cleanupDualStationManagedVllm(dualStationPlan); + if (!cleanup.ok) console.error(` vLLM rollback warning: ${cleanup.reason}`); + }; + + emit("Launching vLLM"); + emit( + `Launch can take 5 minutes to ${String(Math.ceil(runtimeProfile.loadTimeoutSec / 60))} minutes`, + ); + + const ready = await waitForVllmReady(runtimeProfile, start.baseUrl, localDockerEnv); + if (!ready.ok) { + printContainerLogTail(runtimeProfile, localDockerEnv); + await rollbackStartedPair(); + console.error(` vLLM install failed: ${String(ready.reason)}`); + return { ok: false }; + } + + const authBoundary = verifyDualStationVllmAuthBoundary( + start.baseUrl, + dualStationApiKey, + servedModelId, + ); + if (!authBoundary.ok) { + await rollbackStartedPair(); + console.error(` vLLM install failed: ${authBoundary.reason}`); + return { ok: false }; + } + + if (!areDualStationManagedVllmContainersRunning(dualStationPlan)) { + await rollbackStartedPair(); + console.error(" vLLM distributed containers exited unexpectedly after readiness"); + return { ok: false }; + } + + if (start.legacyMigration) { + const commit = await commitDualStationLegacyMigration( + dualStationPlan, + start.legacyMigration, + ); + if (!commit.ok) { + await rollbackStartedPair(); + console.error(` vLLM install failed: ${commit.reason}`); + return { ok: false }; + } + for (const warning of commit.cleanupWarnings) { + console.error(` vLLM cleanup warning: ${warning}`); + } + } + + console.log(` ✓ vLLM ready across two DGX Stations at ${start.baseUrl}`); + return { ok: true }; + }); + } catch (error) { + console.error( + ` vLLM install failed: dual-Station lifecycle lock failed: ${(error as Error).message}`, + ); + return { ok: false }; + } + } + const start = startContainer(runtimeProfile, model); if (!start.ok) { console.error(` vLLM install failed: ${String(start.reason)}`); @@ -1439,9 +1906,9 @@ async function runVllmInstall( `Launch can take 5 minutes to ${String(Math.ceil(runtimeProfile.loadTimeoutSec / 60))} minutes`, ); - const ready = await waitForVllmReady(runtimeProfile); + const ready = await waitForVllmReady(runtimeProfile, undefined, localDockerEnv); if (!ready.ok) { - printContainerLogTail(runtimeProfile); + printContainerLogTail(runtimeProfile, localDockerEnv); dockerStop(runtimeProfile.containerName, { env: buildVllmDockerEnv(), ignoreError: true, @@ -1451,7 +1918,7 @@ async function runVllmInstall( return { ok: false }; } - if (!containerStillRunning(runtimeProfile)) { + if (!containerStillRunning(runtimeProfile, localDockerEnv)) { console.error(" vLLM container exited unexpectedly after readiness"); return { ok: false }; } diff --git a/src/lib/onboard.ts b/src/lib/onboard.ts index 955f67c74a..a237d36ba8 100644 --- a/src/lib/onboard.ts +++ b/src/lib/onboard.ts @@ -1087,7 +1087,7 @@ const { // biome-ignore format: keep src/lib/onboard.ts net-neutral for growth guardrail. const handleVllmSelection = createSetupNimVllmHandler({ VLLM_PORT, runCapture, getLocalProviderBaseUrl, getLocalProviderValidationBaseUrl, - isSafeModelId, requireValue, validateOpenAiLikeSelection, + getManagedVllmProviderBinding: localInference.getManagedDualStationVllmProviderBinding, queryVllmModels: (baseUrl, apiKey) => { const result = localInference.probeVllmModels(baseUrl, apiKey); return result.ok ? result.body : ""; }, isSafeModelId, requireValue, validateOpenAiLikeSelection, applyVllmRuntimeContextWindow: localInference.applyVllmRuntimeContextWindow, isDgxSparkHost: () => nim.detectNvidiaPlatform() === "spark", isNemoClawManagedVllmRunning, exitProcess: (code) => process.exit(code), }); diff --git a/src/lib/onboard/inference-providers/types.ts b/src/lib/onboard/inference-providers/types.ts index c3c312f5df..5ec725290f 100644 --- a/src/lib/onboard/inference-providers/types.ts +++ b/src/lib/onboard/inference-providers/types.ts @@ -226,6 +226,7 @@ export type VllmDeps = CommonDeps & { applyLocalInferenceRoute: (provider: string, model: string) => Promise; run: RunFn; VLLM_LOCAL_CREDENTIAL_ENV: string; + getManagedVllmProviderBinding: () => { baseUrl: string; apiKey: string } | null; }; export type OllamaDeps = CommonDeps & { diff --git a/src/lib/onboard/inference-providers/vllm-local.test.ts b/src/lib/onboard/inference-providers/vllm-local.test.ts new file mode 100644 index 0000000000..ab87a6d7b3 --- /dev/null +++ b/src/lib/onboard/inference-providers/vllm-local.test.ts @@ -0,0 +1,104 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { describe, expect, it, vi } from "vitest"; +import type { VllmDeps } from "./types"; +import { setupVllmLocalInference } from "./vllm-local"; + +const CREDENTIAL_ENV = "NEMOCLAW_VLLM_LOCAL_TOKEN"; + +function deps(overrides: Partial = {}): VllmDeps { + return { + runOpenshell: vi.fn(() => ({ status: 0 })), + upsertProvider: vi.fn(() => ({ ok: true })), + verifyInferenceRoute: vi.fn(), + verifyOnboardInferenceSmoke: vi.fn(), + isNonInteractive: () => true, + registry: { updateSandbox: vi.fn() as VllmDeps["registry"]["updateSandbox"] }, + exitProcess: (code) => { + throw new Error(`exit ${code}`); + }, + error: vi.fn(), + log: vi.fn(), + validateLocalProvider: () => ({ ok: true }), + getLocalProviderHealthCheck: () => ["curl", "-sf", "http://127.0.0.1:8000/v1/models"], + getLocalProviderBaseUrl: () => "http://host.openshell.internal:8000/v1", + applyLocalInferenceRoute: async () => false, + run: vi.fn(() => ({ status: 0 })), + VLLM_LOCAL_CREDENTIAL_ENV: CREDENTIAL_ENV, + getManagedVllmProviderBinding: () => null, + ...overrides, + }; +} + +describe("vLLM local provider credential", () => { + it("preserves the literal dummy credential for legacy single-host vLLM", async () => { + const upsertProvider = vi.fn(() => ({ ok: true })); + + await expect( + setupVllmLocalInference( + { model: "served/model", provider: "vllm-local" }, + deps({ upsertProvider }), + ), + ).resolves.toEqual({ done: false }); + + expect(upsertProvider).toHaveBeenCalledWith( + "vllm-local", + "openai", + CREDENTIAL_ENV, + "http://host.openshell.internal:8000/v1", + { [CREDENTIAL_ENV]: "dummy" }, + ); + }); + + it("registers the persisted managed key through provider env, never as an argv field", async () => { + const apiKey = "c".repeat(64); + const upsertProvider = vi.fn(() => ({ ok: true })); + + await expect( + setupVllmLocalInference( + { model: "served/model", provider: "vllm-local" }, + deps({ + upsertProvider, + getManagedVllmProviderBinding: () => ({ + baseUrl: "http://10.40.0.1:8000/v1", + apiKey, + }), + }), + ), + ).resolves.toEqual({ done: false }); + + expect(upsertProvider).toHaveBeenCalledWith( + "vllm-local", + "openai", + CREDENTIAL_ENV, + "http://10.40.0.1:8000/v1", + { [CREDENTIAL_ENV]: apiKey }, + ); + }); + + it("fails closed without rendering credential-loader details", async () => { + const leaked = "d".repeat(64); + const error = vi.fn(); + const upsertProvider = vi.fn(() => ({ ok: true })); + + await expect( + setupVllmLocalInference( + { model: "served/model", provider: "vllm-local" }, + deps({ + error, + upsertProvider, + getManagedVllmProviderBinding: () => { + throw new Error(`unsafe ${leaked}`); + }, + }), + ), + ).rejects.toThrow("exit 1"); + + expect(upsertProvider).not.toHaveBeenCalled(); + expect(error).toHaveBeenCalledWith( + " Managed vLLM authentication state is unsafe or unreadable.", + ); + expect(error.mock.calls.flat().join("\n")).not.toContain(leaked); + }); +}); diff --git a/src/lib/onboard/inference-providers/vllm-local.ts b/src/lib/onboard/inference-providers/vllm-local.ts index 5dfc4581be..cb1c7d1fcc 100644 --- a/src/lib/onboard/inference-providers/vllm-local.ts +++ b/src/lib/onboard/inference-providers/vllm-local.ts @@ -19,6 +19,7 @@ export async function setupVllmLocalInference( applyLocalInferenceRoute, run, VLLM_LOCAL_CREDENTIAL_ENV, + getManagedVllmProviderBinding, exitProcess, error, } = deps; @@ -49,17 +50,26 @@ export async function setupVllmLocalInference( return exitProcess(1); } } - const baseUrl = getLocalProviderBaseUrl(provider); + let managedBinding: ReturnType; + try { + managedBinding = getManagedVllmProviderBinding(); + } catch { + error(" Managed vLLM authentication state is unsafe or unreadable."); + return exitProcess(1); + } + const baseUrl = managedBinding?.baseUrl ?? getLocalProviderBaseUrl(provider); + const providerToken = managedBinding?.apiKey ?? "dummy"; // Use a dedicated internal credential env so the gateway does not pick // up the user's host OPENAI_API_KEY for local vLLM. vLLM does not enforce - // the bearer at runtime, but a dedicated env name prevents accidental - // hijacking. See GH #2519. + // the bearer for legacy single-host installs; managed dual-Station vLLM + // uses the private persisted key. The dedicated env name prevents + // accidental hijacking by a host OPENAI_API_KEY. See GH #2519. const providerResult = upsertProvider( "vllm-local", "openai", VLLM_LOCAL_CREDENTIAL_ENV, baseUrl, - { [VLLM_LOCAL_CREDENTIAL_ENV]: "dummy" }, + { [VLLM_LOCAL_CREDENTIAL_ENV]: providerToken }, ); if (!providerResult.ok) { error(` ${providerResult.message}`); diff --git a/src/lib/onboard/inference-selection-validation.test.ts b/src/lib/onboard/inference-selection-validation.test.ts index 12eaf250ae..7284d497a7 100644 --- a/src/lib/onboard/inference-selection-validation.test.ts +++ b/src/lib/onboard/inference-selection-validation.test.ts @@ -10,6 +10,45 @@ import { OnboardInferenceCapabilityCache } from "./inference-capability-cache"; import { createInferenceSelectionValidationHelpers } from "./inference-selection-validation"; describe("inference selection validation", () => { + it("uses an explicit managed key without forwarding it as a probe option", async () => { + const apiKey = "f".repeat(64); + const getCredential = vi.fn(() => "ambient-key"); + const probeOpenAiLikeEndpoint = vi.fn(() => ({ + ok: true, + api: "openai-completions", + label: "Chat Completions API", + })); + const helpers = createInferenceSelectionValidationHelpers({ + isNonInteractive: () => false, + agentProductName: () => "OpenClaw", + getCredential, + probeOpenAiLikeEndpoint, + promptValidationRecovery: vi.fn(async () => "selection" as const), + }); + const log = vi.spyOn(console, "log").mockImplementation(() => {}); + + await expect( + helpers.validateOpenAiLikeSelection( + "Local vLLM", + "http://10.40.0.1:8000/v1", + "served/model", + null, + undefined, + undefined, + { apiKey, pinnedAddresses: [] }, + ), + ).resolves.toEqual({ ok: true, api: "openai-completions" }); + expect(getCredential).not.toHaveBeenCalled(); + expect(probeOpenAiLikeEndpoint).toHaveBeenCalledWith( + "http://10.40.0.1:8000/v1", + "served/model", + apiKey, + { pinnedAddresses: [], calibrateTimeouts: true }, + ); + expect(log.mock.calls.flat().join("\n")).not.toContain(apiKey); + log.mockRestore(); + }); + it("records a completed Chat Completions selection for the matching smoke check", async () => { const capabilityCache = new OnboardInferenceCapabilityCache(); const helpers = createInferenceSelectionValidationHelpers({ diff --git a/src/lib/onboard/inference-selection-validation.ts b/src/lib/onboard/inference-selection-validation.ts index a6ecab0788..a8f447a43c 100644 --- a/src/lib/onboard/inference-selection-validation.ts +++ b/src/lib/onboard/inference-selection-validation.ts @@ -82,6 +82,11 @@ export interface InferenceSelectionValidationHelpers { retryMessage?: string, helpUrl?: string | null, options?: { + /** In-memory credential for managed local endpoints; never read from ambient env. */ + apiKey?: string | null; + /** Approved no-DNS endpoint pin; [] also disables ambient proxies for managed IP URLs. */ + pinnedAddresses?: readonly string[]; + trustedPrivateCapability?: TrustedPrivateEndpointCapability; authMode?: "bearer" | "query-param"; extraHeaders?: readonly string[]; requireResponsesToolCalling?: boolean; @@ -245,6 +250,9 @@ export function createInferenceSelectionValidationHelpers( retryMessage = "Please choose a provider/model again.", helpUrl: string | null = null, options: { + apiKey?: string | null; + pinnedAddresses?: readonly string[]; + trustedPrivateCapability?: TrustedPrivateEndpointCapability; authMode?: "bearer" | "query-param"; extraHeaders?: readonly string[]; requireResponsesToolCalling?: boolean; @@ -255,13 +263,19 @@ export function createInferenceSelectionValidationHelpers( capabilityCache?: OnboardInferenceCapabilityCache; } = {}, ): Promise { - const apiKey = credentialEnv ? resolveCredential(credentialEnv) : ""; + const { apiKey: explicitApiKey, ...probeOptions } = options; + const apiKey = + explicitApiKey !== undefined + ? explicitApiKey + : credentialEnv + ? resolveCredential(credentialEnv) + : ""; const probe = await runOpenAiLikeProbe(endpointUrl, model, apiKey, { - ...options, + ...probeOptions, calibrateTimeouts: true, }); if (!probe.ok) { - options.capabilityCache?.invalidate(); + probeOptions.capabilityCache?.invalidate(); printValidationFailure(label, probe); if (deps.isNonInteractive()) { exitNonInteractiveValidationFailure(); @@ -285,12 +299,12 @@ export function createInferenceSelectionValidationHelpers( } const api = probe.api ?? "openai-completions"; if (api === "openai-completions" && probe.validated !== false) { - options.capabilityCache?.rememberCompletedOpenAiChat({ + probeOptions.capabilityCache?.rememberCompletedOpenAiChat({ endpointUrl, model, - authMode: options.authMode, - requireChatCompletionsToolCalling: options.requireChatCompletionsToolCalling, - extraHeaders: options.extraHeaders, + authMode: probeOptions.authMode, + requireChatCompletionsToolCalling: probeOptions.requireChatCompletionsToolCalling, + extraHeaders: probeOptions.extraHeaders, }); } return { ok: true, api }; diff --git a/src/lib/onboard/provider-host-state.test.ts b/src/lib/onboard/provider-host-state.test.ts index 794179a35f..82c995869a 100644 --- a/src/lib/onboard/provider-host-state.test.ts +++ b/src/lib/onboard/provider-host-state.test.ts @@ -148,6 +148,18 @@ describe("detectInferenceProviderHostState", () => { ); }); + it("does not treat curl connection status 000 as a running vLLM", () => { + const state = detectWithDeps( + buildDeps({ + runCapture: vi.fn((command) => + command.join(" ").includes("127.0.0.1:8000/v1/models") ? "000" : "", + ), + }), + ); + + expect(state.vllmRunning).toBe(false); + }); + it("detects a reachable Windows-host Ollama beside WSL-local Ollama and warns outside mirrored networking", () => { const logs: string[] = []; const deps = buildDeps({ diff --git a/src/lib/onboard/provider-host-state.ts b/src/lib/onboard/provider-host-state.ts index fcbd75f02c..4a41d5a8f8 100644 --- a/src/lib/onboard/provider-host-state.ts +++ b/src/lib/onboard/provider-host-state.ts @@ -2,8 +2,13 @@ // SPDX-License-Identifier: Apache-2.0 import { dockerCapture as defaultDockerCapture } from "../adapters/docker"; -import { OLLAMA_PORT, VLLM_PORT } from "../core/ports"; -import { findReachableOllamaHost, OLLAMA_HOST_DOCKER_INTERNAL } from "../inference/local"; +import { OLLAMA_PORT } from "../core/ports"; +import { + findReachableOllamaHost, + getLocalProviderAvailabilityEndpoint, + isLocalProviderProbeOutputHealthy, + OLLAMA_HOST_DOCKER_INTERNAL, +} from "../inference/local"; import type { NvidiaPlatform } from "../inference/nim"; import { detectVllmProfile, type VllmProfile } from "../inference/vllm"; import { buildVllmDockerEnv } from "../inference/vllm-docker-env"; @@ -112,10 +117,18 @@ function buildDeps( } function probeVllmRunning(runCapture: RunCapture): boolean { - return !!runCapture( - ["curl", "-sf", ...LOCAL_PROVIDER_PROBE_CURL_ARGS, `http://127.0.0.1:${VLLM_PORT}/v1/models`], - { ignoreError: true }, + const endpoint = getLocalProviderAvailabilityEndpoint("vllm-local"); + if (!endpoint) return false; + const writeOut = endpoint.endsWith("/health") + ? ["--noproxy", "*", "--write-out", "%{http_code}"] + : []; + const output = runCapture( + ["curl", "-sf", ...LOCAL_PROVIDER_PROBE_CURL_ARGS, ...writeOut, endpoint], + { + ignoreError: true, + }, ); + return isLocalProviderProbeOutputHealthy(endpoint, output); } function probeWindowsOllamaReachable(input: { diff --git a/src/lib/onboard/setup-inference.ts b/src/lib/onboard/setup-inference.ts index 3219ab04a7..ca0fffb653 100644 --- a/src/lib/onboard/setup-inference.ts +++ b/src/lib/onboard/setup-inference.ts @@ -15,6 +15,7 @@ import { isAdvisoryGatewayRouteConflict, } from "../inference/gateway-route-compatibility"; import { withGatewayRouteMutationLock } from "../inference/gateway-route-mutation-lock"; +import { getManagedDualStationVllmProviderBinding } from "../inference/local"; import { assertNoExplicitOpenShellGatewayEndpoint, assertNoOpenShellGatewayEndpointOverride, @@ -117,6 +118,7 @@ export type SetupInferenceDeps = ProviderBranchDeps & { updateSandbox: typeof import("../state/registry").reserveSandboxInferenceRoute; localInferenceTimeoutSecs: number; vllmLocalCredentialEnv: string; + getManagedVllmProviderBinding?: () => { baseUrl: string; apiKey: string } | null; ollamaProxyCredentialEnv: string; isRoutedInferenceProvider: (provider: string) => boolean; applyLocalInferenceRoute?: VllmDeps["applyLocalInferenceRoute"]; @@ -439,6 +441,8 @@ export function createSetupInference( ), run: deps.run, VLLM_LOCAL_CREDENTIAL_ENV: deps.vllmLocalCredentialEnv, + getManagedVllmProviderBinding: + deps.getManagedVllmProviderBinding ?? getManagedDualStationVllmProviderBinding, }, ); if (outcome.done) return outcome.result; diff --git a/src/lib/onboard/setup-nim-flow.test.ts b/src/lib/onboard/setup-nim-flow.test.ts index c8ef099962..79b02cb0db 100644 --- a/src/lib/onboard/setup-nim-flow.test.ts +++ b/src/lib/onboard/setup-nim-flow.test.ts @@ -1068,7 +1068,7 @@ describe("createSetupNim", () => { }), ); - await expect(setupNim(null)).rejects.toThrow("vLLM is already running on localhost:8000"); + await expect(setupNim(null)).rejects.toThrow("vLLM is already running on this host"); expect(error).toHaveBeenCalledWith(expect.stringContaining("Select Local vLLM")); expect(error).toHaveBeenCalledWith(expect.stringContaining("stop the existing server")); diff --git a/src/lib/onboard/setup-nim-flow.ts b/src/lib/onboard/setup-nim-flow.ts index 179d8249c6..55db91ce47 100644 --- a/src/lib/onboard/setup-nim-flow.ts +++ b/src/lib/onboard/setup-nim-flow.ts @@ -568,7 +568,7 @@ export function createSetupNim( } if (vllmRunning) { const message = - `vLLM is already running on localhost:${String(deps.vllmPort)}. ` + + "vLLM is already running on this host. " + "Select Local vLLM, or stop the existing server before selecting the managed install path."; deps.error(` ${message}`); if (deps.isNonInteractive()) { diff --git a/src/lib/onboard/setup-nim-vllm.test.ts b/src/lib/onboard/setup-nim-vllm.test.ts index 4bf4fc7b24..bc7682d0ea 100644 --- a/src/lib/onboard/setup-nim-vllm.test.ts +++ b/src/lib/onboard/setup-nim-vllm.test.ts @@ -30,6 +30,10 @@ function deps(overrides: Partial = {}): SetupNimVllmDeps { runCapture: () => JSON.stringify({ data: [{ id: "served/model" }] }), getLocalProviderBaseUrl: () => "http://host.openshell.internal:8000/v1", getLocalProviderValidationBaseUrl: () => "http://127.0.0.1:8000/v1", + getManagedVllmProviderBinding: () => null, + queryVllmModels: () => { + throw new Error("unexpected authenticated vLLM query"); + }, isSafeModelId: () => true, requireValue, validateOpenAiLikeSelection: async () => ({ ok: true, api: "openai-completions" }), @@ -76,6 +80,140 @@ describe("setupNim vLLM route containment", () => { expect(events).toEqual(["preflight", "probe", "exact", "validate"]); }); + it("authenticates managed model discovery and OpenAI validation without exposing the key", async () => { + const apiKey = "a".repeat(64); + const runCapture = vi.fn(() => ""); + const queryVllmModels = vi.fn(() => JSON.stringify({ data: [{ id: "served/model" }] })); + const validateOpenAiLikeSelection = vi.fn(async () => ({ + ok: true, + api: "openai-completions", + })); + const handler = createSetupNimVllmHandler( + deps({ + runCapture, + getLocalProviderBaseUrl: () => "http://10.40.0.1:8000/v1", + getLocalProviderValidationBaseUrl: () => "http://10.40.0.1:8000/v1", + getManagedVllmProviderBinding: () => ({ + baseUrl: "http://10.40.0.1:8000/v1", + apiKey, + }), + queryVllmModels, + validateOpenAiLikeSelection, + }), + ); + + await expect(handler(state(null))).resolves.toBe("selected"); + expect(runCapture).not.toHaveBeenCalled(); + expect(queryVllmModels).toHaveBeenCalledWith("http://10.40.0.1:8000/v1", apiKey); + expect(validateOpenAiLikeSelection).toHaveBeenCalledWith( + "Local vLLM", + "http://10.40.0.1:8000/v1", + "served/model", + null, + undefined, + undefined, + expect.objectContaining({ + apiKey, + pinnedAddresses: [], + trustedPrivateCapability: expect.objectContaining({ addresses: ["10.40.0.1"] }), + }), + ); + const renderedOutput = [ + ...vi.mocked(console.log).mock.calls, + ...vi.mocked(console.error).mock.calls, + ...vi.mocked(console.warn).mock.calls, + ] + .flat() + .join("\n"); + expect(renderedOutput).not.toContain(apiKey); + expect(renderedOutput).toContain("Using managed dual-Station vLLM endpoint"); + expect(renderedOutput).not.toContain("localhost:8000"); + }); + + it("rejects a root-matched alias with topology-neutral recovery for a managed dual endpoint", async () => { + const selection = state("required/model"); + const validateOpenAiLikeSelection = vi.fn(async () => ({ ok: true })); + const handler = createSetupNimVllmHandler( + deps({ + getLocalProviderBaseUrl: () => "http://10.40.0.1:8000/v1", + getLocalProviderValidationBaseUrl: () => "http://10.40.0.1:8000/v1", + getManagedVllmProviderBinding: () => ({ + baseUrl: "http://10.40.0.1:8000/v1", + apiKey: "a".repeat(64), + }), + queryVllmModels: () => + JSON.stringify({ data: [{ id: "served/model", root: "required/model" }] }), + validateOpenAiLikeSelection, + }), + ); + + await expect(handler(selection)).rejects.toThrow("exit 1"); + expect(validateOpenAiLikeSelection).not.toHaveBeenCalled(); + expect(console.error).toHaveBeenCalledWith( + " To install 'required/model', stop the managed dual-Station vLLM deployment, then rerun the original install/onboard command.", + ); + expect(vi.mocked(console.error).mock.calls.flat().join("\n")).not.toContain("localhost"); + }); + + it("fails closed before probing a managed dual endpoint whose key is missing", async () => { + const runCapture = vi.fn(() => ""); + const queryVllmModels = vi.fn(() => ""); + const handler = createSetupNimVllmHandler( + deps({ + runCapture, + getManagedVllmProviderBinding: () => { + throw new Error("Managed dual-Station vLLM authentication is missing."); + }, + queryVllmModels, + }), + ); + + await expect(handler(state(null))).rejects.toThrow("exit 1"); + expect(runCapture).not.toHaveBeenCalled(); + expect(queryVllmModels).not.toHaveBeenCalled(); + expect(console.error).toHaveBeenCalledWith( + " Managed vLLM authentication state is unsafe or unreadable.", + ); + }); + + it("treats an unexpected undefined managed binding as absent", async () => { + const runCapture = vi.fn(() => JSON.stringify({ data: [{ id: "served/model" }] })); + const queryVllmModels = vi.fn(() => ""); + const handler = createSetupNimVllmHandler( + deps({ + runCapture, + getManagedVllmProviderBinding: () => undefined as never, + queryVllmModels, + }), + ); + + await expect(handler(state(null))).resolves.toBe("selected"); + expect(runCapture).toHaveBeenCalled(); + expect(queryVllmModels).not.toHaveBeenCalled(); + expect(console.log).toHaveBeenCalledWith(" ✓ Using existing vLLM on localhost:8000"); + }); + + it("fails closed before endpoint probes when managed auth state is unsafe", async () => { + const queryVllmModels = vi.fn(() => ""); + const validateOpenAiLikeSelection = vi.fn(async () => ({ ok: true })); + const handler = createSetupNimVllmHandler( + deps({ + getManagedVllmProviderBinding: () => { + throw new Error(`unsafe ${"b".repeat(64)}`); + }, + queryVllmModels, + validateOpenAiLikeSelection, + }), + ); + + await expect(handler(state(null))).rejects.toThrow("exit 1"); + expect(queryVllmModels).not.toHaveBeenCalled(); + expect(validateOpenAiLikeSelection).not.toHaveBeenCalled(); + expect(console.error).toHaveBeenCalledWith( + " Managed vLLM authentication state is unsafe or unreadable.", + ); + }); + it("adopts a served alias when its reported root matches the requested model (#7023)", async () => { const validate = vi.fn(async () => ({ ok: true, api: "openai-completions" })); const selection = state("nvidia/nemotron-3-ultra-550b-a55b"); diff --git a/src/lib/onboard/setup-nim-vllm.ts b/src/lib/onboard/setup-nim-vllm.ts index 216413e6d2..a996156c76 100644 --- a/src/lib/onboard/setup-nim-vllm.ts +++ b/src/lib/onboard/setup-nim-vllm.ts @@ -1,6 +1,10 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +import { + assertEndpointResolvesPublic, + type TrustedPrivateEndpointCapability, +} from "../inference/endpoint-ssrf-preflight"; import { VLLM_MODELS } from "../inference/vllm-models"; import { cliName } from "./branding"; import type { SetupNimSelectionResult, SetupNimSelectionState } from "./setup-nim-flow"; @@ -26,6 +30,8 @@ export interface SetupNimVllmDeps { runCapture(args: string[], options: { ignoreError: boolean }): string; getLocalProviderBaseUrl(provider: string): string | null; getLocalProviderValidationBaseUrl(provider: string): string | null; + getManagedVllmProviderBinding(): { baseUrl: string; apiKey: string } | null; + queryVllmModels(baseUrl: string, apiKey: string): string; isSafeModelId(model: string): boolean; requireValue(value: T | null | undefined, message: string): T; validateOpenAiLikeSelection( @@ -33,6 +39,13 @@ export interface SetupNimVllmDeps { endpointUrl: string, model: string, credentialEnv: string | null, + retryMessage?: string, + helpUrl?: string | null, + options?: { + apiKey?: string | null; + pinnedAddresses?: readonly string[]; + trustedPrivateCapability?: TrustedPrivateEndpointCapability; + }, ): Promise<{ ok: boolean; retry?: string; api?: string | null }>; applyVllmRuntimeContextWindow(models: VllmModels, model: string): void; isDgxSparkHost?: () => boolean; @@ -49,6 +62,21 @@ const NO_QUANTIZATION_VALUES = new Set(["", "false", "none", "null", "unquantize type ModelSizeClass = "large" | "small" | "unknown"; +async function managedVllmValidationOptions(baseUrl: string, apiKey: string) { + const hostname = new URL(baseUrl).hostname.replace(/^\[|\]$/g, ""); + const preflight = await assertEndpointResolvesPublic(baseUrl, undefined, { + trustedPrivateHosts: [hostname], + }); + if (!preflight.ok || !preflight.trustedPrivateCapability) { + throw new Error("Managed dual-Station vLLM endpoint authorization failed."); + } + return { + apiKey, + pinnedAddresses: preflight.addresses ?? [], + trustedPrivateCapability: preflight.trustedPrivateCapability, + }; +} + /** Parse positive integer metadata reported by vLLM model endpoints. */ function parsePositiveInteger(value: unknown): number | null { const normalized = typeof value === "number" ? value : Number(String(value ?? "").trim()); @@ -193,10 +221,16 @@ export function createSetupNimVllmHandler( state: SetupNimSelectionState, options: SetupNimVllmSelectionOptions = {}, ): Promise { - console.log(` ✓ Using existing vLLM on localhost:${deps.VLLM_PORT}`); state.provider = "vllm-local"; state.credentialEnv = null; - state.endpointUrl = deps.getLocalProviderBaseUrl(state.provider); + let managedBinding: ReturnType; + try { + managedBinding = deps.getManagedVllmProviderBinding(); + } catch { + console.error(" Managed vLLM authentication state is unsafe or unreadable."); + deps.exitProcess(1); + } + state.endpointUrl = managedBinding?.baseUrl ?? deps.getLocalProviderBaseUrl(state.provider); if (!state.endpointUrl) { console.error(" Local vLLM base URL could not be determined."); deps.exitProcess(1); @@ -205,15 +239,33 @@ export function createSetupNimVllmHandler( state.assertRouteCompatible?.(); const requiredModel = typeof state.model === "string" ? state.model : null; - const raw = deps.runCapture(["curl", "-sf", `http://127.0.0.1:${deps.VLLM_PORT}/v1/models`], { - ignoreError: true, - }); + const validationBaseUrl = + managedBinding?.baseUrl ?? deps.getLocalProviderValidationBaseUrl(state.provider); + if (!validationBaseUrl) { + console.error(" Local vLLM validation URL could not be determined."); + deps.exitProcess(1); + } + + const apiKey = managedBinding?.apiKey ?? null; + const managedDualEndpoint = managedBinding != null; + console.log( + managedDualEndpoint + ? " ✓ Using managed dual-Station vLLM endpoint" + : ` ✓ Using existing vLLM on localhost:${deps.VLLM_PORT}`, + ); + const raw = apiKey + ? deps.queryVllmModels(validationBaseUrl, apiKey) + : deps.runCapture(["curl", "-sf", `${validationBaseUrl}/models`], { + ignoreError: true, + }); let models: VllmModels; try { models = JSON.parse(raw); } catch { console.error( - ` Could not query vLLM models endpoint. Is vLLM running on localhost:${deps.VLLM_PORT}?`, + managedDualEndpoint + ? " Could not query the managed dual-Station vLLM models endpoint. Is the deployment running and reachable?" + : ` Could not query vLLM models endpoint. Is vLLM running on localhost:${deps.VLLM_PORT}?`, ); deps.exitProcess(1); } @@ -233,13 +285,16 @@ export function createSetupNimVllmHandler( requiredModel && detectedModel !== requiredModel && (options.managedInstall === true || + managedDualEndpoint || !reportedModelMatchesRequest(models, detectedModel, requiredModel)) ) { console.error( ` Detected vLLM model '${detectedModel}' does not match the shared gateway route '${requiredModel}'.`, ); console.error( - ` To install '${requiredModel}', stop the existing vLLM server on localhost:${deps.VLLM_PORT}, then rerun the original install/onboard command.`, + managedDualEndpoint + ? ` To install '${requiredModel}', stop the managed dual-Station vLLM deployment, then rerun the original install/onboard command.` + : ` To install '${requiredModel}', stop the existing vLLM server on localhost:${deps.VLLM_PORT}, then rerun the original install/onboard command.`, ); console.error(` To keep '${detectedModel}' instead, start detailed setup:`); console.error(" unset NEMOCLAW_PROVIDER NEMOCLAW_MODEL NEMOCLAW_VLLM_MODEL"); @@ -264,17 +319,36 @@ export function createSetupNimVllmHandler( } } - const validationBaseUrl = deps.getLocalProviderValidationBaseUrl(state.provider); - if (!validationBaseUrl) { - console.error(" Local vLLM validation URL could not be determined."); - deps.exitProcess(1); + const validationModel = deps.requireValue(state.model, "Expected a detected vLLM model"); + let managedValidationOptions: Awaited> | null = + null; + if (apiKey) { + try { + managedValidationOptions = await managedVllmValidationOptions(validationBaseUrl, apiKey); + } catch { + console.error(" Managed vLLM endpoint authorization could not be verified."); + deps.exitProcess(1); + } } - const validation = await deps.validateOpenAiLikeSelection( - "Local vLLM", - validationBaseUrl, - deps.requireValue(state.model, "Expected a detected vLLM model"), - null, - ); + const validation = apiKey + ? await deps.validateOpenAiLikeSelection( + "Local vLLM", + validationBaseUrl, + validationModel, + null, + undefined, + undefined, + deps.requireValue( + managedValidationOptions, + "Expected managed vLLM validation authorization", + ), + ) + : await deps.validateOpenAiLikeSelection( + "Local vLLM", + validationBaseUrl, + validationModel, + null, + ); if (validation.retry === "selection" || validation.retry === "model" || !validation.ok) { return "retry-selection"; } diff --git a/src/lib/onboard/station-express-resume.test.ts b/src/lib/onboard/station-express-resume.test.ts index 81b0255067..5247a87064 100644 --- a/src/lib/onboard/station-express-resume.test.ts +++ b/src/lib/onboard/station-express-resume.test.ts @@ -49,6 +49,15 @@ function expressEnv(): NodeJS.ProcessEnv { }; } +function dualExpressEnv(): NodeJS.ProcessEnv { + return { + ...expressEnv(), + NEMOCLAW_MODEL: "nemotron-ultra", + NEMOCLAW_DGX_STATION_PEER: "192.168.240.2", + NEMOCLAW_DGX_STATION_SSH_BINDING: "sha256:qualified-pair-binding", + }; +} + function receiptText(generation = receiptGeneration, model = "nemotron-3-ultra-550b-a55b"): string { return `revision=${receiptRevision}\nmodel=${model}\ngeneration=${generation}\n`; } @@ -101,6 +110,28 @@ describe("DGX Station Express resume (#7048)", () => { }); }); + it("captures the qualified dual-Station served alias in the sealed intent", () => { + expect(getStationExpressResumeIntent(dualExpressEnv(), "my-assistant")).toEqual({ + ok: true, + intent: { + ...ultraIntent, + servedModel: "nemotron-ultra", + checkpointModel: "nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-NVFP4", + }, + }); + }); + + it("rejects the dual-Station served alias without both qualified pair signals", () => { + for (const missing of ["NEMOCLAW_DGX_STATION_PEER", "NEMOCLAW_DGX_STATION_SSH_BINDING"]) { + const env = dualExpressEnv(); + delete env[missing]; + expect(getStationExpressResumeIntent(env, "my-assistant")).toEqual({ + ok: false, + message: "DGX Station Express has a conflicting NEMOCLAW_MODEL value.", + }); + } + }); + it("carries the installer receipt generation in the persisted intent", () => { const env = expressEnv(); env[STATION_EXPRESS_RECEIPT_GENERATION_ENV] = receiptGeneration; diff --git a/src/lib/onboard/station-express-resume.ts b/src/lib/onboard/station-express-resume.ts index 2a888a8626..85a2198a28 100644 --- a/src/lib/onboard/station-express-resume.ts +++ b/src/lib/onboard/station-express-resume.ts @@ -99,6 +99,8 @@ const BOUND_RECEIPT_INTENT_KEYS = const SPARK_INTENT_KEYS = "kind,sandboxName,version"; const SPARK_INTENT_KEYS_WITH_MODEL = "kind,model,sandboxName,version"; const SPARK_EXPRESS_PROVIDER = "install-vllm"; +const STATION_ULTRA_ENV_VALUE = "nemotron-3-ultra-550b-a55b"; +const STATION_ULTRA_DUAL_SERVED_MODEL = "nemotron-ultra"; const STATION_EXPRESS_INSTALLER_RESUME_FILE = "station-express-resume"; const STATION_EXPRESS_RETIREMENT_CLAIM_PREFIX = `${STATION_EXPRESS_INSTALLER_RESUME_FILE}.retiring-`; const STATION_EXPRESS_RETIREMENT_CLAIM_RECEIPT = "receipt"; @@ -144,6 +146,21 @@ function servedModel(model: VllmModelDef): string { return model.servedModelId ?? model.id; } +function qualifiedDualStationServedModel( + env: NodeJS.ProcessEnv, + model: VllmModelDef, +): string | null { + if ( + model.envValue !== STATION_ULTRA_ENV_VALUE || + String(env.NEMOCLAW_DGX_STATION_PEER ?? "").trim().length === 0 || + String(env.NEMOCLAW_DGX_STATION_SSH_BINDING ?? "").trim().length === 0 + ) { + return null; + } + const selected = String(env.NEMOCLAW_MODEL ?? "").trim(); + return selected === STATION_ULTRA_DUAL_SERVED_MODEL ? selected : null; +} + function identifiesCheckpoint(model: VllmModelDef, value: string): boolean { const normalized = value.toLowerCase(); return [model.envValue, model.id, model.servedModelId].some( @@ -812,7 +829,7 @@ function expectedEnvironment( if (includeProviderSelection) { expected.NEMOCLAW_PROVIDER = "install-vllm"; expected.NEMOCLAW_VLLM_MODEL = model.envValue; - expected.NEMOCLAW_MODEL = servedModel(model); + expected.NEMOCLAW_MODEL = intent.servedModel ?? servedModel(model); } return expected; } @@ -895,10 +912,12 @@ export function getStationExpressResumeIntent( message: "DGX Station Express requires a registered Station vLLM model and sandbox name.", }; } - const intent: StationExpressResumeIntent = { + const dualServedModel = qualifiedDualStationServedModel(env, model); + const intent: StationResumeIntent = { version: STATION_EXPRESS_INTENT_VERSION, model: model.envValue, sandboxName, + ...(dualServedModel ? { servedModel: dualServedModel, checkpointModel: model.id } : {}), ...(isValidStationExpressReceiptGeneration(env[STATION_EXPRESS_RECEIPT_GENERATION_ENV]) ? { receiptGeneration: env[STATION_EXPRESS_RECEIPT_GENERATION_ENV] } : {}), diff --git a/test/inference-options-docs.test.ts b/test/inference-options-docs.test.ts index 33e4257e11..1ab731490b 100644 --- a/test/inference-options-docs.test.ts +++ b/test/inference-options-docs.test.ts @@ -289,6 +289,29 @@ describe("inference setup navigation", () => { expect(markdown).toContain("only from the OpenShell Docker subnet to its gateway address"); }); + it("documents the dual-Station host-network trust boundary", () => { + const markdown = fs.readFileSync(vllmSetupPath, "utf8"); + + expect(markdown).toContain( + "Existing-server and single-host managed-vLLM paths need port `8000`", + ); + expect(markdown).toContain( + "qualified dual-Station runtime intentionally uses Docker host networking", + ); + expect(markdown).toContain("Neither dual-Station container publishes a Docker port"); + expect(markdown).toContain("all Linux capabilities dropped"); + expect(markdown).toContain("only the selected GPU UUID and exact `uverbs` devices"); + expect(markdown).toContain("worker does not receive the serving key"); + expect(markdown).toContain("`/health` remains unauthenticated for readiness"); + expect(markdown).toContain("deny it on management and LAN interfaces"); + expect(markdown).not.toContain( + "keeps its existing bridge-networked managed-inference topology instead of importing the playbook's host-network setting", + ); + expect(markdown).not.toContain( + "NemoClaw needs port `8000` on host loopback for validation and on the OpenShell Docker bridge", + ); + }); + it("keeps managed image tags, digests, and compressed sizes in sync with source", () => { const markdown = fs.readFileSync(vllmSetupPath, "utf8"); const entries = [ diff --git a/test/install-express-prompt.test.ts b/test/install-express-prompt.test.ts index 28ec0495a3..b93454393a 100644 --- a/test/install-express-prompt.test.ts +++ b/test/install-express-prompt.test.ts @@ -339,14 +339,15 @@ DGX_COMMIT_ID="d0e99cc"\nDGX_PLATFORM="DGX Server for GALAXY-GB300" ); }); - it("uses the Nemotron Ultra recipe without follow-up choices on DGX Station", () => { + it("keeps the Station Ultra default while deferring topology selection", () => { const result = runExpressPromptWithTty("\n", "pipe", "DGX Station"); const output = `${result.stdout}${result.stderr}`; expect(result.status, output).toBe(0); expect(output).toMatch(/Detected DGX Station/); expect(output).toMatch( - /Express install will configure managed local vLLM with NVIDIA Nemotron 3 Ultra 550B/, + /Express install will configure managed local vLLM with NVIDIA Nemotron 3 Ultra 550B and automatic Station topology selection/, ); + expect(output).toMatch(/pretrusted reciprocal dual-Station pair selects distributed serving/); expect(output).toMatch(/approximately 352 GB model/); expect(output).toContain( "Hugging Face authentication is optional for this public model but recommended", @@ -537,7 +538,7 @@ ensure_station_express_host`, ).toBe( `revision=${revision}\nmodel=deepseek-v4-flash\ngeneration=${generation}\n` + "agent=hermes\nsandbox=custom-agent\npolicy_tier=restricted\n" + - "gateway_port=18081\ndashboard_port=18790\nvllm_port=18000\n", + "gateway_port=18081\ndashboard_port=18790\nvllm_port=18000\nmode=express\n", ); expect(output).toContain("A reboot is not required"); expect(output).toContain( @@ -1209,8 +1210,8 @@ printf 'NON_EXPRESS_ALLOWED\n' }); const output = `${result.stdout}${result.stderr}`; expect(result.status, output).toBe(0); - expect(output).toMatch(/managed local vLLM with NVIDIA Nemotron 3 Ultra 550B/); - expect(output).toMatch(/approximately 352 GB model/); + expect(output).toMatch(/automatic Station topology selection/); + expect(output).toMatch(/pretrusted reciprocal dual-Station pair selects distributed serving/); expect(output).toMatch( /RESULT NON_INTERACTIVE=1 SUDO_MODE=prompt PROVIDER=install-vllm MODEL=nvidia\/nemotron-3-ultra-550b-a55b VLLM_MODEL=nemotron-3-ultra-550b-a55b POLICY=suggested YES=1 SANDBOX=my-assistant/, ); @@ -1253,7 +1254,20 @@ detect_express_platform expect(result.stdout).toBe("DGX Station"); }); - it.each(["7.2.0", "7.4.0", "7.5.0"])("recognizes stock DGX OS %s", (version) => { + it("rejects partial and unsupported Station product identifiers", () => { + for (const productName of [ + "Acme XP3830 Workstation", + "Dell Pro Max with Station GB200", + "Dell Pro Max with GB300", + ]) { + const result = detectExpressPlatformForProductName(productName); + + expect(result.status, `${result.stdout}${result.stderr}`).toBe(0); + expect(result.stdout).not.toBe("DGX Station"); + } + }); + + it.each(["7.2.0", "7.4.0", "7.5.0"])("recognizes stock DGX OS %s on Station GB300", (version) => { const result = detectExpressPlatformForStockDgxRelease( "DGX Station GB300", stockDgxRelease(version), diff --git a/test/install-gateway-state-root.test.ts b/test/install-gateway-state-root.test.ts index 5ad91c9488..a40ffe51df 100644 --- a/test/install-gateway-state-root.test.ts +++ b/test/install-gateway-state-root.test.ts @@ -154,6 +154,23 @@ printf 'state=%s\n' "$(nemoclaw_state_dir)"`, } }); + it("normalizes a trailing slash in HOME before selecting its state root", () => { + const home = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-installer-home-slash-")); + try { + const result = runInstallerFunctions( + `${home}/`, + `NEMOCLAW_GATEWAY_PORT=8080 +printf 'state=%s\n' "$(nemoclaw_state_dir)"`, + ); + + expect(result.status, result.output).toBe(0); + expect(result.output).toContain(`state=${home}/.nemoclaw`); + expect(result.output).not.toContain(`${home}//.nemoclaw`); + } finally { + fs.rmSync(home, { recursive: true, force: true }); + } + }); + it("rejects an overlong digit-only gateway port before selecting its state root (#7203)", () => { const home = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-installer-overlong-port-")); try { diff --git a/test/install-station-container-coexistence.test.ts b/test/install-station-container-coexistence.test.ts index 7f5b5d1d63..abcd07a669 100644 --- a/test/install-station-container-coexistence.test.ts +++ b/test/install-station-container-coexistence.test.ts @@ -58,6 +58,54 @@ capture_docker_container_baseline expect(output).toContain("docker_container_baseline_total=0 running=0"); }); + it("uses noninteractive sudo only for read-only peer checks before Docker group access", () => { + const { result, output } = runStationPreparation( + ` +MODE='--check' +ps() { printf '%s %s bash bash prepare-dgx-station-host.sh --check\n' "$$" "$PPID"; } +ss() { :; } +docker() { return 1; } +sudo() { + if [[ "$1" == "-n" ]]; then shift; fi + case "$*" in + 'docker ps -aq --no-trunc'|'docker ps -q --no-trunc') return 0 ;; + *) return 1 ;; + esac +} +systemctl() { return 0; } +capture_docker_container_baseline +`, + { + NEMOCLAW_STATION_PREP_SUDO_NONINTERACTIVE: "1", + PATH: `${path.dirname(process.execPath)}:${TEST_SYSTEM_PATH}`, + }, + ); + + expect(result.status, output).toBe(0); + expect(output).toContain("docker_access=sudo_for_noninteractive_read_only_check"); + expect(output).toContain("docker_container_baseline_total=0 running=0"); + }); + + it("does not hide missing Docker group access during verify", () => { + const { result, output } = runStationPreparation( + ` +MODE='--verify' +docker() { return 1; } +sudo() { return 0; } +systemctl() { return 0; } +capture_docker_container_baseline +`, + { + NEMOCLAW_STATION_PREP_SUDO_NONINTERACTIVE: "1", + PATH: `${path.dirname(process.execPath)}:${TEST_SYSTEM_PATH}`, + }, + ); + + expect(result.status, output).not.toBe(0); + expect(output).toContain("Docker is active but inaccessible to this login"); + expect(output).not.toContain("docker_access=sudo_"); + }); + it("fails closed when Docker is installed but its container state cannot be queried", () => { const { result, output } = runStationPreparation( ` @@ -111,6 +159,7 @@ check_dpkg_database_health() { :; } check_failed_units() { :; } check_vllm_container_conflicts() { :; } check_agent_and_inference_conflicts() { :; } +verify_dual_station_controller_uid_binding() { :; } driver_loaded_exact() { return 0; } package_is_ready() { return 0; } verify_gpu() { :; } diff --git a/test/install-station-controller-binding.test.ts b/test/install-station-controller-binding.test.ts new file mode 100644 index 0000000000..cf8ad53141 --- /dev/null +++ b/test/install-station-controller-binding.test.ts @@ -0,0 +1,442 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { spawnSync } from "node:child_process"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; +import { describe, expect, it } from "vitest"; +import { INSTALLER_PAYLOAD, TEST_SYSTEM_PATH } from "./helpers/installer-sourced-env"; + +const REPO_ROOT = path.resolve(import.meta.dirname, ".."); +const STATION_PREPARE = path.join(REPO_ROOT, "scripts", "prepare-dgx-station-host.sh"); + +function runSourced(body: string) { + const home = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-station-controller-")); + const result = spawnSync( + "bash", + ["--noprofile", "--norc", "-c", `source "$SCRIPT_UNDER_TEST" >/dev/null\n${body}`], + { + cwd: REPO_ROOT, + encoding: "utf8", + env: { HOME: home, PATH: TEST_SYSTEM_PATH, SCRIPT_UNDER_TEST: STATION_PREPARE }, + timeout: 15_000, + killSignal: "SIGKILL", + }, + ); + return { home, result, output: `${result.stdout}${result.stderr}` }; +} + +function runInstallerBody(body: string, extraEnv: Record = {}) { + const home = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-station-migration-")); + const result = spawnSync( + "bash", + ["--noprofile", "--norc", "-c", `source "$INSTALLER_UNDER_TEST" >/dev/null\n${body}`], + { + cwd: REPO_ROOT, + encoding: "utf8", + env: { + HOME: home, + PATH: `${path.dirname(process.execPath)}:${TEST_SYSTEM_PATH}`, + INSTALLER_UNDER_TEST: INSTALLER_PAYLOAD, + ...extraEnv, + }, + timeout: 15_000, + killSignal: "SIGKILL", + }, + ); + return { home, result, output: `${result.stdout}${result.stderr}` }; +} + +function runInstallerOrderHarness(onboardStatus: 0 | 1) { + return runInstallerBody( + ` +record_order() { printf '%s\n' "$1" >>"$HOME/order.trace"; } +resolve_nemoclaw_gateway_port() { printf '18789'; } +preflight_explicit_express_flags() { :; } +print_banner() { :; } +preflight_usage_notice_prompt() { :; } +prepare_installer_host() { _SELECTED_EXPRESS_PLATFORM='DGX Station'; } +bash() { :; } +step() { :; } +install_nodejs() { :; } +ensure_supported_runtime() { :; } +ensure_station_express_pair() { record_order qualify; } +fix_npm_permissions() { :; } +preinstall_backup_and_retire_legacy_gateway() { :; } +install_nemoclaw() { record_order install; } +verify_nemoclaw() { :; } +require_reportable_openshell_version() { :; } +command_exists() { return 0; } +registered_sandbox_count() { printf '0\n'; } +run_installer_host_preflight() { return 0; } +recover_preexisting_sandboxes_before_onboard() { return 0; } +run_onboard() { record_order onboard; return "$ONBOARD_STATUS"; } +restore_onboard_forward_after_post_checks() { return 0; } +finalize_install() { record_order finalize; } +clear_station_dual_pair_resume() { record_order clear; } +clear_station_express_resume() { :; } +main --non-interactive --yes-i-accept-third-party-software +`, + { ONBOARD_STATUS: String(onboardStatus) }, + ); +} + +describe("DGX Station controller UID binding", () => { + it("keeps ordinary Station verification independent of dual-pair binding", () => { + const { home, result, output } = runSourced(` +require_command() { :; } +common_preflight() { STATION_HOST_PROFILE=ai-developer-tools; } +station_uses_factory_runtime() { return 0; } +verify_dgx_os_runtime_user() { printf 'FACTORY_RUNTIME_VERIFIED\\n'; } +verify_dual_station_controller_uid_binding() { + printf 'UNEXPECTED_PAIR_BINDING_CHECK\\n' + return 1 +} +run_verify +`); + try { + expect(result.status, output).toBe(0); + expect(output).toContain("FACTORY_RUNTIME_VERIFIED"); + expect(output).not.toContain("UNEXPECTED_PAIR_BINDING_CHECK"); + } finally { + fs.rmSync(home, { recursive: true, force: true }); + } + }); + + it("runs binding-only preparation without workload inspection and retains sudo acquisition", () => { + const { home, result, output } = runSourced(` +require_command() { :; } +acquire_sudo() { sudo_mode='acquired'; printf 'ACQUIRE_SUDO\\n'; } +check_platform() { printf 'CHECK_PLATFORM\\n'; } +check_no_workloads() { printf 'WORKLOAD_CHECK_MUST_NOT_RUN\\n'; return 1; } +ensure_dual_station_controller_uid_binding() { printf 'ENSURE_CONTROLLER_BINDING sudo=%s\\n' "$sudo_mode"; } +run_bind_controller +`); + try { + expect(result.status, output).toBe(0); + expect(output).toContain("ACQUIRE_SUDO"); + expect(output).toContain("CHECK_PLATFORM"); + expect(output).toContain("ENSURE_CONTROLLER_BINDING sudo=acquired"); + expect(output).toContain("CONTROLLER_UID_BINDING_READY"); + expect(output).not.toContain("WORKLOAD_CHECK_MUST_NOT_RUN"); + } finally { + fs.rmSync(home, { recursive: true, force: true }); + } + }); + + it("creates one exact binding, reuses it, and requires administrator removal to rebind", () => { + const { home, result, output } = runSourced(` +config_dir="$HOME/etc/nemoclaw" +binding_file="$config_dir/dual-station-controller-uid" +mkdir -p "$HOME/etc" +controller_uid=1001 +binding_owned=0 +preparation_controller_uid() { printf '%s\\n' "$controller_uid"; } +ensure_root_directory_safe() { mkdir -p "$1"; } +assert_root_directory_safe() { [[ -d "$1" ]]; } +root_regular_file_is_safe() { ((binding_owned == 1)) && [[ -f "$1" ]]; } +root_directory_is_safe_unprivileged() { [[ -d "$1" ]]; } +root_regular_file_is_safe_unprivileged() { ((binding_owned == 1)) && [[ -f "$1" ]]; } +sudo() { + if [[ "$1" == "chown" && "$2" == "root:root" && "$3" == "\${config_dir}/.dual-station-controller-uid."* ]]; then + binding_owned=1 + return 0 + fi + "$@" +} +ensure_dual_station_controller_uid_binding "$config_dir" "$binding_file" +ensure_dual_station_controller_uid_binding "$config_dir" "$binding_file" +controller_uid=1002 +ensure_dual_station_controller_uid_binding "$config_dir" "$binding_file" +`); + + try { + expect(result.status, output).not.toBe(0); + expect(output.match(/dual_station_controller_uid=installed/g)).toHaveLength(1); + expect(output).toMatch(/administrator must remove .* before rebinding/); + const binding = path.join(home, "etc/nemoclaw/dual-station-controller-uid"); + expect(fs.readFileSync(binding, "utf8")).toBe("1001\n"); + expect(fs.statSync(binding).mode & 0o777).toBe(0o644); + expect( + fs + .readdirSync(path.dirname(binding)) + .filter((name) => name.startsWith(".dual-station-controller-uid.")), + ).toEqual([]); + } finally { + fs.rmSync(home, { recursive: true, force: true }); + } + }); + + it("never lets an atomic publication loser replace the winner", () => { + const { home, result, output } = runSourced(` +config_dir="$HOME/etc/nemoclaw" +binding_file="$config_dir/dual-station-controller-uid" +mkdir -p "$HOME/etc" +preparation_controller_uid() { printf '1001\\n'; } +ensure_root_directory_safe() { mkdir -p "$1"; } +assert_root_directory_safe() { [[ -d "$1" ]]; } +root_regular_file_is_safe() { [[ -f "$1" ]]; } +root_directory_is_safe_unprivileged() { [[ -d "$1" ]]; } +root_regular_file_is_safe_unprivileged() { [[ -f "$1" ]]; } +sudo() { + if [[ "$1" == "chown" ]]; then return 0; fi + if [[ "$1" == "ln" ]]; then + printf '1002\\n' >"$binding_file" + chmod 0644 "$binding_file" + return 1 + fi + "$@" +} +ensure_dual_station_controller_uid_binding "$config_dir" "$binding_file" +`); + + try { + expect(result.status, output).not.toBe(0); + expect(output).toMatch(/administrator must remove .* before rebinding/); + expect( + fs.readFileSync(path.join(home, "etc/nemoclaw/dual-station-controller-uid"), "utf8"), + ).toBe("1002\n"); + } finally { + fs.rmSync(home, { recursive: true, force: true }); + } + }); + + it("rejects a symlink without modifying its target", () => { + const { home, result, output } = runSourced(` +config_dir="$HOME/etc/nemoclaw" +binding_file="$config_dir/dual-station-controller-uid" +mkdir -p "$config_dir" +printf 'preserve\\n' >"$HOME/target" +ln -s "$HOME/target" "$binding_file" +preparation_controller_uid() { printf '1001\\n'; } +ensure_root_directory_safe() { :; } +root_directory_is_safe_unprivileged() { return 0; } +sudo() { "$@"; } +ensure_dual_station_controller_uid_binding "$config_dir" "$binding_file" +`); + + try { + expect(result.status, output).not.toBe(0); + expect(output).toMatch(/must not be a symbolic link/); + expect(fs.readFileSync(path.join(home, "target"), "utf8")).toBe("preserve\n"); + } finally { + fs.rmSync(home, { recursive: true, force: true }); + } + }); + + it("rejects an existing mode-0700 configuration directory before creating a binding", () => { + const { home, result, output } = runSourced(` +config_dir="$HOME/etc/nemoclaw" +binding_file="$config_dir/dual-station-controller-uid" +mkdir -p "$config_dir" +chmod 0700 "$config_dir" +preparation_controller_uid() { printf '1001\\n'; } +ensure_root_directory_safe() { :; } +root_directory_is_safe_unprivileged() { return 1; } +sudo() { printf 'SUDO_AFTER_MODE_CHECK\\n' >&2; return 97; } +ensure_dual_station_controller_uid_binding "$config_dir" "$binding_file" +`); + try { + expect(result.status, output).not.toBe(0); + expect(output).toMatch(/must be root-owned with mode 0755 before binding/); + expect(output).not.toContain("SUDO_AFTER_MODE_CHECK"); + expect(fs.existsSync(path.join(home, "etc/nemoclaw/dual-station-controller-uid"))).toBe( + false, + ); + } finally { + fs.rmSync(home, { recursive: true, force: true }); + } + }); + + it("fails verification on unsafe binding metadata", () => { + const { home, result, output } = runSourced(` +root_directory_is_safe_unprivileged() { return 0; } +root_regular_file_is_safe_unprivileged() { return 1; } +verify_dual_station_controller_uid_binding 1001 /etc/nemoclaw /etc/nemoclaw/dual-station-controller-uid +`); + try { + expect(result.status, output).not.toBe(0); + expect(output).toMatch(/root-owned regular file with mode 0644/); + } finally { + fs.rmSync(home, { recursive: true, force: true }); + } + }); + + it("verifies readable root metadata and exact content without invoking sudo", () => { + const { home, result, output } = runSourced(` +config_dir="$HOME/etc/nemoclaw" +binding_file="$config_dir/dual-station-controller-uid" +mkdir -p "$config_dir" +printf '1001\\n' >"$binding_file" +stat() { + if [[ "\${@: -1}" == "$config_dir" ]]; then printf '0 0 755\\n'; else printf '0 0 644\\n'; fi +} +sudo() { printf 'SUDO_MUST_NOT_RUN\\n' >&2; return 97; } +verify_dual_station_controller_uid_binding 1001 "$config_dir" "$binding_file" +`); + try { + expect(result.status, output).toBe(0); + expect(output).toContain("dual_station_controller_uid=verified uid=1001"); + expect(output).not.toContain("SUDO_MUST_NOT_RUN"); + } finally { + fs.rmSync(home, { recursive: true, force: true }); + } + }); + + it("rejects root preparation without an original non-root sudo UID", () => { + const { home, result, output } = runSourced(`preparation_controller_uid_for 0 ''`); + try { + expect(result.status, output).not.toBe(0); + expect(output).toMatch(/must be run by a non-root controller account/); + } finally { + fs.rmSync(home, { recursive: true, force: true }); + } + }); + + it("recognizes only the frozen legacy head and routes it around workload preparation", () => { + const digest = + "vllm/vllm-openai@sha256:0fec7ec5f3e6bc168e54899935fb0557da908a4832a1dbc88e2debcf2f889416"; + const inspection = ["/nemoclaw-vllm", "true", digest, "true", ...Array(8).fill("-")].join("|"); + const detected = runInstallerBody( + `command_exists() { return 0; }; docker() { printf '%s\\n' "$INSPECTION"; }; station_migratable_legacy_single_head_running`, + { INSPECTION: inspection }, + ); + const routed = runInstallerBody(` +_SELECTED_EXPRESS_PLATFORM='DGX Station' +station_dual_model_requested() { return 0; } +station_managed_dual_head_running() { return 1; } +station_migratable_legacy_single_head_running() { return 0; } +run_station_host_preparation() { printf 'FULL_PREP_MUST_NOT_RUN\\n'; return 1; } +ensure_station_express_host +printf 'MIGRATE=%s REUSE=%s\\n' "$_STATION_EXPRESS_MIGRATING_LEGACY_HEAD" "$_STATION_EXPRESS_DEFERRED_MANAGED_PAIR" +`); + try { + expect(detected.result.status, detected.output).toBe(0); + expect(routed.result.status, routed.output).toBe(0); + expect(routed.output).toContain("MIGRATE=1 REUSE=0"); + expect(routed.output).not.toContain("FULL_PREP_MUST_NOT_RUN"); + } finally { + fs.rmSync(detected.home, { recursive: true, force: true }); + fs.rmSync(routed.home, { recursive: true, force: true }); + } + }); + + it.each([ + [ + "managed dual", + [ + "/nemoclaw-vllm", + "true", + "true", + "head", + "1", + "c".repeat(64), + "d".repeat(64), + "e".repeat(64), + "f".repeat(32), + ].join(" "), + "REUSE=1 MIGRATE=0", + 1, + ], + [ + "legacy single", + [ + "/nemoclaw-vllm", + "true", + "vllm/vllm-openai@sha256:0fec7ec5f3e6bc168e54899935fb0557da908a4832a1dbc88e2debcf2f889416", + "true", + ...Array(8).fill("-"), + ].join("|"), + "REUSE=0 MIGRATE=1", + 2, + ], + ] as const)("uses canonical local Docker to preserve a hidden %s head", (_kind, localInspection, expectedFlags, expectedInspections) => { + const { result, output, home } = runInstallerBody( + ` +command_exists() { return 0; } +docker() { + printf 'host=%s context=%s args=%s,%s\n' "\${DOCKER_HOST-unset}" "\${DOCKER_CONTEXT-unset}" "$1" "$2" >>"$HOME/docker.trace" + if [[ "\${DOCKER_HOST+x}" != x && "\${DOCKER_CONTEXT+x}" != x && "$1" == --context && "$2" == default ]]; then + printf '%s\n' "$LOCAL_DOCKER_INSPECTION" + fi +} +station_dual_model_requested() { return 0; } +run_station_host_preparation() { printf 'FULL_PREP_MUST_NOT_RUN\n'; return 97; } +_SELECTED_EXPRESS_PLATFORM='DGX Station' +ensure_station_express_host +printf 'REUSE=%s MIGRATE=%s\n' "$_STATION_EXPRESS_DEFERRED_MANAGED_PAIR" "$_STATION_EXPRESS_MIGRATING_LEGACY_HEAD" +`, + { + DOCKER_CONTEXT: "ambient-remote", + DOCKER_HOST: "ssh://remote-builder.example.test", + LOCAL_DOCKER_INSPECTION: localInspection, + }, + ); + try { + expect(result.status, output).toBe(0); + expect(output).toContain(expectedFlags); + expect(output).not.toContain("FULL_PREP_MUST_NOT_RUN"); + expect(fs.readFileSync(path.join(home, "docker.trace"), "utf8")).toBe( + "host=unset context=unset args=--context,default\n".repeat(expectedInspections), + ); + } finally { + fs.rmSync(home, { recursive: true, force: true }); + } + }); + + it("passes legacy migration to the coordinator without managed-pair reuse", () => { + const argsFile = path.join(os.tmpdir(), `nemoclaw-legacy-args-${process.pid}-${Date.now()}`); + const { home, result, output } = runInstallerBody( + ` +node() { + if [[ "\${1:-}" == "--no-warnings" ]]; then + printf '%s\\n' "$*" >"$ARGS_FILE" + printf '%s\\n' '{"kind":"single-station","reason":"fixture"}' + return 0 + fi + command node "$@" +} +station_installer_revision() { printf '%040d' 0; } +_SELECTED_EXPRESS_PLATFORM='DGX Station' +_STATION_EXPRESS_MODEL_WAS_EXPLICIT=0 +_STATION_INSTALL_MODE='express' +_STATION_EXPRESS_DEFERRED_MANAGED_PAIR=0 +_STATION_EXPRESS_MIGRATING_LEGACY_HEAD=1 +NEMOCLAW_VLLM_MODEL='nemotron-3-ultra-550b-a55b' +unset NEMOCLAW_DGX_STATION_PEER +ensure_station_express_pair +`, + { ARGS_FILE: argsFile }, + ); + try { + expect(result.status, output).not.toBe(0); + const args = fs.readFileSync(argsFile, "utf8"); + expect(args).toContain("--migrate-legacy-single-head"); + expect(args).not.toContain("--reuse-existing-managed-pair"); + expect(output).toMatch(/legacy single-Station head.*refusing migration/u); + } finally { + fs.rmSync(argsFile, { force: true }); + fs.rmSync(home, { recursive: true, force: true }); + } + }); + + it("qualifies before install/onboarding and clears pair state only after success", () => { + const success = runInstallerOrderHarness(0); + const failure = runInstallerOrderHarness(1); + try { + expect(success.result.status, success.output).toBe(0); + expect(fs.readFileSync(path.join(success.home, "order.trace"), "utf8")).toBe( + "qualify\ninstall\nonboard\nfinalize\nclear\n", + ); + + expect(failure.result.status, failure.output).not.toBe(0); + expect(fs.readFileSync(path.join(failure.home, "order.trace"), "utf8")).toBe( + "qualify\ninstall\nonboard\n", + ); + } finally { + fs.rmSync(success.home, { recursive: true, force: true }); + fs.rmSync(failure.home, { recursive: true, force: true }); + } + }); +}); diff --git a/test/install-station-host-preparation.test.ts b/test/install-station-host-preparation.test.ts index 774cd1a775..9b031ce189 100644 --- a/test/install-station-host-preparation.test.ts +++ b/test/install-station-host-preparation.test.ts @@ -1117,7 +1117,7 @@ ensure_station_express_host expect(fs.readFileSync(stateFile, "utf-8")).toBe( `revision=${STATION_REVISION}\nmodel=nemotron-3-ultra-550b-a55b\ngeneration=${STATION_GENERATION}\n` + "agent=openclaw\nsandbox=my-assistant\npolicy_tier=balanced\n" + - "gateway_port=18081\ndashboard_port=18790\nvllm_port=18000\n", + "gateway_port=18081\ndashboard_port=18790\nvllm_port=18000\nmode=express\n", ); expect(fs.statSync(stateFile).mode & 0o777).toBe(0o600); expect(output).toContain( diff --git a/test/install-station-pair-preparation.test.ts b/test/install-station-pair-preparation.test.ts new file mode 100644 index 0000000000..e782b65ae8 --- /dev/null +++ b/test/install-station-pair-preparation.test.ts @@ -0,0 +1,1499 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { spawnSync } from "node:child_process"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; +import { describe, expect, it } from "vitest"; +import { + type DualStationPreparationDeps, + type DualStationResumeState, + deriveDiscoveryCandidates, + deriveSlash30Counterpart, + type PretrustedSshTarget, + parseDualStationResumeState, + prepareDualStationPair, + type StationDiscoveryHost, + validateResumeFileMetadata, + validateStationPeerTarget, +} from "../scripts/lib/dgx-station-peer.mts"; +import { + buildRemoteHelperCommand, + buildStationPrepSubprocessEnv, + clearDualStationResumeState, + inspectPretrustedSshTarget, + readDualStationResumeState, + strictStationPrepSshTransportArgs, + writeDualStationResumeState, +} from "../scripts/prepare-dual-dgx-station.mts"; +import { + stationKnownHostsDigest, + strictStationSshTransportArgs, +} from "../src/lib/inference/vllm-station-ssh-binding.ts"; +import { INSTALLER_PAYLOAD, TEST_SYSTEM_PATH } from "./helpers/installer-sourced-env"; + +const REPO_ROOT = path.resolve(import.meta.dirname, ".."); +const COORDINATOR = path.join(REPO_ROOT, "scripts", "prepare-dual-dgx-station.mts"); +const STATION_HELPER = path.join(REPO_ROOT, "scripts", "prepare-dgx-station-host.sh"); +const REVISION = "a".repeat(40); +const HELPER_SHA256 = "b".repeat(64); +const HOST_KEY_DATA = "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA"; +const HOST_KEY_DIGEST = stationKnownHostsDigest(`10.10.0.2 ssh-ed25519 ${HOST_KEY_DATA}\n`); +const HOST_KEY_FINGERPRINT = `SHA256:${"A".repeat(43)}`; + +function stationHost(side: "local" | "peer"): StationDiscoveryHost { + const local = side === "local"; + return { + schemaVersion: 1, + hostname: local ? "station-a" : "station-b", + productName: "NVIDIA DGX Station GB300", + architecture: "aarch64", + gpus: [ + { + index: 0, + name: "NVIDIA GB300", + uuid: local ? "GPU-LOCAL-0001" : "GPU-PEER-0002", + }, + ], + rails: [ + { + netdev: "enp1s0f0np0", + macAddress: local ? "02:00:00:00:00:01" : "02:00:00:00:00:02", + pciAddress: "0000:01:00.0", + pciName: "NVIDIA ConnectX-8 Ethernet Controller", + state: "4: ACTIVE", + linkLayer: "Ethernet", + speedMbps: 400_000, + mtu: 9000, + ipv4Addresses: [{ address: local ? "10.10.0.1" : "10.10.0.2", prefixLength: 30 }], + }, + { + netdev: "enp2s0f0np0", + macAddress: local ? "02:00:00:00:00:05" : "02:00:00:00:00:06", + pciAddress: "0000:02:00.0", + pciName: "NVIDIA ConnectX-8 Ethernet Controller", + state: "4: ACTIVE", + linkLayer: "Ethernet", + speedMbps: 400_000, + mtu: 9000, + ipv4Addresses: [{ address: local ? "10.10.0.5" : "10.10.0.6", prefixLength: 30 }], + }, + ], + }; +} + +function stationConnectivity(side: "local" | "peer"): string { + const source = stationHost(side); + const destination = stationHost(side === "local" ? "peer" : "local"); + return JSON.stringify({ + schemaVersion: 1, + checks: source.rails.map((rail, index) => ({ + netdev: rail.netdev, + sourceAddress: rail.ipv4Addresses[0].address, + peerAddress: destination.rails[index].ipv4Addresses[0].address, + routeDevice: rail.netdev, + routeSource: rail.ipv4Addresses[0].address, + routeGateway: null, + routeScope: "link", + peerMac: destination.rails[index].macAddress, + peerNeighborState: "REACHABLE", + jumboPing: true, + })), + }); +} + +function sshBinding(target = "10.10.0.2", keyData = HOST_KEY_DATA): PretrustedSshTarget { + const knownHostsLine = `${target.slice(target.lastIndexOf("@") + 1)} ssh-ed25519 ${keyData}`; + return { + requestedTarget: target, + sshTarget: target, + resolvedHost: target.slice(target.lastIndexOf("@") + 1), + sshUser: "ubuntu", + port: 22, + lookupHost: target.slice(target.lastIndexOf("@") + 1), + hostKeyDigest: stationKnownHostsDigest(`${knownHostsLine}\n`), + keyFingerprints: [HOST_KEY_FINGERPRINT], + knownHostsLines: [knownHostsLine], + }; +} + +function preparationOptions() { + return { revision: REVISION, helperSha256: HELPER_SHA256 }; +} + +function throwFixtureError(error: Error): never { + throw error; +} + +class PreparationHarness { + readonly calls: string[] = []; + readonly statePhases: DualStationResumeState["phase"][] = []; + readonly trusted = new Map(); + readonly trustErrors = new Map(); + readonly localHelperStatus = new Map(); + readonly remoteHelperStatus = new Map(); + local = stationHost("local"); + peer = stationHost("peer"); + resume: DualStationResumeState | null = null; + localConnectivity = true; + peerConnectivity = true; + peerProbeError: Error | null = null; + + readonly deps: DualStationPreparationDeps = { + runLocalHelper: (mode) => { + this.calls.push(`local:${mode}`); + return this.localHelperStatus.get(mode) ?? 0; + }, + probeLocalHost: () => { + this.calls.push("probe:local"); + return structuredClone(this.local); + }, + inspectPretrustedTarget: (target) => { + this.calls.push(`trust:${target}`); + const error = this.trustErrors.get(target); + error && throwFixtureError(error); + return this.trusted.get(target) ?? null; + }, + probePeerHost: (binding) => { + this.calls.push(`probe:peer:${binding.sshTarget}`); + this.peerProbeError && throwFixtureError(this.peerProbeError); + return structuredClone(this.peer); + }, + probeLocalConnectivity: () => { + this.calls.push("connectivity:local"); + return this.localConnectivity; + }, + probePeerConnectivity: (binding) => { + this.calls.push(`connectivity:peer:${binding.sshTarget}`); + return this.peerConnectivity; + }, + runRemoteHelper: (binding, mode) => { + this.calls.push(`remote:${binding.sshTarget}:${mode}`); + return this.remoteHelperStatus.get(mode) ?? 0; + }, + readResumeState: () => { + this.calls.push("state:read"); + return this.resume ? structuredClone(this.resume) : null; + }, + writeResumeState: (state) => { + this.calls.push(`state:write:${state.phase}`); + this.resume = structuredClone(state); + this.statePhases.push(state.phase); + }, + clearResumeState: () => { + this.calls.push("state:clear"); + this.resume = null; + }, + log: (message) => this.calls.push(`log:${message}`), + }; +} + +function trustFirstRail(harness: PreparationHarness): void { + harness.trusted.set("10.10.0.2", sshBinding()); +} + +function readyState(): DualStationResumeState { + return { + schemaVersion: 1, + revision: REVISION, + helperSha256: HELPER_SHA256, + phase: "ready", + peerTarget: "10.10.0.2", + hostKeyDigest: HOST_KEY_DIGEST, + localGpuUuid: "GPU-LOCAL-0001", + peerGpuUuid: "GPU-PEER-0002", + rails: [ + { + localAddress: "10.10.0.1", + localMac: "02:00:00:00:00:01", + peerAddress: "10.10.0.2", + peerMac: "02:00:00:00:00:02", + }, + { + localAddress: "10.10.0.5", + localMac: "02:00:00:00:00:05", + peerAddress: "10.10.0.6", + peerMac: "02:00:00:00:00:06", + }, + ], + }; +} + +function runInstallerBody(body: string, extraEnv: Record = {}) { + const home = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-pair-installer-")); + const result = spawnSync( + "bash", + ["--noprofile", "--norc", "-c", `source "$INSTALLER_UNDER_TEST" >/dev/null\n${body}`], + { + cwd: REPO_ROOT, + encoding: "utf8", + env: { + HOME: home, + INSTALLER_UNDER_TEST: INSTALLER_PAYLOAD, + PATH: `${path.dirname(process.execPath)}:${TEST_SYSTEM_PATH}`, + ...extraEnv, + }, + timeout: 20_000, + killSignal: "SIGKILL", + }, + ); + return { home, result, output: `${result.stdout}${result.stderr}` }; +} + +function coordinatorResult(kind: "ready" | "reboot-required", peer = "10.10.0.2"): string { + const state = readyState(); + state.peerTarget = peer; + const sshBinding = Buffer.from( + JSON.stringify({ + bindingFile: "/tmp/nemoclaw-station-pair/resume.json.ssh-binding/binding.json", + hostKeyDigest: state.hostKeyDigest, + }), + "utf8", + ).toString("base64url"); + return JSON.stringify({ + kind, + peerTarget: peer, + identity: { + peerTarget: peer, + hostKeyDigest: state.hostKeyDigest, + localGpuUuid: state.localGpuUuid, + peerGpuUuid: state.peerGpuUuid, + rails: state.rails, + }, + sshBinding, + }); +} + +describe("deterministic dual-DGX Station peer discovery", () => { + it.each([ + ["10.0.0.1", "10.0.0.2"], + ["10.0.0.2", "10.0.0.1"], + ["172.16.8.5", "172.16.8.6"], + ["172.16.8.6", "172.16.8.5"], + ["192.168.20.1", "192.168.20.2"], + ])("derives only the other usable /30 address: %s -> %s", (address, counterpart) => { + expect(deriveSlash30Counterpart(address)).toBe(counterpart); + }); + + it.each([ + ["10.0.0.0", 30], + ["10.0.0.3", 30], + ["10.0.0.1", 24], + ["8.8.8.1", 30], + ["not-an-ip", 30], + ])("refuses non-usable, non-/30, public, and malformed addresses", (address, prefix) => { + expect(deriveSlash30Counterpart(address, prefix)).toBeNull(); + }); + + it("derives exactly the two reciprocal CX-8 candidates", () => { + expect(deriveDiscoveryCandidates(stationHost("local"))).toEqual(["10.10.0.2", "10.10.0.6"]); + }); + + it.each([ + "DGX-Station", + "P3830", + "NVIDIA Station GB300", + ])("accepts an existing Station firmware product identifier: %s", (productName) => { + const host = stationHost("local"); + host.productName = productName; + + expect(deriveDiscoveryCandidates(host)).toEqual(["10.10.0.2", "10.10.0.6"]); + }); + + it("rejects extra rails, duplicate identities, and non-jumbo links", () => { + const extraRail = stationHost("local"); + extraRail.rails.push({ + ...structuredClone(extraRail.rails[1]), + netdev: "enp3s0f0np0", + pciAddress: "0000:03:00.0", + macAddress: "02:00:00:00:00:09", + ipv4Addresses: [{ address: "10.10.0.9", prefixLength: 30 }], + }); + expect(() => deriveDiscoveryCandidates(extraRail)).toThrow(/exactly two CX-8 rails/); + + const duplicate = stationHost("local"); + duplicate.rails[1].macAddress = duplicate.rails[0].macAddress; + expect(() => deriveDiscoveryCandidates(duplicate)).toThrow(/identity is ambiguous/); + + const nonJumbo = stationHost("local"); + nonJumbo.rails[0].mtu = 1500; + expect(() => deriveDiscoveryCandidates(nonJumbo)).toThrow(/MTU 9000/); + }); + + it("inspects only the two mathematically derived candidates and preserves single-Station behavior", () => { + const harness = new PreparationHarness(); + const result = prepareDualStationPair(preparationOptions(), harness.deps); + + expect(result).toEqual({ + kind: "single-station", + reason: "No derived dual-rail peer address has pre-existing SSH host-key trust", + }); + expect(harness.calls.filter((call) => call.startsWith("trust:"))).toEqual([ + "trust:10.10.0.2", + "trust:10.10.0.6", + ]); + expect(harness.calls.some((call) => call.startsWith("probe:peer"))).toBe(false); + expect(harness.calls.some((call) => call.startsWith("remote:"))).toBe(false); + }); + + it("accepts one pretrusted reciprocal peer and runs preparation in order", () => { + const harness = new PreparationHarness(); + trustFirstRail(harness); + harness.remoteHelperStatus.set("--apply", 11); + const result = prepareDualStationPair(preparationOptions(), harness.deps); + expect(result.kind).toBe("ready"); + expect(result.kind === "ready" && result.peerTarget).toBe("10.10.0.2"); + expect(result.kind === "ready" && result.binding).toEqual(sshBinding()); + expect(harness.statePhases).toEqual(["remote-preparation", "ready"]); + expect(harness.calls.indexOf("local:--verify")).toBeLessThan( + harness.calls.indexOf("probe:peer:10.10.0.2"), + ); + expect(harness.calls.indexOf("state:write:remote-preparation")).toBeLessThan( + harness.calls.indexOf("remote:10.10.0.2:--check"), + ); + expect(harness.calls).toContain("local:--bind-controller"); + expect(harness.calls.filter((call) => call.startsWith("remote:"))).toEqual([ + "remote:10.10.0.2:--check", + "remote:10.10.0.2:--apply", + "remote:10.10.0.2:--bind-controller", + "remote:10.10.0.2:--verify", + ]); + expect(harness.calls).toContain( + "log:Peer Docker access requires a new login; reopening SSH before verification", + ); + }); + it("accepts two rail aliases only when their exact SSH identity is coherent", () => { + const coherent = new PreparationHarness(); + coherent.trusted.set("10.10.0.2", sshBinding("10.10.0.2")); + coherent.trusted.set("10.10.0.6", sshBinding("10.10.0.6")); + expect(prepareDualStationPair(preparationOptions(), coherent.deps).kind).toBe("ready"); + expect(coherent.calls.filter((call) => call.startsWith("probe:peer:"))).toHaveLength(1); + const ambiguous = new PreparationHarness(); + ambiguous.trusted.set("10.10.0.2", sshBinding("10.10.0.2")); + ambiguous.trusted.set("10.10.0.6", sshBinding("10.10.0.6", "AAAAC3NzaChangedKey")); + const result = prepareDualStationPair(preparationOptions(), ambiguous.deps); + expect(result).toMatchObject({ + kind: "single-station", + reason: expect.stringMatching(/different/), + }); + expect(ambiguous.calls.some((call) => call.startsWith("probe:peer:"))).toBe(false); + }); + + it("treats an unusable automatic trust entry as untrusted without contact", () => { + const harness = new PreparationHarness(); + harness.trustErrors.set("10.10.0.2", new Error("unsafe HostKeyAlias")); + + expect(prepareDualStationPair(preparationOptions(), harness.deps)).toMatchObject({ + kind: "single-station", + }); + expect(harness.calls).toContain( + "log:Ignoring derived peer 10.10.0.2: pre-existing SSH trust is unusable (unsafe HostKeyAlias)", + ); + expect(harness.calls.some((call) => call.startsWith("probe:peer:"))).toBe(false); + }); + + it("rejects altered endpoint, port, and known-hosts evidence before peer contact", () => { + const oversizedLines = Array.from( + { length: 5 }, + (_, index) => `10.10.0.${String(index + 2)} ssh-ed25519 A${"B".repeat(14_000)}`, + ); + const scenarios: Array<{ + name: string; + mutate(binding: PretrustedSshTarget): void; + }> = [ + { + name: "requested target substitution", + mutate: (binding) => { + binding.requestedTarget = "10.10.0.6"; + }, + }, + { + name: "fractional port", + mutate: (binding) => { + binding.port = 22.5; + }, + }, + { + name: "blank known-hosts line", + mutate: (binding) => { + binding.knownHostsLines = [""]; + }, + }, + { + name: "untrimmed known-hosts line", + mutate: (binding) => { + binding.knownHostsLines = [`${binding.knownHostsLines[0]} `]; + }, + }, + { + name: "comment known-hosts line", + mutate: (binding) => { + binding.knownHostsLines = ["# no trust evidence"]; + }, + }, + { + name: "oversized known-hosts evidence", + mutate: (binding) => { + binding.knownHostsLines = oversizedLines; + }, + }, + { + name: "known-hosts digest mismatch", + mutate: (binding) => { + binding.knownHostsLines = ["10.10.0.2 ssh-ed25519 AAAAC3NzaChangedKey"]; + }, + }, + ]; + + for (const scenario of scenarios) { + const harness = new PreparationHarness(); + const binding = sshBinding(); + scenario.mutate(binding); + harness.trusted.set("10.10.0.2", binding); + expect( + prepareDualStationPair(preparationOptions(), harness.deps), + scenario.name, + ).toMatchObject({ kind: "single-station" }); + expect( + harness.calls.some((call) => call.startsWith("probe:peer:")), + scenario.name, + ).toBe(false); + expect( + harness.calls.some((call) => call.startsWith("remote:")), + scenario.name, + ).toBe(false); + } + }); + + it("rejects an explicit target whose resolved SSH user changed before peer contact", () => { + const explicit = "ubuntu@station-b"; + const harness = new PreparationHarness(); + const binding = sshBinding(explicit); + binding.sshUser = "root"; + harness.trusted.set(explicit, binding); + + expect(() => + prepareDualStationPair({ ...preparationOptions(), explicitPeer: explicit }, harness.deps), + ).toThrow(/unsafe user or port/); + expect(harness.calls.some((call) => call.startsWith("probe:peer:"))).toBe(false); + expect(harness.calls.some((call) => call.startsWith("remote:"))).toBe(false); + }); + + it("requires reciprocal rail addresses and MACs plus a distinct peer GPU", () => { + const nonreciprocal = new PreparationHarness(); + trustFirstRail(nonreciprocal); + nonreciprocal.peer.rails[1].ipv4Addresses = [{ address: "10.10.0.10", prefixLength: 30 }]; + expect(prepareDualStationPair(preparationOptions(), nonreciprocal.deps)).toMatchObject({ + kind: "single-station", + reason: expect.stringMatching(/not reciprocal/), + }); + expect(nonreciprocal.calls.some((call) => call.startsWith("remote:"))).toBe(false); + + const sameGpu = new PreparationHarness(); + trustFirstRail(sameGpu); + sameGpu.peer.gpus[0].uuid = sameGpu.local.gpus[0].uuid; + expect(prepareDualStationPair(preparationOptions(), sameGpu.deps)).toMatchObject({ + kind: "single-station", + reason: expect.stringMatching(/local Station GPU/), + }); + expect(sameGpu.calls.some((call) => call.startsWith("remote:"))).toBe(false); + }); + + it("keeps an explicit peer authoritative and fail-closed", () => { + const explicit = "ubuntu@station-b"; + const harness = new PreparationHarness(); + harness.trusted.set(explicit, sshBinding(explicit)); + + expect( + prepareDualStationPair({ ...preparationOptions(), explicitPeer: explicit }, harness.deps) + .kind, + ).toBe("ready"); + expect(harness.calls.filter((call) => call.startsWith("trust:"))).toEqual([ + `trust:${explicit}`, + ]); + + const untrusted = new PreparationHarness(); + expect(() => + prepareDualStationPair({ ...preparationOptions(), explicitPeer: explicit }, untrusted.deps), + ).toThrow(/not pretrusted/); + expect(untrusted.calls.some((call) => call.startsWith("probe:peer:"))).toBe(false); + }); + + it.each([ + "root@station;reboot", + "station-b -o ProxyCommand=evil", + "station-b/path", + "user@@station-b", + "[10.10.0.2]", + "$(touch pwned)", + "station-b:2222", + "010.010.000.002", + ])("rejects a malicious or noncanonical peer string: %s", (target) => { + expect(() => validateStationPeerTarget(target)).toThrow(/canonical SSH host/); + }); + + it("fails local verification before trust inspection or remote mutation", () => { + const harness = new PreparationHarness(); + trustFirstRail(harness); + harness.localHelperStatus.set("--verify", 1); + + expect(() => prepareDualStationPair(preparationOptions(), harness.deps)).toThrow( + /verification failed before peer contact/, + ); + expect(harness.calls.some((call) => call.startsWith("trust:"))).toBe(false); + expect(harness.calls.some((call) => call.startsWith("remote:"))).toBe(false); + }); + + it("falls back only before mutation and fails closed for explicit connectivity failure", () => { + const automatic = new PreparationHarness(); + trustFirstRail(automatic); + automatic.localConnectivity = false; + expect(prepareDualStationPair(preparationOptions(), automatic.deps)).toMatchObject({ + kind: "single-station", + reason: expect.stringMatching(/jumbo-frame/), + }); + expect(automatic.statePhases).toEqual([]); + + const explicitTarget = "ubuntu@station-b"; + const explicit = new PreparationHarness(); + explicit.trusted.set(explicitTarget, sshBinding(explicitTarget)); + explicit.peerConnectivity = false; + expect(() => + prepareDualStationPair( + { ...preparationOptions(), explicitPeer: explicitTarget }, + explicit.deps, + ), + ).toThrow(/jumbo-frame/); + }); +}); + +describe("dual-DGX Station reboot resume and reuse", () => { + it("persists the exact pair before remote mutation and resumes remote exit 10", () => { + const first = new PreparationHarness(); + trustFirstRail(first); + first.remoteHelperStatus.set("--apply", 10); + const interrupted = prepareDualStationPair(preparationOptions(), first.deps); + expect(interrupted.kind).toBe("reboot-required"); + expect(interrupted.kind === "reboot-required" && interrupted.binding).toEqual(sshBinding()); + expect(first.resume?.phase).toBe("remote-reboot-required"); + expect(first.resume?.helperSha256).toBe(HELPER_SHA256); + expect(first.calls.some((call) => call.endsWith(":--verify"))).toBe(true); + expect(first.calls).not.toContain("remote:10.10.0.2:--verify"); + const resumed = new PreparationHarness(); + resumed.resume = structuredClone(first.resume); + trustFirstRail(resumed); + expect(prepareDualStationPair(preparationOptions(), resumed.deps).kind).toBe("ready"); + expect(resumed.statePhases).toEqual(["remote-preparation", "ready"]); + expect(resumed.calls.filter((call) => call.startsWith("remote:"))).toEqual([ + "remote:10.10.0.2:--check", + "remote:10.10.0.2:--apply", + "remote:10.10.0.2:--bind-controller", + "remote:10.10.0.2:--verify", + ]); + }); + it("rejects revision, helper, host-key, GPU, and rail substitution on resume", () => { + const scenarios: Array<{ + name: string; + configure(harness: PreparationHarness): void; + options?: ReturnType; + expected: RegExp; + }> = [ + { + name: "revision", + configure: () => undefined, + options: { revision: "d".repeat(40), helperSha256: HELPER_SHA256 }, + expected: /requires NemoClaw revision/, + }, + { + name: "helper", + configure: () => undefined, + options: { revision: REVISION, helperSha256: "d".repeat(64) }, + expected: /helper changed/, + }, + { + name: "host key", + configure: (harness) => { + harness.trusted.set("10.10.0.2", sshBinding("10.10.0.2", "AAAAC3NzaChangedKey")); + }, + expected: /host-key identity changed/, + }, + { + name: "GPU", + configure: (harness) => { + harness.peer.gpus[0].uuid = "GPU-SUBSTITUTED-0003"; + }, + expected: /physical dual-Station pair changed/, + }, + { + name: "rail", + configure: (harness) => { + harness.peer.rails[0].macAddress = "02:00:00:00:00:12"; + }, + expected: /physical dual-Station pair changed/, + }, + ]; + + for (const scenario of scenarios) { + const harness = new PreparationHarness(); + harness.resume = readyState(); + trustFirstRail(harness); + scenario.configure(harness); + expect( + () => prepareDualStationPair(scenario.options ?? preparationOptions(), harness.deps), + scenario.name, + ).toThrow(scenario.expected); + expect( + harness.calls.some((call) => call.startsWith("remote:")), + scenario.name, + ).toBe(false); + } + }); + + it("preserves a remote mismatch as a pinned fail-closed state", () => { + const harness = new PreparationHarness(); + trustFirstRail(harness); + harness.remoteHelperStatus.set("--apply", 1); + + expect(() => prepareDualStationPair(preparationOptions(), harness.deps)).toThrow( + /refusing single-Station fallback/, + ); + expect(harness.resume?.phase).toBe("remote-preparation"); + expect(harness.calls).not.toContain("remote:10.10.0.2:--verify"); + }); + + it("revalidates an exact managed pair and binds both controllers without workload probes", () => { + const harness = new PreparationHarness(); + trustFirstRail(harness); + + const result = prepareDualStationPair( + { ...preparationOptions(), reuseExistingManagedPair: true }, + harness.deps, + ); + expect(result.kind).toBe("ready"); + expect(harness.calls.filter((call) => call.startsWith("local:"))).toEqual([ + "local:--bind-controller", + ]); + expect(harness.calls.filter((call) => call.startsWith("remote:"))).toEqual([ + "remote:10.10.0.2:--bind-controller", + ]); + expect(harness.calls).toContain("connectivity:local"); + expect(harness.calls).toContain("connectivity:peer:10.10.0.2"); + expect(harness.resume?.phase).toBe("ready"); + }); + + it("binds only the active local controller before preparing a peer for legacy migration", () => { + const harness = new PreparationHarness(); + trustFirstRail(harness); + expect( + prepareDualStationPair( + { ...preparationOptions(), migrateLegacySingleStationHead: true }, + harness.deps, + ).kind, + ).toBe("ready"); + expect(harness.calls.filter((call) => call.startsWith("local:"))).toEqual([ + "local:--bind-controller", + ]); + expect(harness.calls.filter((call) => call.startsWith("remote:"))).toEqual([ + "remote:10.10.0.2:--check", + "remote:10.10.0.2:--apply", + "remote:10.10.0.2:--bind-controller", + "remote:10.10.0.2:--verify", + ]); + }); +}); + +describe.sequential("dual-DGX Station trust and resume-state boundaries", () => { + it("validates owner-only regular-file metadata", () => { + expect(() => + validateResumeFileMetadata( + { isFile: true, isSymbolicLink: false, uid: 1000, mode: 0o600, size: 100 }, + 1000, + ), + ).not.toThrow(); + expect(() => + validateResumeFileMetadata( + { isFile: true, isSymbolicLink: false, uid: 1001, mode: 0o600, size: 100 }, + 1000, + ), + ).toThrow(/not owned/); + expect(() => + validateResumeFileMetadata( + { isFile: true, isSymbolicLink: false, uid: 1000, mode: 0o644, size: 100 }, + 1000, + ), + ).toThrow(/0600/); + expect(() => + validateResumeFileMetadata( + { isFile: false, isSymbolicLink: true, uid: 1000, mode: 0o600, size: 100 }, + 1000, + ), + ).toThrow(/symlink/); + }); + + it("writes, fsyncs, reads, and clears canonical owner-only state", () => { + const directory = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-pair-state-")); + fs.chmodSync(directory, 0o700); + const statePath = path.join(directory, "resume.json"); + try { + const state = readyState(); + state.rails.reverse(); + writeDualStationResumeState(statePath, state); + expect(fs.statSync(statePath).mode & 0o777).toBe(0o600); + const loaded = readDualStationResumeState(statePath); + expect(loaded?.rails.map((rail) => rail.localAddress)).toEqual(["10.10.0.1", "10.10.0.5"]); + clearDualStationResumeState(statePath); + expect(fs.existsSync(statePath)).toBe(false); + } finally { + fs.rmSync(directory, { recursive: true, force: true }); + } + }); + + it("clears an owner-only SSH binding orphan even when pair state is absent", () => { + const directory = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-pair-state-")); + fs.chmodSync(directory, 0o700); + const statePath = path.join(directory, "resume.json"); + const bindingDirectory = `${statePath}.ssh-binding`; + try { + fs.mkdirSync(bindingDirectory, { mode: 0o700 }); + fs.writeFileSync(path.join(bindingDirectory, "orphan"), "binding\n", { mode: 0o600 }); + clearDualStationResumeState(statePath); + expect(fs.existsSync(bindingDirectory)).toBe(false); + } finally { + fs.rmSync(directory, { recursive: true, force: true }); + } + }); + + it("rejects malformed, permissive, symlinked, and substitution-prone state", () => { + const directory = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-pair-state-")); + fs.chmodSync(directory, 0o700); + const statePath = path.join(directory, "resume.json"); + try { + fs.writeFileSync(statePath, "not-json\n", { mode: 0o600 }); + expect(() => readDualStationResumeState(statePath)).toThrow(/malformed JSON/); + + fs.writeFileSync(statePath, `${JSON.stringify(readyState())}\n`, { mode: 0o600 }); + fs.chmodSync(statePath, 0o644); + expect(() => readDualStationResumeState(statePath)).toThrow(/0600/); + + fs.rmSync(statePath); + const target = path.join(directory, "target.json"); + fs.writeFileSync(target, `${JSON.stringify(readyState())}\n`, { mode: 0o600 }); + fs.symlinkSync(target, statePath); + expect(() => readDualStationResumeState(statePath)).toThrow(/symlink/); + + const changed = readyState(); + changed.peerGpuUuid = changed.localGpuUuid; + expect(() => parseDualStationResumeState(changed)).toThrow(/GPU identity/); + } finally { + fs.rmSync(directory, { recursive: true, force: true }); + } + }); + + it("requires an existing owner-only resume directory", () => { + const root = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-pair-state-")); + const missingState = path.join(root, "missing", "resume.json"); + try { + expect(() => readDualStationResumeState(missingState)).toThrow(/must already exist/); + fs.chmodSync(root, 0o755); + expect(() => readDualStationResumeState(path.join(root, "resume.json"))).toThrow( + /owner-only/, + ); + } finally { + fs.rmSync(root, { recursive: true, force: true }); + } + }); + + it("uses a strict noninteractive SSH argument boundary and exact-byte helper command", () => { + const args = strictStationPrepSshTransportArgs(); + expect(args).toEqual(strictStationSshTransportArgs()); + expect(args).toContain("BatchMode=yes"); + expect(args).toContain("StrictHostKeyChecking=yes"); + expect(args).toContain("VerifyHostKeyDNS=no"); + expect(args).toContain("NoHostAuthenticationForLocalhost=no"); + expect(args).toContain("ClearAllForwardings=yes"); + expect(args).toContain("ProxyCommand=none"); + expect(args).toContain("ProxyJump=none"); + const command = buildRemoteHelperCommand(HELPER_SHA256, "--apply"); + expect(command).toContain(HELPER_SHA256); + expect(command).toContain("sudo -n true"); + expect(command).toContain("NEMOCLAW_STATION_PREP_SUDO_NONINTERACTIVE=1"); + expect(command).toContain('bash "$f" --apply'); + expect(buildRemoteHelperCommand(HELPER_SHA256, "--bind-controller")).toContain( + 'bash "$f" --bind-controller', + ); + }); + + it("does not expose ambient credentials or shell-loader variables to probes and helpers", () => { + const env = buildStationPrepSubprocessEnv({ + HOME: "/home/operator", + PATH: "/usr/bin:/bin", + SSH_AUTH_SOCK: "/run/user/1000/agent", + HTTPS_PROXY: "http://proxy.example:8080", + LC_CTYPE: "en_US.UTF-8", + NVIDIA_API_KEY: "secret", + HF_TOKEN: "secret", + BASH_ENV: "/tmp/evil", + ENV: "/tmp/evil", + LD_PRELOAD: "/tmp/evil.so", + SSH_ASKPASS: "/tmp/evil", + }); + expect(env).toMatchObject({ + HOME: "/home/operator", + PATH: "/usr/bin:/bin", + SSH_AUTH_SOCK: "/run/user/1000/agent", + HTTPS_PROXY: "http://proxy.example:8080", + LC_ALL: "C", + LC_CTYPE: "en_US.UTF-8", + LANG: "C", + }); + for (const forbidden of [ + "NVIDIA_API_KEY", + "HF_TOKEN", + "BASH_ENV", + "ENV", + "LD_PRELOAD", + "SSH_ASKPASS", + ]) { + expect(env).not.toHaveProperty(forbidden); + } + }); + + it("forces every remote helper sudo call through noninteractive mode", () => { + const root = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-pair-sudo-")); + const fakeSudo = path.join(root, "sudo"); + fs.writeFileSync(fakeSudo, "#!/usr/bin/env bash\nprintf 'SUDO_ARGS=%s\\n' \"$*\"\n", { + mode: 0o700, + }); + try { + const strict = spawnSync( + "bash", + [ + "--noprofile", + "--norc", + "-c", + 'source "$STATION_HELPER" >/dev/null; NEMOCLAW_STATION_PREP_SUDO_NONINTERACTIVE=1; sudo true', + ], + { + encoding: "utf8", + env: { HOME: root, PATH: `${root}:${TEST_SYSTEM_PATH}`, STATION_HELPER }, + }, + ); + const local = spawnSync( + "bash", + ["--noprofile", "--norc", "-c", 'source "$STATION_HELPER" >/dev/null; sudo true'], + { + encoding: "utf8", + env: { HOME: root, PATH: `${root}:${TEST_SYSTEM_PATH}`, STATION_HELPER }, + }, + ); + expect(strict.status, strict.stderr).toBe(0); + expect(strict.stdout).toContain("SUDO_ARGS=-n true"); + expect(local.status, local.stderr).toBe(0); + expect(local.stdout).toContain("SUDO_ARGS=true"); + } finally { + fs.rmSync(root, { recursive: true, force: true }); + } + }); + + it("preserves matching revoked host keys in the pinned trust evidence", () => { + const root = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-pair-trust-")); + const bin = path.join(root, "bin"); + const knownHosts = path.join(root, "known_hosts"); + fs.mkdirSync(bin, { mode: 0o700 }); + fs.writeFileSync(knownHosts, "fixture\n", { mode: 0o600 }); + const ssh = path.join(bin, "ssh"); + const keygen = path.join(bin, "ssh-keygen"); + fs.writeFileSync( + ssh, + `#!/usr/bin/env bash +cat <<'EOF' +user ubuntu +hostname 10.10.0.2 +port 22 +batchmode yes +stricthostkeychecking true +verifyhostkeydns false +nohostauthenticationforlocalhost no +permitlocalcommand no +forwardagent no +forwardx11 no +forwardx11trusted no +tunnel false +updatehostkeys false +controlmaster false +controlpath none +remotecommand none +proxycommand none +proxyjump none +localcommand none +knownhostscommand none +userknownhostsfile ${knownHosts} +globalknownhostsfile none +sendenv LANG +sendenv LC_* +EOF +`, + { mode: 0o700 }, + ); + fs.writeFileSync( + keygen, + `#!/usr/bin/env bash +if [[ " $* " == *" -F "* ]]; then + printf '%s\n' '@revoked 10.10.0.2 ssh-ed25519 AAAAC3NzaRevoked' + printf '%s\n' '10.10.0.2 ssh-ed25519 AAAAC3NzaTrusted' +else + printf '%s\n' '256 ${HOST_KEY_FINGERPRINT} fixture (ED25519)' +fi +`, + { mode: 0o700 }, + ); + const originalPath = process.env.PATH; + process.env.PATH = `${bin}:${originalPath ?? ""}`; + try { + const binding = inspectPretrustedSshTarget("10.10.0.2"); + expect(binding?.knownHostsLines).toContain("@revoked 10.10.0.2 ssh-ed25519 AAAAC3NzaRevoked"); + expect(binding?.knownHostsLines).toContain("10.10.0.2 ssh-ed25519 AAAAC3NzaTrusted"); + fs.chmodSync(knownHosts, 0o666); + expect(inspectPretrustedSshTarget("10.10.0.2")).toBeNull(); + } finally { + process.env.PATH = originalPath; + fs.rmSync(root, { recursive: true, force: true }); + } + }); + + it("uses only deterministic rail candidates without trust enrollment or network discovery", () => { + const root = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-pair-command-boundary-")); + const bin = path.join(root, "bin"); + const stateDirectory = path.join(root, "state"); + const helper = path.join(root, "prepare-dgx-station-host.sh"); + const state = path.join(stateDirectory, "resume.json"); + const forbiddenLog = path.join(root, "forbidden.log"); + fs.mkdirSync(bin, { mode: 0o700 }); + fs.mkdirSync(stateDirectory, { mode: 0o700 }); + fs.writeFileSync(helper, "#!/usr/bin/env bash\nexit 0\n", { mode: 0o700 }); + fs.writeFileSync( + path.join(bin, "python3"), + `#!/usr/bin/env bash\ncat <<'JSON'\n${JSON.stringify(stationHost("local"))}\nJSON\n`, + { mode: 0o700 }, + ); + fs.writeFileSync(path.join(bin, "ssh"), "#!/usr/bin/env bash\nexit 1\n", { mode: 0o700 }); + for (const command of [ + "ssh-keyscan", + "arp-scan", + "avahi-browse", + "dns-sd", + "lldpctl", + "nmap", + "mdns-scan", + ]) { + fs.writeFileSync( + path.join(bin, command), + `#!/usr/bin/env bash\nprintf '%s\\n' ${JSON.stringify(command)} >>${JSON.stringify(forbiddenLog)}\nexit 97\n`, + { mode: 0o700 }, + ); + } + + try { + const result = spawnSync( + process.execPath, + [ + "--no-warnings", + "--experimental-strip-types", + COORDINATOR, + "--helper", + helper, + "--state", + state, + "--revision", + REVISION, + ], + { + cwd: REPO_ROOT, + encoding: "utf8", + env: { + ...process.env, + HOME: root, + PATH: `${bin}:${TEST_SYSTEM_PATH}`, + }, + timeout: 20_000, + killSignal: "SIGKILL", + }, + ); + + expect(result.status, `${result.stdout}${result.stderr}`).toBe(0); + expect(JSON.parse(result.stdout)).toMatchObject({ kind: "single-station" }); + expect(fs.existsSync(forbiddenLog)).toBe(false); + } finally { + fs.rmSync(root, { recursive: true, force: true }); + } + }); + + it("keeps forbidden discovery and trust enrollment unreachable through pair qualification", () => { + const root = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-pair-ready-boundary-")); + const bin = path.join(root, "bin"); + const stateDirectory = path.join(root, "state"); + const helper = path.join(root, "prepare-dgx-station-host.sh"); + const state = path.join(stateDirectory, "resume.json"); + const knownHosts = path.join(root, "known_hosts"); + const forbiddenLog = path.join(root, "forbidden.log"); + fs.mkdirSync(bin, { mode: 0o700 }); + fs.mkdirSync(stateDirectory, { mode: 0o700 }); + fs.writeFileSync(helper, "#!/usr/bin/env bash\nexit 0\n", { mode: 0o700 }); + fs.writeFileSync(knownHosts, "fixture\n", { mode: 0o600 }); + fs.writeFileSync(path.join(bin, "docker"), "#!/usr/bin/env bash\nexit 0\n", { + mode: 0o700, + }); + fs.writeFileSync( + path.join(bin, "python3"), + `#!/usr/bin/env bash +set -Eeuo pipefail +cat >/dev/null +if (($# == 1)); then + cat <<'JSON' +${JSON.stringify(stationHost("local"))} +JSON +else + cat <<'JSON' +${stationConnectivity("local")} +JSON +fi +`, + { mode: 0o700 }, + ); + fs.writeFileSync( + path.join(bin, "ssh-keygen"), + `#!/usr/bin/env bash +set -Eeuo pipefail +if [[ " $* " == *" -F 10.10.0.2 "* ]]; then + printf '%s\n' '10.10.0.2 ssh-ed25519 ${HOST_KEY_DATA}' + exit 0 +fi +if [[ " $* " == *" -F "* ]]; then + exit 1 +fi +printf '%s\n' '256 ${HOST_KEY_FINGERPRINT} fixture (ED25519)' +`, + { mode: 0o700 }, + ); + fs.writeFileSync( + path.join(bin, "ssh"), + `#!/usr/bin/env bash +set -Eeuo pipefail +if [[ " $* " == *" -G "* ]]; then + target='' + for value in "$@"; do target="$value"; done + [[ "$target" == '10.10.0.2' ]] || exit 1 + cat <<'EOF' +user ubuntu +hostname 10.10.0.2 +port 22 +batchmode yes +stricthostkeychecking true +verifyhostkeydns false +nohostauthenticationforlocalhost no +permitlocalcommand no +forwardagent no +forwardx11 no +forwardx11trusted no +tunnel false +updatehostkeys false +controlmaster false +controlpath none +remotecommand none +proxycommand none +proxyjump none +localcommand none +knownhostscommand none +userknownhostsfile ${knownHosts} +globalknownhostsfile none +sendenv LANG +sendenv LC_* +EOF + exit 0 +fi +if [[ " $* " == *'python3 - enp1s0f0np0'* ]]; then + cat >/dev/null + cat <<'JSON' +${stationConnectivity("peer")} +JSON + exit 0 +fi +if [[ " $* " == *'python3 -'* ]]; then + cat >/dev/null + cat <<'JSON' +${JSON.stringify(stationHost("peer"))} +JSON + exit 0 +fi +if [[ " $* " == *'prepare-dgx-station-host.sh'* ]]; then + cat >/dev/null + exit 0 +fi +exit 96 +`, + { mode: 0o700 }, + ); + for (const command of [ + "ssh-keyscan", + "arp-scan", + "avahi-browse", + "dns-sd", + "lldpctl", + "nmap", + "mdns-scan", + ]) { + fs.writeFileSync( + path.join(bin, command), + `#!/usr/bin/env bash\nprintf '%s\\n' ${JSON.stringify(command)} >>${JSON.stringify(forbiddenLog)}\nexit 97\n`, + { mode: 0o700 }, + ); + } + + try { + const result = spawnSync( + process.execPath, + [ + "--no-warnings", + "--experimental-strip-types", + COORDINATOR, + "--helper", + helper, + "--state", + state, + "--revision", + REVISION, + ], + { + cwd: REPO_ROOT, + encoding: "utf8", + env: { + ...process.env, + HOME: root, + PATH: `${bin}:${TEST_SYSTEM_PATH}`, + }, + timeout: 20_000, + killSignal: "SIGKILL", + }, + ); + + expect(result.status, `${result.stdout}${result.stderr}`).toBe(0); + expect(JSON.parse(result.stdout)).toMatchObject({ kind: "ready", peerTarget: "10.10.0.2" }); + expect(fs.existsSync(forbiddenLog)).toBe(false); + } finally { + fs.rmSync(root, { recursive: true, force: true }); + } + }); +}); + +describe("dual-DGX Station installer handoff", () => { + it("changes the Ultra topology only after the coordinator returns a validated ready pair", () => { + const argsFile = path.join(os.tmpdir(), `nemoclaw-pair-args-${process.pid}-${Date.now()}`); + const { result, output, home } = runInstallerBody( + ` +node() { + if [[ "\${1:-}" == "--no-warnings" ]]; then + printf '%s\n' "$PAIR_RESULT" + printf '%s\n' "$*" >"$PAIR_ARGS_FILE" + return 0 + fi + command node "$@" +} +station_installer_revision() { printf '%s' "$PAIR_REVISION"; } +station_dual_pair_resume_pending() { return 0; } +_SELECTED_EXPRESS_PLATFORM='DGX Station' +_STATION_EXPRESS_MODEL_WAS_EXPLICIT=0 +NEMOCLAW_VLLM_MODEL='nemotron-3-ultra-550b-a55b' +NEMOCLAW_MODEL='nvidia/nemotron-3-ultra-550b-a55b' +unset NEMOCLAW_DGX_STATION_PEER NEMOCLAW_DGX_STATION_SSH_BINDING +ensure_station_express_pair +printf 'RESULT peer=%s model=%s selector=%s binding=%s\n' "\${NEMOCLAW_DGX_STATION_PEER:-}" "\${NEMOCLAW_MODEL:-}" "\${NEMOCLAW_VLLM_MODEL:-}" "\${NEMOCLAW_DGX_STATION_SSH_BINDING:-}" +`, + { + PAIR_ARGS_FILE: argsFile, + PAIR_RESULT: coordinatorResult("ready"), + PAIR_REVISION: REVISION, + }, + ); + try { + expect(result.status, output).toBe(0); + expect(output).toContain( + "RESULT peer=10.10.0.2 model=nemotron-ultra selector=nemotron-3-ultra-550b-a55b", + ); + const expectedToken = JSON.parse(coordinatorResult("ready")).sshBinding; + expect(output).toContain(`binding=${expectedToken}`); + const args = fs.readFileSync(argsFile, "utf8"); + expect(args).toContain("--helper"); + expect(args).toContain("prepare-dgx-station-host.sh"); + expect(args).toContain(`--revision ${REVISION}`); + expect(args).not.toContain("--explicit-peer"); + } finally { + fs.rmSync(argsFile, { force: true }); + fs.rmSync(home, { recursive: true, force: true }); + } + }); + + it("passes explicit peer and exact managed-pair reuse without shell interpolation", () => { + const argsFile = path.join(os.tmpdir(), `nemoclaw-pair-args-${process.pid}-${Date.now()}`); + const explicitPeer = "ubuntu@station-b"; + const { result, output, home } = runInstallerBody( + ` +node() { + if [[ "\${1:-}" == "--no-warnings" ]]; then + printf '%s\n' "$PAIR_RESULT" + printf '%s\n' "$*" >"$PAIR_ARGS_FILE" + return 0 + fi + command node "$@" +} +station_installer_revision() { printf '%s' "$PAIR_REVISION"; } +station_dual_pair_resume_pending() { return 0; } +_SELECTED_EXPRESS_PLATFORM='DGX Station' +_STATION_EXPRESS_MODEL_WAS_EXPLICIT=0 +_STATION_EXPRESS_DEFERRED_MANAGED_PAIR=1 +NEMOCLAW_DGX_STATION_PEER="$PAIR_PEER" +NEMOCLAW_VLLM_MODEL='nemotron-3-ultra-550b-a55b' +NEMOCLAW_MODEL='nvidia/nemotron-3-ultra-550b-a55b' +unset NEMOCLAW_DGX_STATION_SSH_BINDING +ensure_station_express_pair +printf 'RESULT peer=%s model=%s binding=%s\n' "$NEMOCLAW_DGX_STATION_PEER" "$NEMOCLAW_MODEL" "$NEMOCLAW_DGX_STATION_SSH_BINDING" +`, + { + PAIR_ARGS_FILE: argsFile, + PAIR_PEER: explicitPeer, + PAIR_RESULT: coordinatorResult("ready", explicitPeer), + PAIR_REVISION: REVISION, + }, + ); + try { + expect(result.status, output).toBe(0); + expect(output).toContain(`RESULT peer=${explicitPeer} model=nemotron-ultra`); + expect(output).toContain(`binding=${JSON.parse(coordinatorResult("ready")).sshBinding}`); + const args = fs.readFileSync(argsFile, "utf8"); + expect(args).toContain(`--explicit-peer ${explicitPeer}`); + expect(args).toContain("--reuse-existing-managed-pair"); + } finally { + fs.rmSync(argsFile, { force: true }); + fs.rmSync(home, { recursive: true, force: true }); + } + }); + + it("preserves the existing single-Station default when discovery finds no peer", () => { + const { result, output, home } = runInstallerBody( + ` +node() { + if [[ "\${1:-}" == "--no-warnings" ]]; then + printf '%s\n' '{"kind":"single-station","reason":"no pretrusted reciprocal peer"}' + return 0 + fi + command node "$@" +} +station_installer_revision() { printf '%s' "$PAIR_REVISION"; } +station_express_resume_generation() { printf '0123456789abcdef0123456789abcdef'; } +_SELECTED_EXPRESS_PLATFORM='DGX Station' +_STATION_EXPRESS_MODEL_WAS_EXPLICIT=0 +NEMOCLAW_VLLM_MODEL='nemotron-3-ultra-550b-a55b' +NEMOCLAW_MODEL='nvidia/nemotron-3-ultra-550b-a55b' +unset NEMOCLAW_DGX_STATION_PEER +NEMOCLAW_DGX_STATION_SSH_BINDING='stale' +ensure_station_express_pair +printf 'RESULT peer=%s model=%s selector=%s binding=%s\n' "\${NEMOCLAW_DGX_STATION_PEER:-}" "\${NEMOCLAW_MODEL:-}" "\${NEMOCLAW_VLLM_MODEL:-}" "\${NEMOCLAW_DGX_STATION_SSH_BINDING:-}" +`, + { PAIR_REVISION: REVISION }, + ); + try { + expect(result.status, output).toBe(0); + expect(output).toContain("No trusted reciprocal dual-DGX Station pair was detected"); + expect(output).toContain( + "RESULT peer= model=nvidia/nemotron-3-ultra-550b-a55b selector=nemotron-3-ultra-550b-a55b binding=", + ); + } finally { + fs.rmSync(home, { recursive: true, force: true }); + } + }); + + it("rejects an explicit peer combined with an explicit non-dual model", () => { + const { result, output, home } = runInstallerBody( + ` +node() { printf 'COORDINATOR_CALLED\n'; return 0; } +_SELECTED_EXPRESS_PLATFORM='DGX Station' +_STATION_EXPRESS_MODEL_WAS_EXPLICIT=1 +NEMOCLAW_VLLM_MODEL='deepseek-v4-flash' +NEMOCLAW_DGX_STATION_PEER='ubuntu@station-b' +ensure_station_express_pair +`, + ); + try { + expect(result.status, output).not.toBe(0); + expect(output).toContain( + "NEMOCLAW_DGX_STATION_PEER requires the DGX Station dual-serving model", + ); + expect(output).not.toContain("COORDINATOR_CALLED"); + } finally { + fs.rmSync(home, { recursive: true, force: true }); + } + }); + + it("rejects conflicting peer and model selections before local host preparation", () => { + const { result, output, home } = runInstallerBody( + ` +maybe_offer_express_install() { + _SELECTED_EXPRESS_PLATFORM='DGX Station' + _STATION_EXPRESS_MODEL_WAS_EXPLICIT=1 + NEMOCLAW_VLLM_MODEL='deepseek-v4-flash' + NEMOCLAW_DGX_STATION_PEER='ubuntu@station-b' +} +ensure_station_express_host() { printf 'LOCAL_HELPER_CALLED\n'; } +ensure_docker() { printf 'DOCKER_CALLED\n'; } +ensure_openshell_build_deps() { printf 'BUILD_DEPS_CALLED\n'; } +prepare_installer_host +`, + ); + try { + expect(result.status, output).not.toBe(0); + expect(output).toContain( + "NEMOCLAW_DGX_STATION_PEER requires the DGX Station dual-serving model", + ); + expect(output).not.toContain("LOCAL_HELPER_CALLED"); + expect(output).not.toContain("DOCKER_CALLED"); + } finally { + fs.rmSync(home, { recursive: true, force: true }); + } + }); + + it("propagates peer exit 10 with manual reboot and exact-revision rerun guidance", () => { + const { result, output, home } = runInstallerBody( + ` +node() { + if [[ "\${1:-}" == "--no-warnings" ]]; then + printf '%s\n' "$PAIR_RESULT" + return 10 + fi + command node "$@" +} +station_installer_revision() { printf '%s' "$PAIR_REVISION"; } +save_station_express_resume() { _STATION_EXPRESS_RESUME_REVISION="$PAIR_REVISION" _STATION_EXPRESS_RESUME_AGENT='openclaw' _STATION_EXPRESS_RESUME_SANDBOX='my-assistant' _STATION_EXPRESS_RESUME_POLICY_TIER='balanced' _STATION_EXPRESS_RESUME_GATEWAY_PORT='8080' _STATION_EXPRESS_RESUME_DASHBOARD_PORT='18789' _STATION_EXPRESS_RESUME_VLLM_PORT='8000'; printf 'SAVED_EXPRESS_RESUME\n'; } +station_dual_pair_resume_pending() { return 0; } +_SELECTED_EXPRESS_PLATFORM='DGX Station' +_STATION_EXPRESS_MODEL_WAS_EXPLICIT=0 +_STATION_INSTALL_MODE='express' +NEMOCLAW_VLLM_MODEL='nemotron-3-ultra-550b-a55b' +NEMOCLAW_MODEL='nvidia/nemotron-3-ultra-550b-a55b' +unset NEMOCLAW_DGX_STATION_PEER +ensure_station_express_pair +`, + { PAIR_RESULT: coordinatorResult("reboot-required"), PAIR_REVISION: REVISION }, + ); + try { + expect(result.status, output).toBe(10); + expect(output).toContain("SAVED_EXPRESS_RESUME"); + expect(output).toContain("requires a manual reboot"); + expect(output).toContain(`NEMOCLAW_INSTALL_TAG=${REVISION}`); + expect(output).not.toMatch(/reboot.*-[a-z]*f/i); + } finally { + fs.rmSync(home, { recursive: true, force: true }); + } + }); + + it("preserves companion resume state when coordinator failure follows pair-state publication", () => { + const { result, output, home } = runInstallerBody( + ` +node() { + if [[ "\${1:-}" == "--no-warnings" ]]; then + while (( $# > 0 )); do + if [[ "$1" == "--state" ]]; then + printf '{}\n' >"$2" + break + fi + shift + done + return 1 + fi + command node "$@" +} +station_installer_revision() { printf '%s' "$PAIR_REVISION"; } +_SELECTED_EXPRESS_PLATFORM='DGX Station' +_STATION_EXPRESS_MODEL_WAS_EXPLICIT=0 +_STATION_INSTALL_MODE='express' +station_express_resume_generation() { printf '0123456789abcdef0123456789abcdef'; } +NEMOCLAW_VLLM_MODEL='nemotron-3-ultra-550b-a55b' +NEMOCLAW_MODEL='nvidia/nemotron-3-ultra-550b-a55b' +unset NEMOCLAW_DGX_STATION_PEER +ensure_station_express_pair +`, + { PAIR_REVISION: REVISION }, + ); + try { + expect(result.status, output).not.toBe(0); + expect(output).toContain("Dual DGX Station preparation failed"); + expect(fs.existsSync(path.join(home, ".nemoclaw", "station-dual-pair-resume.json"))).toBe( + true, + ); + expect(fs.readFileSync(path.join(home, ".nemoclaw", "station-express-resume"), "utf8")).toBe( + `revision=${REVISION}\nmodel=nemotron-3-ultra-550b-a55b\ngeneration=0123456789abcdef0123456789abcdef\n` + + "agent=openclaw\nsandbox=my-assistant\npolicy_tier=balanced\n" + + "gateway_port=8080\ndashboard_port=18789\nvllm_port=8000\nmode=express\n", + ); + } finally { + fs.rmSync(home, { recursive: true, force: true }); + } + }); + + it("defers host preparation only for a complete running managed dual-head candidate", () => { + const valid = [ + "/nemoclaw-vllm", + "true", + "true", + "head", + "1", + "c".repeat(64), + "d".repeat(64), + "e".repeat(64), + "f".repeat(32), + ].join(" "); + const accepted = runInstallerBody( + ` +command_exists() { return 0; } +docker() { printf '%s\n' "$DOCKER_INSPECTION"; } +station_managed_dual_head_running +`, + { DOCKER_INSPECTION: valid }, + ); + const malformed = runInstallerBody( + ` +command_exists() { return 0; } +docker() { printf '%s\n' "$DOCKER_INSPECTION"; } +station_managed_dual_head_running +`, + { DOCKER_INSPECTION: valid.replace(" head ", " worker ") }, + ); + try { + expect(accepted.result.status, accepted.output).toBe(0); + expect(malformed.result.status, malformed.output).not.toBe(0); + } finally { + fs.rmSync(accepted.home, { recursive: true, force: true }); + fs.rmSync(malformed.home, { recursive: true, force: true }); + } + }); + + it("applies Station preparation to an explicitly selected managed-vLLM provider", () => { + const { result, output, home } = runInstallerBody( + ` +detect_express_platform() { printf 'DGX Station'; } +NON_INTERACTIVE='' +NEMOCLAW_NO_EXPRESS='' +NEMOCLAW_PROVIDER='install-vllm' +unset NEMOCLAW_VLLM_MODEL +maybe_offer_express_install +printf 'RESULT selected=%s provider=%s selector=%s\n' "$_SELECTED_EXPRESS_PLATFORM" "$NEMOCLAW_PROVIDER" "\${NEMOCLAW_VLLM_MODEL:-}" +`, + ); + try { + expect(result.status, output).toBe(0); + expect(output).toContain("explicitly selected managed-vLLM provider"); + expect(output).toContain("RESULT selected=DGX Station provider=install-vllm selector="); + } finally { + fs.rmSync(home, { recursive: true, force: true }); + } + }); +}); diff --git a/test/install-station-resume-cleanup.test.ts b/test/install-station-resume-cleanup.test.ts new file mode 100644 index 0000000000..4cf42a4cb8 --- /dev/null +++ b/test/install-station-resume-cleanup.test.ts @@ -0,0 +1,59 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { spawnSync } from "node:child_process"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; + +import { describe, expect, it } from "vitest"; + +import { INSTALLER_PAYLOAD, TEST_SYSTEM_PATH } from "./helpers/installer-sourced-env"; + +describe("DGX Station installer resume cleanup", () => { + it("preserves pair and SSH-binding state when interactive host preflight skips onboarding", () => { + const home = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-station-resume-cleanup-")); + const result = spawnSync( + "bash", + [ + "--noprofile", + "--norc", + "-c", + ` +source "$INSTALLER_UNDER_TEST" >/dev/null +pair_state="$HOME/.nemoclaw/station-dual-pair-resume.json" +binding_state="\${pair_state}.ssh-binding" +mkdir -p "$binding_state" +printf '{}\n' >"$pair_state" +printf 'binding\n' >"$binding_state/token" +printf 'resume\n' >"$HOME/.nemoclaw/station-express-resume" +_SELECTED_EXPRESS_PLATFORM='DGX Station' +ONBOARD_RAN=false +clear_station_resume_after_completed_onboarding +printf 'PAIR=%s BINDING=%s EXPRESS=%s\n' \ + "$([ -f "$pair_state" ] && printf present)" \ + "$([ -f "$binding_state/token" ] && printf present)" \ + "$([ -f "$HOME/.nemoclaw/station-express-resume" ] && printf present)" +`, + ], + { + cwd: path.resolve(import.meta.dirname, ".."), + encoding: "utf8", + env: { + ...process.env, + HOME: home, + INSTALLER_UNDER_TEST: INSTALLER_PAYLOAD, + PATH: TEST_SYSTEM_PATH, + }, + }, + ); + + try { + const output = `${result.stdout}${result.stderr}`; + expect(result.status, output).toBe(0); + expect(output).toContain("PAIR=present BINDING=present EXPRESS=present"); + } finally { + fs.rmSync(home, { recursive: true, force: true }); + } + }); +}); diff --git a/test/install-station-vllm-continuation.test.ts b/test/install-station-vllm-continuation.test.ts index df0918f19c..87a0a9c202 100644 --- a/test/install-station-vllm-continuation.test.ts +++ b/test/install-station-vllm-continuation.test.ts @@ -47,6 +47,25 @@ function runStationPreparationSourced(body: string) { } describe("installer Station Local vLLM continuation", () => { + it("normalizes a trailing slash in HOME before selecting continuation state", () => { + const home = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-vllm-resume-home-slash-")); + try { + const { result, output } = runInstallerSourced( + ` +load_station_vllm_conflict_helpers +printf 'resume=%s\\n' "$(station_local_vllm_resume_file)" +`, + `${home}/`, + ); + + expect(result.status, output).toBe(0); + expect(output).toContain(`resume=${home}/.nemoclaw/station-local-vllm-resume`); + expect(output).not.toContain(`${home}//.nemoclaw`); + } finally { + fs.rmSync(home, { recursive: true, force: true }); + } + }); + it("ignores diagnostic processes that mention vLLM", () => { const { result, output } = runInstallerSourced(` load_station_vllm_conflict_helpers diff --git a/test/prepare-dual-dgx-station-resume-state.test.ts b/test/prepare-dual-dgx-station-resume-state.test.ts new file mode 100644 index 0000000000..bd4c5a9e69 --- /dev/null +++ b/test/prepare-dual-dgx-station-resume-state.test.ts @@ -0,0 +1,224 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { spawnSync } from "node:child_process"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; +import { expect, it, vi } from "vitest"; +import type { DualStationResumeState } from "../scripts/lib/dgx-station-peer.mts"; +import { + CONNECTIVITY_PROBE, + clearDualStationResumeState, + writeDualStationResumeState, +} from "../scripts/prepare-dual-dgx-station.mts"; +import { TEST_SYSTEM_PATH } from "./helpers/installer-sourced-env"; + +function readyState(): DualStationResumeState { + return { + schemaVersion: 1, + revision: "a".repeat(40), + helperSha256: "b".repeat(64), + phase: "ready", + peerTarget: "10.10.0.2", + hostKeyDigest: "c".repeat(64), + localGpuUuid: "GPU-LOCAL-0001", + peerGpuUuid: "GPU-PEER-0002", + rails: [ + { + localAddress: "10.10.0.1", + localMac: "02:00:00:00:00:01", + peerAddress: "10.10.0.2", + peerMac: "02:00:00:00:00:02", + }, + { + localAddress: "10.10.0.5", + localMac: "02:00:00:00:00:05", + peerAddress: "10.10.0.6", + peerMac: "02:00:00:00:00:06", + }, + ], + }; +} + +function captureThrown(operation: () => void): unknown { + try { + operation(); + } catch (error) { + return error; + } + return null; +} + +it("accepts direct routes from iproute2 JSON that omits redundant filtered fields", () => { + const root = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-pair-iproute2-")); + const bin = path.join(root, "bin"); + fs.mkdirSync(bin, { mode: 0o700 }); + fs.writeFileSync( + path.join(bin, "ip"), + `#!/usr/bin/env bash +case "$*" in + "-j route get 10.10.0.2 from 10.10.0.1 oif rail0") + printf '%s\\n' '[{"dst":"10.10.0.2","from":"10.10.0.1","dev":"rail0","flags":[]}]' + ;; + "-j route get 10.10.0.6 from 10.10.0.5 oif rail1") + printf '%s\\n' '[{"dst":"10.10.0.6","from":"10.10.0.5","dev":"rail1","flags":[]}]' + ;; + "-j route show exact 10.10.0.0/30") + printf '%s\\n' '[{"dst":"10.10.0.0/30","dev":"rail0","scope":"link","prefsrc":"10.10.0.1"}]' + ;; + "-j route show exact 10.10.0.4/30") + printf '%s\\n' '[{"dst":"10.10.0.4/30","dev":"rail1","scope":"link","prefsrc":"10.10.0.5"}]' + ;; + "-j neighbor show to 10.10.0.2") + printf '%s\\n' '[{"dst":"10.10.0.2","dev":"rail0","lladdr":"02:00:00:00:00:02","state":["STALE"]}]' + ;; + "-j neighbor show to 10.10.0.6") + printf '%s\\n' '[{"dst":"10.10.0.6","dev":"rail1","lladdr":"02:00:00:00:00:06","state":["REACHABLE"]}]' + ;; + *) + exit 97 + ;; +esac +`, + { mode: 0o700 }, + ); + fs.writeFileSync(path.join(bin, "ping"), "#!/usr/bin/env bash\nexit 0\n", { + mode: 0o700, + }); + + try { + const result = spawnSync( + "python3", + ["-", "rail0", "10.10.0.1", "10.10.0.2", "rail1", "10.10.0.5", "10.10.0.6"], + { + encoding: "utf8", + env: { ...process.env, PATH: `${bin}:${TEST_SYSTEM_PATH}` }, + input: CONNECTIVITY_PROBE, + }, + ); + + expect(result.status, `${result.stdout}${result.stderr}`).toBe(0); + expect(JSON.parse(result.stdout)).toEqual({ + schemaVersion: 1, + checks: [ + expect.objectContaining({ + netdev: "rail0", + routeDevice: "rail0", + routeSource: "10.10.0.1", + routeScope: "link", + peerMac: "02:00:00:00:00:02", + peerNeighborState: "STALE", + jumboPing: true, + }), + expect.objectContaining({ + netdev: "rail1", + routeDevice: "rail1", + routeSource: "10.10.0.5", + routeScope: "link", + peerMac: "02:00:00:00:00:06", + peerNeighborState: "REACHABLE", + jumboPing: true, + }), + ], + }); + } finally { + fs.rmSync(root, { recursive: true, force: true }); + } +}); + +it("preserves the primary resume-state write error when temporary cleanup also fails", () => { + const directory = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-pair-state-cleanup-")); + fs.chmodSync(directory, 0o700); + const statePath = path.join(directory, "resume.json"); + const primaryError = new Error("primary write failure"); + const cleanupError = Object.assign(new Error("temporary unlink failure"), { code: "EACCES" }); + const writeSpy = vi.spyOn(fs, "writeFileSync").mockImplementationOnce(() => { + throw primaryError; + }); + const unlinkSpy = vi.spyOn(fs, "unlinkSync").mockImplementation(() => { + throw cleanupError; + }); + try { + expect(captureThrown(() => writeDualStationResumeState(statePath, readyState()))).toBe( + primaryError, + ); + expect(unlinkSpy).toHaveBeenCalledTimes(1); + } finally { + writeSpy.mockRestore(); + unlinkSpy.mockRestore(); + fs.rmSync(directory, { recursive: true, force: true }); + } +}); + +it("surfaces a non-ENOENT temporary cleanup error when the state write succeeded", () => { + const directory = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-pair-state-cleanup-")); + fs.chmodSync(directory, 0o700); + const statePath = path.join(directory, "resume.json"); + const cleanupError = Object.assign(new Error("temporary unlink failure"), { code: "EACCES" }); + const unlinkSpy = vi.spyOn(fs, "unlinkSync").mockImplementation(() => { + throw cleanupError; + }); + try { + expect(captureThrown(() => writeDualStationResumeState(statePath, readyState()))).toBe( + cleanupError, + ); + expect(fs.existsSync(statePath)).toBe(true); + } finally { + unlinkSpy.mockRestore(); + fs.rmSync(directory, { recursive: true, force: true }); + } +}); + +it("fsyncs the parent directory after deleting resume state", () => { + const directory = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-pair-state-clear-")); + fs.chmodSync(directory, 0o700); + const statePath = path.join(directory, "resume.json"); + writeDualStationResumeState(statePath, readyState()); + const openSpy = vi.spyOn(fs, "openSync"); + const fsyncSpy = vi.spyOn(fs, "fsyncSync"); + const closeSpy = vi.spyOn(fs, "closeSync"); + try { + clearDualStationResumeState(statePath); + + const directoryOpenIndex = openSpy.mock.calls.findIndex(([target]) => target === directory); + expect(directoryOpenIndex).toBeGreaterThanOrEqual(0); + const directoryFd = openSpy.mock.results[directoryOpenIndex]?.value; + expect(fsyncSpy).toHaveBeenCalledWith(directoryFd); + expect(closeSpy).toHaveBeenCalledWith(directoryFd); + expect(fs.existsSync(statePath)).toBe(false); + } finally { + openSpy.mockRestore(); + fsyncSpy.mockRestore(); + closeSpy.mockRestore(); + fs.rmSync(directory, { recursive: true, force: true }); + } +}); + +it("preserves the directory fsync error when closing the directory also fails", () => { + const directory = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-pair-state-clear-")); + fs.chmodSync(directory, 0o700); + const statePath = path.join(directory, "resume.json"); + writeDualStationResumeState(statePath, readyState()); + const primaryError = new Error("directory fsync failure"); + const cleanupError = new Error("directory close failure"); + const closeSync = fs.closeSync.bind(fs); + const fsyncSpy = vi.spyOn(fs, "fsyncSync").mockImplementationOnce(() => { + throw primaryError; + }); + const closeSpy = vi + .spyOn(fs, "closeSync") + .mockImplementationOnce((fd) => closeSync(fd)) + .mockImplementationOnce((fd) => { + closeSync(fd); + throw cleanupError; + }); + try { + expect(captureThrown(() => clearDualStationResumeState(statePath))).toBe(primaryError); + expect(fs.existsSync(statePath)).toBe(false); + } finally { + fsyncSpy.mockRestore(); + closeSpy.mockRestore(); + fs.rmSync(directory, { recursive: true, force: true }); + } +}); diff --git a/test/starter-prompt-docs.test.ts b/test/starter-prompt-docs.test.ts index bb2feb74f7..b42b1ff5b6 100644 --- a/test/starter-prompt-docs.test.ts +++ b/test/starter-prompt-docs.test.ts @@ -30,7 +30,7 @@ const repoRoot = path.resolve(__dirname, ".."); const starterPromptMarkdownSource = path.join(repoRoot, "docs", "resources", "starter-prompt.md"); // CI resolves this Git commit and byte-compares its prompt-asset blobs with // the local files. The digests independently assert those same immutable bytes. -const promptAssetRevision = "ebf29b024c7a24e69ab8386b68e426d31a4c6821"; +const promptAssetRevision = "bf46e62f901825f19e570c17f8c870a0eae04fbc"; type PromptAsset = { path: string; @@ -53,7 +53,7 @@ const promptAssets = { ), dgxStation: definePromptAsset( "docs/resources/prompt-assets/dgx-station.md", - "9f506ece27dcda3cf85735d7c6a80846a7727696b825cf8ca161334ac6925c1f", // gitleaks:allow -- pinned prompt-asset SHA-256 + "9b620ffe898847718fc25e230039c0bf07dd574e26a27479c45297eb17d2cfa4", // gitleaks:allow -- pinned prompt-asset SHA-256 ), windowsWsl: definePromptAsset( "docs/resources/prompt-assets/windows-wsl.md", @@ -817,10 +817,10 @@ describe("starter prompt docs CTA", () => { expect(sparkSource).toContain("nvidia/Qwen3.6-35B-A3B-NVFP4"); expect(sparkSource).toContain("Leave `NEMOCLAW_VLLM_MODEL` and `NEMOCLAW_MODEL` unset"); expect(stationSource).toContain("`nemotron-3-ultra-550b-a55b`"); - expect(stationSource).toContain("`nvidia/nemotron-3-ultra-550b-a55b`"); + expect(stationSource).toContain("`nemotron-ultra`"); expect(stationSource).toContain("`deepseek-v4-flash`"); expect(stationSource).toContain("`deepseek-ai/DeepSeek-V4-Flash`"); - expect(stationSource).toContain("Nemotron 3 Ultra 550B, the ordinary installer default"); + expect(stationSource).toContain("Automatic pair selection"); expect(stationSource).toContain( "DeepSeek V4 Flash, the explicit `--station-deepseek` override", ); @@ -837,8 +837,12 @@ describe("starter prompt docs CTA", () => { expect(stationSource).toContain( "Do not run `scripts/prepare-dgx-station-host.sh --check`, `--verify`, or `--apply` separately", ); - expect(stationSource).toContain("For Nemotron Ultra, run the ordinary installer without"); + expect(stationSource).toContain( + "For automatic pair selection, run the ordinary installer without", + ); expect(stationSource).toContain("For DeepSeek, pass `--station-deepseek`"); + expect(stationSource).toContain("TCP port `6379`"); + expect(stationSource).toContain("shared `/24`"); for (const environmentName of [ "NEMOCLAW_PROVIDER", "NEMOCLAW_VLLM_MODEL", @@ -952,7 +956,7 @@ describe("starter prompt docs CTA", () => { ); const stationConfirmationIndex = stationSource.indexOf("Choices:"); const stationDefaultIndex = stationSource.indexOf( - "For Nemotron Ultra, run the ordinary installer without", + "For automatic pair selection, run the ordinary installer without", ); const stationOverrideIndex = stationSource.indexOf("For DeepSeek, pass `--station-deepseek`"); expect(stationDisclosureIndex).toBeGreaterThan(-1); diff --git a/test/support/vllm-ownership-test-support.ts b/test/support/vllm-ownership-test-support.ts new file mode 100644 index 0000000000..a648f29810 --- /dev/null +++ b/test/support/vllm-ownership-test-support.ts @@ -0,0 +1,30 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +interface DockerCaptureTestOptions { + readonly env?: NodeJS.ProcessEnv; +} + +type OwnershipResponse = () => string; + +/** Route every ownership inspection explicitly and fail closed on fixture drift. */ +export function createStrictVllmOwnershipCapture( + ownershipResponses: readonly OwnershipResponse[], + ambientContext: string, + fallback: (command: string) => string, +): (args: readonly string[], options?: DockerCaptureTestOptions) => string { + const queue = [...ownershipResponses]; + const allowedContexts = new Set(["default", ambientContext]); + return (args, options) => { + const command = args[0] ?? ""; + if (command !== "container") return fallback(command); + + const context = options?.env?.DOCKER_CONTEXT ?? "ambient"; + if (!allowedContexts.has(context)) { + throw new Error(`Unexpected vLLM ownership inspection context: ${context}`); + } + const response = queue.shift(); + if (!response) throw new Error("Unexpected extra vLLM ownership inspection"); + return response(); + }; +} diff --git a/test/support/vllm-station-cluster-lifecycle-test-support.ts b/test/support/vllm-station-cluster-lifecycle-test-support.ts new file mode 100644 index 0000000000..ad30ad30eb --- /dev/null +++ b/test/support/vllm-station-cluster-lifecycle-test-support.ts @@ -0,0 +1,387 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { AsyncLocalStorage } from "node:async_hooks"; +import { vi } from "vitest"; +import { + DUAL_STATION_VLLM_RUNTIME, + type DualStationVllmPlan, +} from "../../src/lib/inference/vllm-station-cluster"; +import { + DUAL_STATION_VLLM_API_KEY_FINGERPRINT_LABEL, + DUAL_STATION_VLLM_CLUSTER_LABEL, + DUAL_STATION_VLLM_ENDPOINT_LABEL, + DUAL_STATION_VLLM_GPU_LABEL, + DUAL_STATION_VLLM_GPU_SMOKE_LABEL, + DUAL_STATION_VLLM_HEAD_CONTAINER_NAME, + DUAL_STATION_VLLM_LAUNCH_CONTRACT_LABEL, + DUAL_STATION_VLLM_LAUNCH_SCHEMA_LABEL, + DUAL_STATION_VLLM_MANAGED_LABEL, + DUAL_STATION_VLLM_ROLE_LABEL, + DUAL_STATION_VLLM_TRANSACTION_LABEL, + DUAL_STATION_VLLM_WORKER_CONTAINER_NAME, + type DualStationDockerOptions, + type DualStationLegacyMigration, + type DualStationVllmLifecycleDeps, + type StartDualStationVllmResult, +} from "../../src/lib/inference/vllm-station-cluster-lifecycle"; +import type { DualStationSshBinding } from "../../src/lib/inference/vllm-station-ssh-binding"; + +export type LifecycleFakeContainer = { + id: string; + name: string; + state: string; + image: string; + labels: Record; +}; + +export type LifecycleHarnessOptions = { + failRole?: "head" | "worker"; + invalidIdRole?: "head" | "worker"; + failSmokeTarget?: "local" | "peer"; + failSmokeCleanupTarget?: "local" | "peer"; + missingImageTarget?: "local" | "peer"; + smokeGpuOutput?: Partial>; + lateCreateRole?: "head" | "worker"; + failedRoleForeignTransaction?: "head" | "worker"; + failFinalInspectionRole?: "head" | "worker"; + failLegacyBackupRemoval?: boolean; +}; + +type LifecycleHarnessFixture = { + apiKey: string; + fakeContainer: ( + role: "head" | "worker", + overrides?: Partial, + ) => LifecycleFakeContainer; + headSmokeId: string; + legacyHeadId: string; + plan: () => DualStationVllmPlan; + workerSmokeId: string; +}; + +export function dualStationDockerValues(args: readonly string[], flag: string): string[] { + return args.flatMap((arg, index) => + arg === flag && index < args.length - 1 ? [args[index + 1]] : [], + ); +} + +export function requireLegacyMigration( + result: StartDualStationVllmResult, +): DualStationLegacyMigration { + if (!result.ok || !result.legacyMigration) { + throw new Error("expected legacy migration handle"); + } + return result.legacyMigration; +} + +function raise(message: string): never { + throw new Error(message); +} + +function row(container: LifecycleFakeContainer): string { + return [ + container.id, + container.name, + container.state, + container.image, + container.labels[DUAL_STATION_VLLM_MANAGED_LABEL] ?? "", + container.labels[DUAL_STATION_VLLM_ROLE_LABEL] ?? "", + container.labels[DUAL_STATION_VLLM_ENDPOINT_LABEL] ?? "", + container.labels[DUAL_STATION_VLLM_CLUSTER_LABEL] ?? "", + container.labels[DUAL_STATION_VLLM_GPU_LABEL] ?? "", + container.labels[DUAL_STATION_VLLM_LAUNCH_SCHEMA_LABEL] ?? "", + container.labels[DUAL_STATION_VLLM_LAUNCH_CONTRACT_LABEL] ?? "", + container.labels[DUAL_STATION_VLLM_API_KEY_FINGERPRINT_LABEL] ?? "", + container.labels[DUAL_STATION_VLLM_TRANSACTION_LABEL] ?? "", + ].join("\t"); +} + +export function createDualStationLifecycleHarness( + fixture: LifecycleHarnessFixture, + options: LifecycleHarnessOptions = {}, +) { + const containers = new Map(); + const operations: Array<{ + kind: "capture" | "rename" | "rm" | "run" | "start" | "stop"; + target: string; + value: string; + }> = []; + const captureOptions: Array = []; + const rmOptions: Array = []; + const runCalls: Array<{ + args: readonly string[]; + options: DualStationDockerOptions | undefined; + }> = []; + const buildRemoteDockerEnv = vi.fn((binding: DualStationSshBinding) => ({ + TARGET: "peer", + DOCKER_HOST: `ssh://${binding.sshUser}@${binding.resolvedHost}`, + VLLM_API_KEY: "ambient-must-be-stripped", + })); + let nonceCounter = 0; + let transactionCounter = 0; + let lifecycleLockActive = 0; + let maxLifecycleLockActive = 0; + let lifecycleLockTail = Promise.resolve(); + const lifecycleLockContext = new AsyncLocalStorage(); + let lateContainer: { targetName: string; container: LifecycleFakeContainer } | null = null; + const launchedRoles = new Set<"head" | "worker">(); + const managedInspectionCounts = { head: 0, worker: 0 }; + let finalInspectionFailureInjected = false; + + async function acquireLifecycleLock(operation: () => Promise | T): Promise { + const previous = lifecycleLockTail; + let release: () => void = () => undefined; + lifecycleLockTail = new Promise((resolve) => { + release = resolve; + }); + await previous; + lifecycleLockActive += 1; + maxLifecycleLockActive = Math.max(maxLifecycleLockActive, lifecycleLockActive); + try { + return await lifecycleLockContext.run(true, operation); + } finally { + lifecycleLockActive -= 1; + release(); + } + } + + function target(optionsArg?: DualStationDockerOptions): string { + return String(optionsArg?.env?.TARGET ?? "unknown"); + } + + function key(targetName: string, name: string): string { + return `${targetName}:${name}`; + } + + function exactContainerById( + targetName: string, + containerId: string, + ): { + containerKey: string; + entries: LifecycleFakeContainer[]; + container: LifecycleFakeContainer; + } | null { + for (const [containerKey, entries] of containers.entries()) { + if (!containerKey.startsWith(`${targetName}:`)) continue; + const container = entries.find((entry) => entry.id === containerId); + if (container) return { containerKey, entries, container }; + } + return null; + } + + const deps: DualStationVllmLifecycleDeps = { + buildLocalDockerEnv: () => ({ + TARGET: "local", + VLLM_API_KEY: "ambient-must-be-stripped", + }), + buildRemoteDockerEnv, + createProbeNonce: () => { + nonceCounter += 1; + return nonceCounter.toString(16).padStart(32, "0"); + }, + createTransactionId: () => { + transactionCounter += 1; + return transactionCounter.toString(16).padStart(32, "0"); + }, + effectiveControllerUid: () => fixture.plan().local.uid, + readControllerUid: () => fixture.plan().local.uid, + loadApiKey: () => fixture.apiKey, + localInterfaceAddresses: () => [fixture.plan().masterAddress], + waitBeforeReconcile: async () => { + const pending = lateContainer; + lateContainer = null; + return pending + ? void containers.set(key(pending.targetName, pending.container.name), [pending.container]) + : undefined; + }, + withLifecycleLock: async (operation: () => Promise | T) => + lifecycleLockContext.getStore() ? await operation() : await acquireLifecycleLock(operation), + dockerCapture: (args, optionsArg) => { + captureOptions.push(optionsArg); + const targetName = target(optionsArg); + if (args[0] === "container" && args[1] === "rename") { + const containerId = args[2]; + const newName = args[3]; + operations.push({ + kind: "rename", + target: targetName, + value: `${containerId}:${newName}`, + }); + const located = exactContainerById(targetName, containerId); + if (!located || (containers.get(key(targetName, newName)) ?? []).length > 0) { + return raise("rename failed"); + } + containers.set( + located.containerKey, + located.entries.filter((entry) => entry.id !== containerId), + ); + located.container.name = newName; + containers.set(key(targetName, newName), [located.container]); + return ""; + } + if (args[0] === "container" && (args[1] === "start" || args[1] === "stop")) { + const action = args[1]; + const containerId = args.at(-1) ?? ""; + operations.push({ kind: action, target: targetName, value: containerId }); + const located = exactContainerById(targetName, containerId); + if (!located) return raise(`${action} failed`); + located.container.state = action === "start" ? "running" : "exited"; + return containerId; + } + switch (args[0]) { + case "image": { + operations.push({ kind: "capture", target: targetName, value: `image:${args.at(-1)}` }); + return options.missingImageTarget === targetName + ? raise("missing image") + : `sha256:${"f".repeat(64)}\n`; + } + case "wait": + operations.push({ kind: "capture", target: targetName, value: `wait:${args[1]}` }); + return "0\n"; + case "logs": { + operations.push({ kind: "capture", target: targetName, value: `logs:${args[1]}` }); + const defaultUuid = + targetName === "local" ? fixture.plan().local.gpu.uuid : fixture.plan().peer.gpu.uuid; + return `${options.smokeGpuOutput?.[targetName as "local" | "peer"] ?? defaultUuid}\n`; + } + default: + break; + } + const filter = dualStationDockerValues(args, "--filter")[0] ?? ""; + const name = filter.replace(/^name=\^\//, "").replace(/\$$/, ""); + operations.push({ kind: "capture", target: targetName, value: name }); + const isSmokeInspection = + dualStationDockerValues(args, "--format")[0]?.includes(DUAL_STATION_VLLM_GPU_SMOKE_LABEL) ?? + false; + let inspected = containers.get(key(targetName, name)) ?? []; + const inspectedRole = + name === DUAL_STATION_VLLM_HEAD_CONTAINER_NAME + ? "head" + : name === DUAL_STATION_VLLM_WORKER_CONTAINER_NAME + ? "worker" + : null; + if (!isSmokeInspection && inspectedRole && launchedRoles.has(inspectedRole)) { + managedInspectionCounts[inspectedRole] += 1; + if ( + options.failFinalInspectionRole === inspectedRole && + managedInspectionCounts[inspectedRole] === 2 && + !finalInspectionFailureInjected + ) { + finalInspectionFailureInjected = true; + inspected = inspected.map((container) => ({ ...container, state: "exited" })); + } + } + return inspected + .map((container) => + isSmokeInspection + ? [ + container.id, + container.name, + container.image, + container.labels[DUAL_STATION_VLLM_GPU_SMOKE_LABEL] ?? "", + container.labels[DUAL_STATION_VLLM_ROLE_LABEL] ?? "", + ].join("\t") + : row(container), + ) + .join("\n"); + }, + dockerRunDetached: (args, optionsArg) => { + const targetName = target(optionsArg); + const name = dualStationDockerValues(args, "--name")[0]; + runCalls.push({ args: [...args], options: optionsArg }); + operations.push({ kind: "run", target: targetName, value: name }); + const labels = Object.fromEntries( + dualStationDockerValues(args, "--label").map((label) => { + const separator = label.indexOf("="); + return [label.slice(0, separator), label.slice(separator + 1)]; + }), + ); + switch (name.startsWith("nemoclaw-vllm-gpu-smoke-")) { + case true: + return options.failSmokeTarget === targetName + ? { status: 1, stdout: "", stderr: "smoke failed" } + : (() => { + const smokeContainer: LifecycleFakeContainer = { + id: targetName === "local" ? fixture.headSmokeId : fixture.workerSmokeId, + name, + state: "exited", + image: DUAL_STATION_VLLM_RUNTIME.image, + labels, + }; + containers.set(key(targetName, name), [smokeContainer]); + return { status: 0, stdout: `${smokeContainer.id}\n` }; + })(); + default: + break; + } + const role = name === DUAL_STATION_VLLM_HEAD_CONTAINER_NAME ? "head" : "worker"; + launchedRoles.add(role); + const imageIndex = args.indexOf("/bin/bash") + 1; + const container = fixture.fakeContainer(role, { + image: args[imageIndex], + labels, + }); + return options.failedRoleForeignTransaction === role + ? (() => { + container.labels[DUAL_STATION_VLLM_TRANSACTION_LABEL] = "f".repeat(32); + containers.set(key(targetName, name), [container]); + return { status: 1, stdout: "", stderr: "ambiguous failed create" }; + })() + : options.lateCreateRole === role + ? (() => { + lateContainer = { targetName, container }; + return { status: null, stdout: "", stderr: "timed out" }; + })() + : options.failRole === role + ? { status: 1, stdout: "", stderr: "failed" } + : (() => { + containers.set(key(targetName, name), [container]); + return { + status: 0, + stdout: + options.invalidIdRole === role ? "not-a-container-id\n" : `${container.id}\n`, + }; + })(); + }, + dockerForceRm: (containerId, optionsArg) => { + rmOptions.push(optionsArg); + const targetName = target(optionsArg); + operations.push({ kind: "rm", target: targetName, value: containerId }); + const shouldFail = + (options.failSmokeCleanupTarget === targetName && + (containerId === fixture.workerSmokeId || containerId === fixture.headSmokeId)) || + (options.failLegacyBackupRemoval && containerId === fixture.legacyHeadId); + const match = [...containers.entries()].find( + ([containerKey, entries]) => + containerKey.startsWith(`${targetName}:`) && + entries.some((entry) => entry.id === containerId), + ); + return shouldFail || !match + ? { status: 1 } + : (() => { + const [containerKey, entries] = match; + const remaining = entries.filter((entry) => entry.id !== containerId); + containers.set(containerKey, remaining); + return { status: 0 }; + })(); + }, + }; + + function seed(targetName: "local" | "peer", container: LifecycleFakeContainer): void { + const containerKey = key(targetName, container.name); + containers.set(containerKey, [...(containers.get(containerKey) ?? []), container]); + } + + return { + buildRemoteDockerEnv, + captureOptions, + containers, + deps, + getMaxLifecycleLockActive: () => maxLifecycleLockActive, + operations, + rmOptions, + runCalls, + seed, + }; +} diff --git a/test/support/vllm-station-model-staging-test-support.ts b/test/support/vllm-station-model-staging-test-support.ts new file mode 100644 index 0000000000..ee9f21cf7b --- /dev/null +++ b/test/support/vllm-station-model-staging-test-support.ts @@ -0,0 +1,209 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { spawnSync } from "node:child_process"; +import fs from "node:fs"; +import path from "node:path"; + +import { vi } from "vitest"; + +import { resolveStationFixturePython } from "../../src/lib/inference/vllm-station-fixture.test-support"; + +import type { + ModelStagingCommandOptions, + ModelStagingCommandResult, +} from "../../src/lib/inference/vllm-station-model-staging"; + +type ModelStagingFixtureCommand = ( + file: string, + args: readonly string[], + options: ModelStagingCommandOptions, +) => Promise; + +function result(stdout = "", status = 0): ModelStagingCommandResult { + return { status, stdout, stderr: "" }; +} + +function runPython( + args: readonly string[], + options: ModelStagingCommandOptions, +): ModelStagingCommandResult { + const completed = spawnSync(resolveStationFixturePython(), [...args], { + encoding: "utf8", + env: options.env, + input: options.input, + timeout: options.timeoutMs, + }); + return { + status: completed.status, + stdout: completed.stdout ?? "", + stderr: completed.stderr ?? "", + error: completed.error?.message, + timedOut: (completed.error as NodeJS.ErrnoException | undefined)?.code === "ETIMEDOUT", + }; +} + +function unexpectedCommand(file: string): never { + throw new Error(`unexpected command: ${file}`); +} + +/** + * Branching in these runners models external command behavior; assertions stay in the test file. + */ +export function createPostAuditMutationRunner(snapshot: string) { + const state = { + configMode: 0, + snapshotMode: 0, + transferSource: "", + transferredConfig: "", + }; + let pythonCall = 0; + let sshCall = 0; + const runCommand = vi.fn(async (file, args, options) => { + if (file === "python3") { + pythonCall += 1; + const audit = runPython(args, options); + if (pythonCall === 2 && audit.status === 0) { + fs.writeFileSync(path.join(snapshot, "config.json"), "changed-after-audit"); + } + return audit; + } + if (file === "ssh") { + sshCall += 1; + return result(sshCall <= 2 ? '{"state":"transfer"}' : '{"state":"ready"}'); + } + if (file === "rsync") { + state.transferSource = String(args.at(-2)).replace(/\/$/, ""); + state.transferredConfig = fs.readFileSync( + path.join(state.transferSource, "config.json"), + "utf8", + ); + state.snapshotMode = fs.statSync(state.transferSource).mode & 0o777; + state.configMode = fs.statSync(path.join(state.transferSource, "config.json")).mode & 0o777; + return result(); + } + return unexpectedCommand(file); + }); + return { runCommand, state }; +} + +export function createBetweenAuditMutationRunner(snapshot: string) { + const state = { materializedSnapshot: "" }; + let pythonCall = 0; + const runCommand = vi.fn(async (file, args, options) => { + if (file === "python3") { + pythonCall += 1; + if (args[3] !== undefined) state.materializedSnapshot = String(args[3]); + const audit = runPython(args, options); + if (pythonCall === 1 && audit.status === 0) { + fs.writeFileSync(path.join(snapshot, "config.json"), "changed-between-audits"); + } + return audit; + } + if (file === "ssh") { + return result(pythonCall === 1 ? '{"state":"transfer"}' : '{"state":"cleaned"}'); + } + return unexpectedCommand(file); + }); + return { runCommand, state }; +} + +export function createPythonOnlyRunner() { + return vi.fn(async (file, args, options) => { + if (file !== "python3") return unexpectedCommand(file); + return runPython(args, options); + }); +} + +export function createManifestPeerPythonRunner(options: { + localManifest: string; + peerHome: string; + peerInputPrefix?: string; +}) { + return vi.fn(async (file, _args, commandOptions) => { + if (file === "python3") return result(options.localManifest); + if (file === "ssh") { + return runPython(["-"], { + ...commandOptions, + env: { ...commandOptions.env, HOME: options.peerHome }, + input: `${options.peerInputPrefix ?? ""}${commandOptions.input ?? ""}`, + }); + } + return unexpectedCommand(file); + }); +} + +export function createPeerIntegrityRunner(options: { + localManifest: string; + peerHome: string; + peerInputPrefix?: string; +}) { + const state = { stagingPath: "" }; + const runCommand = vi.fn(async (file, args, commandOptions) => { + if (file === "python3") return result(options.localManifest); + if (file === "ssh") { + return runPython(["-"], { + ...commandOptions, + env: { ...commandOptions.env, HOME: options.peerHome }, + input: `${options.peerInputPrefix ?? ""}${commandOptions.input ?? ""}`, + }); + } + if (file === "rsync") { + const destination = String(args.at(-1)); + state.stagingPath = destination.slice(destination.indexOf(":") + 1).replace(/\/$/, ""); + fs.mkdirSync(state.stagingPath, { mode: 0o700, recursive: true }); + fs.writeFileSync(path.join(state.stagingPath, "config.json"), "xx"); + return result(); + } + return unexpectedCommand(file); + }); + return { runCommand, state }; +} + +export function createAtomicIdentityReplacementRunner(peerHome: string) { + const state = { materializedSnapshot: "" }; + let sshCall = 0; + const runCommand = vi.fn(async (file, args, options) => { + if (file === "python3") { + if (args[3] !== undefined) state.materializedSnapshot = String(args[3]); + return runPython(args, options); + } + if (file === "ssh") { + sshCall += 1; + const ampleCapacity = `import shutil +class _NemoClawDiskUsage: + free = 1 << 50 +shutil.disk_usage = lambda _path: _NemoClawDiskUsage() +`; + const replaceInstalledIdentity = + sshCall === 3 + ? `import os +_nemoclaw_original_rename = os.rename +def _nemoclaw_replace_after_rename(source, destination): + _nemoclaw_original_rename(source, destination) + _nemoclaw_original_rename(destination, str(destination) + ".nemoclaw-test-original") + os.mkdir(destination, 0o700) +os.rename = _nemoclaw_replace_after_rename +` + : ""; + return runPython(["-"], { + ...options, + env: { ...options.env, HOME: peerHome }, + input: `${ampleCapacity}${replaceInstalledIdentity}${options.input ?? ""}`, + }); + } + if (file === "rsync") { + const source = String(args.at(-2)).replace(/\/$/, ""); + const destination = String(args.at(-1)); + const stagingPath = destination.slice(destination.indexOf(":") + 1).replace(/\/$/, ""); + for (const entry of fs.readdirSync(source)) { + fs.cpSync(path.join(source, entry), path.join(stagingPath, entry), { + recursive: true, + }); + } + return result(); + } + return unexpectedCommand(file); + }); + return { runCommand, state }; +} diff --git a/test/test-boundary-guards.test.ts b/test/test-boundary-guards.test.ts index 6ea5b4b39e..c33ebc3987 100644 --- a/test/test-boundary-guards.test.ts +++ b/test/test-boundary-guards.test.ts @@ -754,6 +754,9 @@ describe("Vitest project membership boundary", () => { ["test/install-openshell-version-check.test.ts", "installer-integration"], ["test/install-preflight-docker-bootstrap.test.ts", "installer-integration"], ["test/install-preflight.test.ts", "installer-integration"], + ["test/install-station-controller-binding.test.ts", "installer-integration"], + ["test/install-station-pair-preparation.test.ts", "installer-integration"], + ["test/install-station-resume-cleanup.test.ts", "installer-integration"], ["test/install-station-dgx-os.test.ts", "installer-integration"], ["test/install-station-docker-repository.test.ts", "installer-integration"], ["test/install-station-host-preparation.test.ts", "installer-integration"], diff --git a/test/vllm-ownership-test-support.test.ts b/test/vllm-ownership-test-support.test.ts new file mode 100644 index 0000000000..02d8e102c2 --- /dev/null +++ b/test/vllm-ownership-test-support.test.ts @@ -0,0 +1,34 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { describe, expect, it } from "vitest"; +import { createStrictVllmOwnershipCapture } from "./support/vllm-ownership-test-support"; + +describe("vLLM ownership test capture", () => { + it("consumes explicit canonical and ambient ownership responses", () => { + const capture = createStrictVllmOwnershipCapture( + [() => "canonical", () => "ambient"], + "builder", + (command) => `fallback:${command}`, + ); + + expect(capture(["container"], { env: { DOCKER_CONTEXT: "default" } })).toBe("canonical"); + expect(capture(["container"], { env: { DOCKER_CONTEXT: "builder" } })).toBe("ambient"); + expect(capture(["ps"])).toBe("fallback:ps"); + }); + + it("preserves explicit absence but rejects ownership response exhaustion", () => { + const capture = createStrictVllmOwnershipCapture([() => ""], "ambient", () => ""); + + expect(capture(["container"], { env: { DOCKER_CONTEXT: "default" } })).toBe(""); + expect(() => capture(["container"])).toThrow("Unexpected extra vLLM ownership inspection"); + }); + + it("rejects ownership inspection through an unregistered Docker context", () => { + const capture = createStrictVllmOwnershipCapture([() => ""], "ambient", () => ""); + + expect(() => capture(["container"], { env: { DOCKER_CONTEXT: "surprise" } })).toThrow( + "Unexpected vLLM ownership inspection context: surprise", + ); + }); +}); diff --git a/vitest.config.ts b/vitest.config.ts index 0888d30880..27981068a1 100644 --- a/vitest.config.ts +++ b/vitest.config.ts @@ -148,6 +148,9 @@ export default defineConfig({ "test/install-clone-ref.test.ts", "test/install-preflight.test.ts", "test/install-preflight-docker-bootstrap.test.ts", + "test/install-station-controller-binding.test.ts", + "test/install-station-pair-preparation.test.ts", + "test/install-station-resume-cleanup.test.ts", "test/install-station-dgx-os.test.ts", "test/install-station-docker-repository.test.ts", "test/install-station-host-preparation.test.ts", @@ -174,6 +177,9 @@ export default defineConfig({ "test/install-clone-ref.test.ts", "test/install-preflight.test.ts", "test/install-preflight-docker-bootstrap.test.ts", + "test/install-station-controller-binding.test.ts", + "test/install-station-pair-preparation.test.ts", + "test/install-station-resume-cleanup.test.ts", "test/install-station-dgx-os.test.ts", "test/install-station-docker-repository.test.ts", "test/install-station-host-preparation.test.ts",