From a21ea1787a7356f00318f349aca57a30a88037bb Mon Sep 17 00:00:00 2001 From: Smitha Jayaram Date: Fri, 9 Oct 2026 13:43:27 +0000 Subject: [PATCH 1/6] Pass NVCRE init containers and a pod-side Megatron-Bridge root Signed-off-by: Smitha Jayaram --- scripts/performance/argument_parser.py | 22 ++++++++++++ scripts/performance/setup_experiment.py | 9 ++++- scripts/performance/utils/executors.py | 6 ++++ .../scripts/performance/test_executors.py | 34 +++++++++++++++++++ 4 files changed, 70 insertions(+), 1 deletion(-) diff --git a/scripts/performance/argument_parser.py b/scripts/performance/argument_parser.py index 89f43a6d1d..15b95dc072 100644 --- a/scripts/performance/argument_parser.py +++ b/scripts/performance/argument_parser.py @@ -752,6 +752,28 @@ def parse_cli_args(): required=False, default=None, ) + nvcre_args.add_argument( + "--nvcre_init_containers_json", + type=str, + help=( + "JSON-encoded list of Kubernetes init container dicts, emitted as the WorkloadRun's " + "spec.initContainers (e.g. to clone a pinned Megatron-Bridge into a shared volume). " + "Volumes they mount must also be given via --nvcre_volumes_json / --nvcre_volume_mounts_json." + ), + required=False, + default=None, + ) + nvcre_args.add_argument( + "--nvcre_bridge_root", + type=str, + help=( + "Megatron-Bridge checkout inside the pod that the job runs from. Defaults to the image's " + "/opt/Megatron-Bridge; set it to the path an init container (--nvcre_init_containers_json) " + "clones a pinned Megatron-Bridge into." + ), + required=False, + default="/opt/Megatron-Bridge", + ) nvcre_args.add_argument( "--nvcre_timeout_per_job", type=str, diff --git a/scripts/performance/setup_experiment.py b/scripts/performance/setup_experiment.py index fd12119cf7..3ce49575cc 100755 --- a/scripts/performance/setup_experiment.py +++ b/scripts/performance/setup_experiment.py @@ -578,6 +578,8 @@ def main( nvcre_node_selector_json: Optional[str] = None, nvcre_volumes_json: Optional[str] = None, nvcre_volume_mounts_json: Optional[str] = None, + nvcre_init_containers_json: Optional[str] = None, + nvcre_bridge_root: str = "/opt/Megatron-Bridge", nvcre_timeout_per_job: str = "24h", nvcre_kubeconfig: Optional[str] = None, nvcre_kube_context: Optional[str] = None, @@ -656,8 +658,10 @@ 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 the root can be redirected to a checkout an init container placed in the pod. if kubeflow_namespace or nvcre_namespace: - in_container_script_dir = "/opt/Megatron-Bridge/scripts/performance" + bridge_root = nvcre_bridge_root.rstrip("/") if nvcre_namespace 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) @@ -732,6 +736,7 @@ def main( 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, + init_containers=json.loads(nvcre_init_containers_json) if nvcre_init_containers_json else None, timeout_per_job=nvcre_timeout_per_job, kubeconfig=nvcre_kubeconfig, kube_context=nvcre_kube_context, @@ -1150,6 +1155,8 @@ 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_init_containers_json=args.nvcre_init_containers_json, + nvcre_bridge_root=args.nvcre_bridge_root, 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..afb2fe2284 100644 --- a/scripts/performance/utils/executors.py +++ b/scripts/performance/utils/executors.py @@ -381,6 +381,7 @@ def nvcre_executor( node_selector: Optional[dict] = None, volumes: Optional[list] = None, volume_mounts: Optional[list] = None, + init_containers: Optional[list] = None, timeout_per_job: str = "24h", kubeconfig: Optional[str] = None, kube_context: Optional[str] = None, @@ -400,6 +401,8 @@ 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. + 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. ``"24h"``). kubeconfig: Path to the kubeconfig file. Defaults to the cluster default. kube_context: Kubernetes context to use from the kubeconfig. @@ -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/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.""" From c1a1e0da669e1908b759ad73db4956e25a610ff3 Mon Sep 17 00:00:00 2001 From: Smitha Jayaram Date: Tue, 13 Oct 2026 09:26:10 +0000 Subject: [PATCH 2/6] Add --nvcre_use_mbridge_commit flag to construct the json required for NVCRE init_container in Run repo Signed-off-by: Smitha Jayaram --- scripts/performance/argument_parser.py | 24 +-- scripts/performance/setup_experiment.py | 64 ++++++-- scripts/performance/utils/nvcre_bridge.py | 100 ++++++++++++ .../scripts/performance/test_nvcre_bridge.py | 142 ++++++++++++++++++ .../performance/test_setup_experiment.py | 23 +++ 5 files changed, 323 insertions(+), 30 deletions(-) create mode 100644 scripts/performance/utils/nvcre_bridge.py create mode 100644 tests/unit_tests/scripts/performance/test_nvcre_bridge.py diff --git a/scripts/performance/argument_parser.py b/scripts/performance/argument_parser.py index 15b95dc072..9df45e0297 100644 --- a/scripts/performance/argument_parser.py +++ b/scripts/performance/argument_parser.py @@ -753,26 +753,14 @@ def parse_cli_args(): default=None, ) nvcre_args.add_argument( - "--nvcre_init_containers_json", - type=str, - help=( - "JSON-encoded list of Kubernetes init container dicts, emitted as the WorkloadRun's " - "spec.initContainers (e.g. to clone a pinned Megatron-Bridge into a shared volume). " - "Volumes they mount must also be given via --nvcre_volumes_json / --nvcre_volume_mounts_json." - ), - required=False, - default=None, - ) - nvcre_args.add_argument( - "--nvcre_bridge_root", - type=str, + "--nvcre_use_mbridge_commit", + action="store_true", help=( - "Megatron-Bridge checkout inside the pod that the job runs from. Defaults to the image's " - "/opt/Megatron-Bridge; set it to the path an init container (--nvcre_init_containers_json) " - "clones a pinned Megatron-Bridge into." + "Run this Megatron-Bridge checkout in the pods instead of the image's /opt/Megatron-Bridge, as the " + "Slurm bind mount does: an init container clones this checkout's origin at its current commit " + "(and its Megatron-LM submodule commit). Commits must be pushed to origin; uncommitted changes " + "are not included." ), - required=False, - default="/opt/Megatron-Bridge", ) nvcre_args.add_argument( "--nvcre_timeout_per_job", diff --git a/scripts/performance/setup_experiment.py b/scripts/performance/setup_experiment.py index 3ce49575cc..c8c1cdf026 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, + MCORE_SUBMODULE, + WORKSPACE_MOUNT, + WORKSPACE_VOLUME, + bridge_init_container, + local_bridge_source, + ) +except (ImportError, ModuleNotFoundError): + from .utils.nvcre_bridge import ( + BRIDGE_POD_ROOT, + MCORE_SUBMODULE, + WORKSPACE_MOUNT, + WORKSPACE_VOLUME, + 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 ( @@ -165,13 +184,20 @@ def _build_nemorun_script( args: List[str], kubeflow_namespace: Optional[str], custom_env_vars: Dict[str, str], + include_mcore: bool = False, ) -> run.Script: - """Build the rank-local task and apply optional Kubeflow NUMA binding.""" + """Build the rank-local task and apply optional Kubeflow NUMA binding. + + ``include_mcore`` puts the checkout's bundled Megatron-LM ahead of the image's on ``PYTHONPATH``. + """ in_container_repo_root = Path(script_dir).parents[1] + python_path = [script_dir, str(in_container_repo_root / "src")] + if include_mcore: + python_path.append(str(in_container_repo_root / MCORE_SUBMODULE)) 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,8 +604,7 @@ def main( nvcre_node_selector_json: Optional[str] = None, nvcre_volumes_json: Optional[str] = None, nvcre_volume_mounts_json: Optional[str] = None, - nvcre_init_containers_json: Optional[str] = None, - nvcre_bridge_root: str = "/opt/Megatron-Bridge", + nvcre_use_mbridge_commit: bool = False, nvcre_timeout_per_job: str = "24h", nvcre_kubeconfig: Optional[str] = None, nvcre_kube_context: Optional[str] = None, @@ -658,9 +683,21 @@ 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 the root can be redirected to a checkout an init container placed in the pod. + # On NVCRE, --nvcre_use_mbridge_commit has an init container clone this checkout's commits into the pod. + local_bridge = None + if nvcre_use_mbridge_commit: + if nvcre_namespace is None: + logger.warning("--nvcre_use_mbridge_commit only applies with --nvcre_namespace; ignoring it.") + else: + local_bridge = local_bridge_source(REPO_ROOT) + logger.info( + "Pods will run Megatron-Bridge %s@%s (Megatron-LM %s) from the init container", + local_bridge.url, + local_bridge.commit, + local_bridge.mcore_commit, + ) if kubeflow_namespace or nvcre_namespace: - bridge_root = nvcre_bridge_root.rstrip("/") if nvcre_namespace else "/opt/Megatron-Bridge" + 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: @@ -724,6 +761,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, @@ -734,9 +774,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, - init_containers=json.loads(nvcre_init_containers_json) if nvcre_init_containers_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, @@ -825,6 +865,7 @@ def main( args=_filter_run_script_args(sys.argv[1:]), kubeflow_namespace=kubeflow_namespace, custom_env_vars=custom_env_vars, + include_mcore=bool(local_bridge and local_bridge.mcore_commit), ) logger.info("Will launch the following command with Nemo-Run: %s", " ".join(nemorun_script.to_command())) @@ -1155,8 +1196,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_init_containers_json=args.nvcre_init_containers_json, - nvcre_bridge_root=args.nvcre_bridge_root, + 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/nvcre_bridge.py b/scripts/performance/utils/nvcre_bridge.py new file mode 100644 index 0000000000..f99b136a79 --- /dev/null +++ b/scripts/performance/utils/nvcre_bridge.py @@ -0,0 +1,100 @@ +# 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 commits (Megatron-Bridge and its Megatron-LM submodule) +into a volume that the training container shares. +""" + +import logging +import re +import shlex +import subprocess +from pathlib import Path +from typing import Any, NamedTuple +from urllib.parse import urlsplit, urlunsplit + + +logger = logging.getLogger(__name__) + +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"} +MCORE_SUBMODULE = "3rdparty/Megatron-LM" + + +class LocalBridge(NamedTuple): + """Where the launcher's checkout can be fetched from, and at which commits.""" + + url: str + commit: str + mcore_commit: str | None + + +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.""" + ssh = re.match(r"^(?:ssh://)?git@([^:/]+)[:/](.+)$", remote_url) + if ssh: + return f"https://{ssh.group(1)}/{ssh.group(2)}" + parts = urlsplit(remote_url) + 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 commits 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") + dirty = _git(repo_root, "status", "--porcelain", "--untracked-files=no") + # A submodule is a gitlink entry ("160000 commit \t") in the tree, checked out or not. + entry = _git(repo_root, "ls-tree", "HEAD", MCORE_SUBMODULE) + 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: + logger.warning("Local Megatron-Bridge has uncommitted changes; pods will run the committed %s only", commit) + return LocalBridge(url, commit, entry.split()[2] if entry.startswith("160000") else None) + + +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", + ] + if source.mcore_commit: + steps.append(f"git submodule update --init --depth 1 {MCORE_SUBMODULE}") + 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_nvcre_bridge.py b/tests/unit_tests/scripts/performance/test_nvcre_bridge.py new file mode 100644 index 0000000000..f85fbd807b --- /dev/null +++ b/tests/unit_tests/scripts/performance/test_nvcre_bridge.py @@ -0,0 +1,142 @@ +# 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 + + +MCORE_SHA = "b" * 40 + + +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 and a Megatron-LM gitlink, 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, "update-index", "--add", "--cacheinfo", f"160000,{MCORE_SHA},3rdparty/Megatron-LM") + _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"), + # credentials must never end up in the WorkloadRun spec + ("https://user:secret-token@github.com/o/r.git", "https://github.com/o/r.git"), + ("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_head_and_submodule_commit(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") + assert source.mcore_commit == MCORE_SHA # read from the tree; the submodule need not be checked out + + +def test_local_bridge_source_without_a_submodule_has_no_mcore_commit(tmp_path: Path) -> None: + 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") + + assert nvcre_bridge.local_bridge_source(repo).mcore_commit is None + + +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_warns_about_uncommitted_changes(checkout: Path, caplog: pytest.LogCaptureFixture) -> None: + (checkout / "README").write_text("changed") + + with caplog.at_level("WARNING"): + nvcre_bridge.local_bridge_source(checkout) + + assert "uncommitted" in caplog.text + + +def test_bridge_init_container_fetches_the_exact_commits() -> None: + source = nvcre_bridge.LocalBridge("https://github.com/o/r.git", "a" * 40, MCORE_SHA) + + 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 "git submodule update --init --depth 1 3rdparty/Megatron-LM" in script + assert nvcre_bridge.BRIDGE_POD_ROOT in script + + +def test_bridge_init_container_skips_the_submodule_when_there_is_none() -> None: + source = nvcre_bridge.LocalBridge("https://github.com/o/r.git", "a" * 40, None) + + assert "submodule" not in nvcre_bridge.bridge_init_container(source, "img")["args"][0] + + +def test_bridge_init_container_quotes_the_url() -> None: + source = nvcre_bridge.LocalBridge("https://github.com/o/r.git; rm -rf /", "a" * 40, None) + + 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..4e1c3dcbc1 100644 --- a/tests/unit_tests/scripts/performance/test_setup_experiment.py +++ b/tests/unit_tests/scripts/performance/test_setup_experiment.py @@ -196,6 +196,29 @@ 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_nemorun_script_puts_bundled_megatron_lm_on_pythonpath_only_when_asked() -> 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" + default = setup_experiment._build_nemorun_script(**kwargs) + assert default.env["PYTHONPATH"] == f"{root}/scripts/performance:{root}/src:$PYTHONPATH" + with_mcore = setup_experiment._build_nemorun_script(**kwargs, include_mcore=True) + assert with_mcore.env["PYTHONPATH"] == ( + f"{root}/scripts/performance:{root}/src:{root}/3rdparty/Megatron-LM:$PYTHONPATH" + ) + + # ── _nvcre_env_vars ─────────────────────────────────────────────────────────── From 9fcc2fd31543202429266e10cba0207db2602678 Mon Sep 17 00:00:00 2001 From: Smitha Jayaram Date: Tue, 13 Oct 2026 10:55:04 +0000 Subject: [PATCH 3/6] Add Unpinned placement for NVCRE jobs Signed-off-by: Smitha Jayaram --- scripts/performance/utils/executors.py | 5 +++++ .../unit_tests/scripts/performance/test_executors.py | 11 +++++++++++ 2 files changed, 16 insertions(+) diff --git a/scripts/performance/utils/executors.py b/scripts/performance/utils/executors.py index afb2fe2284..375be9abc4 100644 --- a/scripts/performance/utils/executors.py +++ b/scripts/performance/utils/executors.py @@ -35,6 +35,10 @@ KUBEFLOW_NUMA_BINDING_ENV = "NEMO_KUBEFLOW_NUMA_BINDING" +# WorkloadRun ``orchestration.placement`` (NVCRE >= v0.7.0): run one job of exactly the requested number of nodes +# instead of covering the whole target. +NVCRE_PLACEMENT = "Unpinned" + def _kubeflow_numa_binding_script(task: run.Script) -> run.Script: """Wrap a task with per-rank GPU-local NUMA binding without dropping task metadata.""" @@ -434,5 +438,6 @@ def nvcre_executor( kubeconfig=kubeconfig, kube_context=kube_context, gang_scheduler_name=gang_scheduler_name, + placement=NVCRE_PLACEMENT, **extra, ) diff --git a/tests/unit_tests/scripts/performance/test_executors.py b/tests/unit_tests/scripts/performance/test_executors.py index 44ea5bcf87..4e58010d87 100644 --- a/tests/unit_tests/scripts/performance/test_executors.py +++ b/tests/unit_tests/scripts/performance/test_executors.py @@ -313,6 +313,17 @@ def test_nvcre_executor_forwards_init_containers(): assert executor.init_containers == [init_container] +@pytest.mark.skipif(not HAS_NVCRE, reason="nemo_run nvcre not available") +def test_nvcre_executor_always_requests_unpinned_placement(): + """Jobs run one job of exactly the requested size, not one per group across the whole target.""" + executor = nvcre_executor(namespace="ns", container_image="img:latest", nodes=2, num_gpus_per_node=4) + executor.job_name = "job1" + executor.experiment_id = "exp_123456789" + + assert executor.placement == "Unpinned" + assert executor.build_workloadrun_yaml(["python"])["spec"]["orchestration"]["placement"] == "Unpinned" + + @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.""" From 6242721d51688a6b358e5927ca004c6cde3a0881 Mon Sep 17 00:00:00 2001 From: Smitha Jayaram Date: Wed, 14 Oct 2026 08:27:32 +0000 Subject: [PATCH 4/6] Clone only the Megatron-Bridge commit in the NVCRE init container. Signed-off-by: Smitha Jayaram --- scripts/performance/argument_parser.py | 9 +- scripts/performance/setup_experiment.py | 37 ++++---- scripts/performance/utils/executors.py | 5 -- scripts/performance/utils/nvcre_bridge.py | 40 ++++----- .../scripts/performance/test_executors.py | 11 --- .../scripts/performance/test_nvcre_bridge.py | 88 ++++++++++++------- .../performance/test_setup_experiment.py | 85 ++++++++++++++++-- 7 files changed, 168 insertions(+), 107 deletions(-) diff --git a/scripts/performance/argument_parser.py b/scripts/performance/argument_parser.py index 9df45e0297..19cb9076d8 100644 --- a/scripts/performance/argument_parser.py +++ b/scripts/performance/argument_parser.py @@ -755,17 +755,12 @@ def parse_cli_args(): 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, as the " - "Slurm bind mount does: an init container clones this checkout's origin at its current commit " - "(and its Megatron-LM submodule commit). Commits must be pushed to origin; uncommitted changes " - "are not included." - ), + 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", ) diff --git a/scripts/performance/setup_experiment.py b/scripts/performance/setup_experiment.py index c8c1cdf026..ca24e568bb 100755 --- a/scripts/performance/setup_experiment.py +++ b/scripts/performance/setup_experiment.py @@ -64,18 +64,18 @@ try: from utils.nvcre_bridge import ( BRIDGE_POD_ROOT, - MCORE_SUBMODULE, WORKSPACE_MOUNT, WORKSPACE_VOLUME, + LocalBridge, bridge_init_container, local_bridge_source, ) except (ImportError, ModuleNotFoundError): from .utils.nvcre_bridge import ( BRIDGE_POD_ROOT, - MCORE_SUBMODULE, WORKSPACE_MOUNT, WORKSPACE_VOLUME, + LocalBridge, bridge_init_container, local_bridge_source, ) @@ -177,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, @@ -184,16 +195,10 @@ def _build_nemorun_script( args: List[str], kubeflow_namespace: Optional[str], custom_env_vars: Dict[str, str], - include_mcore: bool = False, ) -> run.Script: - """Build the rank-local task and apply optional Kubeflow NUMA binding. - - ``include_mcore`` puts the checkout's bundled Megatron-LM ahead of the image's on ``PYTHONPATH``. - """ + """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")] - if include_mcore: - python_path.append(str(in_container_repo_root / MCORE_SUBMODULE)) task = run.Script( path=script_path, entrypoint="python", @@ -684,18 +689,7 @@ def main( # /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 = None - if nvcre_use_mbridge_commit: - if nvcre_namespace is None: - logger.warning("--nvcre_use_mbridge_commit only applies with --nvcre_namespace; ignoring it.") - else: - local_bridge = local_bridge_source(REPO_ROOT) - logger.info( - "Pods will run Megatron-Bridge %s@%s (Megatron-LM %s) from the init container", - local_bridge.url, - local_bridge.commit, - local_bridge.mcore_commit, - ) + local_bridge = _resolve_local_bridge(nvcre_use_mbridge_commit, nvcre_namespace) if kubeflow_namespace or nvcre_namespace: bridge_root = BRIDGE_POD_ROOT if local_bridge else "/opt/Megatron-Bridge" in_container_script_dir = f"{bridge_root}/scripts/performance" @@ -865,7 +859,6 @@ def main( args=_filter_run_script_args(sys.argv[1:]), kubeflow_namespace=kubeflow_namespace, custom_env_vars=custom_env_vars, - include_mcore=bool(local_bridge and local_bridge.mcore_commit), ) logger.info("Will launch the following command with Nemo-Run: %s", " ".join(nemorun_script.to_command())) diff --git a/scripts/performance/utils/executors.py b/scripts/performance/utils/executors.py index 375be9abc4..afb2fe2284 100644 --- a/scripts/performance/utils/executors.py +++ b/scripts/performance/utils/executors.py @@ -35,10 +35,6 @@ KUBEFLOW_NUMA_BINDING_ENV = "NEMO_KUBEFLOW_NUMA_BINDING" -# WorkloadRun ``orchestration.placement`` (NVCRE >= v0.7.0): run one job of exactly the requested number of nodes -# instead of covering the whole target. -NVCRE_PLACEMENT = "Unpinned" - def _kubeflow_numa_binding_script(task: run.Script) -> run.Script: """Wrap a task with per-rank GPU-local NUMA binding without dropping task metadata.""" @@ -438,6 +434,5 @@ def nvcre_executor( kubeconfig=kubeconfig, kube_context=kube_context, gang_scheduler_name=gang_scheduler_name, - placement=NVCRE_PLACEMENT, **extra, ) diff --git a/scripts/performance/utils/nvcre_bridge.py b/scripts/performance/utils/nvcre_bridge.py index f99b136a79..475dff104b 100644 --- a/scripts/performance/utils/nvcre_bridge.py +++ b/scripts/performance/utils/nvcre_bridge.py @@ -14,12 +14,10 @@ """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 commits (Megatron-Bridge and its Megatron-LM submodule) -into a volume that the training container shares. -""" +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 logging import re import shlex import subprocess @@ -28,20 +26,16 @@ from urllib.parse import urlsplit, urlunsplit -logger = logging.getLogger(__name__) - 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"} -MCORE_SUBMODULE = "3rdparty/Megatron-LM" class LocalBridge(NamedTuple): - """Where the launcher's checkout can be fetched from, and at which commits.""" + """Where the launcher's checkout can be fetched from, and at which commit.""" url: str commit: str - mcore_commit: str | None def _git(repo_root: Path, *args: str) -> str: @@ -52,10 +46,14 @@ def _git(repo_root: Path, *args: str) -> str: def pod_clone_url(remote_url: str) -> str: """HTTPS URL a pod can clone anonymously: SSH remotes are rewritten and any credentials dropped.""" - ssh = re.match(r"^(?:ssh://)?git@([^:/]+)[:/](.+)$", remote_url) - if ssh: - return f"https://{ssh.group(1)}/{ssh.group(2)}" + # 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, "", "")) @@ -63,20 +61,22 @@ def pod_clone_url(remote_url: str) -> str: def local_bridge_source(repo_root: Path) -> LocalBridge: - """The origin URL and checked-out commits of the Megatron-Bridge checkout at ``repo_root``.""" + """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") - dirty = _git(repo_root, "status", "--porcelain", "--untracked-files=no") - # A submodule is a gitlink entry ("160000 commit \t") in the tree, checked out or not. - entry = _git(repo_root, "ls-tree", "HEAD", MCORE_SUBMODULE) + # 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: - logger.warning("Local Megatron-Bridge has uncommitted changes; pods will run the committed %s only", commit) - return LocalBridge(url, commit, entry.split()[2] if entry.startswith("160000") else None) + 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]: @@ -89,8 +89,6 @@ def bridge_init_container(source: LocalBridge, image: str) -> dict[str, Any]: f"git fetch -q --depth 1 origin {shlex.quote(source.commit)}", "git checkout -q FETCH_HEAD", ] - if source.mcore_commit: - steps.append(f"git submodule update --init --depth 1 {MCORE_SUBMODULE}") return { "name": "megatron-bridge-clone", "image": image, diff --git a/tests/unit_tests/scripts/performance/test_executors.py b/tests/unit_tests/scripts/performance/test_executors.py index 4e58010d87..44ea5bcf87 100644 --- a/tests/unit_tests/scripts/performance/test_executors.py +++ b/tests/unit_tests/scripts/performance/test_executors.py @@ -313,17 +313,6 @@ def test_nvcre_executor_forwards_init_containers(): assert executor.init_containers == [init_container] -@pytest.mark.skipif(not HAS_NVCRE, reason="nemo_run nvcre not available") -def test_nvcre_executor_always_requests_unpinned_placement(): - """Jobs run one job of exactly the requested size, not one per group across the whole target.""" - executor = nvcre_executor(namespace="ns", container_image="img:latest", nodes=2, num_gpus_per_node=4) - executor.job_name = "job1" - executor.experiment_id = "exp_123456789" - - assert executor.placement == "Unpinned" - assert executor.build_workloadrun_yaml(["python"])["spec"]["orchestration"]["placement"] == "Unpinned" - - @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 index f85fbd807b..84aa68bfa9 100644 --- a/tests/unit_tests/scripts/performance/test_nvcre_bridge.py +++ b/tests/unit_tests/scripts/performance/test_nvcre_bridge.py @@ -30,16 +30,13 @@ from utils import nvcre_bridge # noqa: E402 -MCORE_SHA = "b" * 40 - - 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 and a Megatron-LM gitlink, like a real Megatron-Bridge clone.""" + """A git checkout with an origin remote, like a real Megatron-Bridge clone.""" repo = tmp_path / "bridge" repo.mkdir() _git(repo, "init", "-q") @@ -48,7 +45,6 @@ def checkout(tmp_path: Path) -> Path: _git(repo, "remote", "add", "origin", "https://github.com/example/Megatron-Bridge.git") (repo / "README").write_text("x") _git(repo, "add", "README") - _git(repo, "update-index", "--add", "--cacheinfo", f"160000,{MCORE_SHA},3rdparty/Megatron-LM") _git(repo, "commit", "-q", "-m", "init") return repo @@ -59,6 +55,9 @@ def checkout(tmp_path: Path) -> Path: ("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"), ("https://gitlab.example.com:8443/o/r.git", "https://gitlab.example.com:8443/o/r.git"), @@ -74,26 +73,11 @@ def test_pod_clone_url_rejects_remotes_a_pod_cannot_reach(remote: str) -> None: nvcre_bridge.pod_clone_url(remote) -def test_local_bridge_source_reads_origin_head_and_submodule_commit(checkout: Path) -> None: +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") - assert source.mcore_commit == MCORE_SHA # read from the tree; the submodule need not be checked out - - -def test_local_bridge_source_without_a_submodule_has_no_mcore_commit(tmp_path: Path) -> None: - 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") - - assert nvcre_bridge.local_bridge_source(repo).mcore_commit is None def test_local_bridge_source_requires_a_git_checkout_with_origin(tmp_path: Path) -> None: @@ -101,17 +85,59 @@ def test_local_bridge_source_requires_a_git_checkout_with_origin(tmp_path: Path) nvcre_bridge.local_bridge_source(tmp_path) -def test_local_bridge_source_warns_about_uncommitted_changes(checkout: Path, caplog: pytest.LogCaptureFixture) -> None: +def test_local_bridge_source_rejects_uncommitted_changes(checkout: Path) -> None: (checkout / "README").write_text("changed") - with caplog.at_level("WARNING"): + with pytest.raises(RuntimeError, match="uncommitted"): nvcre_bridge.local_bridge_source(checkout) - assert "uncommitted" in caplog.text + +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_bridge_init_container_fetches_the_exact_commits() -> None: - source = nvcre_bridge.LocalBridge("https://github.com/o/r.git", "a" * 40, MCORE_SHA) +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") @@ -121,18 +147,12 @@ def test_bridge_init_container_fetches_the_exact_commits() -> None: 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 "git submodule update --init --depth 1 3rdparty/Megatron-LM" 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_skips_the_submodule_when_there_is_none() -> None: - source = nvcre_bridge.LocalBridge("https://github.com/o/r.git", "a" * 40, None) - - assert "submodule" not in nvcre_bridge.bridge_init_container(source, "img")["args"][0] - - def test_bridge_init_container_quotes_the_url() -> None: - source = nvcre_bridge.LocalBridge("https://github.com/o/r.git; rm -rf /", "a" * 40, 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] diff --git a/tests/unit_tests/scripts/performance/test_setup_experiment.py b/tests/unit_tests/scripts/performance/test_setup_experiment.py index 4e1c3dcbc1..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( @@ -202,7 +213,71 @@ def test_nvcre_use_mbridge_commit_flag_does_not_swallow_the_next_arg() -> None: assert setup_experiment._filter_run_script_args(argv) == ["--tensor_model_parallel_size", "2"] -def test_nemorun_script_puts_bundled_megatron_lm_on_pythonpath_only_when_asked() -> None: +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", @@ -211,12 +286,8 @@ def test_nemorun_script_puts_bundled_megatron_lm_on_pythonpath_only_when_asked() custom_env_vars={}, ) root = "/mnt/workspace/megatron-bridge" - default = setup_experiment._build_nemorun_script(**kwargs) - assert default.env["PYTHONPATH"] == f"{root}/scripts/performance:{root}/src:$PYTHONPATH" - with_mcore = setup_experiment._build_nemorun_script(**kwargs, include_mcore=True) - assert with_mcore.env["PYTHONPATH"] == ( - f"{root}/scripts/performance:{root}/src:{root}/3rdparty/Megatron-LM:$PYTHONPATH" - ) + script = setup_experiment._build_nemorun_script(**kwargs) + assert script.env["PYTHONPATH"] == f"{root}/scripts/performance:{root}/src:$PYTHONPATH" # ── _nvcre_env_vars ─────────────────────────────────────────────────────────── From 8111854cefe12a95b0e7bb9aec0e90f8892c80c7 Mon Sep 17 00:00:00 2001 From: Smitha Jayaram Date: Wed, 14 Oct 2026 18:10:03 +0000 Subject: [PATCH 5/6] Allowlist the fake credential in the pod clone URL test Default the NVCRE per-job timeout to 12h Signed-off-by: Smitha Jayaram --- scripts/performance/argument_parser.py | 2 +- scripts/performance/setup_experiment.py | 2 +- scripts/performance/utils/executors.py | 4 ++-- tests/unit_tests/scripts/performance/test_nvcre_bridge.py | 2 +- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/scripts/performance/argument_parser.py b/scripts/performance/argument_parser.py index 19cb9076d8..6e6d33561b 100644 --- a/scripts/performance/argument_parser.py +++ b/scripts/performance/argument_parser.py @@ -762,7 +762,7 @@ def parse_cli_args(): type=str, 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 ca24e568bb..519171963e 100755 --- a/scripts/performance/setup_experiment.py +++ b/scripts/performance/setup_experiment.py @@ -610,7 +610,7 @@ def main( nvcre_volumes_json: Optional[str] = None, nvcre_volume_mounts_json: Optional[str] = None, nvcre_use_mbridge_commit: bool = False, - nvcre_timeout_per_job: str = "24h", + nvcre_timeout_per_job: str = "12h", nvcre_kubeconfig: Optional[str] = None, nvcre_kube_context: Optional[str] = None, nvcre_gang_scheduler_name: Optional[str] = None, diff --git a/scripts/performance/utils/executors.py b/scripts/performance/utils/executors.py index afb2fe2284..97e5fee588 100644 --- a/scripts/performance/utils/executors.py +++ b/scripts/performance/utils/executors.py @@ -382,7 +382,7 @@ def nvcre_executor( volumes: Optional[list] = None, volume_mounts: Optional[list] = None, init_containers: Optional[list] = None, - timeout_per_job: str = "24h", + timeout_per_job: str = "12h", kubeconfig: Optional[str] = None, kube_context: Optional[str] = None, gang_scheduler_name: Optional[str] = None, @@ -403,7 +403,7 @@ def nvcre_executor( volume_mounts: Additional Kubernetes volume mount specs. 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. ``"24h"``). + 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"``). diff --git a/tests/unit_tests/scripts/performance/test_nvcre_bridge.py b/tests/unit_tests/scripts/performance/test_nvcre_bridge.py index 84aa68bfa9..7684858f47 100644 --- a/tests/unit_tests/scripts/performance/test_nvcre_bridge.py +++ b/tests/unit_tests/scripts/performance/test_nvcre_bridge.py @@ -59,7 +59,7 @@ def checkout(tmp_path: Path) -> 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"), + ("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"), ], ) From bb0d2611e220802007c0f0f474fc3627e91aa7e9 Mon Sep 17 00:00:00 2001 From: Alex Filby Date: Wed, 14 Oct 2026 18:02:16 -0500 Subject: [PATCH 6/6] Update scripts/performance/setup_experiment.py Signed-off-by: Alex Filby --- scripts/performance/setup_experiment.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/performance/setup_experiment.py b/scripts/performance/setup_experiment.py index 519171963e..8b732a81bf 100755 --- a/scripts/performance/setup_experiment.py +++ b/scripts/performance/setup_experiment.py @@ -1101,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,