feat(perf): support Kubeflow v1 PyTorchJob + Run:ai scheduling

Add --kubeflow_api_version {v1,v2} so the Kubeflow launch path can target the
Training Operator v1 PyTorchJob (kubeflow.org/v1) in addition to the v2
TrainJob. v1 is required on clusters running the v1 operator — notably
NVIDIA Run:ai. kubeflow_executor() picks nemo_run's PyTorchJobExecutor for v1
(clear error if the installed nemo_run predates it); since it inherits the
same dataclass fields, the existing field reconciliation works unchanged.

Extend RunAIPlugin with optional scheduler_name (-> spec.schedulerName, e.g.
runai-scheduler, so raw PyTorchJob/TrainJob submissions gang-schedule under
Run:ai instead of the default scheduler) and a generic labels passthrough for
project/queue membership (key varies by Run:ai version). Both route through
pod_spec_overrides / pod_labels, which apply to v1 and v2 alike. Wire through
build_csp_plugin and add --runai_scheduler_name / --runai_labels_json.

Extend the CSP wiring test to assert the new values reach the plugin.

Note: requires bumping the nemo_run pin to a commit that provides
PyTorchJobExecutor and the pod_annotations/pod_labels/extra_resource_* fields.

Signed-off-by: Rafael M. Koike <koike.rafael@gmail.com>

Author: koiker

scripts/performance/argument_parser.py

@@ -537,6 +537,17 @@ def parse_cli_args():
         help="Kubernetes namespace for Kubeflow TrainJob. When set, uses the Kubeflow executor instead of Slurm.",
         required=False,
     )
+    kubeflow_args.add_argument(
+        "--kubeflow_api_version",
+        type=str,
+        choices=["v1", "v2"],
+        default="v2",
+        help="Kubeflow Training-Operator API to submit against. 'v2' (default) uses the "
+        "TrainJob (trainer.kubeflow.org) via KubeflowExecutor. 'v1' uses the PyTorchJob "
+        "(kubeflow.org/v1) via PyTorchJobExecutor — required for clusters running the v1 "
+        "operator, notably NVIDIA Run:ai. Only applies when --kubeflow_namespace is set.",
+        required=False,
+    )
     kubeflow_args.add_argument(
         "--csp",
         type=str,
@@ -591,6 +602,24 @@ def parse_cli_args():
         required=False,
         default=None,
     )
+    kubeflow_args.add_argument(
+        "--runai_scheduler_name",
+        type=str,
+        help="Kubernetes scheduler to pin Run:ai workload pods to (spec.schedulerName), "
+        "e.g. 'runai-scheduler'. Required for raw PyTorchJob/TrainJob submissions (not via "
+        "the runai CLI) so Run:ai gang-scheduling and quota apply. Omit to use the default scheduler.",
+        required=False,
+        default=None,
+    )
+    kubeflow_args.add_argument(
+        "--runai_labels_json",
+        type=str,
+        help="JSON-encoded dict of extra pod labels for Run:ai project/queue membership "
+        '(e.g. \'{"project": "bench"}\' or \'{"kai.scheduler/queue": "bench"}\'); the exact '
+        "key depends on your Run:ai version.",
+        required=False,
+        default=None,
+    )
     kubeflow_args.add_argument(
         "--kubeflow_workdir_pvc",
         type=str,

scripts/performance/setup_experiment.py

@@ -382,6 +382,43 @@ def maybe_increase_n_attempts_on_flaky_failure(
     return n_attempts
 
 
+def build_csp_plugin(
+    csp: Optional[str],
+    *,
+    runai_extended_resources_json: Optional[str] = None,
+    runai_annotations_json: Optional[str] = None,
+    runai_pvc_claim_name: Optional[str] = None,
+    runai_pvc_mount_path: str = "/nemo-workspace",
+    runai_large_shm: bool = True,
+    runai_env_json: Optional[str] = None,
+    runai_scheduler_name: Optional[str] = None,
+    runai_labels_json: Optional[str] = None,
+):
+    """Map ``--csp`` (+ its ``--runai_*`` args) to a single CSP fabric plugin.
+
+    Returns the plugin instance for the selected CSP, or ``None`` when no CSP is
+    selected. Kept as a small pure function (no executor/cluster state) so the
+    args -> plugin wiring is unit-testable without a live cluster — see
+    ``tests/.../test_setup_experiment_csp.py``.
+    """
+    if csp == "aws":
+        return EKSEnvPlugin()
+    if csp == "gcp":
+        return GKEEnvPlugin()
+    if csp == "runai":
+        return RunAIPlugin(
+            extended_resources=json.loads(runai_extended_resources_json) if runai_extended_resources_json else {},
+            annotations=json.loads(runai_annotations_json) if runai_annotations_json else {},
+            large_shm=runai_large_shm,
+            pvc_claim_name=runai_pvc_claim_name,
+            pvc_mount_path=runai_pvc_mount_path,
+            env_vars=json.loads(runai_env_json) if runai_env_json else {},
+            scheduler_name=runai_scheduler_name,
+            labels=json.loads(runai_labels_json) if runai_labels_json else {},
+        )
+    return None
+
+
 def main(
     use_recipes: bool,
     model_family_name: str,
@@ -459,12 +496,15 @@ def main(
     kubeflow_container_kwargs_json: Optional[str],
     kubeflow_labels_json: Optional[str],
     kubeflow_pod_annotations_json: Optional[str],
+    kubeflow_api_version: str = "v2",
     runai_extended_resources_json: Optional[str] = None,
     runai_annotations_json: Optional[str] = None,
     runai_pvc_claim_name: Optional[str] = None,
     runai_pvc_mount_path: str = "/nemo-workspace",
     runai_large_shm: bool = True,
     runai_env_json: Optional[str] = None,
+    runai_scheduler_name: Optional[str] = None,
+    runai_labels_json: Optional[str] = None,
     deterministic: bool = False,
     config_variant: str = "v1",
     gres: Optional[str] = None,
@@ -621,6 +661,7 @@ def main(
             container_kwargs=json.loads(kubeflow_container_kwargs_json) if kubeflow_container_kwargs_json else None,
             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),
+            api_version=kubeflow_api_version,
         )
     else:
         executor = slurm_executor(
@@ -652,21 +693,19 @@ def main(
     # aws -> EKSEnvPlugin (EFA), gcp -> GKEEnvPlugin (gIB),
     # runai -> RunAIPlugin (RoCE/SR-IOV). Networking/fabric only;
     # arch/recipe/perf env stays in PerfEnvPlugin / the recipe.
-    if csp == "aws":
-        plugins.append(EKSEnvPlugin())
-    elif csp == "gcp":
-        plugins.append(GKEEnvPlugin())
-    elif csp == "runai":
-        plugins.append(
-            RunAIPlugin(
-                extended_resources=json.loads(runai_extended_resources_json) if runai_extended_resources_json else {},
-                annotations=json.loads(runai_annotations_json) if runai_annotations_json else {},
-                large_shm=runai_large_shm,
-                pvc_claim_name=runai_pvc_claim_name,
-                pvc_mount_path=runai_pvc_mount_path,
-                env_vars=json.loads(runai_env_json) if runai_env_json else {},
-            )
-        )
+    csp_plugin = build_csp_plugin(
+        csp,
+        runai_extended_resources_json=runai_extended_resources_json,
+        runai_annotations_json=runai_annotations_json,
+        runai_pvc_claim_name=runai_pvc_claim_name,
+        runai_pvc_mount_path=runai_pvc_mount_path,
+        runai_large_shm=runai_large_shm,
+        runai_env_json=runai_env_json,
+        runai_scheduler_name=runai_scheduler_name,
+        runai_labels_json=runai_labels_json,
+    )
+    if csp_plugin is not None:
+        plugins.append(csp_plugin)
 
     if not use_recipes:
         plugins.append(
@@ -1048,6 +1087,15 @@ 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,
+        kubeflow_api_version=args.kubeflow_api_version,
+        runai_extended_resources_json=args.runai_extended_resources_json,
+        runai_annotations_json=args.runai_annotations_json,
+        runai_pvc_claim_name=args.runai_pvc_claim_name,
+        runai_pvc_mount_path=args.runai_pvc_mount_path,
+        runai_large_shm=args.runai_large_shm,
+        runai_env_json=args.runai_env_json,
+        runai_scheduler_name=args.runai_scheduler_name,
+        runai_labels_json=args.runai_labels_json,
         deterministic=args.deterministic,
         config_variant=config_variant,
         gres=args.gres,

scripts/performance/utils/csp_plugins.py

@@ -133,6 +133,17 @@ class RunAIPlugin(Plugin):
         pvc_mount_path: Container mount path for the PVC.
         env_vars: Additional environment variables injected into the training
             container (e.g. ``TRANSFORMERS_OFFLINE``, ``HF_HOME``).
+        scheduler_name: If set, pin the workload pods to this Kubernetes
+            scheduler (``spec.schedulerName``). Run:ai gang-schedules through its
+            own scheduler — typically ``"runai-scheduler"`` — so raw
+            PyTorchJob/TrainJob submissions (i.e. not via the ``runai`` CLI) must
+            name it explicitly or the default scheduler will place the pods and
+            bypass Run:ai quota/fair-share. Left unset (default) so non-Run:ai
+            paths are unaffected.
+        labels: Extra pod labels merged onto the workload pods. Run:ai expresses
+            project/queue membership through a pod label whose key varies by
+            version (e.g. ``project`` or ``kai.scheduler/queue``); pass the
+            key/value your cluster expects here rather than hardcoding one.
     """
 
     extended_resources: Dict[str, str] = field(default_factory=dict)
@@ -141,12 +152,28 @@ class RunAIPlugin(Plugin):
     pvc_claim_name: Optional[str] = None
     pvc_mount_path: str = "/nemo-workspace"
     env_vars: Dict[str, str] = field(default_factory=dict)
+    scheduler_name: Optional[str] = None
+    labels: Dict[str, str] = field(default_factory=dict)
 
     def setup(self, task: Union["run.Partial", "run.Script"], executor: "run.Executor") -> None:
         """Layer the Run:ai RoCE/SR-IOV fabric onto a Kubeflow executor."""
         if not isinstance(executor, KubeflowExecutor):
             return
 
+        if self.scheduler_name:
+            # pod_spec_overrides merges into the (v2 TrainJob) podTemplateOverrides
+            # spec and the (v1 PyTorchJob) replica pod spec alike.
+            executor.pod_spec_overrides = {
+                **executor.pod_spec_overrides,
+                "schedulerName": self.scheduler_name,
+            }
+
+        if self.labels:
+            if hasattr(executor, "pod_labels"):
+                executor.pod_labels = {**executor.pod_labels, **self.labels}
+            else:
+                executor.labels = {**executor.labels, **self.labels}
+
         if self.extended_resources:
             executor.extra_resource_requests = {**executor.extra_resource_requests, **self.extended_resources}
             executor.extra_resource_limits = {**executor.extra_resource_limits, **self.extended_resources}

scripts/performance/utils/executors.py

@@ -198,6 +198,7 @@ def kubeflow_executor(
     container_kwargs: Optional[Dict[str, Any]] = None,
     labels: Optional[Dict[str, Any]] = None,
     pod_annotations: Optional[Dict[str, Any]] = None,
+    api_version: str = "v2",
 ) -> run.KubeflowExecutor:
     """Build a Kubeflow Training Operator executor.
 
@@ -239,9 +240,16 @@ def kubeflow_executor(
         labels: Pod labels.
         pod_annotations: Annotations applied to the trainer pod template metadata
             (e.g. ``networking.gke.io/interfaces`` for GKE RDMA NIC attachment).
+        api_version: Which Kubeflow Training-Operator API to submit against.
+            ``"v2"`` (default) builds a ``run.KubeflowExecutor`` (TrainJob,
+            ``trainer.kubeflow.org``). ``"v1"`` builds a
+            ``run.PyTorchJobExecutor`` (PyTorchJob, ``kubeflow.org/v1``) for
+            clusters running the v1 operator — notably **NVIDIA Run:ai**. Both
+            share the same configuration surface; only the submitted custom
+            resource and the distributed-launch env contract differ.
 
     Returns:
-        Configured ``run.KubeflowExecutor`` instance.
+        Configured ``run.KubeflowExecutor`` (or ``run.PyTorchJobExecutor`` for v1).
     """
     # K8s/Kubeflow jobs deliberately do NOT inherit PERF_ENV_VARS. That dict was
     # tuned for the Slurm perf-benchmark path; the verified standalone K8s launch
@@ -272,7 +280,7 @@ def kubeflow_executor(
     }
     labels = {**ci_labels, **(labels or {})}
 
-    executor = run.KubeflowExecutor(
+    kf_kwargs = dict(
         # Launch each replica's entrypoint under torchrun so the torch-distributed
         # ClusterTrainingRuntime's rendezvous env (MASTER_ADDR, nnodes, nproc) is
         # consumed and a single WORLD_SIZE = num_nodes * gpus_per_node process
@@ -308,7 +316,8 @@ def kubeflow_executor(
         # `kubectl get trainjob -l` and `kubectl get pods -l` resolve the origin.
         pod_labels=labels,
         # pod_annotations land on the trainer pod template metadata (e.g. GKE
-        # networking.gke.io/interfaces to attach the RDMA NICs for gIB).
+        # networking.gke.io/interfaces to attach the RDMA NICs for gIB, or the
+        # Run:ai Multus k8s.v1.cni.cncf.io/networks RoCE rail attachments).
         pod_annotations=pod_annotations or {},
         # include_submodules=True: KubeflowExecutor.package() ships the packager
         # tarball to <workdir_pvc_path>/<user>/code, which the launcher overlays
@@ -321,4 +330,49 @@ def kubeflow_executor(
         # are harmless cost there.)
         packager=run.GitArchivePackager(include_submodules=True),
     )
+
+    # Select the Training-Operator API. v1 (PyTorchJob) lives in a separate
+    # nemo_run executor subclass; fall back with a clear error if the installed
+    # nemo_run predates it rather than silently launching a v2 TrainJob onto a
+    # v1-only cluster (e.g. Run:ai), which would never schedule.
+    if api_version in ("v1", "pytorchjob", "PyTorchJob"):
+        executor_cls = getattr(run, "PyTorchJobExecutor", None)
+        if executor_cls is None:
+            raise RuntimeError(
+                "kubeflow_api_version='v1' requires a nemo_run build that provides "
+                "PyTorchJobExecutor (Kubeflow Training Operator v1 / PyTorchJob). "
+                "Upgrade nemo_run (NVIDIA-NeMo/Run) to a commit that includes it, or "
+                "use --kubeflow_api_version v2 for the TrainJob (v2) operator."
+            )
+    else:
+        executor_cls = run.KubeflowExecutor
+
+    # Reconcile against the installed executor API. Older releases (e.g.
+    # 0.9.0rc0) predate the granular ``pod_labels`` / ``pod_annotations`` fields
+    # (they expose a single ``annotations``) and ``train_job_basename``. Rather
+    # than pin a newer nemo_run, degrade gracefully so the same recipe launches
+    # across versions — folding the granular values into the legacy fields so
+    # RoCE/Multus pod annotations and CI-origin labels are not silently dropped.
+    # PyTorchJobExecutor inherits the same dataclass fields as KubeflowExecutor,
+    # so this reconciliation works identically for either class.
+    _kf_fields = set(
+        getattr(executor_cls, "model_fields", None)
+        or getattr(executor_cls, "__dataclass_fields__", {})
+        or {}
+    )
+    if _kf_fields:
+        if "pod_annotations" not in _kf_fields and kf_kwargs.get("pod_annotations"):
+            if "annotations" in _kf_fields:
+                kf_kwargs["annotations"] = {
+                    **(kf_kwargs.get("annotations") or {}),
+                    **kf_kwargs["pod_annotations"],
+                }
+        if "pod_labels" not in _kf_fields and kf_kwargs.get("pod_labels"):
+            kf_kwargs["labels"] = {
+                **(kf_kwargs.get("labels") or {}),
+                **kf_kwargs["pod_labels"],
+            }
+        kf_kwargs = {k: v for k, v in kf_kwargs.items() if k in _kf_fields}
+
+    executor = executor_cls(**kf_kwargs)
     return executor