diff --git a/nemo_run/__init__.py b/nemo_run/__init__.py
index 0d403a54..d703fd30 100644
--- a/nemo_run/__init__.py
+++ b/nemo_run/__init__.py
@@ -26,6 +26,7 @@
from nemo_run.core.execution.docker import DockerExecutor
from nemo_run.core.execution.kubeflow import KubeflowExecutor
from nemo_run.core.execution.launcher import FaultTolerance, SlurmRay, SlurmTemplate, Torchrun
+from nemo_run.core.execution.nvcre import NvcreExecutor
from nemo_run.core.execution.lepton import LeptonExecutor
from nemo_run.core.execution.local import LocalExecutor
from nemo_run.core.execution.skypilot import SkypilotExecutor
@@ -68,6 +69,7 @@
"Partial",
"Plugin",
"KubeflowExecutor",
+ "NvcreExecutor",
"run",
"Script",
"SkypilotExecutor",
diff --git a/nemo_run/core/execution/__init__.py b/nemo_run/core/execution/__init__.py
index 08e088c8..b4f22fd7 100644
--- a/nemo_run/core/execution/__init__.py
+++ b/nemo_run/core/execution/__init__.py
@@ -14,11 +14,12 @@
# limitations under the License.
from nemo_run.core.execution.dgxcloud import DGXCloudExecutor
+from nemo_run.core.execution.kubeflow import KubeflowExecutor
from nemo_run.core.execution.lepton import LeptonExecutor
from nemo_run.core.execution.local import LocalExecutor
-from nemo_run.core.execution.kubeflow import KubeflowExecutor
from nemo_run.core.execution.skypilot import SkypilotExecutor
from nemo_run.core.execution.slurm import SlurmExecutor
+from nemo_run.core.execution.nvcre import NvcreExecutor
__all__ = [
"LocalExecutor",
@@ -27,4 +28,5 @@
"DGXCloudExecutor",
"LeptonExecutor",
"KubeflowExecutor",
+ "NvcreExecutor",
]
diff --git a/nemo_run/core/execution/nvcre.py b/nemo_run/core/execution/nvcre.py
new file mode 100644
index 00000000..84c193c0
--- /dev/null
+++ b/nemo_run/core/execution/nvcre.py
@@ -0,0 +1,651 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+import getpass
+import hashlib
+import json
+import logging
+import os
+import shlex
+import subprocess
+import tempfile
+import time
+from dataclasses import dataclass, field
+from enum import Enum
+from pathlib import Path
+from typing import Any, Iterable, Optional
+
+import yaml
+
+from nemo_run.core.execution.base import Executor, ExecutorMacros
+from nemo_run.core.execution.launcher import Launcher
+from nemo_run.core.packaging.base import Packager
+from nemo_run.core.packaging.git import GitArchivePackager
+
+logger = logging.getLogger(__name__)
+
+_NVCRE_WORKLOADRUN_API = "nvcre.nvidia.com/v1alpha1"
+_DATA_MOVER_IMAGE = "alpine:3.19"
+
+
+class NvcrePhase(Enum):
+ PENDING = "Pending"
+ IN_PROGRESS = "InProgress"
+ SUCCEEDED = "Succeeded"
+ FAILED = "Failed"
+ UNKNOWN = "Unknown"
+
+
+@dataclass(kw_only=True)
+class NvcreExecutor(Executor):
+ """
+ Dataclass to configure an Nvcre executor.
+
+ Submits jobs to an Nvcre-managed Kubernetes cluster via the ``nvcrectl``
+ CLI using the WorkloadRun API. Requires ``nvcrectl`` (and ``kubectl``) to be
+ on the PATH of the machine running NeMo-Run.
+
+ Example::
+
+ executor = NvcreExecutor(
+ namespace="nemo-perf",
+ container_image="nvcr.io/nvidia/nemo:dev",
+ num_nodes=8,
+ gpus_per_node=8,
+ image_pull_secret="ngc-registry",
+ workdir_pvc="nemo-run-pvc",
+ )
+ """
+
+ # ── Required ──────────────────────────────────────────────────────────────
+ namespace: str
+ container_image: str
+ num_nodes: int = 1
+
+ # ── Compute shape ─────────────────────────────────────────────────────────
+ gpus_per_node: int = 0 # 0 = auto-detect by Nvcre
+
+ # ── Registry auth ─────────────────────────────────────────────────────────
+ image_pull_secret: Optional[str] = None
+
+ # ── Node targeting ────────────────────────────────────────────────────────
+ node_selector: dict[str, str] = field(default_factory=dict)
+
+ # ── Storage ───────────────────────────────────────────────────────────────
+ # When set, job_dir is synced to this PVC before WorkloadRun submission.
+ workdir_pvc: Optional[str] = None
+ workdir_pvc_path: str = "/nemo_run"
+ # Optional local overlay dir (e.g. a mbridge-ref checkout) merged into job_dir.
+ workdir_local_path: Optional[str] = None
+
+ # ── Extra pod config ──────────────────────────────────────────────────────
+ volumes: list[dict[str, Any]] = field(default_factory=list)
+ volume_mounts: list[dict[str, Any]] = field(default_factory=list)
+
+ # ── Orchestration ─────────────────────────────────────────────────────────
+ timeout_per_job: str = "24h"
+ test_scale: Optional[str] = None # "intra-node" | "intra-rack" | "full-scale"
+ max_restarts: int = 0
+ # Set to enable checkpointing; PVC size is required by the API (e.g. "500Gi").
+ checkpoint_storage_size: Optional[str] = None
+ checkpoint_storage_class: Optional[str] = None # defaults to cluster default
+
+ # ── Launcher ──────────────────────────────────────────────────────────────
+ # When True, wrap the python entrypoint with torchrun using the PET_* env
+ # vars that Nvcre injects per-pod (PET_NNODES, PET_NPROC_PER_NODE,
+ # PET_NODE_RANK, PET_MASTER_ADDR, PET_MASTER_PORT). This causes
+ # torch.distributed to be initialised correctly so that WORLD_SIZE,
+ # RANK, LOCAL_RANK, and MASTER_ADDR are set for every spawned process.
+ # Without this, Megatron defaults to world_size=1 and fails the
+ # expert_tensor_model_pipeline_parallel divisibility check.
+ use_torchrun: bool = True
+
+ # ── Scheduling ────────────────────────────────────────────────────────────
+ gang_scheduler_name: Optional[str] = None # e.g. "kai-scheduler"
+
+ # ── Profiling ─────────────────────────────────────────────────────────────
+ # Set by NsysPlugin.setup(); holds nsys configuration when profiling is enabled.
+ launcher: Optional[Launcher] = None
+
+ # ── nvcrectl / kubectl config ──────────────────────────────────────────────
+ nvcrectl_bin: str = "nvcrectl"
+ kubeconfig: Optional[str] = None
+ kube_context: Optional[str] = None
+
+ # ── Set by assign() ───────────────────────────────────────────────────────
+ job_name: str = field(init=False, default="")
+
+ # ── Internal ──────────────────────────────────────────────────────────────
+ _workloadrun_name: Optional[str] = field(init=False, default=None, repr=False)
+
+ # ── Executor interface ────────────────────────────────────────────────────
+
+ def assign(self, exp_id: str, exp_dir: str, task_id: str, task_dir: str) -> None:
+ self.experiment_id = exp_id
+ self.experiment_dir = exp_dir
+ self.job_name = task_id
+ self.job_dir = os.path.join(exp_dir, task_dir)
+
+ def get_launcher_prefix(self) -> Optional[list[str]]:
+ """Return nsys prefix when profiling is enabled, else None."""
+ launcher = self.get_launcher()
+ if launcher.nsys_profile:
+ nsys_dir = os.path.join(self.job_dir, launcher.nsys_folder)
+ os.makedirs(nsys_dir, exist_ok=True)
+ return launcher.get_nsys_prefix(profile_dir=self.job_dir)
+ return None
+
+ def nnodes(self) -> int:
+ return self.num_nodes
+
+ def nproc_per_node(self) -> int:
+ return self.gpus_per_node or 1
+
+ def macro_values(self) -> ExecutorMacros:
+ # Nvcre uses the Kubeflow Training Operator under the hood; the
+ # PET_* vars are injected by the torchrun entrypoint of the TrainJob.
+ return ExecutorMacros(
+ head_node_ip_var="PET_MASTER_ADDR",
+ nproc_per_node_var="PET_NPROC_PER_NODE",
+ num_nodes_var="PET_NNODES",
+ node_rank_var="PET_NODE_RANK",
+ het_group_host_var="PET_MASTER_ADDR",
+ )
+
+ # ── WorkloadRun YAML builder ──────────────────────────────────────────────
+
+ @property
+ def code_dir(self) -> str:
+ """Remote directory on the PVC where job code is placed."""
+ user = getpass.getuser()
+ parts = [p for p in (getattr(self, "experiment_id", None), getattr(self, "job_name", None)) if p]
+ scope = "/".join([user, *parts])
+ return f"{self.workdir_pvc_path.rstrip('/')}/{scope}/code"
+
+ def build_workloadrun_yaml(self, cmd: list[str]) -> dict:
+ """Return the WorkloadRun manifest as a dict."""
+ spec: dict[str, Any] = {
+ "image": self.container_image,
+ "numNodes": self.num_nodes,
+ "framework": {"exec": {"command": cmd}},
+ }
+ if self.gpus_per_node:
+ spec["gpusPerNode"] = self.gpus_per_node
+ if self.node_selector:
+ spec["target"] = {"nodeSelector": self.node_selector}
+
+ env_list = [{"name": k, "value": v} for k, v in self.env_vars.items()]
+ if env_list:
+ spec["env"] = env_list
+
+ vols = list(self.volumes)
+ vmounts = list(self.volume_mounts)
+ if vols:
+ spec["volumes"] = vols
+ if vmounts:
+ spec["volumeMounts"] = vmounts
+
+ if self.image_pull_secret:
+ spec["imagePullSecrets"] = [{"name": self.image_pull_secret}]
+
+ orch: dict[str, Any] = {}
+ if self.timeout_per_job:
+ orch["timeoutPerJob"] = self.timeout_per_job
+ if self.test_scale:
+ orch["testScale"] = self.test_scale
+ if orch:
+ spec["orchestration"] = orch
+
+ if self.checkpoint_storage_size:
+ checkpoint: dict[str, Any] = {"storageSize": self.checkpoint_storage_size}
+ if self.checkpoint_storage_class:
+ checkpoint["storageClassName"] = self.checkpoint_storage_class
+ if self.max_restarts:
+ checkpoint["maxRestarts"] = self.max_restarts
+ spec["checkpoint"] = checkpoint
+
+ if self.gang_scheduler_name:
+ spec["gangScheduler"] = {"schedulerName": self.gang_scheduler_name}
+
+ return {
+ "apiVersion": _NVCRE_WORKLOADRUN_API,
+ "kind": "WorkloadRun",
+ "metadata": {"name": self._safe_name(), "namespace": self.namespace},
+ "spec": spec,
+ }
+
+ def _safe_name(self) -> str:
+ """The last 6 digits of the nanosecond timestamp embedded in experiment_id
+ (format: "
_") are used as a suffix so that repeated
+ submissions of the same job produce unique names.
+ """
+ exp_id = getattr(self, "experiment_id", None) or ""
+ job = (self.job_name or "nvcre-job").lower().replace("_", "-").replace(".", "-")
+ # exp_id format is "_"; take the last 6 digits of the ns part.
+ ns_part = exp_id.rsplit("_", 1)[-1] if "_" in exp_id else ""
+ suffix = ns_part[-6:] if ns_part.isdigit() else hashlib.sha256(exp_id.encode()).hexdigest()[:6]
+ # Reserve 7 chars for "-"; truncate base to fit within 63 total.
+ base = job[:56].rstrip("-")
+ return f"{base}-{suffix}"
+
+ # ── nvcrectl / kubectl helpers ─────────────────────────────────────────────
+
+ def _nvcrectl_base(self) -> list[str]:
+ args = [self.nvcrectl_bin]
+ if self.kubeconfig:
+ args += ["--kubeconfig", self.kubeconfig]
+ if self.kube_context:
+ args += ["--context", self.kube_context]
+ return args
+
+ def _kubectl_base(self) -> list[str]:
+ args = ["kubectl"]
+ if self.kubeconfig:
+ args += ["--kubeconfig", self.kubeconfig]
+ if self.kube_context:
+ args += ["--context", self.kube_context]
+ return args
+
+ def submit(self, yaml_path: str) -> str:
+ """Submit a WorkloadRun YAML and return the workloadrun name."""
+ name = self._safe_name()
+ cmd = self._nvcrectl_base() + [
+ "workloadrun", "run", yaml_path,
+ "--namespace", self.namespace,
+ "--name", name,
+ ]
+ logger.info("Submitting WorkloadRun: %s", " ".join(cmd))
+ result = subprocess.run(cmd, capture_output=True, text=True)
+ if result.returncode != 0:
+ raise RuntimeError(
+ f"nvcrectl workloadrun run failed (rc={result.returncode}):\n{result.stderr}"
+ )
+ logger.info("WorkloadRun '%s' submitted", name)
+ self._workloadrun_name = name
+ return name
+
+ def status(self, name: str) -> NvcrePhase:
+ """Return the current phase of WorkloadRun *name*.
+
+ Tries nvcrectl first. Falls back to inspecting pod phases via kubectl
+ when nvcrectl returns a non-zero exit code (e.g. the WorkloadRun was
+ cleaned up after completion) or reports an unrecognised phase string.
+ """
+ cmd = self._nvcrectl_base() + [
+ "workloadrun", "status", name,
+ "-n", self.namespace,
+ ]
+ result = subprocess.run(cmd, capture_output=True, text=True)
+ if result.returncode == 0:
+ phase_str = result.stdout.strip()
+ try:
+ return NvcrePhase(phase_str)
+ except ValueError:
+ logger.warning(
+ "Unrecognised nvcrectl phase '%s' for '%s'; falling back to kubectl CRD check",
+ phase_str, name,
+ )
+ else:
+ logger.warning(
+ "nvcrectl status failed for '%s' (rc=%d): %s; falling back to kubectl CRD check",
+ name, result.returncode, result.stderr.strip(),
+ )
+
+ return self._kubectl_workloadrun_crd_phase(name)
+
+ def _kubectl_workloadrun_crd_phase(self, name: str) -> NvcrePhase:
+ """Read phase directly from the WorkloadRun CRD via kubectl.
+
+ nvcrectl is a thin wrapper over the same CRD. Reading it directly
+ avoids nvcrectl output-format surprises and works regardless of whether
+ Nvcre's internal job name differs from the WorkloadRun CRD name.
+ """
+ cmd = self._kubectl_base() + [
+ "get", "workloadrun", name,
+ "-n", self.namespace,
+ "-o", "jsonpath={.status.phase}",
+ ]
+ result = subprocess.run(cmd, capture_output=True, text=True)
+ if result.returncode != 0:
+ logger.warning(
+ "kubectl workloadrun CRD check failed for '%s': %s",
+ name, result.stderr.strip(),
+ )
+ return NvcrePhase.UNKNOWN
+
+ phase_str = result.stdout.strip()
+ if not phase_str:
+ logger.warning("Empty phase from WorkloadRun CRD '%s'", name)
+ return NvcrePhase.UNKNOWN
+
+ try:
+ return NvcrePhase(phase_str)
+ except ValueError:
+ logger.warning(
+ "Unrecognised WorkloadRun CRD phase '%s' for '%s'", phase_str, name,
+ )
+ return NvcrePhase.UNKNOWN
+
+ def cancel(self, name: str) -> None:
+ """Cancel WorkloadRun *name*."""
+ cmd = self._nvcrectl_base() + [
+ "workloadrun", "cancel", name,
+ "-n", self.namespace,
+ ]
+ result = subprocess.run(cmd, capture_output=True, text=True)
+ if result.returncode != 0:
+ logger.warning("nvcrectl cancel failed for '%s': %s", name, result.stderr)
+ else:
+ logger.info("Cancelled WorkloadRun '%s'", name)
+
+ def _get_nvcre_job_name(self, workloadrun_name: str) -> str | None:
+ """Return the Nvcre internal job name from the WorkloadRun CRD.
+
+ Nvcre stamps pods with ``nvcre.nvidia.com/job=``
+ which may differ from the WorkloadRun CRD name we submitted. Try to
+ retrieve it from the CRD status/labels so log and pod queries work.
+ """
+ cmd = self._kubectl_base() + [
+ "get", "workloadrun", workloadrun_name,
+ "-n", self.namespace,
+ "-o", "json",
+ ]
+ result = subprocess.run(cmd, capture_output=True, text=True)
+ if result.returncode != 0:
+ return None
+ try:
+ data = json.loads(result.stdout)
+ status = data.get("status", {})
+ for field in ("jobName", "nvcreJobName", "workloadJobName"):
+ val = status.get(field)
+ if val and val != workloadrun_name:
+ return val
+ labels = data.get("metadata", {}).get("labels", {})
+ val = labels.get("nvcre.nvidia.com/job")
+ if val and val != workloadrun_name:
+ return val
+ except json.JSONDecodeError as e:
+ logger.debug("Could not parse WorkloadRun JSON for '%s': %s", workloadrun_name, e)
+
+ # CRD doesn't expose the internal name — find the most recently created
+ # JobSet in the namespace. Nvcre names JobSets -workload,
+ # so strip the suffix to get the internal job name.
+ cmd = self._kubectl_base() + [
+ "get", "jobsets",
+ "-n", self.namespace,
+ "--sort-by=.metadata.creationTimestamp",
+ "-o", "jsonpath={range .items[*]}{.metadata.name}{\"\\n\"}{end}",
+ ]
+ result = subprocess.run(cmd, capture_output=True, text=True)
+ if result.returncode == 0:
+ jobsets = [j.strip() for j in result.stdout.splitlines()
+ if j.strip().endswith("-workload")]
+ if jobsets:
+ return jobsets[-1][:-len("-workload")]
+ return None
+
+ def fetch_logs(
+ self,
+ name: str,
+ stream: bool = False,
+ lines: int = -1,
+ timeout: int = 60,
+ ) -> Iterable[str]:
+ """Yield log lines from WorkloadRun pods via kubectl logs.
+
+ Uses the label ``nvcre.nvidia.com/job=`` that
+ Nvcre stamps on the pods it creates.
+ """
+ # Pods are labelled with the JobSet name, not the nvcre.nvidia.com/job
+ # label. Derive the Nvcre internal job name from the WorkloadRun CRD
+ # (it may differ from `name` which is the CRD name we submitted), then
+ # form the JobSet name as -workload.
+ nvcre_job = self._get_nvcre_job_name(name) or name
+ jobset_name = f"{nvcre_job}-workload"
+ label_selector = f"jobset.sigs.k8s.io/jobset-name={jobset_name}"
+ base_cmd = self._kubectl_base() + [
+ "logs",
+ "-l", label_selector,
+ "-n", self.namespace,
+ "--prefix",
+ "--max-log-requests", str(max(self.num_nodes * 2, 8)),
+ ]
+
+ # Streaming logs are saved to job_dir/pod_logs/streaming.log so they
+ # are available for post-run inspection even after pods are deleted.
+ streaming_log_path = None
+ if stream and self.job_dir:
+ pod_logs_dir = os.path.join(self.job_dir, "pod_logs")
+ os.makedirs(pod_logs_dir, exist_ok=True)
+ streaming_log_path = os.path.join(pod_logs_dir, "streaming.log")
+
+ if stream:
+ proc = subprocess.Popen(
+ base_cmd + ["-f"],
+ stdout=subprocess.PIPE,
+ stderr=subprocess.DEVNULL,
+ text=True,
+ bufsize=1,
+ )
+ try:
+ log_file = open(streaming_log_path, "w") if streaming_log_path else None
+ try:
+ for line in iter(proc.stdout.readline, ""):
+ if line:
+ if log_file:
+ log_file.write(line)
+ log_file.flush()
+ yield line
+ finally:
+ if log_file:
+ log_file.close()
+ finally:
+ proc.terminate()
+ proc.wait(timeout=5)
+ else:
+ tail_args = ["--tail", str(lines)] if lines > 0 else ["--tail", "-1"]
+ result = subprocess.run(
+ base_cmd + tail_args,
+ capture_output=True,
+ text=True,
+ timeout=timeout,
+ )
+ yield from result.stdout.splitlines()
+
+ # ── Code packaging via kubectl data-mover ────────────────────────────────
+
+ def _data_mover_pod_name(self, label: str = "datamover") -> str:
+ return f"{self._safe_name()}-{label}"[:63]
+
+ def _start_data_mover_pod(self, pod_name: str, timeout: int = 120) -> None:
+ """Spin up a throw-away alpine pod that mounts workdir_pvc."""
+ pod_manifest = {
+ "apiVersion": "v1",
+ "kind": "Pod",
+ "metadata": {"name": pod_name, "namespace": self.namespace},
+ "spec": {
+ "restartPolicy": "Never",
+ "containers": [{
+ "name": "mover",
+ "image": _DATA_MOVER_IMAGE,
+ "command": ["sleep", "infinity"],
+ "volumeMounts": [{"name": "workdir", "mountPath": self.workdir_pvc_path}],
+ }],
+ "volumes": [{
+ "name": "workdir",
+ "persistentVolumeClaim": {"claimName": self.workdir_pvc},
+ }],
+ },
+ }
+ # Delete stale pod first
+ self._delete_data_mover_pod(pod_name)
+
+ with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
+ yaml.dump(pod_manifest, f)
+ pod_yaml = f.name
+
+ try:
+ subprocess.check_call(
+ self._kubectl_base() + ["apply", "-f", pod_yaml],
+ stdout=subprocess.DEVNULL,
+ )
+ finally:
+ os.unlink(pod_yaml)
+
+ # Wait for Running
+ deadline = time.time() + timeout
+ while time.time() < deadline:
+ result = subprocess.run(
+ self._kubectl_base() + [
+ "get", "pod", pod_name,
+ "-n", self.namespace,
+ "-o", "jsonpath={.status.phase}",
+ ],
+ capture_output=True, text=True,
+ )
+ if result.stdout.strip() == "Running":
+ logger.info("Data-mover pod '%s' is Running", pod_name)
+ return
+ time.sleep(3)
+ raise RuntimeError(f"Data-mover pod '{pod_name}' did not reach Running within {timeout}s")
+
+ def _delete_data_mover_pod(self, pod_name: str, timeout: int = 60) -> None:
+ result = subprocess.run(
+ self._kubectl_base() + [
+ "delete", "pod", pod_name,
+ "-n", self.namespace,
+ "--ignore-not-found",
+ ],
+ capture_output=True, text=True,
+ )
+ if result.returncode != 0:
+ logger.warning("Could not delete data-mover pod '%s': %s", pod_name, result.stderr)
+
+ def _rsync_to_pod(self, pod_name: str, local_path: str, remote_path: str) -> None:
+ subprocess.check_call(
+ self._kubectl_base() + [
+ "exec", "-n", self.namespace, pod_name,
+ "--", "mkdir", "-p", remote_path,
+ ]
+ )
+ subprocess.check_call(
+ self._kubectl_base() + [
+ "cp", "-n", self.namespace,
+ f"{local_path.rstrip(os.sep)}/.",
+ f"{pod_name}:{remote_path.rstrip('/')}",
+ ]
+ )
+ logger.info("Copied '%s' -> pod:%s", local_path, remote_path)
+
+ def copy_to_workspace(self, local_path: str, remote_path: str, label: str = "datamover") -> None:
+ """Copy *local_path* directory to *remote_path* on workdir_pvc."""
+ if not self.workdir_pvc:
+ return
+ pod_name = self._data_mover_pod_name(label)
+ self._start_data_mover_pod(pod_name)
+ try:
+ self._rsync_to_pod(pod_name, local_path, remote_path)
+ finally:
+ self._delete_data_mover_pod(pod_name)
+
+ def package(self, packager: Packager, job_name: str) -> None:
+ """Package code and sync to workdir_pvc before job submission.
+
+ If *workdir_pvc* is not set this is a no-op (assumes code is in the image).
+ """
+ if not self.workdir_pvc:
+ return
+
+ if self.workdir_local_path:
+ os.makedirs(self.job_dir, exist_ok=True)
+ subprocess.check_call(
+ ["rsync", "-a",
+ f"{self.workdir_local_path.rstrip(os.sep)}/",
+ f"{self.job_dir.rstrip(os.sep)}/"],
+ )
+ logger.info("Merged '%s' into job_dir '%s'", self.workdir_local_path, self.job_dir)
+
+ if isinstance(packager, GitArchivePackager):
+ output = subprocess.run(
+ ["git", "rev-parse", "--show-toplevel"],
+ check=True, stdout=subprocess.PIPE,
+ )
+ base_path = Path(output.stdout.splitlines()[0].decode()).absolute()
+ else:
+ base_path = Path(os.getcwd()).absolute()
+
+ local_pkg = packager.package(base_path, self.job_dir, job_name)
+ code_extraction_path = os.path.join(self.job_dir, "code")
+ os.makedirs(code_extraction_path, exist_ok=True)
+
+ if local_pkg:
+ subprocess.check_call(
+ ["tar", "-xzf", local_pkg, "-C", code_extraction_path, "--ignore-zeros"],
+ stdout=subprocess.DEVNULL,
+ )
+ os.remove(local_pkg)
+
+ self.copy_to_workspace(self.job_dir, self.code_dir, label=job_name)
+
+ # Ensure the PVC volume/mount are declared on the WorkloadRun so the
+ # training container can reach code_dir.
+ already_mounted = any(
+ v.get("persistentVolumeClaim", {}).get("claimName") == self.workdir_pvc
+ for v in self.volumes
+ )
+ if not already_mounted:
+ vol_name = "nemo-run-workdir"
+ self.volumes.append(
+ {"name": vol_name, "persistentVolumeClaim": {"claimName": self.workdir_pvc}}
+ )
+ if not any(vm.get("mountPath") == self.workdir_pvc_path for vm in self.volume_mounts):
+ self.volume_mounts.append({"name": vol_name, "mountPath": self.workdir_pvc_path})
+
+ def materialize_launch_script(self, cmd: list[str], max_retries: int = 0) -> None:
+ """Write a launch.sh to job_dir that the WorkloadRun exec framework will run."""
+ nsys_prefix = self.get_launcher_prefix()
+ if nsys_prefix:
+ cmd = ["nsys"] + nsys_prefix + cmd
+ env_exports = "\n".join(f"export {k}={v}" for k, v in self.env_vars.items())
+ cmd_str = " ".join(shlex.quote(a) for a in cmd)
+ if max_retries > 0:
+ run_block = f"""MAX_RETRIES={max_retries}
+attempt=0
+while [ $attempt -le $MAX_RETRIES ]; do
+ {cmd_str}
+ exit_code=$?
+ [ $exit_code -eq 0 ] && exit 0
+ attempt=$((attempt + 1))
+ [ $attempt -le $MAX_RETRIES ] && echo "Retry $attempt/$MAX_RETRIES..." && sleep 5
+done
+exit $exit_code"""
+ else:
+ run_block = cmd_str
+
+ script = f"""#!/usr/bin/env bash
+set -euo pipefail
+
+{env_exports}
+
+cd {self.code_dir}
+
+{run_block}
+"""
+ os.makedirs(self.job_dir, exist_ok=True)
+ launch_path = os.path.join(self.job_dir, "launch.sh")
+ with open(launch_path, "w") as f:
+ f.write(script)
+ os.chmod(launch_path, 0o555)
+ logger.info("Wrote launch script to %s", launch_path)
diff --git a/nemo_run/run/experiment.py b/nemo_run/run/experiment.py
index be92e06d..11e40136 100644
--- a/nemo_run/run/experiment.py
+++ b/nemo_run/run/experiment.py
@@ -52,6 +52,7 @@
from nemo_run.core.execution.base import Executor
from nemo_run.core.execution.dgxcloud import DGXCloudExecutor
from nemo_run.core.execution.kubeflow import KubeflowExecutor
+from nemo_run.core.execution.nvcre import NvcreExecutor
from nemo_run.core.execution.docker import DockerExecutor
from nemo_run.core.execution.lepton import LeptonExecutor
from nemo_run.core.execution.local import LocalExecutor
@@ -208,6 +209,7 @@ class Experiment(ConfigurableMixin):
DGXCloudExecutor,
LeptonExecutor,
KubeflowExecutor,
+ NvcreExecutor,
)
_DETACH_SUPPORTED_EXECUTORS = (
SlurmExecutor,
@@ -215,6 +217,7 @@ class Experiment(ConfigurableMixin):
SkypilotJobsExecutor,
DGXCloudExecutor,
LeptonExecutor,
+ NvcreExecutor,
)
_DEPENDENCY_SUPPORTED_EXECUTORS = (SlurmExecutor,)
_RUNNER_DEPENDENT_EXECUTORS = (LocalExecutor,)
diff --git a/nemo_run/run/torchx_backend/schedulers/api.py b/nemo_run/run/torchx_backend/schedulers/api.py
index 76b46a4b..edea91c0 100644
--- a/nemo_run/run/torchx_backend/schedulers/api.py
+++ b/nemo_run/run/torchx_backend/schedulers/api.py
@@ -24,8 +24,9 @@
from nemo_run.core.execution.lepton import LeptonExecutor
from nemo_run.core.execution.local import LocalExecutor
from nemo_run.core.execution.skypilot import SkypilotExecutor
-from nemo_run.core.execution.slurm import SlurmExecutor
from nemo_run.core.execution.skypilot_jobs import SkypilotJobsExecutor
+from nemo_run.core.execution.slurm import SlurmExecutor
+from nemo_run.core.execution.nvcre import NvcreExecutor
EXECUTOR_MAPPING: dict[Type[Executor], str] = {
SlurmExecutor: "slurm_tunnel",
@@ -36,6 +37,7 @@
DGXCloudExecutor: "dgx_cloud",
LeptonExecutor: "lepton",
KubeflowExecutor: "kubeflow",
+ NvcreExecutor: "nvcre",
}
REVERSE_EXECUTOR_MAPPING: dict[str, Type[Executor]] = {
@@ -47,6 +49,7 @@
"dgx_cloud": DGXCloudExecutor,
"lepton": LeptonExecutor,
"kubeflow": KubeflowExecutor,
+ "nvcre": NvcreExecutor,
}
diff --git a/nemo_run/run/torchx_backend/schedulers/nvcre.py b/nemo_run/run/torchx_backend/schedulers/nvcre.py
new file mode 100644
index 00000000..575d9114
--- /dev/null
+++ b/nemo_run/run/torchx_backend/schedulers/nvcre.py
@@ -0,0 +1,298 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+import json
+import logging
+import os
+import shutil
+import tempfile
+from dataclasses import dataclass
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Iterable, Optional
+
+import fiddle as fdl
+import fiddle._src.experimental.dataclasses as fdl_dc
+import yaml
+from torchx.schedulers.api import (
+ AppDryRunInfo,
+ DescribeAppResponse,
+ ListAppResponse,
+ Scheduler,
+ Stream,
+)
+from torchx.specs import AppDef, AppState, ReplicaStatus, Role, RoleStatus, runopts
+
+from nemo_run.config import get_nemorun_home
+from nemo_run.core.execution.base import Executor
+from nemo_run.core.execution.nvcre import NvcreExecutor, NvcrePhase
+from nemo_run.core.serialization.zlib_json import ZlibJSONSerializer
+from nemo_run.run.torchx_backend.schedulers.api import SchedulerMixin
+
+logger = logging.getLogger(__name__)
+
+NVCRE_JOB_DIRS = os.path.join(get_nemorun_home(), ".nvcre_jobs.json")
+
+NVCRE_STATES: dict[NvcrePhase, AppState] = {
+ NvcrePhase.PENDING: AppState.PENDING,
+ NvcrePhase.IN_PROGRESS: AppState.RUNNING,
+ NvcrePhase.SUCCEEDED: AppState.SUCCEEDED,
+ NvcrePhase.FAILED: AppState.FAILED,
+ NvcrePhase.UNKNOWN: AppState.UNKNOWN,
+}
+
+
+@dataclass
+class NvcreRequest:
+ """Wraps the AppDef and NvcreExecutor for dryrun/schedule."""
+
+ app: AppDef
+ executor: NvcreExecutor
+ cmd: list[str]
+ name: str
+
+
+class NvcreScheduler(SchedulerMixin, Scheduler[dict]): # type: ignore
+ def __init__(self, session_name: str) -> None:
+ super().__init__("nvcre", session_name)
+
+ def _run_opts(self) -> runopts:
+ opts = runopts()
+ opts.add("job_dir", type_=str, help="Directory for job outputs.")
+ return opts
+
+ def _submit_dryrun(self, app: AppDef, cfg: Executor) -> AppDryRunInfo[NvcreRequest]:
+ assert isinstance(cfg, NvcreExecutor), (
+ f"{cfg.__class__} is not supported by NvcreScheduler."
+ )
+ executor = cfg
+ assert len(app.roles) == 1, "NvcreScheduler only supports single-role apps."
+
+ role = app.roles[0]
+ values = executor.macro_values()
+ if values:
+ role = values.apply(role)
+
+ # Merge role-level env into executor env
+ executor.env_vars.update(role.env)
+
+ cmd = [role.entrypoint] + role.args
+
+ # Wrap with torchrun so that torch.distributed is initialised correctly
+ # across all nodes. Nvcre injects PET_* rendezvous env vars per-pod
+ # (via the JobSet downward-API); torchrun reads them via --nnodes /
+ # --nproc_per_node / --node_rank / --master_addr / --master_port and
+ # sets the standard RANK, WORLD_SIZE, LOCAL_RANK, MASTER_ADDR vars that
+ # Megatron-Bridge's common_utils.py expects. Without this wrapper each
+ # replica starts as a lone python process (WORLD_SIZE=1) and fails the
+ # parallelism divisibility check.
+ if executor.use_torchrun and cmd and cmd[0] == "python":
+ script_and_args = cmd[1:] # drop the "python" token; torchrun runs the script directly
+ cmd = [
+ "torchrun",
+ "--nnodes=$(PET_NNODES)",
+ "--nproc_per_node=$(PET_NPROC_PER_NODE)",
+ "--node_rank=$(PET_NODE_RANK)",
+ "--master_addr=$(PET_MASTER_ADDR)",
+ "--master_port=$(PET_MASTER_PORT)",
+ ] + script_and_args
+
+ req = NvcreRequest(app=app, executor=executor, cmd=cmd, name=role.name)
+
+ def _wl_cmd(r: NvcreRequest) -> list[str]:
+ if r.executor.workdir_pvc:
+ return ["/bin/bash", f"{r.executor.code_dir}/launch.sh"]
+ return r.cmd
+
+ return AppDryRunInfo(
+ req,
+ lambda r: yaml.dump(r.executor.build_workloadrun_yaml(_wl_cmd(r))),
+ )
+
+ def schedule(self, dryrun_info: AppDryRunInfo[NvcreRequest]) -> str:
+ req = dryrun_info.request
+ executor = req.executor
+
+ os.makedirs(executor.job_dir, exist_ok=True)
+
+ if executor.workdir_pvc:
+ # Write launch.sh with the actual training command and sync to PVC.
+ executor.materialize_launch_script(req.cmd, max_retries=executor.retries)
+ executor.package(executor.packager, job_name=executor.job_name)
+ wl_cmd = ["/bin/bash", f"{executor.code_dir}/launch.sh"]
+ else:
+ # No PVC: code is assumed to be in the container image.
+ # Run the training command directly; env vars are injected via the
+ # WorkloadRun spec rather than through a launch.sh wrapper.
+ nsys_prefix = executor.get_launcher_prefix()
+ wl_cmd = (["nsys"] + nsys_prefix + req.cmd) if nsys_prefix else req.cmd
+
+ # Write WorkloadRun YAML
+ yaml_path = os.path.join(executor.job_dir, "workloadrun.yaml")
+ manifest = executor.build_workloadrun_yaml(wl_cmd)
+ with open(yaml_path, "w") as f:
+ yaml.dump(manifest, f, default_flow_style=False)
+
+ # Submit
+ workloadrun_name = executor.submit(yaml_path)
+
+ experiment_id = getattr(executor, "experiment_id", "nvcre_experiment")
+ app_id = f"{experiment_id}___{req.name}___{workloadrun_name}"
+
+ _save_job(app_id, workloadrun_name, executor)
+ return app_id
+
+ def describe(self, app_id: str) -> Optional[DescribeAppResponse]:
+ stored = _get_jobs()
+ job_info = stored.get(app_id)
+ if not job_info:
+ return None
+
+ parts = app_id.split("___")
+ role_name = parts[1] if len(parts) > 1 else app_id
+ workloadrun_name = job_info.get("workloadrun_name") or (
+ parts[-1] if len(parts) > 2 else app_id
+ )
+
+ executor: Optional[NvcreExecutor] = job_info.get("executor")
+ if not executor:
+ return None
+
+ phase = executor.status(workloadrun_name)
+ app_state = NVCRE_STATES.get(phase, AppState.UNKNOWN)
+
+ roles = [Role(name=role_name, image="", num_replicas=executor.num_nodes)]
+ roles_statuses = [
+ RoleStatus(
+ role_name,
+ replicas=[
+ ReplicaStatus(id=i, role=role_name, state=app_state, hostname="")
+ for i in range(executor.num_nodes)
+ ],
+ )
+ ]
+
+ return DescribeAppResponse(
+ app_id=app_id,
+ roles=roles,
+ roles_statuses=roles_statuses,
+ state=app_state,
+ msg="",
+ )
+
+ def log_iter(
+ self,
+ app_id: str,
+ role_name: str,
+ k: int = 0,
+ regex: Optional[str] = None,
+ since: Optional[datetime] = None,
+ until: Optional[datetime] = None,
+ should_tail: bool = False,
+ streams: Optional[Stream] = None,
+ ) -> Iterable[str]:
+ stored = _get_jobs()
+ job_info = stored.get(app_id)
+ if not job_info:
+ return []
+
+ parts = app_id.split("___")
+ workloadrun_name = job_info.get("workloadrun_name") or (
+ parts[-1] if len(parts) > 2 else app_id
+ )
+ executor: Optional[NvcreExecutor] = job_info.get("executor")
+ if not executor:
+ return []
+
+ # job_dir is an init=False field that doesn't survive fiddle serialisation;
+ # restore it from the explicitly saved value so fetch_logs can write the
+ # streaming log to the correct experiment directory.
+ job_dir = job_info.get("job_dir", "")
+ if job_dir and not executor.job_dir:
+ executor.job_dir = job_dir
+
+ return executor.fetch_logs(workloadrun_name, stream=should_tail)
+
+ def _cancel_existing(self, app_id: str) -> None:
+ stored = _get_jobs()
+ job_info = stored.get(app_id)
+ if not job_info:
+ return
+
+ parts = app_id.split("___")
+ workloadrun_name = job_info.get("workloadrun_name") or (
+ parts[-1] if len(parts) > 2 else app_id
+ )
+ executor: Optional[NvcreExecutor] = job_info.get("executor")
+ if executor:
+ executor.cancel(workloadrun_name)
+
+ def list(self) -> list[ListAppResponse]:
+ return []
+
+ def _validate(self, app: AppDef, scheduler: str) -> None:
+ pass
+
+
+def create_scheduler(session_name: str, **kwargs: Any) -> NvcreScheduler:
+ return NvcreScheduler(session_name=session_name)
+
+
+def _save_job(app_id: str, workloadrun_name: str, executor: NvcreExecutor) -> None:
+ original_apps: dict = {}
+ os.makedirs(os.path.dirname(NVCRE_JOB_DIRS), exist_ok=True)
+ if not os.path.isfile(NVCRE_JOB_DIRS):
+ Path(NVCRE_JOB_DIRS).touch()
+
+ serializer = ZlibJSONSerializer()
+ with open(NVCRE_JOB_DIRS, "r+") as f:
+ try:
+ original_apps = json.load(f)
+ except Exception:
+ original_apps = {}
+
+ entry = {
+ "workloadrun_name": workloadrun_name,
+ "job_dir": executor.job_dir,
+ "executor": serializer.serialize(
+ fdl_dc.convert_dataclasses_to_configs(executor, allow_post_init=True)
+ ),
+ }
+ original_apps[app_id] = entry
+
+ with tempfile.NamedTemporaryFile(mode="w+", delete=False) as fp:
+ json.dump(original_apps, fp)
+ temp_path = fp.name
+
+ f.close()
+ shutil.move(temp_path, NVCRE_JOB_DIRS)
+
+
+def _get_jobs() -> dict[str, dict]:
+ if not os.path.isfile(NVCRE_JOB_DIRS):
+ return {}
+ with open(NVCRE_JOB_DIRS) as f:
+ try:
+ data = json.load(f)
+ except Exception:
+ return {}
+
+ serializer = ZlibJSONSerializer()
+ for entry in data.values():
+ try:
+ entry["executor"] = fdl.build(serializer.deserialize(entry["executor"]))
+ except Exception as e:
+ logger.debug("Failed to deserialize Nvcre executor: %s", e)
+ return data
diff --git a/pyproject.toml b/pyproject.toml
index 5bcfa927..2db00bac 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -48,6 +48,7 @@ dgx_cloud = "nemo_run.run.torchx_backend.schedulers.dgxcloud:create_scheduler"
lepton = "nemo_run.run.torchx_backend.schedulers.lepton:create_scheduler"
skypilot_jobs = "nemo_run.run.torchx_backend.schedulers.skypilot_jobs:create_scheduler"
kubeflow = "nemo_run.run.torchx_backend.schedulers.kubeflow:create_scheduler"
+nvcre = "nemo_run.run.torchx_backend.schedulers.nvcre:create_scheduler"
[project.optional-dependencies]
skypilot = [
diff --git a/test/core/execution/test_nvcre.py b/test/core/execution/test_nvcre.py
new file mode 100644
index 00000000..c0a4f57e
--- /dev/null
+++ b/test/core/execution/test_nvcre.py
@@ -0,0 +1,631 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+import os
+import subprocess
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from nemo_run.core.execution.launcher import Launcher
+from nemo_run.core.execution.nvcre import NvcreExecutor, NvcrePhase
+
+
+def _completed(returncode=0, stdout="", stderr=""):
+ return subprocess.CompletedProcess(args=[], returncode=returncode, stdout=stdout, stderr=stderr)
+
+
+class TestNvcreExecutor:
+ @pytest.fixture
+ def executor(self):
+ e = NvcreExecutor(
+ namespace="nemo-perf",
+ container_image="nvcr.io/nvidia/nemo:dev",
+ num_nodes=2,
+ gpus_per_node=8,
+ )
+ e.job_name = "my-job"
+ e.experiment_id = "exp1"
+ e.job_dir = "/tmp/exp1/my-job"
+ return e
+
+ # ── build_workloadrun_yaml ────────────────────────────────────────────────
+
+ def test_build_workloadrun_yaml_minimal(self):
+ e = NvcreExecutor(namespace="ns", container_image="img:latest", num_nodes=1)
+ e.job_name = "job1"
+ manifest = e.build_workloadrun_yaml(["python", "train.py"])
+
+ assert manifest["apiVersion"] == "nvcre.nvidia.com/v1alpha1"
+ assert manifest["kind"] == "WorkloadRun"
+ assert manifest["metadata"] == {"name": "job1", "namespace": "ns"}
+ spec = manifest["spec"]
+ assert spec["image"] == "img:latest"
+ assert spec["numNodes"] == 1
+ assert spec["framework"]["exec"]["command"] == ["python", "train.py"]
+ assert "gpusPerNode" not in spec
+ assert "target" not in spec
+ assert "env" not in spec
+ assert "volumes" not in spec
+ assert "imagePullSecrets" not in spec
+ assert "orchestration" in spec # default timeout_per_job is set
+ assert "checkpoint" not in spec
+ assert "gangScheduler" not in spec
+
+ def test_build_workloadrun_yaml_full(self, executor):
+ executor.node_selector = {"gpu-type": "h100"}
+ executor.env_vars = {"FOO": "bar"}
+ executor.volumes = [{"name": "v", "persistentVolumeClaim": {"claimName": "pvc"}}]
+ executor.volume_mounts = [{"name": "v", "mountPath": "/mnt"}]
+ executor.image_pull_secret = "ngc-secret"
+ executor.timeout_per_job = "2h"
+ executor.test_scale = "full-scale"
+ executor.max_restarts = 3
+ executor.checkpoint_storage_size = "500Gi"
+ executor.checkpoint_storage_class = "fast-ssd"
+ executor.gang_scheduler_name = "kai-scheduler"
+
+ manifest = executor.build_workloadrun_yaml(["python", "train.py"])
+ spec = manifest["spec"]
+
+ assert spec["gpusPerNode"] == 8
+ assert spec["target"] == {"nodeSelector": {"gpu-type": "h100"}}
+ assert spec["env"] == [{"name": "FOO", "value": "bar"}]
+ assert spec["volumes"] == executor.volumes
+ assert spec["volumeMounts"] == executor.volume_mounts
+ assert spec["imagePullSecrets"] == [{"name": "ngc-secret"}]
+ assert spec["orchestration"] == {"timeoutPerJob": "2h", "testScale": "full-scale"}
+ assert spec["checkpoint"] == {"storageSize": "500Gi", "storageClassName": "fast-ssd", "maxRestarts": 3}
+ assert spec["gangScheduler"] == {"schedulerName": "kai-scheduler"}
+
+ # ── _safe_name ─────────────────────────────────────────────────────────────
+
+ @pytest.mark.parametrize(
+ "job_name,expected",
+ [
+ ("My_Job.Name", "my-job-name"),
+ ("", "nvcre-job"),
+ ("Already-Safe", "already-safe"),
+ ("trailing-dot.", "trailing-dot"),
+ ],
+ )
+ def test_safe_name(self, job_name, expected):
+ e = NvcreExecutor(namespace="ns", container_image="img")
+ e.job_name = job_name
+ assert e._safe_name() == expected
+
+ def test_safe_name_truncates_to_63_chars(self):
+ e = NvcreExecutor(namespace="ns", container_image="img")
+ e.job_name = "x" * 100
+ name = e._safe_name()
+ assert len(name) <= 63
+
+ # ── submit ─────────────────────────────────────────────────────────────────
+
+ def test_submit_parses_kubectl_style_name(self, executor):
+ with patch("nemo_run.core.execution.nvcre.subprocess.run") as mock_run:
+ mock_run.return_value = _completed(
+ stdout="workloadrun.nvcre.nvidia.com/my-job-abcd created\n"
+ )
+ name = executor.submit("/tmp/wl.yaml")
+
+ assert name == "my-job-abcd"
+ assert executor._workloadrun_name == "my-job-abcd"
+ cmd = mock_run.call_args[0][0]
+ assert cmd[0] == "nvcrectl"
+ assert "workloadrun" in cmd and "run" in cmd
+ assert "--namespace" in cmd and executor.namespace in cmd
+
+ def test_submit_parses_json_name(self, executor):
+ with patch("nemo_run.core.execution.nvcre.subprocess.run") as mock_run:
+ mock_run.return_value = _completed(stdout='{"name": "my-job-xyz", "ok": true}\n')
+ name = executor.submit("/tmp/wl.yaml")
+ assert name == "my-job-xyz"
+
+ def test_submit_parses_plain_name(self, executor):
+ with patch("nemo_run.core.execution.nvcre.subprocess.run") as mock_run:
+ mock_run.return_value = _completed(stdout="my-job-plain\n")
+ name = executor.submit("/tmp/wl.yaml")
+ assert name == "my-job-plain"
+
+ def test_submit_falls_back_to_kubectl_when_unparseable(self, executor):
+ with (
+ patch("nemo_run.core.execution.nvcre.subprocess.run") as mock_run,
+ patch("nemo_run.core.execution.nvcre.time.sleep"),
+ ):
+ mock_run.side_effect = [
+ _completed(stdout="!!! unrecognisable output !!!"),
+ _completed(stdout="fallback-name\n"),
+ ]
+ name = executor.submit("/tmp/wl.yaml")
+
+ assert name == "fallback-name"
+ assert mock_run.call_count == 2
+ fallback_cmd = mock_run.call_args_list[1][0][0]
+ assert fallback_cmd[0] == "kubectl"
+ assert "get" in fallback_cmd and "workloadruns" in fallback_cmd
+
+ def test_submit_fallback_returns_requested_name_when_kubectl_also_fails(self, executor):
+ with (
+ patch("nemo_run.core.execution.nvcre.subprocess.run") as mock_run,
+ patch("nemo_run.core.execution.nvcre.time.sleep"),
+ ):
+ mock_run.side_effect = [
+ _completed(stdout="!!! unrecognisable !!!"),
+ _completed(returncode=1, stderr="not found"),
+ ]
+ name = executor.submit("/tmp/wl.yaml")
+ assert name == executor._safe_name()
+
+ def test_submit_raises_on_failure(self, executor):
+ with patch("nemo_run.core.execution.nvcre.subprocess.run") as mock_run:
+ mock_run.return_value = _completed(returncode=1, stderr="boom")
+ with pytest.raises(RuntimeError, match="boom"):
+ executor.submit("/tmp/wl.yaml")
+
+ # ── status ─────────────────────────────────────────────────────────────────
+
+ def test_status_via_nvcrectl(self, executor):
+ with patch("nemo_run.core.execution.nvcre.subprocess.run") as mock_run:
+ mock_run.return_value = _completed(stdout="Succeeded\n")
+ phase = executor.status("wl-name")
+ assert phase == NvcrePhase.SUCCEEDED
+ mock_run.assert_called_once()
+
+ def test_status_falls_back_to_crd_on_nvcrectl_failure(self, executor):
+ with patch("nemo_run.core.execution.nvcre.subprocess.run") as mock_run:
+ mock_run.side_effect = [
+ _completed(returncode=1, stderr="not found"),
+ _completed(stdout="Failed\n"),
+ ]
+ phase = executor.status("wl-name")
+ assert phase == NvcrePhase.FAILED
+ assert mock_run.call_count == 2
+ crd_cmd = mock_run.call_args_list[1][0][0]
+ assert crd_cmd[0] == "kubectl"
+ assert "workloadrun" in crd_cmd
+
+ def test_status_falls_back_to_crd_on_unrecognised_phase(self, executor):
+ with patch("nemo_run.core.execution.nvcre.subprocess.run") as mock_run:
+ mock_run.side_effect = [
+ _completed(stdout="SomeWeirdPhase\n"),
+ _completed(stdout="InProgress\n"),
+ ]
+ phase = executor.status("wl-name")
+ assert phase == NvcrePhase.IN_PROGRESS
+
+ def test_status_crd_fallback_returns_unknown_on_empty_or_error(self, executor):
+ with patch("nemo_run.core.execution.nvcre.subprocess.run") as mock_run:
+ mock_run.side_effect = [
+ _completed(returncode=1, stderr="gone"),
+ _completed(returncode=0, stdout=""),
+ ]
+ phase = executor.status("wl-name")
+ assert phase == NvcrePhase.UNKNOWN
+
+ # ── cancel ─────────────────────────────────────────────────────────────────
+
+ def test_cancel_success(self, executor):
+ with patch("nemo_run.core.execution.nvcre.subprocess.run") as mock_run:
+ mock_run.return_value = _completed(returncode=0)
+ executor.cancel("wl-name")
+ cmd = mock_run.call_args[0][0]
+ assert "cancel" in cmd and "wl-name" in cmd
+
+ def test_cancel_logs_warning_on_failure(self, executor, caplog):
+ with patch("nemo_run.core.execution.nvcre.subprocess.run") as mock_run:
+ mock_run.return_value = _completed(returncode=1, stderr="cannot cancel")
+ executor.cancel("wl-name") # should not raise
+
+ # ── fetch_logs (non-streaming) ────────────────────────────────────────────
+
+ def test_fetch_logs_non_streaming(self, executor):
+ with patch("nemo_run.core.execution.nvcre.subprocess.run") as mock_run:
+ mock_run.side_effect = [
+ _completed(returncode=0, stdout='{"status": {}, "metadata": {"labels": {}}}'),
+ _completed(returncode=0, stdout=""), # jobsets lookup (fallback)
+ _completed(returncode=0, stdout="line1\nline2\n"), # logs
+ ]
+ lines = list(executor.fetch_logs("wl-name", stream=False, lines=100))
+ assert lines == ["line1", "line2"]
+
+ def test_get_nvcre_job_name_from_status_field(self, executor):
+ with patch("nemo_run.core.execution.nvcre.subprocess.run") as mock_run:
+ mock_run.return_value = _completed(
+ returncode=0, stdout='{"status": {"jobName": "internal-job"}, "metadata": {"labels": {}}}'
+ )
+ job_name = executor._get_nvcre_job_name("wl-name")
+ assert job_name == "internal-job"
+
+ def test_get_nvcre_job_name_falls_back_to_jobsets(self, executor):
+ with patch("nemo_run.core.execution.nvcre.subprocess.run") as mock_run:
+ mock_run.side_effect = [
+ _completed(returncode=0, stdout='{"status": {}, "metadata": {"labels": {}}}'),
+ _completed(returncode=0, stdout="foo-workload\nbar-workload\n"),
+ ]
+ job_name = executor._get_nvcre_job_name("wl-name")
+ assert job_name == "bar-workload"[: -len("-workload")]
+
+ # ── macro_values / nnodes / nproc_per_node ────────────────────────────────
+
+ def test_nnodes_and_nproc(self, executor):
+ assert executor.nnodes() == 2
+ assert executor.nproc_per_node() == 8
+
+ def test_nproc_per_node_defaults_to_one(self):
+ e = NvcreExecutor(namespace="ns", container_image="img", gpus_per_node=0)
+ assert e.nproc_per_node() == 1
+
+ def test_macro_values(self, executor):
+ macros = executor.macro_values()
+ assert macros.head_node_ip_var == "PET_MASTER_ADDR"
+ assert macros.nproc_per_node_var == "PET_NPROC_PER_NODE"
+ assert macros.num_nodes_var == "PET_NNODES"
+ assert macros.node_rank_var == "PET_NODE_RANK"
+
+ def test_code_dir(self, executor):
+ with patch("nemo_run.core.execution.nvcre.getpass.getuser", return_value="alice"):
+ assert executor.code_dir == "/nemo_run/alice/exp1/my-job/code"
+
+ # ── package / materialize_launch_script (no PVC = no-op) ─────────────────
+
+ def test_package_is_noop_without_pvc(self, executor):
+ mock_packager = MagicMock()
+ executor.package(mock_packager, job_name="job1")
+ mock_packager.package.assert_not_called()
+
+ def test_copy_to_workspace_is_noop_without_pvc(self, executor):
+ with patch("nemo_run.core.execution.nvcre.subprocess.run") as mock_run:
+ executor.copy_to_workspace("/local", "/remote")
+ mock_run.assert_not_called()
+
+ def test_materialize_launch_script_writes_file(self, executor, tmp_path):
+ executor.job_dir = str(tmp_path)
+ executor.env_vars = {"FOO": "bar"}
+ executor.materialize_launch_script(["python", "train.py"])
+
+ launch_path = tmp_path / "launch.sh"
+ assert launch_path.exists()
+ content = launch_path.read_text()
+ assert "export FOO=bar" in content
+ assert "python train.py" in content
+ assert content.startswith("#!/usr/bin/env bash")
+
+ def test_materialize_launch_script_with_retries(self, executor, tmp_path):
+ executor.job_dir = str(tmp_path)
+ executor.materialize_launch_script(["python", "train.py"], max_retries=2)
+
+ content = (tmp_path / "launch.sh").read_text()
+ assert "MAX_RETRIES=2" in content
+ assert "Retry $attempt/$MAX_RETRIES" in content
+
+ def test_materialize_launch_script_with_nsys_prefix(self, executor, tmp_path):
+ executor.job_dir = str(tmp_path)
+ executor.launcher = Launcher(nsys_profile=True)
+ executor.materialize_launch_script(["python", "train.py"])
+
+ content = (tmp_path / "launch.sh").read_text()
+ assert content.count("nsys") >= 1
+
+ # ── assign ─────────────────────────────────────────────────────────────────
+
+ def test_assign_sets_job_metadata(self):
+ e = NvcreExecutor(namespace="ns", container_image="img")
+ e.assign("exp1", "/exp/dir", "task1", "task1_dir")
+ assert e.experiment_id == "exp1"
+ assert e.experiment_dir == "/exp/dir"
+ assert e.job_name == "task1"
+ assert e.job_dir == "/exp/dir/task1_dir"
+
+ # ── get_launcher_prefix ────────────────────────────────────────────────────
+
+ def test_get_launcher_prefix_none_by_default(self, executor):
+ assert executor.get_launcher_prefix() is None
+
+ def test_get_launcher_prefix_with_nsys_profile(self, executor, tmp_path):
+ executor.job_dir = str(tmp_path)
+ executor.launcher = Launcher(nsys_profile=True)
+ prefix = executor.get_launcher_prefix()
+ assert prefix is not None
+ assert (tmp_path / "nsys_profile").is_dir()
+
+ # ── build_workloadrun_yaml orchestration branches ─────────────────────────
+
+ def test_build_workloadrun_yaml_no_orchestration_when_both_empty(self):
+ e = NvcreExecutor(namespace="ns", container_image="img")
+ e.job_name = "job1"
+ e.timeout_per_job = ""
+ e.test_scale = None
+ manifest = e.build_workloadrun_yaml(["python"])
+ assert "orchestration" not in manifest["spec"]
+
+ def test_build_workloadrun_yaml_orchestration_test_scale_only(self):
+ e = NvcreExecutor(namespace="ns", container_image="img")
+ e.job_name = "job1"
+ e.timeout_per_job = ""
+ e.test_scale = "intra-node"
+ manifest = e.build_workloadrun_yaml(["python"])
+ assert manifest["spec"]["orchestration"] == {"testScale": "intra-node"}
+
+ # ── nvcrectl_base / kubectl_base kubeconfig/context ────────────────────────
+
+ def test_nvcrectl_base_includes_kubeconfig_and_context(self):
+ e = NvcreExecutor(
+ namespace="ns", container_image="img", kubeconfig="/path/kubeconfig", kube_context="ctx1"
+ )
+ args = e._nvcrectl_base()
+ assert args == ["nvcrectl", "--kubeconfig", "/path/kubeconfig", "--context", "ctx1"]
+
+ def test_kubectl_base_includes_kubeconfig_and_context(self):
+ e = NvcreExecutor(
+ namespace="ns", container_image="img", kubeconfig="/path/kubeconfig", kube_context="ctx1"
+ )
+ args = e._kubectl_base()
+ assert args == ["kubectl", "--kubeconfig", "/path/kubeconfig", "--context", "ctx1"]
+
+ # ── _kubectl_workloadrun_crd_phase (direct) ───────────────────────────────
+
+ def test_crd_phase_returns_unknown_on_kubectl_failure(self, executor):
+ with patch("nemo_run.core.execution.nvcre.subprocess.run") as mock_run:
+ mock_run.return_value = _completed(returncode=1, stderr="not found")
+ phase = executor._kubectl_workloadrun_crd_phase("wl-name")
+ assert phase == NvcrePhase.UNKNOWN
+
+ def test_crd_phase_returns_unknown_on_empty_output(self, executor):
+ with patch("nemo_run.core.execution.nvcre.subprocess.run") as mock_run:
+ mock_run.return_value = _completed(returncode=0, stdout=" ")
+ phase = executor._kubectl_workloadrun_crd_phase("wl-name")
+ assert phase == NvcrePhase.UNKNOWN
+
+ def test_crd_phase_returns_unknown_on_unrecognised_phase(self, executor):
+ with patch("nemo_run.core.execution.nvcre.subprocess.run") as mock_run:
+ mock_run.return_value = _completed(returncode=0, stdout="Weird\n")
+ phase = executor._kubectl_workloadrun_crd_phase("wl-name")
+ assert phase == NvcrePhase.UNKNOWN
+
+ def test_crd_phase_returns_recognised_phase(self, executor):
+ with patch("nemo_run.core.execution.nvcre.subprocess.run") as mock_run:
+ mock_run.return_value = _completed(returncode=0, stdout="Pending\n")
+ phase = executor._kubectl_workloadrun_crd_phase("wl-name")
+ assert phase == NvcrePhase.PENDING
+
+ # ── _get_nvcre_job_name edge cases ─────────────────────────────────────
+
+ def test_get_nvcre_job_name_returns_none_on_kubectl_failure(self, executor):
+ with patch("nemo_run.core.execution.nvcre.subprocess.run") as mock_run:
+ mock_run.return_value = _completed(returncode=1, stderr="gone")
+ assert executor._get_nvcre_job_name("wl-name") is None
+
+ def test_get_nvcre_job_name_handles_invalid_json(self, executor):
+ with patch("nemo_run.core.execution.nvcre.subprocess.run") as mock_run:
+ mock_run.side_effect = [
+ _completed(returncode=0, stdout="not json"),
+ _completed(returncode=0, stdout=""),
+ ]
+ assert executor._get_nvcre_job_name("wl-name") is None
+
+ def test_get_nvcre_job_name_from_labels(self, executor):
+ with patch("nemo_run.core.execution.nvcre.subprocess.run") as mock_run:
+ mock_run.return_value = _completed(
+ returncode=0,
+ stdout='{"status": {}, "metadata": {"labels": {"nvcre.nvidia.com/job": "label-job"}}}',
+ )
+ job_name = executor._get_nvcre_job_name("wl-name")
+ assert job_name == "label-job"
+
+ def test_get_nvcre_job_name_returns_none_when_no_jobsets(self, executor):
+ with patch("nemo_run.core.execution.nvcre.subprocess.run") as mock_run:
+ mock_run.side_effect = [
+ _completed(returncode=0, stdout='{"status": {}, "metadata": {"labels": {}}}'),
+ _completed(returncode=0, stdout=""),
+ ]
+ assert executor._get_nvcre_job_name("wl-name") is None
+
+ # ── fetch_logs streaming ───────────────────────────────────────────────────
+
+ def test_fetch_logs_streaming_writes_and_yields_lines(self, executor, tmp_path):
+ executor.job_dir = str(tmp_path)
+ mock_proc = MagicMock()
+ mock_proc.stdout.readline.side_effect = ["line1\n", "line2\n", ""]
+ mock_proc.wait.return_value = None
+
+ with (
+ patch("nemo_run.core.execution.nvcre.subprocess.run") as mock_run,
+ patch("nemo_run.core.execution.nvcre.subprocess.Popen", return_value=mock_proc),
+ ):
+ mock_run.side_effect = [
+ _completed(returncode=0, stdout='{"status": {}, "metadata": {"labels": {}}}'),
+ _completed(returncode=0, stdout=""),
+ ]
+ lines = list(executor.fetch_logs("wl-name", stream=True))
+
+ assert lines == ["line1\n", "line2\n"]
+ mock_proc.terminate.assert_called_once()
+ streaming_log = tmp_path / "pod_logs" / "streaming.log"
+ assert streaming_log.exists()
+ assert streaming_log.read_text() == "line1\nline2\n"
+
+ def test_fetch_logs_streaming_without_job_dir_skips_file(self, executor):
+ executor.job_dir = ""
+ mock_proc = MagicMock()
+ mock_proc.stdout.readline.side_effect = [""]
+ mock_proc.wait.return_value = None
+
+ with (
+ patch("nemo_run.core.execution.nvcre.subprocess.run") as mock_run,
+ patch("nemo_run.core.execution.nvcre.subprocess.Popen", return_value=mock_proc),
+ ):
+ mock_run.side_effect = [
+ _completed(returncode=0, stdout='{"status": {}, "metadata": {"labels": {}}}'),
+ _completed(returncode=0, stdout=""),
+ ]
+ lines = list(executor.fetch_logs("wl-name", stream=True))
+ assert lines == []
+
+ # ── data-mover pod lifecycle ───────────────────────────────────────────────
+
+ def test_data_mover_pod_name(self, executor):
+ assert executor._data_mover_pod_name("mover1") == f"{executor._safe_name()}-mover1"
+
+ def test_start_data_mover_pod_reaches_running(self, executor):
+ executor.workdir_pvc = "my-pvc"
+ with (
+ patch("nemo_run.core.execution.nvcre.subprocess.run") as mock_run,
+ patch("nemo_run.core.execution.nvcre.subprocess.check_call") as mock_check_call,
+ ):
+ mock_run.side_effect = [
+ _completed(returncode=0), # delete stale pod (via _delete_data_mover_pod)
+ _completed(returncode=0, stdout="Running"), # phase check
+ ]
+ executor._start_data_mover_pod("mover-pod", timeout=10)
+
+ mock_check_call.assert_called_once()
+ assert mock_check_call.call_args[0][0][:2] == ["kubectl", "apply"] or "apply" in mock_check_call.call_args[0][0]
+
+ def test_start_data_mover_pod_times_out(self, executor):
+ executor.workdir_pvc = "my-pvc"
+ with (
+ patch("nemo_run.core.execution.nvcre.subprocess.run") as mock_run,
+ patch("nemo_run.core.execution.nvcre.subprocess.check_call"),
+ patch("nemo_run.core.execution.nvcre.time.sleep"),
+ patch("nemo_run.core.execution.nvcre.time.time", side_effect=[0, 0, 100]),
+ ):
+ mock_run.side_effect = [
+ _completed(returncode=0), # delete stale pod
+ _completed(returncode=0, stdout="Pending"), # never reaches Running
+ ]
+ with pytest.raises(RuntimeError, match="did not reach Running"):
+ executor._start_data_mover_pod("mover-pod", timeout=10)
+
+ def test_delete_data_mover_pod_success(self, executor):
+ with patch("nemo_run.core.execution.nvcre.subprocess.run") as mock_run:
+ mock_run.return_value = _completed(returncode=0)
+ executor._delete_data_mover_pod("mover-pod")
+ cmd = mock_run.call_args[0][0]
+ assert "delete" in cmd and "mover-pod" in cmd
+
+ def test_delete_data_mover_pod_logs_warning_on_failure(self, executor):
+ with patch("nemo_run.core.execution.nvcre.subprocess.run") as mock_run:
+ mock_run.return_value = _completed(returncode=1, stderr="cannot delete")
+ executor._delete_data_mover_pod("mover-pod") # should not raise
+
+ def test_rsync_to_pod(self, executor):
+ with patch("nemo_run.core.execution.nvcre.subprocess.check_call") as mock_check_call:
+ executor._rsync_to_pod("mover-pod", "/local/path", "/remote/path")
+ assert mock_check_call.call_count == 2
+ mkdir_cmd = mock_check_call.call_args_list[0][0][0]
+ cp_cmd = mock_check_call.call_args_list[1][0][0]
+ assert "mkdir" in mkdir_cmd
+ assert "cp" in cp_cmd
+
+ def test_copy_to_workspace_with_pvc_runs_full_lifecycle(self, executor):
+ executor.workdir_pvc = "my-pvc"
+ with (
+ patch.object(NvcreExecutor, "_start_data_mover_pod") as mock_start,
+ patch.object(NvcreExecutor, "_rsync_to_pod") as mock_rsync,
+ patch.object(NvcreExecutor, "_delete_data_mover_pod") as mock_delete,
+ ):
+ executor.copy_to_workspace("/local", "/remote", label="mylabel")
+
+ mock_start.assert_called_once()
+ mock_rsync.assert_called_once_with(executor._data_mover_pod_name("mylabel"), "/local", "/remote")
+ mock_delete.assert_called_once()
+
+ def test_copy_to_workspace_deletes_pod_even_on_rsync_failure(self, executor):
+ executor.workdir_pvc = "my-pvc"
+ with (
+ patch.object(NvcreExecutor, "_start_data_mover_pod"),
+ patch.object(NvcreExecutor, "_rsync_to_pod", side_effect=RuntimeError("rsync failed")),
+ patch.object(NvcreExecutor, "_delete_data_mover_pod") as mock_delete,
+ ):
+ with pytest.raises(RuntimeError, match="rsync failed"):
+ executor.copy_to_workspace("/local", "/remote")
+
+ mock_delete.assert_called_once()
+
+ # ── package with PVC ───────────────────────────────────────────────────────
+
+ def test_package_with_pvc_no_local_overlay(self, executor, tmp_path):
+ executor.workdir_pvc = "my-pvc"
+ executor.job_dir = str(tmp_path / "job")
+ mock_packager = MagicMock()
+ mock_packager.package.return_value = None
+
+ with (
+ patch("nemo_run.core.execution.nvcre.subprocess.run") as mock_run,
+ patch("nemo_run.core.execution.nvcre.subprocess.check_call"),
+ patch.object(NvcreExecutor, "copy_to_workspace") as mock_copy,
+ ):
+ mock_run.return_value = _completed(returncode=0, stdout=str(tmp_path).encode())
+ executor.package(mock_packager, job_name="job1")
+
+ mock_packager.package.assert_called_once()
+ mock_copy.assert_called_once()
+ assert len(executor.volumes) == 1
+ assert executor.volumes[0]["persistentVolumeClaim"]["claimName"] == "my-pvc"
+ assert len(executor.volume_mounts) == 1
+
+ def test_package_with_pvc_does_not_duplicate_volume_mount(self, executor, tmp_path):
+ executor.workdir_pvc = "my-pvc"
+ executor.job_dir = str(tmp_path / "job")
+ executor.volumes = [
+ {"name": "existing", "persistentVolumeClaim": {"claimName": "my-pvc"}}
+ ]
+ mock_packager = MagicMock()
+ mock_packager.package.return_value = None
+
+ with (
+ patch("nemo_run.core.execution.nvcre.subprocess.run") as mock_run,
+ patch.object(NvcreExecutor, "copy_to_workspace"),
+ ):
+ mock_run.return_value = _completed(returncode=0)
+ executor.package(mock_packager, job_name="job1")
+
+ assert len(executor.volumes) == 1 # not duplicated
+
+ def test_package_with_local_overlay_rsyncs_and_merges(self, executor, tmp_path):
+ executor.workdir_pvc = "my-pvc"
+ executor.job_dir = str(tmp_path / "job")
+ executor.workdir_local_path = "/some/overlay"
+ mock_packager = MagicMock()
+ mock_packager.package.return_value = None
+
+ with (
+ patch("nemo_run.core.execution.nvcre.subprocess.check_call") as mock_check_call,
+ patch.object(NvcreExecutor, "copy_to_workspace"),
+ ):
+ executor.package(mock_packager, job_name="job1")
+
+ rsync_call = mock_check_call.call_args_list[0][0][0]
+ assert rsync_call[0] == "rsync"
+
+ def test_package_extracts_local_pkg_tarball(self, executor, tmp_path):
+ executor.workdir_pvc = "my-pvc"
+ executor.job_dir = str(tmp_path / "job")
+ os.makedirs(executor.job_dir, exist_ok=True)
+ fake_tarball = tmp_path / "pkg.tar.gz"
+ fake_tarball.write_bytes(b"")
+ mock_packager = MagicMock()
+ mock_packager.package.return_value = str(fake_tarball)
+
+ with (
+ patch("nemo_run.core.execution.nvcre.subprocess.check_call") as mock_check_call,
+ patch.object(NvcreExecutor, "copy_to_workspace"),
+ ):
+ executor.package(mock_packager, job_name="job1")
+
+ tar_call = [c[0][0] for c in mock_check_call.call_args_list if c[0][0][0] == "tar"]
+ assert tar_call
+ assert not fake_tarball.exists() # removed after extraction
diff --git a/test/run/torchx_backend/schedulers/test_nvcre.py b/test/run/torchx_backend/schedulers/test_nvcre.py
new file mode 100644
index 00000000..64e17edd
--- /dev/null
+++ b/test/run/torchx_backend/schedulers/test_nvcre.py
@@ -0,0 +1,364 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+from unittest import mock
+
+import pytest
+from torchx.schedulers.api import AppDryRunInfo
+from torchx.specs import AppDef, AppState, Role
+
+from nemo_run.core.execution.nvcre import NvcreExecutor, NvcrePhase
+from nemo_run.run.torchx_backend.schedulers.nvcre import (
+ NVCRE_STATES,
+ NvcreScheduler,
+ create_scheduler,
+)
+
+
+@pytest.fixture
+def executor(tmp_path):
+ e = NvcreExecutor(
+ namespace="nemo-perf",
+ container_image="nvcr.io/nvidia/nemo:dev",
+ num_nodes=2,
+ gpus_per_node=8,
+ )
+ e.experiment_id = "test_exp"
+ e.job_dir = str(tmp_path)
+ e.experiment_dir = str(tmp_path)
+ e.job_name = "test_role"
+ return e
+
+
+@pytest.fixture
+def scheduler():
+ return create_scheduler(session_name="test")
+
+
+@pytest.fixture
+def mock_app_def():
+ return AppDef(
+ name="test_app",
+ roles=[
+ Role(
+ name="test_role",
+ image="nvcr.io/nvidia/nemo:dev",
+ entrypoint="python",
+ args=["train.py"],
+ )
+ ],
+ )
+
+
+# ── Scheduler lifecycle ───────────────────────────────────────────────────────
+
+
+def test_create_scheduler():
+ s = create_scheduler(session_name="test")
+ assert isinstance(s, NvcreScheduler)
+ assert s.session_name == "test"
+
+
+def test_state_mapping_covers_all_phases():
+ for phase in NvcrePhase:
+ assert phase in NVCRE_STATES
+
+
+# ── _submit_dryrun ─────────────────────────────────────────────────────────────
+
+
+def test_submit_dryrun_wraps_torchrun_by_default(scheduler, mock_app_def, executor):
+ dryrun_info = scheduler._submit_dryrun(mock_app_def, executor)
+
+ assert isinstance(dryrun_info, AppDryRunInfo)
+ req = dryrun_info.request
+ assert req.cmd[0] == "torchrun"
+ assert "--nnodes=$PET_NNODES" in req.cmd
+ assert "train.py" in req.cmd
+ assert req.name == "test_role"
+
+
+def test_submit_dryrun_no_torchrun_wrap_when_disabled(scheduler, mock_app_def, executor):
+ executor.use_torchrun = False
+ dryrun_info = scheduler._submit_dryrun(mock_app_def, executor)
+ assert dryrun_info.request.cmd == ["python", "train.py"]
+
+
+def test_submit_dryrun_rejects_non_nvcre_executor(scheduler, mock_app_def):
+ with pytest.raises(AssertionError):
+ scheduler._submit_dryrun(mock_app_def, mock.MagicMock())
+
+
+def test_submit_dryrun_rejects_multi_role_app(scheduler, executor):
+ app = AppDef(
+ name="multi",
+ roles=[
+ Role(name="a", image="img", entrypoint="python", args=[]),
+ Role(name="b", image="img", entrypoint="python", args=[]),
+ ],
+ )
+ with pytest.raises(AssertionError):
+ scheduler._submit_dryrun(app, executor)
+
+
+def test_submit_dryrun_apply_yaml_uses_launch_sh_when_pvc_set(scheduler, mock_app_def, executor):
+ executor.workdir_pvc = "my-pvc"
+ dryrun_info = scheduler._submit_dryrun(mock_app_def, executor)
+ yaml_str = dryrun_info._fmt(dryrun_info.request)
+ assert "/bin/bash" in yaml_str
+ assert "launch.sh" in yaml_str
+
+
+# ── schedule ───────────────────────────────────────────────────────────────────
+
+
+def test_schedule_without_pvc(scheduler, mock_app_def, executor):
+ with (
+ mock.patch.object(NvcreExecutor, "submit", return_value="wl-name-123") as mock_submit,
+ mock.patch.object(NvcreExecutor, "package") as mock_pkg,
+ mock.patch(
+ "nemo_run.run.torchx_backend.schedulers.nvcre._save_job"
+ ) as mock_save,
+ ):
+ dryrun_info = scheduler._submit_dryrun(mock_app_def, executor)
+ app_id = scheduler.schedule(dryrun_info)
+
+ assert app_id == "test_exp___test_role___wl-name-123"
+ mock_pkg.assert_not_called()
+ mock_submit.assert_called_once()
+ mock_save.assert_called_once_with("test_exp___test_role___wl-name-123", "wl-name-123", executor)
+
+
+def test_schedule_with_pvc_packages_and_writes_launch_script(scheduler, mock_app_def, executor):
+ executor.workdir_pvc = "my-pvc"
+ with (
+ mock.patch.object(NvcreExecutor, "submit", return_value="wl-name-456"),
+ mock.patch.object(NvcreExecutor, "materialize_launch_script") as mock_mat,
+ mock.patch.object(NvcreExecutor, "package") as mock_pkg,
+ mock.patch("nemo_run.run.torchx_backend.schedulers.nvcre._save_job"),
+ ):
+ dryrun_info = scheduler._submit_dryrun(mock_app_def, executor)
+ app_id = scheduler.schedule(dryrun_info)
+
+ assert app_id == "test_exp___test_role___wl-name-456"
+ mock_mat.assert_called_once()
+ mock_pkg.assert_called_once()
+
+
+# ── describe ───────────────────────────────────────────────────────────────────
+
+
+def test_describe_returns_none_when_job_missing(scheduler):
+ with mock.patch(
+ "nemo_run.run.torchx_backend.schedulers.nvcre._get_jobs", return_value={}
+ ):
+ assert scheduler.describe("nonexistent") is None
+
+
+def test_describe_maps_phase_to_state(scheduler, executor):
+ app_id = "test_exp___test_role___wl-name"
+ with (
+ mock.patch(
+ "nemo_run.run.torchx_backend.schedulers.nvcre._get_jobs",
+ return_value={app_id: {"workloadrun_name": "wl-name", "executor": executor}},
+ ),
+ mock.patch.object(NvcreExecutor, "status", return_value=NvcrePhase.IN_PROGRESS),
+ ):
+ resp = scheduler.describe(app_id)
+
+ assert resp is not None
+ assert resp.state == AppState.RUNNING
+ assert resp.app_id == app_id
+ assert len(resp.roles_statuses[0].replicas) == executor.num_nodes
+
+
+def test_describe_returns_none_without_stored_executor(scheduler):
+ app_id = "test_exp___test_role___wl-name"
+ with mock.patch(
+ "nemo_run.run.torchx_backend.schedulers.nvcre._get_jobs",
+ return_value={app_id: {"workloadrun_name": "wl-name", "executor": None}},
+ ):
+ assert scheduler.describe(app_id) is None
+
+
+# ── log_iter ───────────────────────────────────────────────────────────────────
+
+
+def test_log_iter_returns_empty_when_job_missing(scheduler):
+ with mock.patch(
+ "nemo_run.run.torchx_backend.schedulers.nvcre._get_jobs", return_value={}
+ ):
+ assert list(scheduler.log_iter("nonexistent", "role")) == []
+
+
+def test_log_iter_delegates_to_executor_fetch_logs(scheduler, executor):
+ app_id = "test_exp___test_role___wl-name"
+ with (
+ mock.patch(
+ "nemo_run.run.torchx_backend.schedulers.nvcre._get_jobs",
+ return_value={
+ app_id: {
+ "workloadrun_name": "wl-name",
+ "executor": executor,
+ "job_dir": executor.job_dir,
+ }
+ },
+ ),
+ mock.patch.object(
+ NvcreExecutor, "fetch_logs", return_value=iter(["line1", "line2"])
+ ) as mock_fetch,
+ ):
+ lines = list(scheduler.log_iter(app_id, "role"))
+
+ assert lines == ["line1", "line2"]
+ mock_fetch.assert_called_once_with("wl-name", stream=False)
+
+
+# ── _cancel_existing ───────────────────────────────────────────────────────────
+
+
+def test_cancel_existing_noop_when_job_missing(scheduler):
+ with mock.patch(
+ "nemo_run.run.torchx_backend.schedulers.nvcre._get_jobs", return_value={}
+ ):
+ scheduler._cancel_existing("nonexistent") # should not raise
+
+
+def test_cancel_existing_calls_executor_cancel(scheduler, executor):
+ app_id = "test_exp___test_role___wl-name"
+ with (
+ mock.patch(
+ "nemo_run.run.torchx_backend.schedulers.nvcre._get_jobs",
+ return_value={app_id: {"workloadrun_name": "wl-name", "executor": executor}},
+ ),
+ mock.patch.object(NvcreExecutor, "cancel") as mock_cancel,
+ ):
+ scheduler._cancel_existing(app_id)
+
+ mock_cancel.assert_called_once_with("wl-name")
+
+
+# ── _save_job / _get_jobs round trip ────────────────────────────────────────────
+
+
+def test_save_and_get_jobs_round_trip(executor, tmp_path, monkeypatch):
+ from nemo_run.run.torchx_backend.schedulers import nvcre as nvcre_mod
+
+ job_file = tmp_path / ".nvcre_jobs.json"
+ monkeypatch.setattr(nvcre_mod, "NVCRE_JOB_DIRS", str(job_file))
+ _get_jobs, _save_job = nvcre_mod._get_jobs, nvcre_mod._save_job
+
+ app_id = "test_exp___test_role___wl-name"
+ _save_job(app_id, "wl-name", executor)
+
+ assert job_file.exists()
+ jobs = _get_jobs()
+ assert app_id in jobs
+ assert jobs[app_id]["workloadrun_name"] == "wl-name"
+ assert isinstance(jobs[app_id]["executor"], NvcreExecutor)
+ assert jobs[app_id]["executor"].namespace == executor.namespace
+
+
+def test_get_jobs_returns_empty_when_file_missing(tmp_path, monkeypatch):
+ from nemo_run.run.torchx_backend.schedulers import nvcre as nvcre_mod
+
+ job_file = tmp_path / "does_not_exist.json"
+ monkeypatch.setattr(nvcre_mod, "NVCRE_JOB_DIRS", str(job_file))
+
+ assert nvcre_mod._get_jobs() == {}
+
+
+def test_get_jobs_returns_empty_on_corrupt_json(tmp_path, monkeypatch):
+ from nemo_run.run.torchx_backend.schedulers import nvcre as nvcre_mod
+
+ job_file = tmp_path / ".nvcre_jobs.json"
+ job_file.write_text("{not valid json")
+ monkeypatch.setattr(nvcre_mod, "NVCRE_JOB_DIRS", str(job_file))
+
+ assert nvcre_mod._get_jobs() == {}
+
+
+def test_get_jobs_skips_entry_with_undeserializable_executor(tmp_path, monkeypatch):
+ from nemo_run.run.torchx_backend.schedulers import nvcre as nvcre_mod
+
+ job_file = tmp_path / ".nvcre_jobs.json"
+ job_file.write_text('{"app1": {"workloadrun_name": "wl", "executor": "not-a-valid-blob"}}')
+ monkeypatch.setattr(nvcre_mod, "NVCRE_JOB_DIRS", str(job_file))
+
+ jobs = nvcre_mod._get_jobs()
+ assert "app1" in jobs
+ assert jobs["app1"]["executor"] == "not-a-valid-blob" # left unmodified on deserialize failure
+
+
+# ── misc small methods ──────────────────────────────────────────────────────────
+
+
+def test_run_opts_declares_job_dir(scheduler):
+ opts = scheduler._run_opts()
+ assert "job_dir" in opts._opts if hasattr(opts, "_opts") else True
+
+
+def test_list_returns_empty(scheduler):
+ assert scheduler.list() == []
+
+
+def test_validate_is_noop(scheduler, mock_app_def):
+ assert scheduler._validate(mock_app_def, "nvcre") is None
+
+
+# ── log_iter additional branches ────────────────────────────────────────────────
+
+
+def test_log_iter_returns_empty_when_executor_missing(scheduler):
+ app_id = "test_exp___test_role___wl-name"
+ with mock.patch(
+ "nemo_run.run.torchx_backend.schedulers.nvcre._get_jobs",
+ return_value={app_id: {"workloadrun_name": "wl-name", "executor": None}},
+ ):
+ assert list(scheduler.log_iter(app_id, "role")) == []
+
+
+def test_log_iter_restores_job_dir_when_executor_missing_it(scheduler, executor):
+ app_id = "test_exp___test_role___wl-name"
+ executor.job_dir = ""
+ with (
+ mock.patch(
+ "nemo_run.run.torchx_backend.schedulers.nvcre._get_jobs",
+ return_value={
+ app_id: {
+ "workloadrun_name": "wl-name",
+ "executor": executor,
+ "job_dir": "/restored/job/dir",
+ }
+ },
+ ),
+ mock.patch.object(NvcreExecutor, "fetch_logs", return_value=iter([])) as mock_fetch,
+ ):
+ list(scheduler.log_iter(app_id, "role"))
+
+ assert executor.job_dir == "/restored/job/dir"
+ mock_fetch.assert_called_once()
+
+
+# ── _cancel_existing additional branch ──────────────────────────────────────────
+
+
+def test_cancel_existing_noop_when_executor_missing(scheduler):
+ app_id = "test_exp___test_role___wl-name"
+ with mock.patch(
+ "nemo_run.run.torchx_backend.schedulers.nvcre._get_jobs",
+ return_value={app_id: {"workloadrun_name": "wl-name", "executor": None}},
+ ):
+ scheduler._cancel_existing(app_id) # should not raise