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

Signed-off-by: Nils Matteson <nilsmatteson@icloud.com>
2026-10-03 03:59:17 +00:00
parent faa9860dbd
commit 9fdb174678
4 changed files with 227 additions and 9 deletions
@@ -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 (