[Bugfix] Avoid InfiniBand state in TP1 snapshots (#59699)

Signed-off-by: Nils Matteson <nilsmatteson@icloud.com>

.buildkite/scripts/initialized-snapshot-e2e.sh

@@ -34,6 +34,7 @@ readonly DEADLINE_EPOCH="$(( $(date +%s) + RUN_LIMIT_S ))"
 RUN_ID="" CONTAINER_NAME="" CONTAINER_ID=""
 GPU_UUID="" PORT_ONE="" PORT_TWO=""
 LINK_REMAP_BASELINE=""
+CHECK_CLEANUP_STATE=0
 SNAPSHOT_PIDS=()
 die() {
     echo "initialized snapshot E2E: $*" >&2
@@ -113,7 +114,7 @@ cleanup() {
         names="$(container_names 2>/dev/null)" || return 1
         grep -Fxq "$CONTAINER_NAME" <<< "$names" && failed=1
     fi
-    if [[ -n "$PORT_TWO" ]] && ! wait_clean; then
+    if (( CHECK_CLEANUP_STATE )) && ! wait_clean; then
         echo "process, GPU, port, or link-remap residue remained" >&2
         failed=1
     fi
@@ -218,7 +219,8 @@ GPU_PIDS="$(gpu_pids)" || die "failed to inspect the selected GPU"
 RUN_ID="snapshot-e2e-$(date +%s)-$$-$RANDOM"
 CONTAINER_NAME="vllm-$RUN_ID"
 read -r PORT_ONE PORT_TWO < <(python3 -c 'import socket; s=[socket.socket(),socket.socket()]; [x.bind(("127.0.0.1",0)) for x in s]; print(*(x.getsockname()[1] for x in s))')
-LINK_REMAP_BASELINE="$(link_remaps)"
+LINK_REMAP_BASELINE="$(link_remaps)" \
+    || die "cannot inventory /dev/shm; check CI runner permissions"
 
 run "pull exact candidate image" 900 docker pull --platform linux/amd64 \
     "$IMAGE_TAG" >/dev/null
@@ -231,6 +233,8 @@ IFS='|' read -r IMAGE_ID IMAGE_PLATFORM OCI_COMMIT VLLM_COMMIT < <(
     || die "candidate image identity is invalid"
 [[ "$OCI_COMMIT" == "$EXPECTED_COMMIT" && "$VLLM_COMMIT" == "$EXPECTED_COMMIT" ]] \
     || die "candidate image commit labels do not match"
+# A failed preflight has not started any workload whose cleanup we can verify.
+CHECK_CLEANUP_STATE=1
 CONTAINER_ID="$(run "start exact candidate container" 90 docker run --detach \
     --name "$CONTAINER_NAME" --label "ai.vllm.snapshot.e2e.run=$RUN_ID" \
     --gpus "device=$GPU_UUID" --user 0 --privileged --pid=host --ipc=host \

docs/features/initialized_snapshots.md

@@ -31,6 +31,11 @@ Snapshots currently require:
   Current Hugging Face hub clients hold their connections for the process
   lifetime, so download the model in a separate step and run create with
   `HF_HUB_OFFLINE=1`, as the quickstart below does.
+- Snapshot creation defaults `NCCL_IB_DISABLE=1` for its singleton donor and
+  inherited workers because CRIU cannot capture live InfiniBand/RDMA state.
+  An explicit caller value is retained, but creation rejects an open
+  `/dev/infiniband/` descriptor before CRIU. Close non-NCCL RDMA clients before
+  capture.
 - A remote model ID and an immutable 40-character `--revision`. Local model
   directories and mutable revisions are not supported.
 - Enough disk for the artifact, with the same installed vLLM package, model
@@ -138,9 +143,13 @@ snapshot or external CRIU operation may use a shared `/dev/shm` mount at a time.
   and directories. Recreate the artifact after an unclean host shutdown.
 - Restore currently requires the same host, GPU, driver, kernel, Python,
   PyTorch, installed vLLM version, model revision, engine arguments,
-  selected environment variables, and CRIU plugin binaries.
+  selected effective environment variables, and CRIU plugin binaries. An unset
+  `NCCL_IB_DISABLE` therefore matches the snapshot donor default of `1`, while
+  an explicit different value does not.
 - Only dense float16 TP1 has been validated. Other model formats depend on their
-  existing sleep level 2 reload support; NCCL state is not restored.
+  existing sleep level 2 reload support; distributed and RDMA snapshots are not
+  supported. The snapshot-only NCCL default does not change ordinary serve or
+  multi-GPU defaults.
 - CRIU support varies by kernel and driver. Preserve package, library, model,
   and generated-cache paths for the artifact lifetime.
 - A snapshot can include application secrets or request state present in the

tests/entrypoints/launchers/test_launch_cli.py

@@ -237,6 +237,46 @@ def test_snapshot_environment_contract(
     )
 
 
+@pytest.mark.parametrize(
+    ("configured", "expected"),
+    [(None, "1"), ("0", "0"), ("1", "1")],
+)
+def test_snapshot_child_defaults_nccl_ib_without_mutating_parent(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+    configured: str | None,
+    expected: str,
+):
+    if configured is None:
+        monkeypatch.delenv("NCCL_IB_DISABLE", raising=False)
+    else:
+        monkeypatch.setenv("NCCL_IB_DISABLE", configured)
+    parent_environment = os.environ.copy()
+    process = create_autospec(subprocess.Popen, instance=True)
+    process.pid = 100
+    popen = create_autospec(subprocess.Popen, return_value=process)
+    monkeypatch.setattr(snapshot_runtime.subprocess, "Popen", popen)
+
+    LocalSnapshotTools().launch_child(tmp_path, ("model",))
+
+    assert popen.call_args.kwargs["env"]["NCCL_IB_DISABLE"] == expected
+    assert os.environ == parent_environment
+
+
+def test_snapshot_environment_identity_applies_effective_nccl_ib_default(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    monkeypatch.delenv("NCCL_IB_DISABLE", raising=False)
+    implicit = LocalSnapshotTools()._environment_identity()
+    monkeypatch.setenv("NCCL_IB_DISABLE", "1")
+    explicit_disabled = LocalSnapshotTools()._environment_identity()
+    monkeypatch.setenv("NCCL_IB_DISABLE", "0")
+    explicit_enabled = LocalSnapshotTools()._environment_identity()
+
+    assert implicit == explicit_disabled
+    assert explicit_enabled != implicit
+
+
 @pytest.mark.parametrize(
     "case",
     [
@@ -369,6 +409,33 @@ ss() { :; }
     assert result.returncode == (0 if probe is None else 1), result.stderr
 
 
+@pytest.mark.parametrize("container_attempted", [False, True])
+def test_snapshot_e2e_cleanup_checks_state_only_after_container_attempt(
+    container_attempted: bool,
+):
+    script = (
+        Path(__file__).resolve().parents[3]
+        / ".buildkite/scripts/initialized-snapshot-e2e.sh"
+    )
+    if not script.exists():
+        pytest.skip("snapshot E2E script is not packaged in this test image")
+    source = script.read_text()
+    cleanup = source[source.index("cleanup() {") : source.index("on_exit() {")]
+    shell = (
+        "set -Eeuo pipefail\n"
+        "CONTAINER_NAME='' PORT_TWO=18002\n"
+        f"CHECK_CLEANUP_STATE={int(container_attempted)}\n"
+        "wait_clean() { echo checked-state; return 1; }\n"
+        + cleanup
+        + "\nif cleanup; then exit 0; else exit 1; fi\n"
+    )
+
+    result = subprocess.run(["bash", "-c", shell], capture_output=True, timeout=10)
+
+    assert result.returncode == int(container_attempted), result.stderr
+    assert (b"checked-state" in result.stdout) == container_attempted
+
+
 def test_snapshot_create_cli_accepts_only_pinned_compact_mode(
     monkeypatch: pytest.MonkeyPatch,
 ):
@@ -868,6 +935,96 @@ def test_snapshot_private_path_and_link_remap_security(
         tools.restore(artifact, _manifest())
 
 
+@pytest.mark.parametrize("inventory_error", [None, PermissionError("denied")])
+def test_snapshot_dump_rejects_unavailable_link_remap_inventory_before_criu(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+    inventory_error: PermissionError | None,
+):
+    tools = LocalSnapshotTools()
+    workdir = tmp_path / "snapshot"
+    workdir.mkdir()
+    tools.shm_dir = tmp_path / "missing-shm"
+    if inventory_error is not None:
+        tools.shm_dir.mkdir()
+
+        def fail_inventory(_path: Path):
+            raise inventory_error
+
+        monkeypatch.setattr(Path, "iterdir", fail_inventory)
+    monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called"))
+
+    with pytest.raises(SnapshotCreateError, match="inventory CRIU link remaps"):
+        tools.dump(workdir, ProcessInventory(100, (100,), (100,), "GPU"))
+
+
+def test_snapshot_restore_rejects_saved_remap_without_prefix(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+):
+    _mark_snapshot_pids_free(monkeypatch)
+    artifact = _local_restore_artifact(tmp_path)
+    saved = artifact / "link-remaps"
+    saved.mkdir()
+    (saved / "270").write_bytes(b"invalid remap name")
+    tools = LocalSnapshotTools()
+    tools.shm_dir = tmp_path / "shm"
+    tools.shm_dir.mkdir()
+    monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called"))
+
+    with pytest.raises(SnapshotRestoreError, match="invalid saved link remap"):
+        tools.restore(artifact, _manifest())
+
+    assert not tuple(tools.shm_dir.iterdir())
+
+
+def test_snapshot_link_remap_inventory_matches_only_numeric_suffixes(tmp_path: Path):
+    for name in (
+        "link_remap.0",
+        "link_remap.270",
+        "link_remap.",
+        "link_remap.12a",
+        "link-remap.12",
+        "prefix.link_remap.12",
+        "270",
+    ):
+        (tmp_path / name).touch()
+    tools = LocalSnapshotTools()
+    tools.shm_dir = tmp_path
+
+    assert tools._link_remap_names() == {"link_remap.0", "link_remap.270"}
+
+
+def test_snapshot_dump_preserves_primary_failure_when_cleanup_inventory_fails(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+):
+    tools = LocalSnapshotTools()
+    workdir = tmp_path / "snapshot"
+    workdir.mkdir()
+    inventory_calls = 0
+
+    def inventory_names() -> set[str]:
+        nonlocal inventory_calls
+        inventory_calls += 1
+        if inventory_calls > 1:
+            raise SnapshotCreateError(
+                "could not inventory CRIU link remaps: PermissionError: denied"
+            )
+        return set()
+
+    monkeypatch.setattr(tools, "_link_remap_names", inventory_names)
+
+    def fail_dump(*_args):
+        raise RuntimeError("dump failed")
+
+    monkeypatch.setattr(tools, "_criu", fail_dump)
+
+    with pytest.raises(SnapshotCreateError) as failure:
+        tools.dump(workdir, ProcessInventory(100, (100,), (100,), "GPU"))
+
+    assert "RuntimeError: dump failed" in str(failure.value)
+    assert "cleanup inventory failed" in str(failure.value)
+
+
 @pytest.mark.parametrize(
     ("blocked_state", "message"),
     [("tcp", "external established TCP"), ("io_uring", "kernel.io_uring_disabled=1")],
@@ -900,6 +1057,28 @@ def test_snapshot_rejects_unsafe_process_state_before_criu(
         assert "1.1.1.1:443" in str(excinfo.value)
 
 
+def test_snapshot_rejects_open_infiniband_device_before_criu(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+):
+    tools = LocalSnapshotTools()
+    monkeypatch.setattr(tools, "_tree_pids", lambda _root_pid: (100, 101))
+    monkeypatch.setattr(tools, "_cuda_process_rows", lambda: ("101, GPU-abc",))
+    monkeypatch.setattr(
+        tools,
+        "_descriptor_targets",
+        lambda pid: ("/dev/infiniband/uverbs7",) if pid == 101 else (),
+    )
+    monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called"))
+
+    with pytest.raises(SnapshotCreateError) as failure:
+        tools.inventory(100, tmp_path)
+
+    message = str(failure.value)
+    assert "pid 101" in message
+    assert "/dev/infiniband/uverbs7" in message
+    assert "NCCL_IB_DISABLE=1" in message
+
+
 def test_decode_endpoint_families():
     assert _decode_endpoint("AF_INET", "3C00000A:D6F6") == "10.0.0.60:55030"
     assert (

vllm/snapshot/runtime.py

@@ -254,6 +254,7 @@ class LocalSnapshotTools:
         ]
         child_environment = os.environ.copy()
         child_environment.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+        child_environment.setdefault("NCCL_IB_DISABLE", "1")
         with (workdir / "child.log").open("wb") as log_file:
             process = subprocess.Popen(
                 command,
@@ -357,6 +358,13 @@ class LocalSnapshotTools:
             if "anon_inode:[io_uring]" in targets:
                 io_uring_pids.append(pid)
             for target in targets:
+                if target.startswith("/dev/infiniband/"):
+                    raise SnapshotCreateError(
+                        "snapshot process has open InfiniBand/RDMA state that "
+                        f"CRIU cannot capture: pid {pid}, {target}. If NCCL owns "
+                        "it, use NCCL_IB_DISABLE=1. Close other RDMA clients "
+                        "before snapshot creation."
+                    )
                 if target.startswith("socket:[") and target.endswith("]"):
                     # Shared descriptors name one holder; any of them locates
                     # the connection for the operator.
@@ -512,10 +520,18 @@ class LocalSnapshotTools:
             os.fsync(output.fileno())
 
     def _link_remap_names(self) -> set[str]:
+        try:
+            entries = tuple(self.shm_dir.iterdir())
+        except OSError as error:
+            raise SnapshotCreateError(
+                f"could not inventory CRIU link remaps in {self.shm_dir}: "
+                f"{_error_detail(error)}"
+            ) from error
         return {
             path.name
-            for path in self.shm_dir.glob("link_remap.*")
-            if path.name.removeprefix("link_remap.").isdigit()
+            for path in entries
+            if path.name.startswith("link_remap.")
+            and path.name.removeprefix("link_remap.").isdigit()
         }
 
     def _capture_link_remaps(self, artifact: Path, names: set[str]) -> None:
@@ -565,7 +581,8 @@ class LocalSnapshotTools:
         try:
             for source in sorted(source_dir.iterdir()):
                 if (
-                    not source.name.removeprefix("link_remap.").isdigit()
+                    not source.name.startswith("link_remap.")
+                    or not source.name.removeprefix("link_remap.").isdigit()
                     or source.is_symlink()
                     or not source.is_file()
                 ):
@@ -662,8 +679,16 @@ class LocalSnapshotTools:
                 workdir,
                 self._link_remap_names() - remaps_before,
             )
-        except BaseException:
-            for name in self._link_remap_names() - remaps_before:
+        except BaseException as primary:
+            try:
+                remaps_after = self._link_remap_names()
+            except BaseException as cleanup_error:
+                raise SnapshotCreateError(
+                    f"snapshot dump failed: {_error_detail(primary)}; "
+                    "remap cleanup inventory failed: "
+                    f"{_error_detail(cleanup_error)}"
+                ) from primary
+            for name in remaps_after - remaps_before:
                 (self.shm_dir / name).unlink(missing_ok=True)
             raise
         self._record_child_log_size(workdir)
@@ -747,6 +772,7 @@ class LocalSnapshotTools:
         prefixes = ("VLLM_", "CUDA_", "NCCL_", "TORCH_", "TRITON_")
         environment = os.environ.copy()
         environment.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+        environment.setdefault("NCCL_IB_DISABLE", "1")
         selected = tuple(
             sorted(
                 (key, hashlib.sha256(key.encode() + b"\0" + value.encode()).hexdigest())