sjayaram/k8s excalibur

scripts/performance/argument_parser.py

@@ -681,6 +681,92 @@ 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,
+    )
+
     # For performance
     performance_args = parser.add_argument_group("Performance arguments")
     performance_args.add_argument(

scripts/performance/setup_experiment.py

@@ -36,6 +36,7 @@
         _kubeflow_numa_binding_script,
         kubeflow_executor,
         slurm_executor,
+        xcalibur_executor,
     )
     from utils.utils import configure_slurm_gpu_tuning, select_config_variant_interactive
 except (ImportError, ModuleNotFoundError):
@@ -45,9 +46,15 @@
         _kubeflow_numa_binding_script,
         kubeflow_executor,
         slurm_executor,
+        xcalibur_executor,
     )
     from .utils.utils import configure_slurm_gpu_tuning, select_config_variant_interactive
 
+try:
+    from nemo_run.core.execution.xcalibur import XCaliburExecutor as _XCaliburExecutor
+except ImportError:
+    _XCaliburExecutor = None  # type: ignore[assignment,misc]
+
 try:
     import wandb
 
@@ -86,10 +93,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 (
@@ -101,7 +112,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
@@ -110,6 +121,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
@@ -540,6 +553,19 @@ 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,
+    deterministic: bool = False,
     config_variant: str | None = None,
     gres: Optional[str] = None,
     packager: str = "git",
@@ -587,7 +613,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.
@@ -630,7 +657,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:
@@ -686,6 +713,32 @@ 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:
+        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,
+        )
+        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,
@@ -1008,7 +1061,7 @@ 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,
         dryrun=args.dryrun,
@@ -1087,6 +1140,19 @@ 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,
+        deterministic=args.deterministic,
         config_variant=config_variant,
         gres=args.gres,
         packager=args.packager,

scripts/performance/utils/executors.py

@@ -21,6 +21,7 @@
 import nemo_run as run
 from nemo_run.config import get_nemorun_home, set_nemorun_home
 from nemo_run.core.execution.launcher import SlurmTemplate
+from nemo_run.core.execution.xcalibur import XCaliburExecutor
 
 
 DEFAULT_NEMO_CACHE_HOME = Path.home() / ".cache" / "nemo"
@@ -361,3 +362,39 @@ 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,
+) -> 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,
+    )