feat(perf): add RunAIPlugin for NVIDIA Run:ai (RoCE/SR-IOV) Kubeflow fabric

koiker · koiker · commit cbb0dfd11d5f · 2026-06-23T21:41:42.000-04:00
Add a new CSP fabric plugin (`--csp runai`) for on-prem and colocation
clusters managed by NVIDIA Run:ai. The plugin injects RoCE/GDR rail
extended resources, Multus network-attachment annotations, a
memory-backed /dev/shm, and an optional workspace PVC onto the
KubeflowExecutor — following the same pattern as EKSEnvPlugin (EFA)
and GKEEnvPlugin (gIB).

Changes:
- csp_plugins.py: new RunAIPlugin dataclass with setup() method
- argument_parser.py: add "runai" to --csp choices; add --runai_*
  args (extended_resources_json, annotations_json, pvc_claim_name,
  pvc_mount_path, large_shm, env_json)
- setup_experiment.py: import RunAIPlugin, wire --runai_* args to
  main() signature and CSP plugin selection, filter --runai_* from
  rank-local script args

Validated on an 8-node B300 NVL72 cluster (56 GPUs, RoCE fabric)
running Nemotron 3 30B BF16 pretraining at ~698 TFLOP/s/GPU.

Signed-off-by: Rafael M. Koike &lt;koike.rafael@gmail.com&gt;
diff --git a/scripts/performance/argument_parser.py b/scripts/performance/argument_parser.py
@@ -540,13 +540,57 @@ def parse_cli_args():
     kubeflow_args.add_argument(
         "--csp",
         type=str,
-        choices=["aws", "gcp"],
+        choices=["aws", "gcp", "runai"],
         default=None,
-        help="Cloud provider of the Kubeflow cluster. Selects the CSP fabric plugin "
-        "(aws -> EKSEnvPlugin/EFA, gcp -> GKEEnvPlugin/gIB). Omit on Slurm or when no "
-        "CSP-specific fabric config is needed.",
+        help="Cloud / platform provider of the Kubeflow cluster. Selects the CSP fabric plugin "
+        "(aws -> EKSEnvPlugin/EFA, gcp -> GKEEnvPlugin/gIB, runai -> RunAIPlugin/RoCE). "
+        "Omit on Slurm or when no CSP-specific fabric config is needed.",
         required=False,
     )
+    kubeflow_args.add_argument(
+        "--runai_extended_resources_json",
+        type=str,
+        help="JSON-encoded dict of Run:ai extended resource requests per pod "
+        '(e.g. \'{"nvidia.com/r0-p0": "1", "nvidia.com/r1-p0": "1"}\').',
+        required=False,
+        default=None,
+    )
+    kubeflow_args.add_argument(
+        "--runai_annotations_json",
+        type=str,
+        help="JSON-encoded dict of pod annotations for Run:ai networking "
+        '(e.g. \'{"k8s.v1.cni.cncf.io/networks": "default/r0-p0,..."}\').',
+        required=False,
+        default=None,
+    )
+    kubeflow_args.add_argument(
+        "--runai_pvc_claim_name",
+        type=str,
+        help="PVC claim name for the shared workspace on Run:ai clusters.",
+        required=False,
+        default=None,
+    )
+    kubeflow_args.add_argument(
+        "--runai_pvc_mount_path",
+        type=str,
+        help="Container mount path for the Run:ai workspace PVC. Defaults to '/nemo-workspace'.",
+        required=False,
+        default="/nemo-workspace",
+    )
+    kubeflow_args.add_argument(
+        "--runai_large_shm",
+        type=bool_arg,
+        help="Mount a memory-backed /dev/shm for NCCL shared-memory collectives. Defaults to true.",
+        required=False,
+        default=True,
+    )
+    kubeflow_args.add_argument(
+        "--runai_env_json",
+        type=str,
+        help="JSON-encoded dict of additional environment variables for the Run:ai training container.",
+        required=False,
+        default=None,
+    )
     kubeflow_args.add_argument(
         "--kubeflow_workdir_pvc",
         type=str,
diff --git a/scripts/performance/setup_experiment.py b/scripts/performance/setup_experiment.py
@@ -54,9 +54,9 @@
     from .perf_plugins import NsysPlugin, PerfEnvPlugin, PyTorchProfilerPlugin
 
 try:
-    from utils.csp_plugins import EKSEnvPlugin, GKEEnvPlugin
+    from utils.csp_plugins import EKSEnvPlugin, GKEEnvPlugin, RunAIPlugin
 except (ImportError, ModuleNotFoundError):
-    from .utils.csp_plugins import EKSEnvPlugin, GKEEnvPlugin
+    from .utils.csp_plugins import EKSEnvPlugin, GKEEnvPlugin, RunAIPlugin
 
 
 SCRIPT_DIR = Path(__file__).parent.resolve()
@@ -88,7 +88,7 @@ def _filter_run_script_args(argv: List[str]) -> List[str]:
     """
 
     def _is_launcher_only(flag: str) -> bool:
-        return flag in ("--additional_slurm_params", "--csp") or flag.startswith("--kubeflow_")
+        return flag in ("--additional_slurm_params", "--csp") or flag.startswith(("--kubeflow_", "--runai_"))
 
     filtered_args = []
     skip_next = False
@@ -459,6 +459,12 @@ def main(
     kubeflow_container_kwargs_json: Optional[str],
     kubeflow_labels_json: Optional[str],
     kubeflow_pod_annotations_json: 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,
     deterministic: bool = False,
     config_variant: str = "v1",
     gres: Optional[str] = None,
@@ -643,12 +649,24 @@ def main(
     plugins = []
 
     # CSP fabric plugins (Kubeflow only; inert on Slurm via their isinstance guard):
-    # aws -> EKSEnvPlugin (EFA), gcp -> GKEEnvPlugin (gIB). Networking/fabric only;
+    # 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 {},
+            )
+        )
 
     if not use_recipes:
         plugins.append(
diff --git a/scripts/performance/utils/csp_plugins.py b/scripts/performance/utils/csp_plugins.py
@@ -27,7 +27,7 @@
 
 import json
 from dataclasses import dataclass, field
-from typing import List, Union
+from typing import Dict, List, Optional, Union
 
 import nemo_run as run
 from nemo_run import Plugin
@@ -100,3 +100,71 @@ def setup(self, task: Union["run.Partial", "run.Script"], executor: "run.Executo
             "networking.gke.io/interfaces": json.dumps(interfaces, separators=(",", ":")),
         }
         executor.env_vars.setdefault("NCCL_NET", "gIB")
+
+
+@dataclass(kw_only=True)
+class RunAIPlugin(Plugin):
+    """NVIDIA Run:ai (RoCE / SR-IOV) fabric.
+
+    Attaches RoCE/GDR rails as Kubernetes extended resources and Multus
+    network-attachment annotations, and optionally enlarges ``/dev/shm`` via an
+    ``emptyDir`` volume.  This is the on-prem / colocation equivalent of the
+    cloud CSP plugins: the networking topology is expressed through Multus
+    NetworkAttachmentDefinitions and SR-IOV device-plugin resources instead of
+    cloud-provider device plugins (EFA, gIB).
+
+    Typical B300 NVL72 topology exposes 8 RoCE rails (``r0-p0`` … ``r7-p0``)
+    each with one SR-IOV VF. A single Multus annotation attaches all rails.
+
+    Attributes:
+        extended_resources: Kubernetes extended-resource requests per pod,
+            e.g. ``{"nvidia.com/r0-p0": "1", "nvidia.com/r1-p0": "1", …}``.
+            Maps directly to ``extra_resource_requests`` / ``limits``.
+        annotations: Pod annotations dict, typically a single
+            ``k8s.v1.cni.cncf.io/networks`` key listing comma-separated Multus
+            NetworkAttachmentDefinitions.
+        large_shm: Mount a memory-backed ``/dev/shm`` (``emptyDir.medium:
+            Memory``) to avoid the default 64 MiB limit, which is too small for
+            NCCL shared-memory collectives on multi-GPU nodes.
+        pvc_claim_name: If set, attach the named PersistentVolumeClaim to the
+            pod and mount it at ``pvc_mount_path``.  Typically the shared
+            workspace PVC that holds model assets, HuggingFace cache, and
+            experiment logs.
+        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``).
+    """
+
+    extended_resources: Dict[str, str] = field(default_factory=dict)
+    annotations: Dict[str, str] = field(default_factory=dict)
+    large_shm: bool = True
+    pvc_claim_name: Optional[str] = None
+    pvc_mount_path: str = "/nemo-workspace"
+    env_vars: 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.extended_resources:
+            executor.extra_resource_requests = {**executor.extra_resource_requests, **self.extended_resources}
+            executor.extra_resource_limits = {**executor.extra_resource_limits, **self.extended_resources}
+
+        if self.annotations:
+            executor.pod_annotations = {**executor.pod_annotations, **self.annotations}
+
+        if self.large_shm and not any(v.get("name") == "dshm" for v in executor.volumes):
+            executor.volumes.append({"name": "dshm", "emptyDir": {"medium": "Memory"}})
+            executor.volume_mounts.append({"name": "dshm", "mountPath": "/dev/shm"})
+
+        if self.pvc_claim_name:
+            vol_name = "runai-workspace"
+            if not any(v.get("name") == vol_name for v in executor.volumes):
+                executor.volumes.append(
+                    {"name": vol_name, "persistentVolumeClaim": {"claimName": self.pvc_claim_name}}
+                )
+                executor.volume_mounts.append({"name": vol_name, "mountPath": self.pvc_mount_path})
+
+        if self.env_vars:
+            executor.env_vars.update(self.env_vars)
