Allow to configure Nebius InfiniBand fabrics

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):