Support for NVCRE init containers (#6375)

Signed-off-by: Smitha Jayaram <sjayaram@nvidia.com>
Signed-off-by: Alex Filby <afilby@nvidia.com>

Author: sjayaram-nv

scripts/performance/argument_parser.py

@@ -738,12 +738,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(

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,

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,
     )

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
+    }

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."""