Support Nebius InfiniBand clusters (#2604)

An InfiniBand cluster is now created automatically
when provisioning 8xH100 or 8xH200 instances using
a fleet configuration with `placement: cluster`.

Nebius clusters were supported using `dstack`
placement groups. Several changes were made to the
placement group management logic:
- The offer for the master instance of the fleet
  is passed to `Compute.create_placement_group`,
  which allows to set different placement group
  settings based on the offer. Nebius requires
  different settings for H100 and H200 clusters.
- `Compute.is_suitable_placement_group` is
  introduced to allow choosing an appropriate
  placement group when creating the master
  instance and filtering offers for non-master
  instances based on backend-specific placement
  group properties. Nebius currently only provides
  homogeneous clusters, so offers need to be
  filtered based on the placement group.
- The placement group object is passed to
  `Compute.create_instance` to allow adding the
  instance to the placement group using its
  backend-specific properties, such as cluster ID
  on Nebius.
- The placement group name is generated at master
  instance provisioning time, not at fleet
  creation time. This allows to have different
  placement group names for the same fleet and
  avoid name conflicts, since multiple placement
  groups can be created while `dstack` is trying
  different offers for the master instance.
- Placement groups that were created during master
  instance provisioning but didn't end up being
  used are now cleaned up. Nebius quotas limit the
  number of clusters, so unused clusters need to
  be cleaned up quickly, without waiting for fleet
  deletion.
- If all offers failed for the master instance,
  `dstack` will no longer attempt to provision
  other fleet instances to avoid them being
  provisioned without a placement group or without
  connectivity at all.
- Placement group creation errors are now handled
  gracefully, so that `dstack` can move on to
  other master instance offers, which may lead to
  creating different placement groups. For
  example, if `dstack` cannot create a cluster in
  one Nebius region because of a missing quota, it
  may attempt to create a cluster in another
  region.

Author: jvstme

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

@@ -159,6 +159,7 @@ def create_instance(
         self,
         instance_offer: InstanceOfferWithAvailability,
         instance_config: InstanceConfiguration,
+        placement_group: Optional[PlacementGroup],
     ) -> JobProvisioningData:
         project_name = instance_config.project_name
         ec2_resource = self.session.resource("ec2", region_name=instance_offer.region)
@@ -248,7 +249,7 @@ def create_instance(
                         spot=instance_offer.instance.resources.spot,
                         subnet_id=subnet_id,
                         allocate_public_ip=allocate_public_ip,
-                        placement_group_name=instance_config.placement_group_name,
+                        placement_group_name=placement_group.name if placement_group else None,
                         enable_efa=enable_efa,
                         max_efa_interfaces=max_efa_interfaces,
                         reservation_id=instance_config.reservation,
@@ -291,6 +292,7 @@ def create_instance(
     def create_placement_group(
         self,
         placement_group: PlacementGroup,
+        master_instance_offer: InstanceOffer,
     ) -> PlacementGroupProvisioningData:
         ec2_client = self.session.client("ec2", region_name=placement_group.configuration.region)
         logger.debug("Creating placement group %s...", placement_group.name)
@@ -323,6 +325,16 @@ def delete_placement_group(
                 raise e
         logger.debug("Deleted placement group %s", placement_group.name)
 
+    def is_suitable_placement_group(
+        self,
+        placement_group: PlacementGroup,
+        instance_offer: InstanceOffer,
+    ) -> bool:
+        return (
+            placement_group.configuration.backend == BackendType.AWS
+            and placement_group.configuration.region == instance_offer.region
+        )
+
     def create_gateway(
         self,
         configuration: GatewayComputeConfiguration,

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

@@ -62,6 +62,7 @@
     InstanceOfferWithAvailability,
     InstanceType,
 )
+from dstack._internal.core.models.placement import PlacementGroup
 from dstack._internal.core.models.resources import Memory, Range
 from dstack._internal.core.models.runs import JobProvisioningData, Requirements
 from dstack._internal.utils.logging import get_logger
@@ -109,6 +110,7 @@ def create_instance(
         self,
         instance_offer: InstanceOfferWithAvailability,
         instance_config: InstanceConfiguration,
+        placement_group: Optional[PlacementGroup],
     ) -> JobProvisioningData:
         instance_name = generate_unique_instance_name(
             instance_config, max_length=azure_resources.MAX_RESOURCE_NAME_LEN

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

@@ -25,6 +25,7 @@
 )
 from dstack._internal.core.models.instances import (
     InstanceConfiguration,
+    InstanceOffer,
     InstanceOfferWithAvailability,
     SSHKey,
 )
@@ -144,6 +145,7 @@ def create_instance(
         self,
         instance_offer: InstanceOfferWithAvailability,
         instance_config: InstanceConfiguration,
+        placement_group: Optional[PlacementGroup],
     ) -> JobProvisioningData:
         """
         Launches a new instance. It should return `JobProvisioningData` ASAP.
@@ -176,7 +178,7 @@ def run_job(
         )
         instance_offer = instance_offer.copy()
         self._restrict_instance_offer_az_to_volumes_az(instance_offer, volumes)
-        return self.create_instance(instance_offer, instance_config)
+        return self.create_instance(instance_offer, instance_config, placement_group=None)
 
     def _restrict_instance_offer_az_to_volumes_az(
         self,
@@ -225,9 +227,15 @@ class ComputeWithPlacementGroupSupport(ABC):
     def create_placement_group(
         self,
         placement_group: PlacementGroup,
+        master_instance_offer: InstanceOffer,
     ) -> PlacementGroupProvisioningData:
         """
         Creates a placement group.
+
+        Args:
+            placement_group: details about the placement group to be created
+            master_instance_offer: the first instance dstack will attempt to add
+                                   to the placement group
         """
         pass
 
@@ -242,10 +250,27 @@ def delete_placement_group(
         """
         pass
 
+    @abstractmethod
+    def is_suitable_placement_group(
+        self,
+        placement_group: PlacementGroup,
+        instance_offer: InstanceOffer,
+    ) -> bool:
+        """
+        Checks if the instance offer can be provisioned in the placement group.
+
+        Should return immediately, without performing API calls.
+
+        Can be called with an offer originating from a different backend, because some backends
+        (BackendType.DSTACK) produce offers on behalf of other backends. Should return `False`
+        in that case.
+        """
+        pass
+
 
 class ComputeWithGatewaySupport(ABC):
     """
-    Must be subclassed and imlemented to support gateways.
+    Must be subclassed and implemented to support gateways.
     """
 
     @abstractmethod
@@ -418,6 +443,21 @@ def generate_unique_volume_name(
     )
 
 
+def generate_unique_placement_group_name(
+    project_name: str,
+    fleet_name: str,
+    max_length: int = _DEFAULT_MAX_RESOURCE_NAME_LEN,
+) -> str:
+    """
+    Generates a unique placement group name valid across all backends.
+    """
+    return generate_unique_backend_name(
+        resource_name=fleet_name,
+        project_name=project_name,
+        max_length=max_length,
+    )
+
+
 def generate_unique_backend_name(
     resource_name: str,
     project_name: Optional[str],

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

@@ -18,6 +18,7 @@
     InstanceConfiguration,
     InstanceOfferWithAvailability,
 )
+from dstack._internal.core.models.placement import PlacementGroup
 from dstack._internal.core.models.runs import JobProvisioningData, Requirements
 from dstack._internal.utils.logging import get_logger
 
@@ -58,6 +59,7 @@ def create_instance(
         self,
         instance_offer: InstanceOfferWithAvailability,
         instance_config: InstanceConfiguration,
+        placement_group: Optional[PlacementGroup],
     ) -> JobProvisioningData:
         vm_id = generate_unique_instance_name(instance_config, max_length=MAX_RESOURCE_NAME_LEN)
         public_keys = instance_config.get_public_keys()

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

@@ -20,6 +20,7 @@
     InstanceOffer,
     InstanceOfferWithAvailability,
 )
+from dstack._internal.core.models.placement import PlacementGroup
 from dstack._internal.core.models.resources import Memory, Range
 from dstack._internal.core.models.runs import JobProvisioningData, Requirements
 from dstack._internal.utils.logging import get_logger
@@ -85,6 +86,7 @@ def create_instance(
         self,
         instance_offer: InstanceOfferWithAvailability,
         instance_config: InstanceConfiguration,
+        placement_group: Optional[PlacementGroup],
     ) -> JobProvisioningData:
         instance_name = generate_unique_instance_name(
             instance_config, max_length=MAX_INSTANCE_NAME_LEN

src/dstack/_internal/core/backends/gcp/auth.py

@@ -19,7 +19,7 @@ def authenticate(creds: AnyGCPCreds, project_id: Optional[str] = None) -> Tuple[
     credentials, credentials_project_id = get_credentials(creds)
     if project_id is None:
         # If project_id is not specified explicitly, try using credentials' project_id.
-        # Explicit project_id takes precedence bacause credentials' project_id may be irrelevant.
+        # Explicit project_id takes precedence because credentials' project_id may be irrelevant.
         # For example, with Workload Identity Federation for GKE, it's cluster project_id.
         project_id = credentials_project_id
     if project_id is None:

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

@@ -166,6 +166,7 @@ def create_instance(
         self,
         instance_offer: InstanceOfferWithAvailability,
         instance_config: InstanceConfiguration,
+        placement_group: Optional[PlacementGroup],
     ) -> JobProvisioningData:
         instance_name = generate_unique_instance_name(
             instance_config, max_length=gcp_resources.MAX_RESOURCE_NAME_LEN
@@ -199,11 +200,11 @@ def create_instance(
             instance_type_name=instance_offer.instance.name,
         )
         placement_policy = None
-        if instance_config.placement_group_name is not None:
+        if placement_group is not None:
             placement_policy = gcp_resources.get_placement_policy_resource_name(
                 project_id=self.config.project_id,
                 region=instance_offer.region,
-                placement_policy=instance_config.placement_group_name,
+                placement_policy=placement_group.name,
             )
         labels = {
             "owner": "dstack",
@@ -406,6 +407,7 @@ def update_provisioning_data(
     def create_placement_group(
         self,
         placement_group: PlacementGroup,
+        master_instance_offer: InstanceOffer,
     ) -> PlacementGroupProvisioningData:
         policy = compute_v1.ResourcePolicy(
             name=placement_group.name,
@@ -440,6 +442,16 @@ def delete_placement_group(
                 raise PlacementGroupInUseError()
             raise
 
+    def is_suitable_placement_group(
+        self,
+        placement_group: PlacementGroup,
+        instance_offer: InstanceOffer,
+    ) -> bool:
+        return (
+            placement_group.configuration.backend == BackendType.GCP
+            and placement_group.configuration.region == instance_offer.region
+        )
+
     def create_gateway(
         self,
         configuration: GatewayComputeConfiguration,

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

@@ -20,6 +20,7 @@
     InstanceOffer,
     InstanceOfferWithAvailability,
 )
+from dstack._internal.core.models.placement import PlacementGroup
 from dstack._internal.core.models.runs import JobProvisioningData, Requirements
 
 MAX_INSTANCE_NAME_LEN = 60
@@ -46,7 +47,10 @@ def get_offers(
         return offers_with_availability
 
     def create_instance(
-        self, instance_offer: InstanceOfferWithAvailability, instance_config: InstanceConfiguration
+        self,
+        instance_offer: InstanceOfferWithAvailability,
+        instance_config: InstanceConfiguration,
+        placement_group: Optional[PlacementGroup],
     ) -> JobProvisioningData:
         instance_name = generate_unique_instance_name(
             instance_config, max_length=MAX_INSTANCE_NAME_LEN

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

@@ -15,6 +15,7 @@
     InstanceType,
     Resources,
 )
+from dstack._internal.core.models.placement import PlacementGroup
 from dstack._internal.core.models.runs import Job, JobProvisioningData, Requirements, Run
 from dstack._internal.core.models.volumes import Volume, VolumeProvisioningData
 from dstack._internal.utils.logging import get_logger
@@ -53,6 +54,7 @@ def create_instance(
         self,
         instance_offer: InstanceOfferWithAvailability,
         instance_config: InstanceConfiguration,
+        placement_group: Optional[PlacementGroup],
     ) -> JobProvisioningData:
         return JobProvisioningData(
             backend=instance_offer.backend,