|
| 1 | +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | + |
| 15 | +"""CPU-only checks for examples/models/deepseek_v4/inference.sh. |
| 16 | +
|
| 17 | +The launcher (``uv``) and the Slurm tools (``srun``, ``scontrol``) are replaced by stubs, so the |
| 18 | +tests only exercise the shell logic: default parallelism per variant, the multi-node torchrun |
| 19 | +arguments, the divisibility guard and the Slurm command documented in the README. |
| 20 | +""" |
| 21 | + |
| 22 | +import os |
| 23 | +import re |
| 24 | +import subprocess |
| 25 | +from pathlib import Path |
| 26 | + |
| 27 | +import pytest |
| 28 | + |
| 29 | + |
| 30 | +pytestmark = pytest.mark.unit |
| 31 | + |
| 32 | +REPO_ROOT = Path(__file__).resolve().parents[3] |
| 33 | +SCRIPT = "examples/models/deepseek_v4/inference.sh" |
| 34 | +README = REPO_ROOT / "examples/models/deepseek_v4/README.md" |
| 35 | +ENTRYPOINT = "examples/conversion/hf_to_megatron_generate_text.py" |
| 36 | +MULTINODE_ARGS = {"--nproc_per_node=4", "--nnodes=8", "--node_rank=3", "--master_addr=node0"} |
| 37 | + |
| 38 | + |
| 39 | +def _write_stub(path: Path, body: str) -> None: |
| 40 | + path.write_text("#!/bin/bash\n" + body) |
| 41 | + path.chmod(0o755) |
| 42 | + |
| 43 | + |
| 44 | +@pytest.fixture |
| 45 | +def launcher_env(tmp_path): |
| 46 | + """Environment whose uv/srun/scontrol are stubs; every uv call appends its arguments to UV_LOG.""" |
| 47 | + fake_bin = tmp_path / "bin" |
| 48 | + fake_bin.mkdir() |
| 49 | + _write_stub(fake_bin / "uv", 'printf "%s\\n" "$*" >> "$UV_LOG"\n') |
| 50 | + # srun: drop the srun options and run the payload once, as a single task would. |
| 51 | + _write_stub(fake_bin / "srun", 'while [[ "${1:-}" == --* ]]; do shift; done\nexec "$@"\n') |
| 52 | + _write_stub(fake_bin / "scontrol", 'printf "node0\\nnode1\\n"\n') |
| 53 | + |
| 54 | + env = os.environ.copy() |
| 55 | + for key in ("NNODES", "NODE_RANK", "MASTER_ADDR", "MASTER_PORT", "GPUS_PER_NODE", "EP", "PP", "MODEL_VARIANT"): |
| 56 | + env.pop(key, None) |
| 57 | + env["PATH"] = f"{fake_bin}:{env['PATH']}" |
| 58 | + env["UV_LOG"] = str(tmp_path / "uv.log") |
| 59 | + env["WORKSPACE"] = str(tmp_path) |
| 60 | + return env |
| 61 | + |
| 62 | + |
| 63 | +def _run(args, env): |
| 64 | + return subprocess.run(args, cwd=REPO_ROOT, env=env, capture_output=True, text=True) |
| 65 | + |
| 66 | + |
| 67 | +def _uv_calls(env): |
| 68 | + log = Path(env["UV_LOG"]) |
| 69 | + return [line.split() for line in log.read_text().splitlines()] if log.exists() else [] |
| 70 | + |
| 71 | + |
| 72 | +def _documented_slurm_command() -> str: |
| 73 | + match = re.search(r"`(srun --ntasks-per-node=1 bash -c '[^`]*')`", README.read_text()) |
| 74 | + assert match, "README no longer documents the multi-node srun command" |
| 75 | + return match.group(1) |
| 76 | + |
| 77 | + |
| 78 | +@pytest.mark.parametrize(("variant", "nproc"), [("DeepSeek-V4-Flash", 4), ("DeepSeek-V4-Pro", 32)]) |
| 79 | +def test_single_node_defaults(launcher_env, variant, nproc): |
| 80 | + result = _run(["bash", SCRIPT], {**launcher_env, "MODEL_VARIANT": variant}) |
| 81 | + |
| 82 | + assert result.returncode == 0, result.stderr |
| 83 | + calls = _uv_calls(launcher_env) |
| 84 | + assert len(calls) == 1, "no imported checkpoint, so only the HF path should launch" |
| 85 | + args = calls[0] |
| 86 | + assert args[:4] == ["run", "python", "-m", "torch.distributed.run"] |
| 87 | + assert f"--nproc_per_node={nproc}" in args |
| 88 | + assert not any(a.startswith(("--nnodes", "--node_rank", "--master_addr", "--master_port")) for a in args) |
| 89 | + assert ENTRYPOINT in args |
| 90 | + assert (REPO_ROOT / ENTRYPOINT).is_file() |
| 91 | + |
| 92 | + |
| 93 | +def test_multinode_both_loading_paths(launcher_env, tmp_path): |
| 94 | + (tmp_path / "models" / "DeepSeek-V4-Pro" / "iter_0000000").mkdir(parents=True) |
| 95 | + env = {**launcher_env, "MODEL_VARIANT": "DeepSeek-V4-Pro", "NNODES": "8", "NODE_RANK": "3", "MASTER_ADDR": "node0"} |
| 96 | + |
| 97 | + result = _run(["bash", SCRIPT], env) |
| 98 | + |
| 99 | + assert result.returncode == 0, result.stderr |
| 100 | + calls = _uv_calls(env) |
| 101 | + assert len(calls) == 2, "HF path and imported-checkpoint path must both launch" |
| 102 | + for args in calls: |
| 103 | + assert MULTINODE_ARGS | {"--master_port=29500", "--legacy-full-prefix", ENTRYPOINT} <= set(args) |
| 104 | + assert "--megatron_model_path" not in calls[0] |
| 105 | + assert "--megatron_model_path" in calls[1] |
| 106 | + |
| 107 | + |
| 108 | +def test_nondivisible_nodes_exit_before_launch(launcher_env): |
| 109 | + result = _run(["bash", SCRIPT], {**launcher_env, "PP": "4", "EP": "8", "NNODES": "3"}) |
| 110 | + |
| 111 | + assert result.returncode != 0 |
| 112 | + assert "not divisible by NNODES=3" in result.stderr |
| 113 | + assert _uv_calls(launcher_env) == [] |
| 114 | + |
| 115 | + |
| 116 | +def test_documented_slurm_command_runs_from_repo_root(launcher_env): |
| 117 | + command = _documented_slurm_command() |
| 118 | + assert SCRIPT in command, "the documented command must use the root-relative script path" |
| 119 | + assert "./inference.sh" not in (REPO_ROOT / SCRIPT).read_text(), "header example must match the README" |
| 120 | + env = { |
| 121 | + **launcher_env, |
| 122 | + "MODEL_VARIANT": "DeepSeek-V4-Pro", |
| 123 | + "SLURM_JOB_NUM_NODES": "8", |
| 124 | + "SLURM_NODEID": "3", |
| 125 | + "SLURM_JOB_NODELIST": "node[0-7]", |
| 126 | + } |
| 127 | + |
| 128 | + result = _run(["bash", "-c", command], env) |
| 129 | + |
| 130 | + assert result.returncode == 0, result.stderr |
| 131 | + calls = _uv_calls(env) |
| 132 | + assert len(calls) == 1 |
| 133 | + assert MULTINODE_ARGS | {ENTRYPOINT} <= set(calls[0]) |
0 commit comments