Merge pull request #2392 from dstackai/issue_2372_backend_features

Introduce ComputeWith classes to detect compute features

Author: r4victor

contributing/BACKENDS.md

@@ -115,13 +115,8 @@ Refer to examples:
 ##### 2.4.4. Create the backend compute class
 
 Under the backend directory you've created, create the `compute.py` file and define the
-backend compute class there (should extend `dstack._internal.core.backends.base.compute.Compute`).
-
-You'll have to implement `get_offers`, `run_job` and `terminate_instance`.
-You may need to implement `update_provisioning_data`, see its docstring for details.
-
-For VM-based backends, also implement the `create_instance` method and add the backend name to
-[`BACKENDS_WITH_CREATE_INSTANCE_SUPPORT`](`https://github.com/dstackai/dstack/blob/master/src/dstack/_internal/core/backends/__init__.py`).
+backend compute class that extends the `dstack._internal.core.backends.base.compute.Compute` class.
+It can also extend and implement `ComputeWith*` classes to support additional features such as fleets, volumes, gateways, placement groups, etc. For example, it should extend `ComputeWithCreateInstanceSupport` to support fleets.
 
 Refer to examples:
 [datacrunch](https://github.com/dstackai/dstack/blob/master/src/dstack/_internal/core/backends/datacrunch/compute.py),

src/dstack/_internal/core/backends/__init__.py

@@ -1,42 +1,59 @@
+from dstack._internal.core.backends.base.compute import (
+    ComputeWithCreateInstanceSupport,
+    ComputeWithGatewaySupport,
+    ComputeWithMultinodeSupport,
+    ComputeWithPlacementGroupSupport,
+    ComputeWithPrivateGatewaySupport,
+    ComputeWithReservationSupport,
+    ComputeWithVolumeSupport,
+)
+from dstack._internal.core.backends.base.configurator import Configurator
+from dstack._internal.core.backends.configurators import list_available_configurator_classes
 from dstack._internal.core.models.backends.base import BackendType
 
-BACKENDS_WITH_MULTINODE_SUPPORT = [
-    BackendType.AWS,
-    BackendType.AZURE,
-    BackendType.GCP,
-    BackendType.REMOTE,
-    BackendType.OCI,
-    BackendType.VULTR,
-]
-BACKENDS_WITH_CREATE_INSTANCE_SUPPORT = [
-    BackendType.AWS,
-    BackendType.DSTACK,
-    BackendType.AZURE,
-    BackendType.CUDO,
-    BackendType.DATACRUNCH,
-    BackendType.GCP,
-    BackendType.LAMBDA,
-    BackendType.OCI,
-    BackendType.TENSORDOCK,
-    BackendType.VULTR,
-]
-BACKENDS_WITH_PLACEMENT_GROUPS_SUPPORT = [
-    BackendType.AWS,
-]
-BACKENDS_WITH_RESERVATION_SUPPORT = [
-    BackendType.AWS,
-]
 
-BACKENDS_WITH_GATEWAY_SUPPORT = [
-    BackendType.AWS,
-    BackendType.AZURE,
-    BackendType.GCP,
-    BackendType.KUBERNETES,
-]
-BACKENDS_WITH_PRIVATE_GATEWAY_SUPPORT = [BackendType.AWS]
-BACKENDS_WITH_VOLUMES_SUPPORT = [
-    BackendType.AWS,
-    BackendType.GCP,
-    BackendType.LOCAL,
-    BackendType.RUNPOD,
-]
+def _get_backends_with_compute_feature(
+    configurator_classes: list[type[Configurator]],
+    compute_feature_class: type,
+) -> list[BackendType]:
+    backend_types = []
+    for configurator_class in configurator_classes:
+        compute_class = configurator_class.BACKEND_CLASS.COMPUTE_CLASS
+        if issubclass(compute_class, compute_feature_class):
+            backend_types.append(configurator_class.TYPE)
+    return backend_types
+
+
+_configurator_classes = list_available_configurator_classes()
+
+
+# The following backend lists do not include unavailable backends (i.e. backends missing deps).
+# TODO: Add LocalBackend to lists if it's enabled
+BACKENDS_WITH_CREATE_INSTANCE_SUPPORT = _get_backends_with_compute_feature(
+    configurator_classes=_configurator_classes,
+    compute_feature_class=ComputeWithCreateInstanceSupport,
+)
+BACKENDS_WITH_MULTINODE_SUPPORT = [BackendType.REMOTE] + _get_backends_with_compute_feature(
+    configurator_classes=_configurator_classes,
+    compute_feature_class=ComputeWithMultinodeSupport,
+)
+BACKENDS_WITH_PLACEMENT_GROUPS_SUPPORT = _get_backends_with_compute_feature(
+    configurator_classes=_configurator_classes,
+    compute_feature_class=ComputeWithPlacementGroupSupport,
+)
+BACKENDS_WITH_RESERVATION_SUPPORT = _get_backends_with_compute_feature(
+    configurator_classes=_configurator_classes,
+    compute_feature_class=ComputeWithReservationSupport,
+)
+BACKENDS_WITH_GATEWAY_SUPPORT = _get_backends_with_compute_feature(
+    configurator_classes=_configurator_classes,
+    compute_feature_class=ComputeWithGatewaySupport,
+)
+BACKENDS_WITH_PRIVATE_GATEWAY_SUPPORT = _get_backends_with_compute_feature(
+    configurator_classes=_configurator_classes,
+    compute_feature_class=ComputeWithPrivateGatewaySupport,
+)
+BACKENDS_WITH_VOLUMES_SUPPORT = _get_backends_with_compute_feature(
+    configurator_classes=_configurator_classes,
+    compute_feature_class=ComputeWithVolumeSupport,
+)

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

@@ -8,7 +8,8 @@
 
 
 class AWSBackend(Backend):
-    TYPE: BackendType = BackendType.AWS
+    TYPE = BackendType.AWS
+    COMPUTE_CLASS = AWSCompute
 
     def __init__(self, config: AWSConfig):
         self.config = config

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

@@ -12,6 +12,13 @@
 from dstack._internal.core.backends.aws.models import AWSAccessKeyCreds
 from dstack._internal.core.backends.base.compute import (
     Compute,
+    ComputeWithCreateInstanceSupport,
+    ComputeWithGatewaySupport,
+    ComputeWithMultinodeSupport,
+    ComputeWithPlacementGroupSupport,
+    ComputeWithPrivateGatewaySupport,
+    ComputeWithReservationSupport,
+    ComputeWithVolumeSupport,
     generate_unique_gateway_instance_name,
     generate_unique_instance_name,
     generate_unique_volume_name,
@@ -62,7 +69,16 @@ class AWSVolumeBackendData(CoreModel):
     iops: int
 
 
-class AWSCompute(Compute):
+class AWSCompute(
+    Compute,
+    ComputeWithCreateInstanceSupport,
+    ComputeWithMultinodeSupport,
+    ComputeWithReservationSupport,
+    ComputeWithPlacementGroupSupport,
+    ComputeWithGatewaySupport,
+    ComputeWithPrivateGatewaySupport,
+    ComputeWithVolumeSupport,
+):
     def __init__(self, config: AWSConfig):
         super().__init__()
         self.config = config

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

@@ -53,7 +53,8 @@
 
 
 class AWSConfigurator(Configurator):
-    TYPE: BackendType = BackendType.AWS
+    TYPE = BackendType.AWS
+    BACKEND_CLASS = AWSBackend
 
     def validate_config(self, config: AWSBackendConfigWithCreds, default_creds_enabled: bool):
         if is_core_model_instance(config.creds, AWSDefaultCreds) and not default_creds_enabled:

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

@@ -6,7 +6,8 @@
 
 
 class AzureBackend(Backend):
-    TYPE: BackendType = BackendType.AZURE
+    TYPE = BackendType.AZURE
+    COMPUTE_CLASS = AzureCompute
 
     def __init__(self, config: AzureConfig):
         self.config = config

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

@@ -39,6 +39,9 @@
 from dstack._internal.core.backends.azure.config import AzureConfig
 from dstack._internal.core.backends.base.compute import (
     Compute,
+    ComputeWithCreateInstanceSupport,
+    ComputeWithGatewaySupport,
+    ComputeWithMultinodeSupport,
     generate_unique_gateway_instance_name,
     generate_unique_instance_name,
     get_gateway_user_data,
@@ -71,7 +74,12 @@
 CONFIGURABLE_DISK_SIZE = Range[Memory](min=Memory.parse("30GB"), max=Memory.parse("4095GB"))
 
 
-class AzureCompute(Compute):
+class AzureCompute(
+    Compute,
+    ComputeWithCreateInstanceSupport,
+    ComputeWithMultinodeSupport,
+    ComputeWithGatewaySupport,
+):
     def __init__(self, config: AzureConfig, credential: TokenCredential):
         super().__init__()
         self.config = config

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

@@ -72,7 +72,8 @@
 
 
 class AzureConfigurator(Configurator):
-    TYPE: BackendType = BackendType.AZURE
+    TYPE = BackendType.AZURE
+    BACKEND_CLASS = AzureBackend
 
     def validate_config(self, config: AzureBackendConfigWithCreds, default_creds_enabled: bool):
         if is_core_model_instance(config.creds, AzureDefaultCreds) and not default_creds_enabled:

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

@@ -1,12 +1,18 @@
 from abc import ABC, abstractmethod
+from typing import ClassVar
 
 from dstack._internal.core.backends.base.compute import Compute
 from dstack._internal.core.models.backends.base import BackendType
 
 
 class Backend(ABC):
-    TYPE: BackendType
+    TYPE: ClassVar[BackendType]
+    # `COMPUTE_CLASS` is used to introspect compute features without initializing it.
+    COMPUTE_CLASS: ClassVar[type[Compute]]
 
     @abstractmethod
     def compute(self) -> Compute:
+        """
+        Returns Compute instance.
+        """
         pass