Support Runpod Instant Clusters (#3214)

* Fix runpod type annotations

* Add _generate_create_cluster_mutation

* feat: add create_cluster method to RunpodApiClient

* feat: add delete_cluster method to RunpodApiClient

* Use keyword arguments

* Implement run_jobs and terminate_compute_group

* Prototype compute.run_jobs calling

* Prototype compute group provisioning for multinode tasks

* Add JobModel.waiting_master_job

* Add ComputeGroupModel

* Implement process_compute_groups to terminate compute groups

* Remove todo

* Fix comments

* Set internal_ip

* Support Runpod Clusters offers

* Respect supported pod_counts

* Support registry_auth

* Fix tests

* Add feature flag DSTACK_FF_RUNPOD_CLUSTER_OFFERS_ENABLED

* Remove hardcoded template_id

* Use BACKENDS_WITH_GROUP_PROVISIONING_SUPPORT

Author: r4victor

pyproject.toml

@@ -84,6 +84,7 @@ include = [
     "src/dstack/_internal/server",
     "src/dstack/_internal/core/services",
     "src/dstack/_internal/core/backends/kubernetes",
+    "src/dstack/_internal/core/backends/runpod",
     "src/dstack/_internal/cli/services/configurators",
     "src/dstack/_internal/cli/commands",
 ]

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

@@ -17,13 +17,15 @@
 from gpuhunt import CPUArchitecture
 
 from dstack._internal import settings
+from dstack._internal.core.backends.base.models import JobConfiguration
 from dstack._internal.core.backends.base.offers import OfferModifier, filter_offers_by_requirements
 from dstack._internal.core.consts import (
     DSTACK_RUNNER_HTTP_PORT,
     DSTACK_RUNNER_SSH_PORT,
     DSTACK_SHIM_HTTP_PORT,
 )
 from dstack._internal.core.models.backends.base import BackendType
+from dstack._internal.core.models.compute_groups import ComputeGroup, ComputeGroupProvisioningData
 from dstack._internal.core.models.configurations import LEGACY_REPO_DIR
 from dstack._internal.core.models.gateways import (
     GatewayComputeConfiguration,
@@ -324,6 +326,23 @@ def _restrict_instance_offer_az_to_volumes_az(
             ]
 
 
+class ComputeWithGroupProvisioningSupport(ABC):
+    @abstractmethod
+    def run_jobs(
+        self,
+        run: Run,
+        job_configurations: List[JobConfiguration],
+        instance_offer: InstanceOfferWithAvailability,
+        project_ssh_public_key: str,
+        project_ssh_private_key: str,
+    ) -> ComputeGroupProvisioningData:
+        pass
+
+    @abstractmethod
+    def terminate_compute_group(self, compute_group: ComputeGroup):
+        pass
+
+
 class ComputeWithPrivilegedSupport:
     """
     Must be subclassed to support runs with `privileged: true`.

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

@@ -1,4 +1,14 @@
 from pathlib import Path
+from typing import List
+
+from dstack._internal.core.models.common import CoreModel
+from dstack._internal.core.models.runs import Job
+from dstack._internal.core.models.volumes import Volume
+
+
+class JobConfiguration(CoreModel):
+    job: Job
+    volumes: List[Volume]
 
 
 def fill_data(values: dict, filename_field: str = "filename", data_field: str = "data") -> dict:

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

@@ -25,6 +25,7 @@
     "gcp-a4",
     "gcp-g4",
     "gcp-dws-calendar-mode",
+    "runpod-cluster",
 ]
 
 

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

@@ -1,6 +1,7 @@
 from dstack._internal.core.backends.base.compute import (
     ComputeWithCreateInstanceSupport,
     ComputeWithGatewaySupport,
+    ComputeWithGroupProvisioningSupport,
     ComputeWithMultinodeSupport,
     ComputeWithPlacementGroupSupport,
     ComputeWithPrivateGatewaySupport,
@@ -39,6 +40,10 @@ def _get_backends_with_compute_feature(
     configurator_classes=_configurator_classes,
     compute_feature_class=ComputeWithCreateInstanceSupport,
 )
+BACKENDS_WITH_GROUP_PROVISIONING_SUPPORT = _get_backends_with_compute_feature(
+    configurator_classes=_configurator_classes,
+    compute_feature_class=ComputeWithGroupProvisioningSupport,
+)
 BACKENDS_WITH_PRIVILEGED_SUPPORT = _get_backends_with_compute_feature(
     configurator_classes=_configurator_classes,
     compute_feature_class=ComputeWithPrivilegedSupport,