sjayaram/k8s excalibur

scripts/performance/argument_parser.py

@@ -681,6 +681,99 @@ def parse_cli_args():
         default=None,
     )
 
+    # XCalibur
+    xcalibur_args = parser.add_argument_group("XCalibur arguments")
+    xcalibur_args.add_argument(
+        "--xcalibur_namespace",
+        type=str,
+        help="Kubernetes namespace for XCalibur WorkloadRun. When set, uses the XCalibur executor instead of Slurm.",
+        required=False,
+    )
+    xcalibur_args.add_argument(
+        "--xcalibur_image_pull_secret",
+        type=str,
+        help="Kubernetes image pull secret name for pulling the container image.",
+        required=False,
+        default=None,
+    )
+    xcalibur_args.add_argument(
+        "--xcalibur_workdir_pvc",
+        type=str,
+        help="PVC name for syncing job workdir to the cluster before launch.",
+        required=False,
+        default=None,
+    )
+    xcalibur_args.add_argument(
+        "--xcalibur_workdir_pvc_path",
+        type=str,
+        help="Mount path for the workdir PVC inside the training pod.",
+        default="/nemo_run",
+        required=False,
+    )
+    xcalibur_args.add_argument(
+        "--xcalibur_workdir_local_path",
+        type=str,
+        help="Local directory rsynced into the workdir PVC before launch.",
+        required=False,
+        default=None,
+    )
+    xcalibur_args.add_argument(
+        "--xcalibur_node_selector_json",
+        type=str,
+        help="JSON-encoded dict of node selector labels for targeting specific nodes.",
+        required=False,
+        default=None,
+    )
+    xcalibur_args.add_argument(
+        "--xcalibur_volumes_json",
+        type=str,
+        help="JSON-encoded list of Kubernetes Volume dicts.",
+        required=False,
+        default=None,
+    )
+    xcalibur_args.add_argument(
+        "--xcalibur_volume_mounts_json",
+        type=str,
+        help="JSON-encoded list of Kubernetes VolumeMount dicts.",
+        required=False,
+        default=None,
+    )
+    xcalibur_args.add_argument(
+        "--xcalibur_timeout_per_job",
+        type=str,
+        help="Per-job timeout passed to XCalibur orchestration (e.g. '24h').",
+        required=False,
+        default="24h",
+    )
+    xcalibur_args.add_argument(
+        "--xcalibur_test_scale",
+        type=str,
+        help="XCalibur test scale: 'intra-node', 'intra-rack', or 'full-scale'.",
+        required=False,
+        default=None,
+    )
+    xcalibur_args.add_argument(
+        "--xcalibur_kubeconfig",
+        type=str,
+        help="Path to kubeconfig file for xcalctl/kubectl.",
+        required=False,
+        default=None,
+    )
+    xcalibur_args.add_argument(
+        "--xcalibur_kube_context",
+        type=str,
+        help="Kubernetes context to use with xcalctl/kubectl.",
+        required=False,
+        default=None,
+    )
+    xcalibur_args.add_argument(
+        "--xcalibur_gang_scheduler_name",
+        type=str,
+        help="Gang scheduler name for the WorkloadRun (e.g. 'kai-scheduler').",
+        required=False,
+        default=None,
+    )
+
     # For performance
     performance_args = parser.add_argument_group("Performance arguments")
     performance_args.add_argument(

scripts/performance/perf_plugins.py

@@ -33,6 +33,10 @@
 
 import nemo_run as run
 from nemo_run import Plugin, Script, SlurmExecutor
+try:
+    from nemo_run.core.execution.xcalibur import XCaliburExecutor
+except ImportError:
+    XCaliburExecutor = None  # type: ignore[assignment,misc]
 
 
 logger: logging.Logger = logging.getLogger(__name__)
@@ -185,6 +189,9 @@ def setup(self, task: Union["run.Partial", "run.Script"], executor: "run.Executo
         if isinstance(executor, SlurmExecutor):
             # NOTE: DO NOT change to f-string, `%q{}` is Slurm placeholder
             launcher.nsys_filename = "profile_%p_%q{SLURM_JOB_ID}_node%q{SLURM_NODEID}_rank%q{SLURM_PROCID}"
+        elif XCaliburExecutor is not None and isinstance(executor, XCaliburExecutor):
+            # XCalibur pods use PET_* env vars for distributed rank info.
+            launcher.nsys_filename = "profile_node${PET_NODE_RANK}_rank%p"
 
         if self.nsys_gpu_metrics:
             if hasattr(launcher, "nsys_gpu_metrics"):

scripts/performance/setup_experiment.py

@@ -87,10 +87,14 @@ def _filter_run_script_args(argv: List[str]) -> List[str]:
       carry JSON values whose ``{}`` / ``[]`` are brace/glob-expanded by the
       shell in the generated launch command, corrupting argv and leaking tokens
       into the training entrypoint's Hydra override parser.
+    * ``--offline`` — controls HF Hub access on the launcher node only; offline
+      behaviour in the rank-local script is governed by the ``HF_HUB_OFFLINE``
+      environment variable, not by this flag.
 
-    All of these take a value, passed either as ``--flag value`` (two tokens) or
-    ``--flag=value`` (one token).
+    Value-taking flags are passed as ``--flag value`` or ``--flag=value``.
+    Boolean flags (no following value) are listed in ``_LAUNCHER_ONLY_BOOL``.
     """
+    _LAUNCHER_ONLY_BOOL = {"--offline", "--dryrun"}
 
     def _is_launcher_only(flag: str) -> bool:
         return flag in (
@@ -102,7 +106,7 @@ def _is_launcher_only(flag: str) -> bool:
             "--enable_vboost",
             "--lock_gpu_freq",
             "--peak_mem_clk",
-        ) or flag.startswith("--kubeflow_")
+        ) or flag.startswith("--kubeflow_") or flag.startswith("--xcalibur_")
 
     filtered_args = []
     skip_next = False
@@ -111,6 +115,8 @@ def _is_launcher_only(flag: str) -> bool:
         if skip_next:
             skip_next = False
             continue
+        if arg in _LAUNCHER_ONLY_BOOL:
+            continue
         if _is_launcher_only(arg.split("=", 1)[0]):
             skip_next = "=" not in arg
             continue
@@ -541,6 +547,20 @@ def main(
     kubeflow_container_kwargs_json: Optional[str],
     kubeflow_labels_json: Optional[str],
     kubeflow_pod_annotations_json: Optional[str],
+    xcalibur_namespace: Optional[str] = None,
+    xcalibur_image_pull_secret: Optional[str] = None,
+    xcalibur_workdir_pvc: Optional[str] = None,
+    xcalibur_workdir_pvc_path: str = "/nemo_run",
+    xcalibur_workdir_local_path: Optional[str] = None,
+    xcalibur_node_selector_json: Optional[str] = None,
+    xcalibur_volumes_json: Optional[str] = None,
+    xcalibur_volume_mounts_json: Optional[str] = None,
+    xcalibur_timeout_per_job: str = "24h",
+    xcalibur_test_scale: Optional[str] = None,
+    xcalibur_kubeconfig: Optional[str] = None,
+    xcalibur_kube_context: Optional[str] = None,
+    xcalibur_gang_scheduler_name: Optional[str] = None,
+    deterministic: bool = False,
     config_variant: str | None = None,
     gres: Optional[str] = None,
     packager: str = "git",
@@ -588,7 +608,8 @@ def main(
     if export_nsys_sqlite and not enable_nsys:
         logger.warning("--export_nsys_sqlite was set without --enable_nsys; no Nsys SQLite export will be generated.")
 
-    script_name = ENTRYPOINT_BOOTSTRAP
+    # XCalibur uses run_script.py directly; all other executors use bootstrap.py.
+    script_name = "run_script.py" if xcalibur_namespace else ENTRYPOINT_BOOTSTRAP
     # Keep the historical W&B-name behavior for CI. The lightweight fallback
     # deliberately avoids resolving a recipe: effective parallelism, batches,
     # and process environment are finalized by bootstrap.py in the container.
@@ -612,7 +633,7 @@ def main(
         # runs. Creating the dir from the launcher would either fail (PVC
         # not present) or create a useless dir on the launcher's local FS.
         # Let the trainer container create its own dirs on first write.
-        if kubeflow_namespace is None:
+        if kubeflow_namespace is None and xcalibur_namespace is None:
             save_dir_path.mkdir(parents=True, exist_ok=True)
             save_dir_mount = f"{save_dir_path}:{save_dir_path}"
             if save_dir_mount not in custom_mounts:
@@ -631,7 +652,7 @@ def main(
     # Kubeflow the trainer pod runs the image — which ships Megatron-Bridge at
     # /opt/Megatron-Bridge — and custom_mounts do not apply, so the launcher's
     # /tmp path does not exist in the pod; use the image's script path instead.
-    if kubeflow_namespace:
+    if kubeflow_namespace or xcalibur_namespace:
         in_container_script_dir = "/opt/Megatron-Bridge/scripts/performance"
         in_container_script_path = f"{in_container_script_dir}/{script_name}"
     else:
@@ -688,6 +709,37 @@ def main(
             labels=json.loads(kubeflow_labels_json) if kubeflow_labels_json else None,
             pod_annotations=(json.loads(kubeflow_pod_annotations_json) if kubeflow_pod_annotations_json else None),
         )
+    elif xcalibur_namespace is not None:
+        try:
+            from utils.executors import xcalibur_executor
+        except ImportError:
+            from .utils.executors import xcalibur_executor
+        executor = xcalibur_executor(
+            namespace=xcalibur_namespace,
+            image=container_image,
+            num_nodes=-(num_gpus // -gpus_per_node),
+            gpus_per_node=gpus_per_node,
+            image_pull_secret=xcalibur_image_pull_secret,
+            workdir_pvc=xcalibur_workdir_pvc,
+            workdir_pvc_path=xcalibur_workdir_pvc_path,
+            workdir_local_path=xcalibur_workdir_local_path,
+            node_selector=json.loads(xcalibur_node_selector_json) if xcalibur_node_selector_json else None,
+            volumes=json.loads(xcalibur_volumes_json) if xcalibur_volumes_json else None,
+            volume_mounts=json.loads(xcalibur_volume_mounts_json) if xcalibur_volume_mounts_json else None,
+            timeout_per_job=xcalibur_timeout_per_job,
+            test_scale=xcalibur_test_scale,
+            kubeconfig=xcalibur_kubeconfig,
+            kube_context=xcalibur_kube_context,
+            gang_scheduler_name=xcalibur_gang_scheduler_name,
+        )
+        xcal_env = custom_env_vars.copy()
+        if hf_token:
+            # Always allow the pod to reach HF to download gated model files
+            # (tokenizer configs, etc.) — the pod has no access to the host
+            # HF cache so offline mode must not be forced here even when
+            # --offline was passed for the launcher-side setup.
+            xcal_env.update({"HF_TOKEN": hf_token, "HF_HUB_OFFLINE": "0", "TRANSFORMERS_OFFLINE": "0"})
+        executor.env_vars = xcal_env
     else:
         executor = slurm_executor(
             gpu=gpu,
@@ -808,15 +860,17 @@ def main(
                 logger.info("dryrun requested: exiting")
                 return
 
-            job_dir, job_status = get_job_dir_and_status_from_run(exp_name)
-
-            terminal_failure = job_status not in ["SUCCEEDED", "SUBMITTED", "PENDING", "RUNNING"]
-
             if detach:
+                # For detached runs (e.g. XCalibur), skip status polling — the
+                # caller (llmb-run) polls for completion via xcalctl.
                 is_finished_experiment = True
                 is_testing_passed = True
                 break
 
+            job_dir, job_status = get_job_dir_and_status_from_run(exp_name)
+
+            terminal_failure = job_status not in ["SUCCEEDED", "SUBMITTED", "PENDING", "RUNNING"]
+
             log_file_paths = list(Path(f"{job_dir}").glob("log*.out"))
             ensure_logs_where_written(log_file_paths)
 
@@ -1010,9 +1064,10 @@ def main(
         task=args.task,
         compute_dtype=args.compute_dtype,
         gpu=args.gpu,
-        hf_token=args.hf_token,
+        hf_token=args.hf_token or os.environ.get('HF_TOKEN'),
         offline=args.offline,
-        detach=args.detach,
+        # Force detach for XCalibur — llmb-run polls for completion via xcalctl
+        detach=True if args.xcalibur_namespace else args.detach,
         dryrun=args.dryrun,
         enable_vboost=args.enable_vboost,
         lock_gpu_freq=args.lock_gpu_freq,
@@ -1089,6 +1144,20 @@ def main(
         kubeflow_container_kwargs_json=args.kubeflow_container_kwargs_json,
         kubeflow_labels_json=args.kubeflow_labels_json,
         kubeflow_pod_annotations_json=args.kubeflow_pod_annotations_json,
+        xcalibur_namespace=args.xcalibur_namespace,
+        xcalibur_image_pull_secret=args.xcalibur_image_pull_secret,
+        xcalibur_workdir_pvc=args.xcalibur_workdir_pvc,
+        xcalibur_workdir_pvc_path=args.xcalibur_workdir_pvc_path,
+        xcalibur_workdir_local_path=args.xcalibur_workdir_local_path,
+        xcalibur_node_selector_json=args.xcalibur_node_selector_json,
+        xcalibur_volumes_json=args.xcalibur_volumes_json,
+        xcalibur_volume_mounts_json=args.xcalibur_volume_mounts_json,
+        xcalibur_timeout_per_job=args.xcalibur_timeout_per_job,
+        xcalibur_test_scale=args.xcalibur_test_scale,
+        xcalibur_kubeconfig=args.xcalibur_kubeconfig,
+        xcalibur_kube_context=args.xcalibur_kube_context,
+        xcalibur_gang_scheduler_name=args.xcalibur_gang_scheduler_name,
+        deterministic=args.deterministic,
         config_variant=config_variant,
         gres=args.gres,
         packager=args.packager,

scripts/performance/utils/executors.py

@@ -361,3 +361,42 @@ def kubeflow_executor(
         packager=run.GitArchivePackager(include_submodules=True),
     )
     return executor
+
+
+def xcalibur_executor(
+    namespace: str,
+    image: str,
+    num_nodes: int,
+    gpus_per_node: int,
+    image_pull_secret: Optional[str] = None,
+    workdir_pvc: Optional[str] = None,
+    workdir_pvc_path: str = "/nemo_run",
+    workdir_local_path: Optional[str] = None,
+    node_selector: Optional[dict] = None,
+    volumes: Optional[list] = None,
+    volume_mounts: Optional[list] = None,
+    timeout_per_job: str = "24h",
+    test_scale: Optional[str] = None,
+    kubeconfig: Optional[str] = None,
+    kube_context: Optional[str] = None,
+    gang_scheduler_name: Optional[str] = None,
+):
+    from nemo_run.core.execution.xcalibur import XCaliburExecutor
+    return XCaliburExecutor(
+        namespace=namespace,
+        container_image=image,
+        num_nodes=num_nodes,
+        gpus_per_node=gpus_per_node,
+        image_pull_secret=image_pull_secret,
+        node_selector=node_selector or {},
+        workdir_pvc=workdir_pvc,
+        workdir_pvc_path=workdir_pvc_path,
+        workdir_local_path=workdir_local_path,
+        volumes=volumes or [],
+        volume_mounts=volume_mounts or [],
+        timeout_per_job=timeout_per_job,
+        test_scale=test_scale,
+        kubeconfig=kubeconfig,
+        kube_context=kube_context,
+        gang_scheduler_name=gang_scheduler_name,
+    )