dstack 0.18.44__py3-none-any.whl → 0.19.0__py3-none-any.whl

dstack/_internal/core/backends/cudo/configurator.py

@@ -0,0 +1,72 @@
+import json
+
+from dstack._internal.core.backends.base.configurator import (
+    BackendRecord,
+    Configurator,
+    raise_invalid_credentials_error,
+)
+from dstack._internal.core.backends.cudo import api_client
+from dstack._internal.core.backends.cudo.backend import CudoBackend
+from dstack._internal.core.backends.cudo.models import (
+    AnyCudoBackendConfig,
+    CudoBackendConfig,
+    CudoBackendConfigWithCreds,
+    CudoConfig,
+    CudoCreds,
+    CudoStoredConfig,
+)
+from dstack._internal.core.models.backends.base import BackendType
+
+REGIONS = [
+    "no-luster-1",
+    "se-smedjebacken-1",
+    "gb-london-1",
+    "se-stockholm-1",
+    "us-newyork-1",
+    "us-santaclara-1",
+]
+
+DEFAULT_REGION = "no-luster-1"
+
+
+class CudoConfigurator(Configurator):
+    TYPE = BackendType.CUDO
+    BACKEND_CLASS = CudoBackend
+
+    def validate_config(self, config: CudoBackendConfigWithCreds, default_creds_enabled: bool):
+        self._validate_cudo_api_key(config.creds.api_key)
+
+    def create_backend(
+        self, project_name: str, config: CudoBackendConfigWithCreds
+    ) -> BackendRecord:
+        if config.regions is None:
+            config.regions = REGIONS
+        return BackendRecord(
+            config=CudoStoredConfig(
+                **CudoBackendConfig.__response__.parse_obj(config).dict()
+            ).json(),
+            auth=CudoCreds.parse_obj(config.creds).json(),
+        )
+
+    def get_backend_config(
+        self, record: BackendRecord, include_creds: bool
+    ) -> AnyCudoBackendConfig:
+        config = self._get_config(record)
+        if include_creds:
+            return CudoBackendConfigWithCreds.__response__.parse_obj(config)
+        return CudoBackendConfig.__response__.parse_obj(config)
+
+    def get_backend(self, record: BackendRecord) -> CudoBackend:
+        config = self._get_config(record)
+        return CudoBackend(config=config)
+
+    def _get_config(self, record: BackendRecord) -> CudoConfig:
+        return CudoConfig.__response__(
+            **json.loads(record.config),
+            creds=CudoCreds.parse_raw(record.auth),
+        )
+
+    def _validate_cudo_api_key(self, api_key: str):
+        client = api_client.CudoApiClient(api_key=api_key)
+        if not client.validate_api_key():
+            raise_invalid_credentials_error(fields=[["creds", "api_key"]])

dstack/_internal/core/backends/cudo/models.py

@@ -0,0 +1,37 @@
+from typing import Annotated, List, Literal, Optional, Union
+
+from pydantic import Field
+
+from dstack._internal.core.models.common import CoreModel
+
+
+class CudoAPIKeyCreds(CoreModel):
+    type: Annotated[Literal["api_key"], Field(description="The type of credentials")] = "api_key"
+    api_key: Annotated[str, Field(description="The API key")]
+
+
+AnyCudoCreds = CudoAPIKeyCreds
+CudoCreds = AnyCudoCreds
+
+
+class CudoBackendConfig(CoreModel):
+    type: Annotated[Literal["cudo"], Field(description="The type of backend")] = "cudo"
+    regions: Annotated[
+        Optional[List[str]], Field(description="The list of Cudo regions. Omit to use all regions")
+    ] = None
+    project_id: Annotated[str, Field(description="The project ID")]
+
+
+class CudoBackendConfigWithCreds(CudoBackendConfig):
+    creds: Annotated[AnyCudoCreds, Field(description="The credentials")]
+
+
+AnyCudoBackendConfig = Union[CudoBackendConfig, CudoBackendConfigWithCreds]
+
+
+class CudoStoredConfig(CudoBackendConfig):
+    pass
+
+
+class CudoConfig(CudoStoredConfig):
+    creds: AnyCudoCreds

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

@@ -1,15 +0,0 @@
-from dstack._internal.core.backends.base import Backend
-from dstack._internal.core.backends.datacrunch.compute import DataCrunchCompute
-from dstack._internal.core.backends.datacrunch.config import DataCrunchConfig
-from dstack._internal.core.models.backends.base import BackendType
-
-
-class DataCrunchBackend(Backend):
-    TYPE: BackendType = BackendType.DATACRUNCH
-
-    def __init__(self, config: DataCrunchConfig):
-        self.config = config
-        self._compute = DataCrunchCompute(self.config)
-
-    def compute(self) -> DataCrunchCompute:
-        return self._compute

dstack/_internal/core/backends/datacrunch/backend.py

@@ -0,0 +1,16 @@
+from dstack._internal.core.backends.base.backend import Backend
+from dstack._internal.core.backends.datacrunch.compute import DataCrunchCompute
+from dstack._internal.core.backends.datacrunch.models import DataCrunchConfig
+from dstack._internal.core.models.backends.base import BackendType
+
+
+class DataCrunchBackend(Backend):
+    TYPE = BackendType.DATACRUNCH
+    COMPUTE_CLASS = DataCrunchCompute
+
+    def __init__(self, config: DataCrunchConfig):
+        self.config = config
+        self._compute = DataCrunchCompute(self.config)
+
+    def compute(self) -> DataCrunchCompute:
+        return self._compute

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

@@ -1,24 +1,23 @@
 from typing import Dict, List, Optional
 
-from dstack._internal.core.backends.base import Compute
+from dstack._internal.core.backends.base.backend import Compute
 from dstack._internal.core.backends.base.compute import (
+    ComputeWithCreateInstanceSupport,
     generate_unique_instance_name,
     get_shim_commands,
 )
 from dstack._internal.core.backends.base.offers import get_catalog_offers
 from dstack._internal.core.backends.datacrunch.api_client import DataCrunchAPIClient
-from dstack._internal.core.backends.datacrunch.config import DataCrunchConfig
+from dstack._internal.core.backends.datacrunch.models import DataCrunchConfig
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.instances import (
     InstanceAvailability,
     InstanceConfiguration,
     InstanceOffer,
     InstanceOfferWithAvailability,
-    SSHKey,
 )
 from dstack._internal.core.models.resources import Memory, Range
-from dstack._internal.core.models.runs import Job, JobProvisioningData, Requirements, Run
-from dstack._internal.core.models.volumes import Volume
+from dstack._internal.core.models.runs import JobProvisioningData, Requirements
 from dstack._internal.utils.logging import get_logger
 
 logger = get_logger("datacrunch.compute")
@@ -33,7 +32,10 @@ IMAGE_SIZE = Memory.parse("50GB")
 CONFIGURABLE_DISK_SIZE = Range[Memory](min=IMAGE_SIZE, max=None)
 
 
-class DataCrunchCompute(Compute):
+class DataCrunchCompute(
+    ComputeWithCreateInstanceSupport,
+    Compute,
+):
     def __init__(self, config: DataCrunchConfig):
         super().__init__()
         self.config = config
@@ -148,25 +150,6 @@ class DataCrunchCompute(Compute):
             backend_data=None,
         )
 
-    def run_job(
-        self,
-        run: Run,
-        job: Job,
-        instance_offer: InstanceOfferWithAvailability,
-        project_ssh_public_key: str,
-        project_ssh_private_key: str,
-        volumes: List[Volume],
-    ) -> JobProvisioningData:
-        instance_config = InstanceConfiguration(
-            project_name=run.project_name,
-            instance_name=job.job_spec.job_name,  # TODO: generate name
-            ssh_keys=[
-                SSHKey(public=project_ssh_public_key.strip()),
-            ],
-            user=run.user,
-        )
-        return self.create_instance(instance_offer, instance_config)
-
     def terminate_instance(
         self, instance_id: str, region: str, backend_data: Optional[str] = None
     ) -> None:

dstack/_internal/core/backends/datacrunch/configurator.py

@@ -0,0 +1,66 @@
+import json
+
+from dstack._internal.core.backends.base.configurator import (
+    BackendRecord,
+    Configurator,
+)
+from dstack._internal.core.backends.datacrunch.backend import DataCrunchBackend
+from dstack._internal.core.backends.datacrunch.models import (
+    AnyDataCrunchBackendConfig,
+    DataCrunchBackendConfig,
+    DataCrunchBackendConfigWithCreds,
+    DataCrunchConfig,
+    DataCrunchCreds,
+    DataCrunchStoredConfig,
+)
+from dstack._internal.core.models.backends.base import (
+    BackendType,
+)
+
+REGIONS = [
+    "FIN-01",
+    "ICE-01",
+]
+
+DEFAULT_REGION = "FIN-01"
+
+
+class DataCrunchConfigurator(Configurator):
+    TYPE = BackendType.DATACRUNCH
+    BACKEND_CLASS = DataCrunchBackend
+
+    def validate_config(
+        self, config: DataCrunchBackendConfigWithCreds, default_creds_enabled: bool
+    ):
+        # FIXME: validate datacrunch creds
+        return
+
+    def create_backend(
+        self, project_name: str, config: DataCrunchBackendConfigWithCreds
+    ) -> BackendRecord:
+        if config.regions is None:
+            config.regions = REGIONS
+        return BackendRecord(
+            config=DataCrunchStoredConfig(
+                **DataCrunchBackendConfig.__response__.parse_obj(config).dict()
+            ).json(),
+            auth=DataCrunchCreds.parse_obj(config.creds).json(),
+        )
+
+    def get_backend_config(
+        self, record: BackendRecord, include_creds: bool
+    ) -> AnyDataCrunchBackendConfig:
+        config = self._get_config(record)
+        if include_creds:
+            return DataCrunchBackendConfigWithCreds.__response__.parse_obj(config)
+        return DataCrunchBackendConfig.__response__.parse_obj(config)
+
+    def get_backend(self, record: BackendRecord) -> DataCrunchBackend:
+        config = self._get_config(record)
+        return DataCrunchBackend(config=config)
+
+    def _get_config(self, record: BackendRecord) -> DataCrunchConfig:
+        return DataCrunchConfig.__response__(
+            **json.loads(record.config),
+            creds=DataCrunchCreds.parse_raw(record.auth),
+        )

dstack/_internal/core/backends/datacrunch/models.py

@@ -0,0 +1,38 @@
+from typing import Annotated, List, Literal, Optional, Union
+
+from pydantic import Field
+
+from dstack._internal.core.models.common import CoreModel
+
+
+class DataCrunchAPIKeyCreds(CoreModel):
+    type: Annotated[Literal["api_key"], Field(description="The type of credentials")] = "api_key"
+    client_id: Annotated[str, Field(description="The client ID")]
+    client_secret: Annotated[str, Field(description="The client secret")]
+
+
+AnyDataCrunchCreds = DataCrunchAPIKeyCreds
+DataCrunchCreds = AnyDataCrunchCreds
+
+
+class DataCrunchBackendConfig(CoreModel):
+    type: Annotated[Literal["datacrunch"], Field(description="The type of backend")] = "datacrunch"
+    regions: Annotated[
+        Optional[List[str]],
+        Field(description="The list of DataCrunch regions. Omit to use all regions"),
+    ] = None
+
+
+class DataCrunchBackendConfigWithCreds(DataCrunchBackendConfig):
+    creds: Annotated[AnyDataCrunchCreds, Field(description="The credentials")]
+
+
+AnyDataCrunchBackendConfig = Union[DataCrunchBackendConfig, DataCrunchBackendConfigWithCreds]
+
+
+class DataCrunchStoredConfig(DataCrunchBackendConfig):
+    pass
+
+
+class DataCrunchConfig(DataCrunchStoredConfig):
+    creds: AnyDataCrunchCreds

dstack/_internal/core/{models/backends/dstack.py → backends/dstack/models.py}

@@ -1,15 +1,15 @@
-from typing import List
+from typing import Annotated, List, Literal
 
-from typing_extensions import Literal
+from pydantic import Field
 
 from dstack._internal.core.models.common import CoreModel
 
-# The OSS is currently aware of some of the DstackBackend internals (DstackConfigInfo) to be able to
+# The OSS is currently aware of some of the DstackBackend internals (DstackBackendConfig) to be able to
 # show DstackBackend base backends as regular backends.
 # Consider designing an API that would allow DstackBackend to do the same without exposing its internals.
 
 
-class DstackConfigInfo(CoreModel):
+class DstackBackendConfig(CoreModel):
     """
     This is a config model of DstackBackend stored in BackendModel.config and used by DstackConfigurator.
     """
@@ -18,9 +18,9 @@ class DstackConfigInfo(CoreModel):
     base_backends: List[str]
 
 
-class DstackBaseBackendConfigInfo(CoreModel):
+class DstackBaseBackendConfig(CoreModel):
     type: str
 
 
-class DstackConfigValues(CoreModel):
-    type: Literal["dstack"] = "dstack"
+class DstackConfig(CoreModel):
+    type: Annotated[Literal["dstack"], Field(description="The type of backend")] = "dstack"

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

@@ -1,16 +0,0 @@
-from dstack._internal.core.backends.base import Backend
-from dstack._internal.core.backends.gcp.compute import GCPCompute
-from dstack._internal.core.backends.gcp.config import GCPConfig
-from dstack._internal.core.models.backends.base import BackendType
-
-
-class GCPBackend(Backend):
-    TYPE: BackendType = BackendType.GCP
-
-    def __init__(self, config: GCPConfig):
-        self.config = config
-        self._compute = GCPCompute(self.config)
-        # self._check_credentials()
-
-    def compute(self) -> GCPCompute:
-        return self._compute

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

@@ -8,12 +8,11 @@ from google.auth.credentials import Credentials
 from google.auth.exceptions import DefaultCredentialsError
 from google.oauth2 import service_account
 
-from dstack._internal.core.errors import BackendAuthError
-from dstack._internal.core.models.backends.gcp import (
+from dstack._internal.core.backends.gcp.models import (
     AnyGCPCreds,
-    GCPDefaultCreds,
     GCPServiceAccountCreds,
 )
+from dstack._internal.core.errors import BackendAuthError
 from dstack._internal.core.models.common import is_core_model_instance
 
 
@@ -57,11 +56,3 @@ def validate_credentials(credentials: Credentials, project_id: str):
         raise BackendAuthError(f"project_id {project_id} not found")
     except Exception:
         raise BackendAuthError("Insufficient permissions")
-
-
-def default_creds_available() -> bool:
-    try:
-        authenticate(GCPDefaultCreds())
-    except BackendAuthError:
-        return False
-    return True

dstack/_internal/core/backends/gcp/backend.py

@@ -0,0 +1,17 @@
+from dstack._internal.core.backends.base.backend import Backend
+from dstack._internal.core.backends.gcp.compute import GCPCompute
+from dstack._internal.core.backends.gcp.models import GCPConfig
+from dstack._internal.core.models.backends.base import BackendType
+
+
+class GCPBackend(Backend):
+    TYPE = BackendType.GCP
+    COMPUTE_CLASS = GCPCompute
+
+    def __init__(self, config: GCPConfig):
+        self.config = config
+        self._compute = GCPCompute(self.config)
+        # self._check_credentials()
+
+    def compute(self) -> GCPCompute:
+        return self._compute

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

@@ -12,17 +12,20 @@ import dstack._internal.core.backends.gcp.auth as auth
 import dstack._internal.core.backends.gcp.resources as gcp_resources
 from dstack._internal.core.backends.base.compute import (
     Compute,
+    ComputeWithCreateInstanceSupport,
+    ComputeWithGatewaySupport,
+    ComputeWithMultinodeSupport,
+    ComputeWithVolumeSupport,
     generate_unique_gateway_instance_name,
     generate_unique_instance_name,
     generate_unique_volume_name,
     get_gateway_user_data,
-    get_job_instance_name,
     get_shim_commands,
     get_user_data,
     merge_tags,
 )
 from dstack._internal.core.backends.base.offers import get_catalog_offers
-from dstack._internal.core.backends.gcp.config import GCPConfig
+from dstack._internal.core.backends.gcp.models import GCPConfig
 from dstack._internal.core.errors import (
     ComputeError,
     ComputeResourceNotFoundError,
@@ -42,10 +45,9 @@ from dstack._internal.core.models.instances import (
     InstanceOfferWithAvailability,
     InstanceType,
     Resources,
-    SSHKey,
 )
 from dstack._internal.core.models.resources import Memory, Range
-from dstack._internal.core.models.runs import Job, JobProvisioningData, Requirements, Run
+from dstack._internal.core.models.runs import JobProvisioningData, Requirements
 from dstack._internal.core.models.volumes import (
     Volume,
     VolumeAttachmentData,
@@ -69,7 +71,13 @@ class GCPVolumeDiskBackendData(CoreModel):
     disk_type: str
 
 
-class GCPCompute(Compute):
+class GCPCompute(
+    ComputeWithCreateInstanceSupport,
+    ComputeWithMultinodeSupport,
+    ComputeWithGatewaySupport,
+    ComputeWithVolumeSupport,
+    Compute,
+):
     def __init__(self, config: GCPConfig):
         super().__init__()
         self.config = config
@@ -363,44 +371,6 @@ class GCPCompute(Compute):
             f"Failed to get instance IP address. Instance status: {instance.status}"
         )
 
-    def run_job(
-        self,
-        run: Run,
-        job: Job,
-        instance_offer: InstanceOfferWithAvailability,
-        project_ssh_public_key: str,
-        project_ssh_private_key: str,
-        volumes: List[Volume],
-    ) -> JobProvisioningData:
-        # TODO: run_job is the same for vm-based backends, refactor
-        instance_config = InstanceConfiguration(
-            project_name=run.project_name,
-            instance_name=get_job_instance_name(run, job),  # TODO: generate name
-            ssh_keys=[
-                SSHKey(public=project_ssh_public_key.strip()),
-            ],
-            user=run.user,
-            volumes=volumes,
-            reservation=run.run_spec.configuration.reservation,
-        )
-        instance_offer = instance_offer.copy()
-        if len(volumes) > 0:
-            volume = volumes[0]
-            if (
-                volume.provisioning_data is not None
-                and volume.provisioning_data.availability_zone is not None
-            ):
-                if instance_offer.availability_zones is None:
-                    instance_offer.availability_zones = [
-                        volume.provisioning_data.availability_zone
-                    ]
-                instance_offer.availability_zones = [
-                    z
-                    for z in instance_offer.availability_zones
-                    if z == volume.provisioning_data.availability_zone
-                ]
-        return self.create_instance(instance_offer, instance_config)
-
     def create_gateway(
         self,
         configuration: GatewayComputeConfiguration,