Allow to configure Nebius InfiniBand fabrics (#2607)

Add an option in Nebius backend settings to limit
the list of allowed fabrics for InfiniBand
clusters. This can be useful for larger customers
that have capacity reservations tied to a specific
fabric.

Author: jvstme

src/dstack/_internal/core/backends/nebius/compute.py

@@ -2,7 +2,6 @@
 import random
 import shlex
 import time
-from dataclasses import dataclass
 from functools import cached_property
 from typing import List, Optional
 
@@ -21,6 +20,7 @@
 )
 from dstack._internal.core.backends.base.offers import get_catalog_offers
 from dstack._internal.core.backends.nebius import resources
+from dstack._internal.core.backends.nebius.fabrics import get_suitable_infiniband_fabrics
 from dstack._internal.core.backends.nebius.models import NebiusConfig, NebiusServiceAccountCreds
 from dstack._internal.core.errors import (
     BackendError,
@@ -81,24 +81,6 @@
 ]
 
 
-@dataclass(frozen=True)
-class InfinibandFabric:
-    name: str
-    platform: str
-    region: str
-
-
-# https://docs.nebius.com/compute/clusters/gpu#fabrics
-INFINIBAND_FABRICS = [
-    InfinibandFabric("fabric-2", "gpu-h100-sxm", "eu-north1"),
-    InfinibandFabric("fabric-3", "gpu-h100-sxm", "eu-north1"),
-    InfinibandFabric("fabric-4", "gpu-h100-sxm", "eu-north1"),
-    InfinibandFabric("fabric-5", "gpu-h200-sxm", "eu-west1"),
-    InfinibandFabric("fabric-6", "gpu-h100-sxm", "eu-north1"),
-    InfinibandFabric("fabric-7", "gpu-h200-sxm", "eu-north1"),
-]
-
-
 class NebiusCompute(
     ComputeWithCreateInstanceSupport,
     ComputeWithMultinodeSupport,
@@ -280,7 +262,9 @@ def create_placement_group(
         backend_data = NebiusPlacementGroupBackendData(cluster=None)
         # Only create a Nebius cluster if the instance supports it.
         # For other instances, return dummy PlacementGroupProvisioningData.
-        if fabrics := _get_suitable_infiniband_fabrics(master_instance_offer):
+        if fabrics := get_suitable_infiniband_fabrics(
+            master_instance_offer, allowed_fabrics=self.config.fabrics
+        ):
             fabric = random.choice(fabrics)
             op = resources.create_cluster(
                 self._sdk,
@@ -319,7 +303,11 @@ def is_suitable_placement_group(
         )
         return (
             backend_data.cluster is None
-            or backend_data.cluster.fabric in _get_suitable_infiniband_fabrics(instance_offer)
+            or backend_data.cluster.fabric
+            in get_suitable_infiniband_fabrics(
+                instance_offer,
+                allowed_fabrics=None,  # enforced at cluster creation time, no need to enforce here
+            )
         )
 
 
@@ -380,15 +368,3 @@ def _wait_for_instance(sdk: SDK, op: SDKOperation[Operation]) -> None:
 def _supported_instances(offer: InstanceOffer) -> bool:
     platform, _ = offer.instance.name.split()
     return platform in SUPPORTED_PLATFORMS and not offer.instance.resources.spot
-
-
-def _get_suitable_infiniband_fabrics(offer: InstanceOffer) -> list[str]:
-    if len(offer.instance.resources.gpus) < 8:
-        # From the create VM page in the Nebius Console:
-        # > Only virtual machines with at least 8 NVIDIA® Hopper® H100 or H200 GPUs
-        # > can be added to the cluster
-        return []
-    platform, _ = offer.instance.name.split()
-    return [
-        f.name for f in INFINIBAND_FABRICS if f.platform == platform and f.region == offer.region
-    ]

src/dstack/_internal/core/backends/nebius/configurator.py

@@ -9,6 +9,7 @@
 )
 from dstack._internal.core.backends.nebius import resources
 from dstack._internal.core.backends.nebius.backend import NebiusBackend
+from dstack._internal.core.backends.nebius.fabrics import get_all_infiniband_fabrics
 from dstack._internal.core.backends.nebius.models import (
     AnyNebiusBackendConfig,
     NebiusBackendConfig,
@@ -38,6 +39,16 @@ def validate_config(self, config: NebiusBackendConfigWithCreds, default_creds_en
                 fields=[["creds"]],
                 details=str(e),
             )
+        valid_fabrics = get_all_infiniband_fabrics()
+        if invalid_fabrics := set(config.fabrics or []) - valid_fabrics:
+            raise_invalid_credentials_error(
+                fields=[["fabrics"]],
+                details=(
+                    "These InfiniBand fabrics do not exist or are not known to dstack:"
+                    f" {sorted(invalid_fabrics)}. Omit `fabrics` to allow all fabrics or select"
+                    f" some of the valid options: {sorted(valid_fabrics)}"
+                ),
+            )
 
     def create_backend(
         self, project_name: str, config: NebiusBackendConfigWithCreds

src/dstack/_internal/core/backends/nebius/fabrics.py

@@ -0,0 +1,47 @@
+from collections.abc import Container
+from dataclasses import dataclass
+from typing import Optional
+
+from dstack._internal.core.models.instances import InstanceOffer
+
+
+@dataclass(frozen=True)
+class InfinibandFabric:
+    name: str
+    platform: str
+    region: str
+
+
+# https://docs.nebius.com/compute/clusters/gpu#fabrics
+INFINIBAND_FABRICS = [
+    InfinibandFabric("fabric-2", "gpu-h100-sxm", "eu-north1"),
+    InfinibandFabric("fabric-3", "gpu-h100-sxm", "eu-north1"),
+    InfinibandFabric("fabric-4", "gpu-h100-sxm", "eu-north1"),
+    InfinibandFabric("fabric-5", "gpu-h200-sxm", "eu-west1"),
+    InfinibandFabric("fabric-6", "gpu-h100-sxm", "eu-north1"),
+    InfinibandFabric("fabric-7", "gpu-h200-sxm", "eu-north1"),
+]
+
+
+def get_suitable_infiniband_fabrics(
+    offer: InstanceOffer, allowed_fabrics: Optional[Container[str]]
+) -> list[str]:
+    if len(offer.instance.resources.gpus) < 8:
+        # From the create VM page in the Nebius Console:
+        # > Only virtual machines with at least 8 NVIDIA® Hopper® H100 or H200 GPUs
+        # > can be added to the cluster
+        return []
+    platform, _ = offer.instance.name.split()
+    return [
+        f.name
+        for f in INFINIBAND_FABRICS
+        if (
+            f.platform == platform
+            and f.region == offer.region
+            and (allowed_fabrics is None or f.name in allowed_fabrics)
+        )
+    ]
+
+
+def get_all_infiniband_fabrics() -> set[str]:
+    return {f.name for f in INFINIBAND_FABRICS}

src/dstack/_internal/core/backends/nebius/models.py

@@ -87,6 +87,14 @@ class NebiusBackendConfig(CoreModel):
         Optional[list[str]],
         Field(description="The list of allowed Nebius regions. Omit to allow all regions"),
     ] = None
+    fabrics: Annotated[
+        Optional[list[str]],
+        Field(
+            description=(
+                "The list of allowed fabrics for InfiniBand clusters. Omit to allow all fabrics"
+            )
+        ),
+    ] = None
 
 
 class NebiusBackendConfigWithCreds(NebiusBackendConfig):