diff --git a/scripts/performance/argument_parser.py b/scripts/performance/argument_parser.py index 2edd4a683e..d2a2fc19b2 100644 --- a/scripts/performance/argument_parser.py +++ b/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( diff --git a/scripts/performance/perf_plugins.py b/scripts/performance/perf_plugins.py index 2c0f12451b..68e75322fc 100644 --- a/scripts/performance/perf_plugins.py +++ b/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"): diff --git a/scripts/performance/setup_experiment.py b/scripts/performance/setup_experiment.py index aaa89b46dd..668323e37b 100755 --- a/scripts/performance/setup_experiment.py +++ b/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, diff --git a/scripts/performance/utils/executors.py b/scripts/performance/utils/executors.py index 072bf3b56d..e748e3479d 100644 --- a/scripts/performance/utils/executors.py +++ b/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, + )