diff --git a/nemo_run/__init__.py b/nemo_run/__init__.py index 0d403a54..6bc19c1a 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.xcalibur import XCaliburExecutor 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", + "XCaliburExecutor", "run", "Script", "SkypilotExecutor", diff --git a/nemo_run/core/execution/__init__.py b/nemo_run/core/execution/__init__.py index 08e088c8..4f68076f 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.xcalibur import XCaliburExecutor __all__ = [ "LocalExecutor", @@ -27,4 +28,5 @@ "DGXCloudExecutor", "LeptonExecutor", "KubeflowExecutor", + "XCaliburExecutor", ] diff --git a/nemo_run/core/execution/xcalibur.py b/nemo_run/core/execution/xcalibur.py new file mode 100644 index 00000000..012bebc3 --- /dev/null +++ b/nemo_run/core/execution/xcalibur.py @@ -0,0 +1,695 @@ +# 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 json +import logging +import os +import re +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__) + +_XCALIBUR_WORKLOADRUN_API = "excalibur.nvidia.com/v1alpha1" +_DATA_MOVER_IMAGE = "alpine:3.19" + + +class XCaliburPhase(Enum): + PENDING = "Pending" + IN_PROGRESS = "InProgress" + SUCCEEDED = "Succeeded" + FAILED = "Failed" + UNKNOWN = "Unknown" + + +@dataclass(kw_only=True) +class XCaliburExecutor(Executor): + """ + Dataclass to configure an XCalibur executor. + + Submits jobs to an XCalibur-managed Kubernetes cluster via the ``xcalctl`` + CLI using the WorkloadRun API. Requires ``xcalctl`` (and ``kubectl``) to be + on the PATH of the machine running NeMo-Run. + + Example:: + + executor = XCaliburExecutor( + 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 XCalibur + + # ── 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 + + # ── Launcher ────────────────────────────────────────────────────────────── + # When True, wrap the python entrypoint with torchrun using the PET_* env + # vars that XCalibur 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 + + # ── xcalctl / kubectl config ────────────────────────────────────────────── + xcalctl_bin: str = "xcalctl" + 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: + # XCalibur 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.max_restarts: + spec["checkpoint"] = {"maxRestarts": self.max_restarts} + + if self.gang_scheduler_name: + spec["gangScheduler"] = {"schedulerName": self.gang_scheduler_name} + + return { + "apiVersion": _XCALIBUR_WORKLOADRUN_API, + "kind": "WorkloadRun", + "metadata": {"name": self._safe_name(), "namespace": self.namespace}, + "spec": spec, + } + + def _safe_name(self) -> str: + """RFC-1123 safe WorkloadRun name derived from job_name.""" + name = (self.job_name or "xcalibur-job").lower().replace("_", "-").replace(".", "-") + return name[:63].rstrip("-") + + # ── xcalctl / kubectl helpers ───────────────────────────────────────────── + + def _xcalctl_base(self) -> list[str]: + args = [self.xcalctl_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. + + xcalctl generates its own WorkloadRun name and does not necessarily + use the ``--name`` flag we pass. We parse the actual name from + xcalctl's stdout so that subsequent ``status()`` and ``fetch_logs()`` + calls use the right resource name. + """ + name = self._safe_name() + cmd = self._xcalctl_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"xcalctl workloadrun run failed (rc={result.returncode}):\n{result.stderr}" + ) + + actual_name = self._parse_submitted_name(result.stdout) + if not actual_name: + # xcalctl output format not recognised — ask kubectl for the most + # recently created WorkloadRun in our namespace as a fallback. + actual_name = self._latest_workloadrun_name() or name + if actual_name != name: + logger.info( + "WorkloadRun submitted: xcalctl used name '%s' (we requested '%s')", + actual_name, name, + ) + else: + logger.info("WorkloadRun '%s' submitted", actual_name) + self._workloadrun_name = actual_name + return actual_name + + def _latest_workloadrun_name(self) -> str | None: + """Return the name of the most recently created WorkloadRun in our namespace. + + Used as a last-resort fallback when xcalctl output cannot be parsed. + A short sleep is applied first to allow the API server to reflect the + newly created resource. + """ + time.sleep(2) + cmd = self._kubectl_base() + [ + "get", "workloadruns", + "-n", self.namespace, + "--sort-by=.metadata.creationTimestamp", + "-o", "jsonpath={.items[-1].metadata.name}", + ] + result = subprocess.run(cmd, capture_output=True, text=True) + if result.returncode == 0 and result.stdout.strip(): + return result.stdout.strip() + logger.warning("Could not retrieve latest WorkloadRun via kubectl: %s", result.stderr.strip()) + return None + + def _parse_submitted_name(self, output: str) -> str | None: + """Extract the WorkloadRun name xcalctl actually assigned from its output. + + xcalctl may output the name in several formats, e.g.: + - kubectl-style: ``workloadrun.excalibur.nvidia.com/name created`` + - plain: ``name`` + - JSON: ``{"name": "name", ...}`` + Returns None if no recognisable name is found. + """ + output = output.strip() + # kubectl-style: "workloadrun.*/name created|configured|unchanged" + m = re.search(r'workloadrun[^/]*/([a-z0-9][a-z0-9-]{2,61})', output, re.IGNORECASE) + if m: + return m.group(1) + # JSON: {"name": "value"} or {"workloadrun": {"name": "value"}} + m = re.search(r'"name"\s*:\s*"([a-z0-9][a-z0-9-]{2,61})"', output, re.IGNORECASE) + if m: + return m.group(1) + # Plain: a single token that looks like a k8s name on its own line + m = re.search(r'^([a-z][a-z0-9-]{2,61})\s*$', output, re.MULTILINE) + if m: + return m.group(1) + return None + + def status(self, name: str) -> XCaliburPhase: + """Return the current phase of WorkloadRun *name*. + + Tries xcalctl first. Falls back to inspecting pod phases via kubectl + when xcalctl returns a non-zero exit code (e.g. the WorkloadRun was + cleaned up after completion) or reports an unrecognised phase string. + """ + cmd = self._xcalctl_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 XCaliburPhase(phase_str) + except ValueError: + logger.warning( + "Unrecognised xcalctl phase '%s' for '%s'; falling back to kubectl CRD check", + phase_str, name, + ) + else: + logger.warning( + "xcalctl 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) -> XCaliburPhase: + """Read phase directly from the WorkloadRun CRD via kubectl. + + xcalctl is a thin wrapper over the same CRD. Reading it directly + avoids xcalctl output-format surprises and works regardless of whether + XCalibur'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 XCaliburPhase.UNKNOWN + + phase_str = result.stdout.strip() + if not phase_str: + logger.warning("Empty phase from WorkloadRun CRD '%s'", name) + return XCaliburPhase.UNKNOWN + + try: + return XCaliburPhase(phase_str) + except ValueError: + logger.warning( + "Unrecognised WorkloadRun CRD phase '%s' for '%s'", phase_str, name, + ) + return XCaliburPhase.UNKNOWN + + def cancel(self, name: str) -> None: + """Cancel WorkloadRun *name*.""" + cmd = self._xcalctl_base() + [ + "workloadrun", "cancel", name, + "-n", self.namespace, + ] + result = subprocess.run(cmd, capture_output=True, text=True) + if result.returncode != 0: + logger.warning("xcalctl cancel failed for '%s': %s", name, result.stderr) + else: + logger.info("Cancelled WorkloadRun '%s'", name) + + def _get_xcalibur_job_name(self, workloadrun_name: str) -> str | None: + """Return the XCalibur internal job name from the WorkloadRun CRD. + + XCalibur stamps pods with ``excalibur.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", "xcaliburJobName", "workloadJobName"): + val = status.get(field) + if val and val != workloadrun_name: + return val + labels = data.get("metadata", {}).get("labels", {}) + val = labels.get("excalibur.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. XCalibur 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 ``excalibur.nvidia.com/job=`` that + XCalibur stamps on the pods it creates. + """ + # Pods are labelled with the JobSet name, not the excalibur.nvidia.com/job + # label. Derive the XCalibur 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. + xcalibur_job = self._get_xcalibur_job_name(name) or name + jobset_name = f"{xcalibur_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()) + if max_retries > 0: + cmd_str = " ".join(cmd) + 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 = " ".join(cmd) + + 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..2cdc4073 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.xcalibur import XCaliburExecutor 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, + XCaliburExecutor, ) _DETACH_SUPPORTED_EXECUTORS = ( SlurmExecutor, @@ -215,6 +217,7 @@ class Experiment(ConfigurableMixin): SkypilotJobsExecutor, DGXCloudExecutor, LeptonExecutor, + XCaliburExecutor, ) _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..f62b1f27 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.xcalibur import XCaliburExecutor EXECUTOR_MAPPING: dict[Type[Executor], str] = { SlurmExecutor: "slurm_tunnel", @@ -36,6 +37,7 @@ DGXCloudExecutor: "dgx_cloud", LeptonExecutor: "lepton", KubeflowExecutor: "kubeflow", + XCaliburExecutor: "xcalibur", } REVERSE_EXECUTOR_MAPPING: dict[str, Type[Executor]] = { @@ -47,6 +49,7 @@ "dgx_cloud": DGXCloudExecutor, "lepton": LeptonExecutor, "kubeflow": KubeflowExecutor, + "xcalibur": XCaliburExecutor, } diff --git a/nemo_run/run/torchx_backend/schedulers/xcalibur.py b/nemo_run/run/torchx_backend/schedulers/xcalibur.py new file mode 100644 index 00000000..3e20b08f --- /dev/null +++ b/nemo_run/run/torchx_backend/schedulers/xcalibur.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.xcalibur import XCaliburExecutor, XCaliburPhase +from nemo_run.core.serialization.zlib_json import ZlibJSONSerializer +from nemo_run.run.torchx_backend.schedulers.api import SchedulerMixin + +logger = logging.getLogger(__name__) + +XCALIBUR_JOB_DIRS = os.path.join(get_nemorun_home(), ".xcalibur_jobs.json") + +XCALIBUR_STATES: dict[XCaliburPhase, AppState] = { + XCaliburPhase.PENDING: AppState.PENDING, + XCaliburPhase.IN_PROGRESS: AppState.RUNNING, + XCaliburPhase.SUCCEEDED: AppState.SUCCEEDED, + XCaliburPhase.FAILED: AppState.FAILED, + XCaliburPhase.UNKNOWN: AppState.PENDING, +} + + +@dataclass +class XCaliburRequest: + """Wraps the AppDef and XCaliburExecutor for dryrun/schedule.""" + + app: AppDef + executor: XCaliburExecutor + cmd: list[str] + name: str + + +class XCaliburScheduler(SchedulerMixin, Scheduler[dict]): # type: ignore + def __init__(self, session_name: str) -> None: + super().__init__("xcalibur", 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[XCaliburRequest]: + assert isinstance(cfg, XCaliburExecutor), ( + f"{cfg.__class__} is not supported by XCaliburScheduler." + ) + executor = cfg + assert len(app.roles) == 1, "XCaliburScheduler 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. XCalibur 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 = XCaliburRequest(app=app, executor=executor, cmd=cmd, name=role.name) + + def _wl_cmd(r: XCaliburRequest) -> 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[XCaliburRequest]) -> 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", "xcalibur_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[XCaliburExecutor] = job_info.get("executor") + if not executor: + return None + + phase = executor.status(workloadrun_name) + app_state = XCALIBUR_STATES.get(phase, AppState.PENDING) + + 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[XCaliburExecutor] = 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[XCaliburExecutor] = 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) -> XCaliburScheduler: + return XCaliburScheduler(session_name=session_name) + + +def _save_job(app_id: str, workloadrun_name: str, executor: XCaliburExecutor) -> None: + original_apps: dict = {} + os.makedirs(os.path.dirname(XCALIBUR_JOB_DIRS), exist_ok=True) + if not os.path.isfile(XCALIBUR_JOB_DIRS): + Path(XCALIBUR_JOB_DIRS).touch() + + serializer = ZlibJSONSerializer() + with open(XCALIBUR_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, XCALIBUR_JOB_DIRS) + + +def _get_jobs() -> dict[str, dict]: + if not os.path.isfile(XCALIBUR_JOB_DIRS): + return {} + with open(XCALIBUR_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 XCalibur executor: %s", e) + return data diff --git a/pyproject.toml b/pyproject.toml index 6b5e4a99..270a956c 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" +xcalibur = "nemo_run.run.torchx_backend.schedulers.xcalibur:create_scheduler" [project.optional-dependencies] skypilot = [ diff --git a/test/core/execution/test_xcalibur.py b/test/core/execution/test_xcalibur.py new file mode 100644 index 00000000..0bc2dd80 --- /dev/null +++ b/test/core/execution/test_xcalibur.py @@ -0,0 +1,629 @@ +# 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.xcalibur import XCaliburExecutor, XCaliburPhase + + +def _completed(returncode=0, stdout="", stderr=""): + return subprocess.CompletedProcess(args=[], returncode=returncode, stdout=stdout, stderr=stderr) + + +class TestXCaliburExecutor: + @pytest.fixture + def executor(self): + e = XCaliburExecutor( + 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 = XCaliburExecutor(namespace="ns", container_image="img:latest", num_nodes=1) + e.job_name = "job1" + manifest = e.build_workloadrun_yaml(["python", "train.py"]) + + assert manifest["apiVersion"] == "excalibur.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.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"] == {"maxRestarts": 3} + assert spec["gangScheduler"] == {"schedulerName": "kai-scheduler"} + + # ── _safe_name ───────────────────────────────────────────────────────────── + + @pytest.mark.parametrize( + "job_name,expected", + [ + ("My_Job.Name", "my-job-name"), + ("", "xcalibur-job"), + ("Already-Safe", "already-safe"), + ("trailing-dot.", "trailing-dot"), + ], + ) + def test_safe_name(self, job_name, expected): + e = XCaliburExecutor(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 = XCaliburExecutor(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.xcalibur.subprocess.run") as mock_run: + mock_run.return_value = _completed( + stdout="workloadrun.excalibur.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] == "xcalctl" + 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.xcalibur.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.xcalibur.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.xcalibur.subprocess.run") as mock_run, + patch("nemo_run.core.execution.xcalibur.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.xcalibur.subprocess.run") as mock_run, + patch("nemo_run.core.execution.xcalibur.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.xcalibur.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_xcalctl(self, executor): + with patch("nemo_run.core.execution.xcalibur.subprocess.run") as mock_run: + mock_run.return_value = _completed(stdout="Succeeded\n") + phase = executor.status("wl-name") + assert phase == XCaliburPhase.SUCCEEDED + mock_run.assert_called_once() + + def test_status_falls_back_to_crd_on_xcalctl_failure(self, executor): + with patch("nemo_run.core.execution.xcalibur.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 == XCaliburPhase.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.xcalibur.subprocess.run") as mock_run: + mock_run.side_effect = [ + _completed(stdout="SomeWeirdPhase\n"), + _completed(stdout="InProgress\n"), + ] + phase = executor.status("wl-name") + assert phase == XCaliburPhase.IN_PROGRESS + + def test_status_crd_fallback_returns_unknown_on_empty_or_error(self, executor): + with patch("nemo_run.core.execution.xcalibur.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 == XCaliburPhase.UNKNOWN + + # ── cancel ───────────────────────────────────────────────────────────────── + + def test_cancel_success(self, executor): + with patch("nemo_run.core.execution.xcalibur.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.xcalibur.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.xcalibur.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_xcalibur_job_name_from_status_field(self, executor): + with patch("nemo_run.core.execution.xcalibur.subprocess.run") as mock_run: + mock_run.return_value = _completed( + returncode=0, stdout='{"status": {"jobName": "internal-job"}, "metadata": {"labels": {}}}' + ) + job_name = executor._get_xcalibur_job_name("wl-name") + assert job_name == "internal-job" + + def test_get_xcalibur_job_name_falls_back_to_jobsets(self, executor): + with patch("nemo_run.core.execution.xcalibur.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_xcalibur_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 = XCaliburExecutor(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.xcalibur.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.xcalibur.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 = XCaliburExecutor(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 = XCaliburExecutor(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 = XCaliburExecutor(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"} + + # ── xcalctl_base / kubectl_base kubeconfig/context ──────────────────────── + + def test_xcalctl_base_includes_kubeconfig_and_context(self): + e = XCaliburExecutor( + namespace="ns", container_image="img", kubeconfig="/path/kubeconfig", kube_context="ctx1" + ) + args = e._xcalctl_base() + assert args == ["xcalctl", "--kubeconfig", "/path/kubeconfig", "--context", "ctx1"] + + def test_kubectl_base_includes_kubeconfig_and_context(self): + e = XCaliburExecutor( + 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.xcalibur.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 == XCaliburPhase.UNKNOWN + + def test_crd_phase_returns_unknown_on_empty_output(self, executor): + with patch("nemo_run.core.execution.xcalibur.subprocess.run") as mock_run: + mock_run.return_value = _completed(returncode=0, stdout=" ") + phase = executor._kubectl_workloadrun_crd_phase("wl-name") + assert phase == XCaliburPhase.UNKNOWN + + def test_crd_phase_returns_unknown_on_unrecognised_phase(self, executor): + with patch("nemo_run.core.execution.xcalibur.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 == XCaliburPhase.UNKNOWN + + def test_crd_phase_returns_recognised_phase(self, executor): + with patch("nemo_run.core.execution.xcalibur.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 == XCaliburPhase.PENDING + + # ── _get_xcalibur_job_name edge cases ───────────────────────────────────── + + def test_get_xcalibur_job_name_returns_none_on_kubectl_failure(self, executor): + with patch("nemo_run.core.execution.xcalibur.subprocess.run") as mock_run: + mock_run.return_value = _completed(returncode=1, stderr="gone") + assert executor._get_xcalibur_job_name("wl-name") is None + + def test_get_xcalibur_job_name_handles_invalid_json(self, executor): + with patch("nemo_run.core.execution.xcalibur.subprocess.run") as mock_run: + mock_run.side_effect = [ + _completed(returncode=0, stdout="not json"), + _completed(returncode=0, stdout=""), + ] + assert executor._get_xcalibur_job_name("wl-name") is None + + def test_get_xcalibur_job_name_from_labels(self, executor): + with patch("nemo_run.core.execution.xcalibur.subprocess.run") as mock_run: + mock_run.return_value = _completed( + returncode=0, + stdout='{"status": {}, "metadata": {"labels": {"excalibur.nvidia.com/job": "label-job"}}}', + ) + job_name = executor._get_xcalibur_job_name("wl-name") + assert job_name == "label-job" + + def test_get_xcalibur_job_name_returns_none_when_no_jobsets(self, executor): + with patch("nemo_run.core.execution.xcalibur.subprocess.run") as mock_run: + mock_run.side_effect = [ + _completed(returncode=0, stdout='{"status": {}, "metadata": {"labels": {}}}'), + _completed(returncode=0, stdout=""), + ] + assert executor._get_xcalibur_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.xcalibur.subprocess.run") as mock_run, + patch("nemo_run.core.execution.xcalibur.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.xcalibur.subprocess.run") as mock_run, + patch("nemo_run.core.execution.xcalibur.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.xcalibur.subprocess.run") as mock_run, + patch("nemo_run.core.execution.xcalibur.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.xcalibur.subprocess.run") as mock_run, + patch("nemo_run.core.execution.xcalibur.subprocess.check_call"), + patch("nemo_run.core.execution.xcalibur.time.sleep"), + patch("nemo_run.core.execution.xcalibur.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.xcalibur.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.xcalibur.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.xcalibur.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(XCaliburExecutor, "_start_data_mover_pod") as mock_start, + patch.object(XCaliburExecutor, "_rsync_to_pod") as mock_rsync, + patch.object(XCaliburExecutor, "_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(XCaliburExecutor, "_start_data_mover_pod"), + patch.object(XCaliburExecutor, "_rsync_to_pod", side_effect=RuntimeError("rsync failed")), + patch.object(XCaliburExecutor, "_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.xcalibur.subprocess.run") as mock_run, + patch("nemo_run.core.execution.xcalibur.subprocess.check_call"), + patch.object(XCaliburExecutor, "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.xcalibur.subprocess.run") as mock_run, + patch.object(XCaliburExecutor, "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.xcalibur.subprocess.check_call") as mock_check_call, + patch.object(XCaliburExecutor, "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.xcalibur.subprocess.check_call") as mock_check_call, + patch.object(XCaliburExecutor, "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_xcalibur.py b/test/run/torchx_backend/schedulers/test_xcalibur.py new file mode 100644 index 00000000..392c1849 --- /dev/null +++ b/test/run/torchx_backend/schedulers/test_xcalibur.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.xcalibur import XCaliburExecutor, XCaliburPhase +from nemo_run.run.torchx_backend.schedulers.xcalibur import ( + XCALIBUR_STATES, + XCaliburScheduler, + create_scheduler, +) + + +@pytest.fixture +def executor(tmp_path): + e = XCaliburExecutor( + 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, XCaliburScheduler) + assert s.session_name == "test" + + +def test_state_mapping_covers_all_phases(): + for phase in XCaliburPhase: + assert phase in XCALIBUR_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_xcalibur_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(XCaliburExecutor, "submit", return_value="wl-name-123") as mock_submit, + mock.patch.object(XCaliburExecutor, "package") as mock_pkg, + mock.patch( + "nemo_run.run.torchx_backend.schedulers.xcalibur._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(XCaliburExecutor, "submit", return_value="wl-name-456"), + mock.patch.object(XCaliburExecutor, "materialize_launch_script") as mock_mat, + mock.patch.object(XCaliburExecutor, "package") as mock_pkg, + mock.patch("nemo_run.run.torchx_backend.schedulers.xcalibur._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.xcalibur._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.xcalibur._get_jobs", + return_value={app_id: {"workloadrun_name": "wl-name", "executor": executor}}, + ), + mock.patch.object(XCaliburExecutor, "status", return_value=XCaliburPhase.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.xcalibur._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.xcalibur._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.xcalibur._get_jobs", + return_value={ + app_id: { + "workloadrun_name": "wl-name", + "executor": executor, + "job_dir": executor.job_dir, + } + }, + ), + mock.patch.object( + XCaliburExecutor, "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.xcalibur._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.xcalibur._get_jobs", + return_value={app_id: {"workloadrun_name": "wl-name", "executor": executor}}, + ), + mock.patch.object(XCaliburExecutor, "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 xcalibur as xcalibur_mod + + job_file = tmp_path / ".xcalibur_jobs.json" + monkeypatch.setattr(xcalibur_mod, "XCALIBUR_JOB_DIRS", str(job_file)) + _get_jobs, _save_job = xcalibur_mod._get_jobs, xcalibur_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"], XCaliburExecutor) + 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 xcalibur as xcalibur_mod + + job_file = tmp_path / "does_not_exist.json" + monkeypatch.setattr(xcalibur_mod, "XCALIBUR_JOB_DIRS", str(job_file)) + + assert xcalibur_mod._get_jobs() == {} + + +def test_get_jobs_returns_empty_on_corrupt_json(tmp_path, monkeypatch): + from nemo_run.run.torchx_backend.schedulers import xcalibur as xcalibur_mod + + job_file = tmp_path / ".xcalibur_jobs.json" + job_file.write_text("{not valid json") + monkeypatch.setattr(xcalibur_mod, "XCALIBUR_JOB_DIRS", str(job_file)) + + assert xcalibur_mod._get_jobs() == {} + + +def test_get_jobs_skips_entry_with_undeserializable_executor(tmp_path, monkeypatch): + from nemo_run.run.torchx_backend.schedulers import xcalibur as xcalibur_mod + + job_file = tmp_path / ".xcalibur_jobs.json" + job_file.write_text('{"app1": {"workloadrun_name": "wl", "executor": "not-a-valid-blob"}}') + monkeypatch.setattr(xcalibur_mod, "XCALIBUR_JOB_DIRS", str(job_file)) + + jobs = xcalibur_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, "xcalibur") 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.xcalibur._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.xcalibur._get_jobs", + return_value={ + app_id: { + "workloadrun_name": "wl-name", + "executor": executor, + "job_dir": "/restored/job/dir", + } + }, + ), + mock.patch.object(XCaliburExecutor, "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.xcalibur._get_jobs", + return_value={app_id: {"workloadrun_name": "wl-name", "executor": None}}, + ): + scheduler._cancel_existing(app_id) # should not raise