diff --git a/scripts/performance/argument_parser.py b/scripts/performance/argument_parser.py index 89f43a6d1d..6e6d33561b 100644 --- a/scripts/performance/argument_parser.py +++ b/scripts/performance/argument_parser.py @@ -752,12 +752,17 @@ def parse_cli_args(): required=False, default=None, ) + nvcre_args.add_argument( + "--nvcre_use_mbridge_commit", + action="store_true", + help=("Run this Megatron-Bridge checkout in the pods instead of the image's /opt/Megatron-Bridge"), + ) nvcre_args.add_argument( "--nvcre_timeout_per_job", type=str, - help="Per-job timeout passed to Nvcre orchestration (e.g. '24h').", + help="Per-job timeout passed to Nvcre orchestration (e.g. '12h').", required=False, - default="24h", + default="12h", ) nvcre_args.add_argument( diff --git a/scripts/performance/setup_experiment.py b/scripts/performance/setup_experiment.py index fd12119cf7..8b732a81bf 100755 --- a/scripts/performance/setup_experiment.py +++ b/scripts/performance/setup_experiment.py @@ -61,6 +61,25 @@ except (ImportError, ModuleNotFoundError): from .perf_plugins import NsysPlugin, PreemptionPlugin, PyTorchProfilerPlugin +try: + from utils.nvcre_bridge import ( + BRIDGE_POD_ROOT, + WORKSPACE_MOUNT, + WORKSPACE_VOLUME, + LocalBridge, + bridge_init_container, + local_bridge_source, + ) +except (ImportError, ModuleNotFoundError): + from .utils.nvcre_bridge import ( + BRIDGE_POD_ROOT, + WORKSPACE_MOUNT, + WORKSPACE_VOLUME, + LocalBridge, + bridge_init_container, + local_bridge_source, + ) + SCRIPT_DIR = Path(__file__).parent.resolve() REPO_ROOT = SCRIPT_DIR.parents[1] BRIDGE_SOURCE_DIR = REPO_ROOT / "src" / "megatron" / "bridge" @@ -95,7 +114,7 @@ def _filter_run_script_args(argv: List[str]) -> List[str]: Value-taking flags are passed as ``--flag value`` or ``--flag=value``. Boolean flags (no following value) are listed in ``_LAUNCHER_ONLY_BOOL``. """ - _LAUNCHER_ONLY_BOOL = {"--offline", "--dryrun", "-d"} + _LAUNCHER_ONLY_BOOL = {"--offline", "--dryrun", "-d", "--nvcre_use_mbridge_commit"} def _is_launcher_only(flag: str) -> bool: return ( @@ -158,6 +177,17 @@ def _bridge_source_mount(in_container_script_dir: str) -> str: return f"{BRIDGE_SOURCE_DIR}:{target}" +def _resolve_local_bridge(nvcre_use_mbridge_commit: bool, nvcre_namespace: Optional[str]) -> Optional[LocalBridge]: + """The checkout the pods should run for ``--nvcre_use_mbridge_commit``, or None when the flag is off.""" + if not nvcre_use_mbridge_commit: + return None + if nvcre_namespace is None: + raise ValueError("--nvcre_use_mbridge_commit requires --nvcre_namespace.") + local_bridge = local_bridge_source(REPO_ROOT) + logger.info("Pods will run Megatron-Bridge %s@%s from the init container", local_bridge.url, local_bridge.commit) + return local_bridge + + def _build_nemorun_script( *, script_path: str, @@ -168,10 +198,11 @@ def _build_nemorun_script( ) -> run.Script: """Build the rank-local task and apply optional Kubeflow NUMA binding.""" in_container_repo_root = Path(script_dir).parents[1] + python_path = [script_dir, str(in_container_repo_root / "src")] task = run.Script( path=script_path, entrypoint="python", - env={"PYTHONPATH": f"{script_dir}:{in_container_repo_root / 'src'}:$PYTHONPATH"}, + env={"PYTHONPATH": ":".join([*python_path, "$PYTHONPATH"])}, args=args, ) if kubeflow_namespace and _kubeflow_numa_binding_enabled(custom_env_vars): @@ -578,7 +609,8 @@ def main( nvcre_node_selector_json: Optional[str] = None, nvcre_volumes_json: Optional[str] = None, nvcre_volume_mounts_json: Optional[str] = None, - nvcre_timeout_per_job: str = "24h", + nvcre_use_mbridge_commit: bool = False, + nvcre_timeout_per_job: str = "12h", nvcre_kubeconfig: Optional[str] = None, nvcre_kube_context: Optional[str] = None, nvcre_gang_scheduler_name: Optional[str] = None, @@ -656,8 +688,11 @@ def main( # Kubeflow the trainer pod runs the image — which ships Megatron-Bridge at # /opt/Megatron-Bridge — and custom_mounts do not apply, so the launcher's # /tmp path does not exist in the pod; use the image's script path instead. + # On NVCRE, --nvcre_use_mbridge_commit has an init container clone this checkout's commits into the pod. + local_bridge = _resolve_local_bridge(nvcre_use_mbridge_commit, nvcre_namespace) if kubeflow_namespace or nvcre_namespace: - in_container_script_dir = "/opt/Megatron-Bridge/scripts/performance" + bridge_root = BRIDGE_POD_ROOT if local_bridge else "/opt/Megatron-Bridge" + in_container_script_dir = f"{bridge_root}/scripts/performance" in_container_script_path = f"{in_container_script_dir}/{script_name}" else: in_container_script_dir = str(SCRIPT_DIR) @@ -720,6 +755,9 @@ def main( from utils.executors import nvcre_executor except ImportError: from .utils.executors import nvcre_executor + init_containers = [bridge_init_container(local_bridge, container_image)] if local_bridge else None + extra_volumes = [dict(WORKSPACE_VOLUME)] if local_bridge else [] + extra_mounts = [dict(WORKSPACE_MOUNT)] if local_bridge else [] executor = nvcre_executor( namespace=nvcre_namespace, container_image=container_image, @@ -730,8 +768,9 @@ def main( workdir_pvc_path=nvcre_workdir_pvc_path, workdir_local_path=nvcre_workdir_local_path, node_selector=json.loads(nvcre_node_selector_json) if nvcre_node_selector_json else None, - volumes=json.loads(nvcre_volumes_json) if nvcre_volumes_json else None, - volume_mounts=json.loads(nvcre_volume_mounts_json) if nvcre_volume_mounts_json else None, + volumes=(json.loads(nvcre_volumes_json) if nvcre_volumes_json else []) + extra_volumes, + volume_mounts=(json.loads(nvcre_volume_mounts_json) if nvcre_volume_mounts_json else []) + extra_mounts, + init_containers=init_containers, timeout_per_job=nvcre_timeout_per_job, kubeconfig=nvcre_kubeconfig, kube_context=nvcre_kube_context, @@ -1062,7 +1101,7 @@ def main( task=args.task, compute_dtype=args.compute_dtype, gpu=args.gpu, - hf_token=args.hf_token or os.environ.get("HF_TOKEN"), + hf_token=args.hf_token, offline=args.offline, # Force detach for Nvcre — llmb-run polls for completion via nvcrectl detach=True if args.nvcre_namespace else args.detach, @@ -1150,6 +1189,7 @@ def main( nvcre_node_selector_json=args.nvcre_node_selector_json, nvcre_volumes_json=args.nvcre_volumes_json, nvcre_volume_mounts_json=args.nvcre_volume_mounts_json, + nvcre_use_mbridge_commit=args.nvcre_use_mbridge_commit, nvcre_timeout_per_job=args.nvcre_timeout_per_job, nvcre_kubeconfig=args.nvcre_kubeconfig, nvcre_kube_context=args.nvcre_kube_context, diff --git a/scripts/performance/utils/executors.py b/scripts/performance/utils/executors.py index f229eef2db..97e5fee588 100644 --- a/scripts/performance/utils/executors.py +++ b/scripts/performance/utils/executors.py @@ -381,7 +381,8 @@ def nvcre_executor( node_selector: Optional[dict] = None, volumes: Optional[list] = None, volume_mounts: Optional[list] = None, - timeout_per_job: str = "24h", + init_containers: Optional[list] = None, + timeout_per_job: str = "12h", kubeconfig: Optional[str] = None, kube_context: Optional[str] = None, gang_scheduler_name: Optional[str] = None, @@ -400,7 +401,9 @@ def nvcre_executor( node_selector: Kubernetes node selector labels for targeting specific nodes. volumes: Additional Kubernetes volume specs to attach to the pod. volume_mounts: Additional Kubernetes volume mount specs. - timeout_per_job: Maximum wall-clock time per job (e.g. ``"24h"``). + init_containers: Kubernetes init container specs emitted as the WorkloadRun's ``spec.initContainers`` + (e.g. to clone a pinned Megatron-Bridge into a shared volume that replaces the image's copy). + timeout_per_job: Maximum wall-clock time per job (e.g. ``"12h"``). kubeconfig: Path to the kubeconfig file. Defaults to the cluster default. kube_context: Kubernetes context to use from the kubeconfig. gang_scheduler_name: Name of the gang scheduler (e.g. ``"kai-scheduler"``). @@ -413,6 +416,8 @@ def nvcre_executor( "NvcreExecutor is unavailable: the installed nemo_run does not provide nemo_run.core.execution.nvcre" ) + # Only passed when requested, so a nemo_run without ``init_containers`` still works for jobs that don't use it. + extra = {"init_containers": init_containers} if init_containers else {} return NvcreExecutor( namespace=namespace, container_image=container_image, @@ -429,4 +434,5 @@ def nvcre_executor( kubeconfig=kubeconfig, kube_context=kube_context, gang_scheduler_name=gang_scheduler_name, + **extra, ) diff --git a/scripts/performance/utils/nvcre_bridge.py b/scripts/performance/utils/nvcre_bridge.py new file mode 100644 index 0000000000..475dff104b --- /dev/null +++ b/scripts/performance/utils/nvcre_bridge.py @@ -0,0 +1,98 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Run the launcher's Megatron-Bridge checkout in NVCRE pods instead of the image's copy. + +On Slurm the launcher's checkout is bind-mounted over the image's ``/opt/Megatron-Bridge``. +Pods have no such mount, so an init container clones the same Megatron-Bridge commit into +a volume that the training container shares.""" + +import re +import shlex +import subprocess +from pathlib import Path +from typing import Any, NamedTuple +from urllib.parse import urlsplit, urlunsplit + + +BRIDGE_POD_ROOT = "/mnt/workspace/megatron-bridge" +WORKSPACE_VOLUME: dict[str, Any] = {"name": "workspace", "emptyDir": {}} +WORKSPACE_MOUNT: dict[str, Any] = {"name": "workspace", "mountPath": "/mnt/workspace"} + + +class LocalBridge(NamedTuple): + """Where the launcher's checkout can be fetched from, and at which commit.""" + + url: str + commit: str + + +def _git(repo_root: Path, *args: str) -> str: + return subprocess.run( + ["git", "-C", str(repo_root), *args], check=True, capture_output=True, text=True + ).stdout.strip() + + +def pod_clone_url(remote_url: str) -> str: + """HTTPS URL a pod can clone anonymously: SSH remotes are rewritten and any credentials dropped.""" + # scp-style ``git@host:path``: the colon separates host from path. + scp = re.match(r"^git@([^:/]+):(.+)$", remote_url) + if scp: + return f"https://{scp.group(1)}/{scp.group(2)}" + parts = urlsplit(remote_url) + if parts.scheme == "ssh" and parts.hostname: + # The port of an ssh:// URL is the SSH port, which does not apply over HTTPS. + return urlunsplit(("https", parts.hostname, parts.path, "", "")) + if parts.scheme in ("http", "https") and parts.hostname: + host = parts.hostname + (f":{parts.port}" if parts.port else "") + return urlunsplit((parts.scheme, host, parts.path, "", "")) + raise ValueError(f"Cannot derive a clone URL a pod can reach from remote {remote_url!r}") + + +def local_bridge_source(repo_root: Path) -> LocalBridge: + """The origin URL and checked-out commit of the Megatron-Bridge checkout at ``repo_root``.""" + try: + url = pod_clone_url(_git(repo_root, "config", "--get", "remote.origin.url")) + commit = _git(repo_root, "rev-parse", "HEAD") + # Submodule pointers are ignored: only this repository's own commit is fetched. + dirty = _git(repo_root, "status", "--porcelain", "--untracked-files=no", "--ignore-submodules=all") + except subprocess.CalledProcessError as exc: + raise RuntimeError( + f"--nvcre_use_mbridge_commit needs {repo_root} to be a git checkout with an 'origin' remote: {exc.stderr}" + ) from exc + if dirty: + raise RuntimeError( + f"--nvcre_use_mbridge_commit runs commit {commit} in the pods, but {repo_root} has uncommitted changes " + "to tracked files. Commit and push them, or stash them, and resubmit." + ) + return LocalBridge(url, commit) + + +def bridge_init_container(source: LocalBridge, image: str) -> dict[str, Any]: + """Init container that clones ``source`` into ``BRIDGE_POD_ROOT``; ``image`` is the training image (has git).""" + steps = [ + "set -ex", + f"git init -q {BRIDGE_POD_ROOT}", + f"cd {BRIDGE_POD_ROOT}", + f"git remote add origin {shlex.quote(source.url)}", + f"git fetch -q --depth 1 origin {shlex.quote(source.commit)}", + "git checkout -q FETCH_HEAD", + ] + return { + "name": "megatron-bridge-clone", + "image": image, + "command": ["/bin/bash", "-c"], + "args": ["\n".join(steps)], + "volumeMounts": [dict(WORKSPACE_MOUNT)], # a copy: shared objects would be dumped as YAML anchors + } diff --git a/tests/unit_tests/scripts/performance/test_executors.py b/tests/unit_tests/scripts/performance/test_executors.py index 036b4d470e..44ea5bcf87 100644 --- a/tests/unit_tests/scripts/performance/test_executors.py +++ b/tests/unit_tests/scripts/performance/test_executors.py @@ -279,6 +279,40 @@ def test_nvcre_executor_optional_fields(): assert executor.volume_mounts == [] +@pytest.mark.skipif(not HAS_NVCRE, reason="nemo_run nvcre not available") +def test_nvcre_executor_init_containers_default_to_none(): + """Without init containers the executor is built exactly as before.""" + executor = nvcre_executor(namespace="ns", container_image="img:latest", nodes=1, num_gpus_per_node=8) + + assert not getattr(executor, "init_containers", []) + + +@pytest.mark.skipif( + not HAS_NVCRE or "init_containers" not in getattr(NvcreExecutor, "__dataclass_fields__", {}), + reason="nemo_run NvcreExecutor without init_containers", +) +def test_nvcre_executor_forwards_init_containers(): + """init_containers must reach NvcreExecutor unchanged.""" + init_container = { + "name": "megatron-bridge-clone", + "image": "nvcr.io/nvidia/nemo:26.08.01", + "command": ["/bin/bash", "-c"], + "args": ["git clone https://github.com/NVIDIA-NeMo/Megatron-Bridge.git /mnt/workspace/megatron-bridge"], + "volumeMounts": [{"name": "workspace", "mountPath": "/mnt/workspace"}], + } + executor = nvcre_executor( + namespace="ns", + container_image="img:latest", + nodes=1, + num_gpus_per_node=8, + volumes=[{"name": "workspace", "emptyDir": {}}], + volume_mounts=[{"name": "workspace", "mountPath": "/mnt/workspace"}], + init_containers=[init_container], + ) + + assert executor.init_containers == [init_container] + + @pytest.mark.skipif(not HAS_NVCRE, reason="nemo_run nvcre not available") def test_nvcre_executor_has_no_test_scale(): """test_scale must not be wired up — it was intentionally removed.""" diff --git a/tests/unit_tests/scripts/performance/test_nvcre_bridge.py b/tests/unit_tests/scripts/performance/test_nvcre_bridge.py new file mode 100644 index 0000000000..7684858f47 --- /dev/null +++ b/tests/unit_tests/scripts/performance/test_nvcre_bridge.py @@ -0,0 +1,162 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for scripts/performance/utils/nvcre_bridge.py — running the launcher's checkout in NVCRE pods.""" + +import subprocess +import sys +from pathlib import Path + +import pytest + + +pytestmark = pytest.mark.unit + +_PERF_SCRIPTS_DIR = Path(__file__).resolve().parents[4] / "scripts" / "performance" +if str(_PERF_SCRIPTS_DIR) not in sys.path: + sys.path.insert(0, str(_PERF_SCRIPTS_DIR)) + +from utils import nvcre_bridge # noqa: E402 + + +def _git(repo: Path, *args: str) -> str: + return subprocess.run(["git", "-C", str(repo), *args], check=True, capture_output=True, text=True).stdout.strip() + + +@pytest.fixture +def checkout(tmp_path: Path) -> Path: + """A git checkout with an origin remote, like a real Megatron-Bridge clone.""" + repo = tmp_path / "bridge" + repo.mkdir() + _git(repo, "init", "-q") + _git(repo, "config", "user.email", "t@example.com") + _git(repo, "config", "user.name", "t") + _git(repo, "remote", "add", "origin", "https://github.com/example/Megatron-Bridge.git") + (repo / "README").write_text("x") + _git(repo, "add", "README") + _git(repo, "commit", "-q", "-m", "init") + return repo + + +@pytest.mark.parametrize( + "remote, expected", + [ + ("https://github.com/NVIDIA-NeMo/Megatron-Bridge.git", "https://github.com/NVIDIA-NeMo/Megatron-Bridge.git"), + ("git@github.com:NVIDIA-NeMo/Megatron-Bridge.git", "https://github.com/NVIDIA-NeMo/Megatron-Bridge.git"), + ("ssh://git@github.com/NVIDIA-NeMo/Megatron-Bridge.git", "https://github.com/NVIDIA-NeMo/Megatron-Bridge.git"), + # the port of an ssh:// URL is the SSH port; it must not leak into the HTTPS URL or its path + ("ssh://git@gitlab.example.com:2222/o/r.git", "https://gitlab.example.com/o/r.git"), + ("git@gitlab.example.com:o/sub/r.git", "https://gitlab.example.com/o/sub/r.git"), + # credentials must never end up in the WorkloadRun spec + ("https://user:secret-token@github.com/o/r.git", "https://github.com/o/r.git"), # pragma: allowlist secret + ("https://gitlab.example.com:8443/o/r.git", "https://gitlab.example.com:8443/o/r.git"), + ], +) +def test_pod_clone_url(remote: str, expected: str) -> None: + assert nvcre_bridge.pod_clone_url(remote) == expected + + +@pytest.mark.parametrize("remote", ["/srv/git/bridge.git", "file:///srv/git/bridge.git", "../bridge"]) +def test_pod_clone_url_rejects_remotes_a_pod_cannot_reach(remote: str) -> None: + with pytest.raises(ValueError, match="clone URL"): + nvcre_bridge.pod_clone_url(remote) + + +def test_local_bridge_source_reads_origin_and_head(checkout: Path) -> None: + source = nvcre_bridge.local_bridge_source(checkout) + + assert source.url == "https://github.com/example/Megatron-Bridge.git" + assert source.commit == _git(checkout, "rev-parse", "HEAD") + + +def test_local_bridge_source_requires_a_git_checkout_with_origin(tmp_path: Path) -> None: + with pytest.raises(RuntimeError, match="origin"): + nvcre_bridge.local_bridge_source(tmp_path) + + +def test_local_bridge_source_rejects_uncommitted_changes(checkout: Path) -> None: + (checkout / "README").write_text("changed") + + with pytest.raises(RuntimeError, match="uncommitted"): + nvcre_bridge.local_bridge_source(checkout) + + +def test_local_bridge_source_rejects_staged_changes(checkout: Path) -> None: + (checkout / "README").write_text("changed") + _git(checkout, "add", "README") + + with pytest.raises(RuntimeError, match="uncommitted"): + nvcre_bridge.local_bridge_source(checkout) + + +def test_local_bridge_source_ignores_untracked_files(checkout: Path) -> None: + (checkout / "scratch.txt").write_text("not tracked") + + assert nvcre_bridge.local_bridge_source(checkout).commit == _git(checkout, "rev-parse", "HEAD") + + +def test_local_bridge_source_ignores_a_drifted_submodule(checkout: Path, tmp_path: Path) -> None: + """Only this repository's commit is fetched, so a submodule checked out at another commit is not a dirty tree.""" + sub_repo = tmp_path / "sub_repo" + sub_repo.mkdir() + _git(sub_repo, "init", "-q") + _git(sub_repo, "config", "user.email", "t@example.com") + _git(sub_repo, "config", "user.name", "t") + (sub_repo / "f").write_text("1") + _git(sub_repo, "add", "f") + _git(sub_repo, "commit", "-q", "-m", "one") + recorded = _git(sub_repo, "rev-parse", "HEAD") + _git(checkout, "update-index", "--add", "--cacheinfo", f"160000,{recorded},3rdparty/sub") + _git(checkout, "commit", "-q", "-m", "add submodule") + _git(checkout, "clone", "-q", str(sub_repo), "3rdparty/sub") + (checkout / "3rdparty" / "sub" / "f").write_text("2") + _git( + checkout / "3rdparty" / "sub", + "-c", + "user.email=t@example.com", + "-c", + "user.name=t", + "commit", + "-qam", + "two", + ) + assert "3rdparty/sub" in _git(checkout, "status", "--porcelain", "--untracked-files=no") # git sees drift + + assert nvcre_bridge.local_bridge_source(checkout).commit == _git(checkout, "rev-parse", "HEAD") + + +def test_bridge_init_container_fetches_the_exact_commit() -> None: + source = nvcre_bridge.LocalBridge("https://github.com/o/r.git", "a" * 40) + + container = nvcre_bridge.bridge_init_container(source, "nvcr.io/nvidia/nemo:26.08.01") + + assert container["name"] == "megatron-bridge-clone" + assert container["image"] == "nvcr.io/nvidia/nemo:26.08.01" # the training image, which has git + assert container["volumeMounts"] == [nvcre_bridge.WORKSPACE_MOUNT] + script = container["args"][0] + assert f"git fetch -q --depth 1 origin {'a' * 40}" in script + assert "git remote add origin https://github.com/o/r.git" in script + assert "submodule" not in script # only the Megatron-Bridge commit is fetched + assert nvcre_bridge.BRIDGE_POD_ROOT in script + + +def test_bridge_init_container_quotes_the_url() -> None: + source = nvcre_bridge.LocalBridge("https://github.com/o/r.git; rm -rf /", "a" * 40) + + assert "'https://github.com/o/r.git; rm -rf /'" in nvcre_bridge.bridge_init_container(source, "img")["args"][0] + + +def test_workspace_volume_and_mount_agree() -> None: + assert nvcre_bridge.WORKSPACE_VOLUME["name"] == nvcre_bridge.WORKSPACE_MOUNT["name"] + assert nvcre_bridge.BRIDGE_POD_ROOT.startswith(nvcre_bridge.WORKSPACE_MOUNT["mountPath"] + "/") diff --git a/tests/unit_tests/scripts/performance/test_setup_experiment.py b/tests/unit_tests/scripts/performance/test_setup_experiment.py index 56f3cebd6c..300fe14348 100644 --- a/tests/unit_tests/scripts/performance/test_setup_experiment.py +++ b/tests/unit_tests/scripts/performance/test_setup_experiment.py @@ -16,7 +16,9 @@ from __future__ import annotations +import json import os +import runpy import subprocess import sys from pathlib import Path @@ -34,6 +36,15 @@ import setup_experiment +try: + import nemo_run + from nemo_run.core.execution.nvcre import NvcreExecutor # noqa: F401 + + HAS_NVCRE = True +except (ImportError, ModuleNotFoundError): + HAS_NVCRE = False + + def test_default_experiment_name_uses_only_cli_metadata() -> None: assert ( setup_experiment._default_experiment_name( @@ -196,6 +207,89 @@ def test_nvcre_flags_are_filtered_from_recipe_args() -> None: assert filtered == ["--tensor_model_parallel_size", "2"] +def test_nvcre_use_mbridge_commit_flag_does_not_swallow_the_next_arg() -> None: + """--nvcre_use_mbridge_commit takes no value, so the arg after it must still reach the recipe.""" + argv = ["--nvcre_use_mbridge_commit", "--tensor_model_parallel_size", "2"] + assert setup_experiment._filter_run_script_args(argv) == ["--tensor_model_parallel_size", "2"] + + +def test_use_mbridge_commit_without_nvcre_namespace_is_an_error() -> None: + with pytest.raises(ValueError, match="--nvcre_namespace"): + setup_experiment._resolve_local_bridge(True, None) + + +def test_resolve_local_bridge_is_none_when_the_flag_is_off() -> None: + assert setup_experiment._resolve_local_bridge(False, "ns") is None + assert setup_experiment._resolve_local_bridge(False, None) is None + + +def test_resolve_local_bridge_reads_this_checkout(monkeypatch: pytest.MonkeyPatch) -> None: + source = setup_experiment.LocalBridge("https://github.com/o/r.git", "a" * 40) + monkeypatch.setattr(setup_experiment, "local_bridge_source", lambda root: source) + + assert setup_experiment._resolve_local_bridge(True, "ns") is source + + +@pytest.mark.skipif(not HAS_NVCRE, reason="nemo_run nvcre not available") +def test_main_wires_the_pod_checkout_into_the_workloadrun(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + """--nvcre_use_mbridge_commit runs the scripts from the pod clone, via an init container and its volume.""" + from utils import nvcre_bridge + + source = nvcre_bridge.LocalBridge("https://github.com/o/r.git", "a" * 40) + monkeypatch.setattr(nvcre_bridge, "local_bridge_source", lambda root: source) + user_volume = {"name": "data", "persistentVolumeClaim": {"claimName": "pvc"}} + user_mount = {"name": "data", "mountPath": "/data"} + submitted = {} + + def fake_run(script, **kwargs): + submitted.update(script=script, **kwargs) + + monkeypatch.setattr(nemo_run, "run", fake_run) + monkeypatch.setenv("NEMORUN_HOME", str(tmp_path / "nemorun")) + monkeypatch.setattr( + sys, + "argv", + [ + "setup_experiment.py", + *("--model_family_name", "llama", "--model_recipe_name", "llama3_8b"), + *("--num_gpus", "8", "--gpus_per_node", "4", "--gpu", "gb300"), + *("--container_image", "example.invalid/nemo:test", "--dryrun"), + *("--nvcre_namespace", "ns", "--nvcre_use_mbridge_commit"), + *("--nvcre_volumes_json", json.dumps([user_volume])), + *("--nvcre_volume_mounts_json", json.dumps([user_mount])), + ], + ) + + runpy.run_path(str(_PERF_SCRIPTS_DIR / "setup_experiment.py"), run_name="__main__") + + executor = submitted["executor"] + executor.job_name = "job1" + executor.experiment_id = "exp_123456789" + command = submitted["script"].to_command() + spec = executor.build_workloadrun_yaml(command)["spec"] + pod_scripts = f"{nvcre_bridge.BRIDGE_POD_ROOT}/scripts/performance/" + + assert any(pod_scripts in part for part in command) + assert [c["name"] for c in spec["initContainers"]] == ["megatron-bridge-clone"] + assert f"git fetch -q --depth 1 origin {'a' * 40}" in spec["initContainers"][0]["args"][0] + assert spec["volumes"] == [user_volume, nvcre_bridge.WORKSPACE_VOLUME] + assert spec["volumeMounts"] == [user_mount, nvcre_bridge.WORKSPACE_MOUNT] + assert submitted["script"].env["PYTHONPATH"].startswith(f"{pod_scripts.rstrip('/')}:") + + +def test_nemorun_script_pythonpath_uses_the_pod_checkout_without_its_megatron_lm() -> None: + kwargs = dict( + script_path="/mnt/workspace/megatron-bridge/scripts/performance/run.py", + script_dir="/mnt/workspace/megatron-bridge/scripts/performance", + args=[], + kubeflow_namespace=None, + custom_env_vars={}, + ) + root = "/mnt/workspace/megatron-bridge" + script = setup_experiment._build_nemorun_script(**kwargs) + assert script.env["PYTHONPATH"] == f"{root}/scripts/performance:{root}/src:$PYTHONPATH" + + # ── _nvcre_env_vars ───────────────────────────────────────────────────────────