dstack 0.18.43__py3-none-any.whl → 0.19.0rc1__py3-none-any.whl

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

@@ -0,0 +1,82 @@
+import json
+
+from dstack._internal.core.backends.base.configurator import (
+    BackendRecord,
+    Configurator,
+    raise_invalid_credentials_error,
+)
+from dstack._internal.core.backends.lambdalabs import api_client
+from dstack._internal.core.backends.lambdalabs.backend import LambdaBackend
+from dstack._internal.core.backends.lambdalabs.models import (
+    AnyLambdaBackendConfig,
+    LambdaBackendConfig,
+    LambdaBackendConfigWithCreds,
+    LambdaConfig,
+    LambdaCreds,
+    LambdaStoredConfig,
+)
+from dstack._internal.core.models.backends.base import (
+    BackendType,
+)
+
+REGIONS = [
+    "us-south-1",
+    "us-south-2",
+    "us-south-3",
+    "us-west-2",
+    "us-west-1",
+    "us-midwest-1",
+    "us-west-3",
+    "us-east-1",
+    "us-east-2",
+    "europe-central-1",
+    "asia-south-1",
+    "me-west-1",
+    "asia-northeast-1",
+    "asia-northeast-2",
+]
+
+DEFAULT_REGION = "us-east-1"
+
+
+class LambdaConfigurator(Configurator):
+    TYPE = BackendType.LAMBDA
+    BACKEND_CLASS = LambdaBackend
+
+    def validate_config(self, config: LambdaBackendConfigWithCreds, default_creds_enabled: bool):
+        self._validate_lambda_api_key(config.creds.api_key)
+
+    def create_backend(
+        self, project_name: str, config: LambdaBackendConfigWithCreds
+    ) -> BackendRecord:
+        if config.regions is None:
+            config.regions = REGIONS
+        return BackendRecord(
+            config=LambdaStoredConfig(
+                **LambdaBackendConfig.__response__.parse_obj(config).dict()
+            ).json(),
+            auth=LambdaCreds.parse_obj(config.creds).json(),
+        )
+
+    def get_backend_config(
+        self, record: BackendRecord, include_creds: bool
+    ) -> AnyLambdaBackendConfig:
+        config = self._get_config(record)
+        if include_creds:
+            return LambdaBackendConfigWithCreds.__response__.parse_obj(config)
+        return LambdaBackendConfig.__response__.parse_obj(config)
+
+    def get_backend(self, record: BackendRecord) -> LambdaBackend:
+        config = self._get_config(record)
+        return LambdaBackend(config=config)
+
+    def _get_config(self, record: BackendRecord) -> LambdaConfig:
+        return LambdaConfig.__response__(
+            **json.loads(record.config),
+            creds=LambdaCreds.parse_raw(record.auth),
+        )
+
+    def _validate_lambda_api_key(self, api_key: str):
+        client = api_client.LambdaAPIClient(api_key=api_key)
+        if not client.validate_api_key():
+            raise_invalid_credentials_error(fields=[["creds", "api_key"]])

dstack/_internal/core/backends/lambdalabs/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 LambdaAPIKeyCreds(CoreModel):
+    type: Annotated[Literal["api_key"], Field(description="The type of credentials")] = "api_key"
+    api_key: Annotated[str, Field(description="The API key")]
+
+
+AnyLambdaCreds = LambdaAPIKeyCreds
+LambdaCreds = AnyLambdaCreds
+
+
+class LambdaBackendConfig(CoreModel):
+    type: Annotated[Literal["lambda"], Field(description="The type of backend")] = "lambda"
+    regions: Annotated[
+        Optional[List[str]],
+        Field(description="The list of Lambda regions. Omit to use all regions"),
+    ] = None
+
+
+class LambdaBackendConfigWithCreds(LambdaBackendConfig):
+    creds: Annotated[AnyLambdaCreds, Field(description="The credentials")]
+
+
+AnyLambdaBackendConfig = Union[LambdaBackendConfig, LambdaBackendConfigWithCreds]
+
+
+class LambdaStoredConfig(LambdaBackendConfig):
+    pass
+
+
+class LambdaConfig(LambdaStoredConfig):
+    creds: AnyLambdaCreds

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

@@ -1,13 +0,0 @@
-from dstack._internal.core.backends.base import Backend
-from dstack._internal.core.backends.local.compute import LocalCompute
-from dstack._internal.core.models.backends.base import BackendType
-
-
-class LocalBackend(Backend):
-    TYPE: BackendType = BackendType.LOCAL
-
-    def __init__(self):
-        self._compute = LocalCompute()
-
-    def compute(self) -> LocalCompute:
-        return self._compute

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

@@ -0,0 +1,14 @@
+from dstack._internal.core.backends.base.backend import Backend
+from dstack._internal.core.backends.local.compute import LocalCompute
+from dstack._internal.core.models.backends.base import BackendType
+
+
+class LocalBackend(Backend):
+    TYPE = BackendType.LOCAL
+    COMPUTE_CLASS = LocalCompute
+
+    def __init__(self):
+        self._compute = LocalCompute()
+
+    def compute(self) -> LocalCompute:
+        return self._compute

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

@@ -1,6 +1,10 @@
 from typing import List, Optional
 
-from dstack._internal.core.backends.base.compute import Compute
+from dstack._internal.core.backends.base.compute import (
+    Compute,
+    ComputeWithCreateInstanceSupport,
+    ComputeWithVolumeSupport,
+)
 from dstack._internal.core.consts import DSTACK_RUNNER_SSH_PORT
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.instances import (
@@ -18,7 +22,11 @@ from dstack._internal.utils.logging import get_logger
 logger = get_logger(__name__)
 
 
-class LocalCompute(Compute):
+class LocalCompute(
+    ComputeWithCreateInstanceSupport,
+    ComputeWithVolumeSupport,
+    Compute,
+):
     def get_offers(
         self, requirements: Optional[Requirements] = None
     ) -> List[InstanceOfferWithAvailability]:
@@ -85,6 +93,12 @@ class LocalCompute(Compute):
             backend_data=None,
         )
 
+    def register_volume(self, volume: Volume) -> VolumeProvisioningData:
+        return VolumeProvisioningData(
+            volume_id=volume.volume_id,
+            size_gb=volume.configuration.size_gb,
+        )
+
     def create_volume(self, volume: Volume) -> VolumeProvisioningData:
         return VolumeProvisioningData(
             volume_id=volume.name,

dstack/_internal/core/backends/models.py

@@ -0,0 +1,128 @@
+from typing import Union
+
+from dstack._internal.core.backends.aws.models import (
+    AWSBackendConfig,
+    AWSBackendConfigWithCreds,
+)
+from dstack._internal.core.backends.azure.models import (
+    AzureBackendConfig,
+    AzureBackendConfigWithCreds,
+)
+from dstack._internal.core.backends.cudo.models import (
+    CudoBackendConfig,
+    CudoBackendConfigWithCreds,
+)
+from dstack._internal.core.backends.datacrunch.models import (
+    DataCrunchBackendConfig,
+    DataCrunchBackendConfigWithCreds,
+)
+from dstack._internal.core.backends.dstack.models import (
+    DstackBackendConfig,
+    DstackBaseBackendConfig,
+)
+from dstack._internal.core.backends.gcp.models import (
+    GCPBackendConfig,
+    GCPBackendConfigWithCreds,
+    GCPBackendFileConfigWithCreds,
+)
+from dstack._internal.core.backends.kubernetes.models import (
+    KubernetesBackendConfig,
+    KubernetesBackendConfigWithCreds,
+    KubernetesBackendFileConfigWithCreds,
+)
+from dstack._internal.core.backends.lambdalabs.models import (
+    LambdaBackendConfig,
+    LambdaBackendConfigWithCreds,
+)
+from dstack._internal.core.backends.oci.models import (
+    OCIBackendConfig,
+    OCIBackendConfigWithCreds,
+)
+from dstack._internal.core.backends.runpod.models import (
+    RunpodBackendConfig,
+    RunpodBackendConfigWithCreds,
+)
+from dstack._internal.core.backends.tensordock.models import (
+    TensorDockBackendConfig,
+    TensorDockBackendConfigWithCreds,
+)
+from dstack._internal.core.backends.vastai.models import (
+    VastAIBackendConfig,
+    VastAIBackendConfigWithCreds,
+)
+from dstack._internal.core.backends.vultr.models import (
+    VultrBackendConfig,
+    VultrBackendConfigWithCreds,
+)
+from dstack._internal.core.models.common import CoreModel
+
+# Backend config returned by the API
+AnyBackendConfigWithoutCreds = Union[
+    AWSBackendConfig,
+    AzureBackendConfig,
+    CudoBackendConfig,
+    DataCrunchBackendConfig,
+    GCPBackendConfig,
+    KubernetesBackendConfig,
+    LambdaBackendConfig,
+    OCIBackendConfig,
+    RunpodBackendConfig,
+    TensorDockBackendConfig,
+    VastAIBackendConfig,
+    VultrBackendConfig,
+    DstackBackendConfig,
+    DstackBaseBackendConfig,
+]
+
+# Same as AnyBackendConfigWithoutCreds but also includes creds.
+# Used to create/update backend.
+# Also returned by the API to project admins so that they can see/update backend creds.
+AnyBackendConfigWithCreds = Union[
+    AWSBackendConfigWithCreds,
+    AzureBackendConfigWithCreds,
+    CudoBackendConfigWithCreds,
+    DataCrunchBackendConfigWithCreds,
+    GCPBackendConfigWithCreds,
+    KubernetesBackendConfigWithCreds,
+    LambdaBackendConfigWithCreds,
+    OCIBackendConfigWithCreds,
+    RunpodBackendConfigWithCreds,
+    TensorDockBackendConfigWithCreds,
+    VastAIBackendConfigWithCreds,
+    VultrBackendConfigWithCreds,
+    DstackBackendConfig,
+]
+
+# Backend config accepted in server/config.yaml.
+# This can be different from the API config.
+# For example, it can make creds data optional and resolve it by filename.
+AnyBackendFileConfigWithCreds = Union[
+    AWSBackendConfigWithCreds,
+    AzureBackendConfigWithCreds,
+    CudoBackendConfigWithCreds,
+    DataCrunchBackendConfigWithCreds,
+    GCPBackendFileConfigWithCreds,
+    KubernetesBackendFileConfigWithCreds,
+    LambdaBackendConfigWithCreds,
+    OCIBackendConfigWithCreds,
+    RunpodBackendConfigWithCreds,
+    TensorDockBackendConfigWithCreds,
+    VastAIBackendConfigWithCreds,
+    VultrBackendConfigWithCreds,
+]
+
+
+# The API can return backend config with or without creds
+AnyBackendConfig = Union[AnyBackendConfigWithoutCreds, AnyBackendConfigWithCreds]
+
+
+# In case we'll support multiple backends of the same type,
+# this adds backend name to backend config.
+class BackendInfo(CoreModel):
+    name: str
+    config: AnyBackendConfigWithoutCreds
+
+
+class BackendInfoYAML(CoreModel):
+    name: str
+    config_yaml: str

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

@@ -1,15 +0,0 @@
-from dstack._internal.core.backends.base import Backend
-from dstack._internal.core.backends.oci.compute import OCICompute
-from dstack._internal.core.backends.oci.config import OCIConfig
-from dstack._internal.core.models.backends.base import BackendType
-
-
-class OCIBackend(Backend):
-    TYPE: BackendType = BackendType.OCI
-
-    def __init__(self, config: OCIConfig):
-        self.config = config
-        self._compute = OCICompute(self.config)
-
-    def compute(self) -> OCICompute:
-        return self._compute

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

@@ -2,7 +2,7 @@ import oci
 from typing_extensions import Any, Mapping
 
 from dstack._internal.core.backends.oci.exceptions import any_oci_exception
-from dstack._internal.core.models.backends.oci import AnyOCICreds, OCIDefaultCreds
+from dstack._internal.core.backends.oci.models import AnyOCICreds, OCIDefaultCreds
 from dstack._internal.core.models.common import is_core_model_instance
 
 
@@ -20,7 +20,3 @@ def creds_valid(creds: AnyOCICreds) -> bool:
     except any_oci_exception:
         return False
     return True
-
-
-def default_creds_available() -> bool:
-    return creds_valid(OCIDefaultCreds())

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

@@ -0,0 +1,16 @@
+from dstack._internal.core.backends.base.backend import Backend
+from dstack._internal.core.backends.oci.compute import OCICompute
+from dstack._internal.core.backends.oci.models import OCIConfig
+from dstack._internal.core.models.backends.base import BackendType
+
+
+class OCIBackend(Backend):
+    TYPE = BackendType.OCI
+    COMPUTE_CLASS = OCICompute
+
+    def __init__(self, config: OCIConfig):
+        self.config = config
+        self._compute = OCICompute(self.config)
+
+    def compute(self) -> OCICompute:
+        return self._compute

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

@@ -6,13 +6,14 @@ import oci
 
 from dstack._internal.core.backends.base.compute import (
     Compute,
+    ComputeWithCreateInstanceSupport,
+    ComputeWithMultinodeSupport,
     generate_unique_instance_name,
-    get_job_instance_name,
     get_user_data,
 )
 from dstack._internal.core.backends.base.offers import get_catalog_offers
 from dstack._internal.core.backends.oci import resources
-from dstack._internal.core.backends.oci.config import OCIConfig
+from dstack._internal.core.backends.oci.models import OCIConfig
 from dstack._internal.core.backends.oci.region import make_region_clients_map
 from dstack._internal.core.errors import NoCapacityError
 from dstack._internal.core.models.backends.base import BackendType
@@ -21,11 +22,9 @@ from dstack._internal.core.models.instances import (
     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
 
 SUPPORTED_SHAPE_FAMILIES = [
     "VM.Standard2.",
@@ -46,7 +45,11 @@ SUPPORTED_SHAPE_FAMILIES = [
 CONFIGURABLE_DISK_SIZE = Range[Memory](min=Memory.parse("50GB"), max=Memory.parse("32TB"))
 
 
-class OCICompute(Compute):
+class OCICompute(
+    ComputeWithCreateInstanceSupport,
+    ComputeWithMultinodeSupport,
+    Compute,
+):
     def __init__(self, config: OCIConfig):
         super().__init__()
         self.config = config
@@ -92,23 +95,6 @@ class OCICompute(Compute):
 
         return offers_with_availability
 
-    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=get_job_instance_name(run, job),
-            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/{server/services/backends/configurators/oci.py → core/backends/oci/configurator.py}

@@ -1,9 +1,23 @@
 import json
 from typing import Dict, Iterable, List, Set, Tuple
 
-from dstack._internal.core.backends.oci import OCIBackend, auth, resources
-from dstack._internal.core.backends.oci.config import OCIConfig
+from dstack._internal.core.backends.base.configurator import (
+    BackendRecord,
+    Configurator,
+    raise_invalid_credentials_error,
+)
+from dstack._internal.core.backends.oci import resources
+from dstack._internal.core.backends.oci.backend import OCIBackend
 from dstack._internal.core.backends.oci.exceptions import any_oci_exception
+from dstack._internal.core.backends.oci.models import (
+    AnyOCIBackendConfig,
+    OCIBackendConfig,
+    OCIBackendConfigWithCreds,
+    OCIConfig,
+    OCICreds,
+    OCIDefaultCreds,
+    OCIStoredConfig,
+)
 from dstack._internal.core.backends.oci.region import (
     get_subscribed_regions,
     make_region_client,
@@ -12,26 +26,8 @@ from dstack._internal.core.backends.oci.region import (
 from dstack._internal.core.errors import ServerClientError
 from dstack._internal.core.models.backends.base import (
     BackendType,
-    ConfigElementValue,
-    ConfigMultiElement,
-)
-from dstack._internal.core.models.backends.oci import (
-    AnyOCIConfigInfo,
-    OCIConfigInfo,
-    OCIConfigInfoWithCreds,
-    OCIConfigInfoWithCredsPartial,
-    OCIConfigValues,
-    OCICreds,
-    OCIDefaultCreds,
-    OCIStoredConfig,
 )
 from dstack._internal.core.models.common import is_core_model_instance
-from dstack._internal.server import settings
-from dstack._internal.server.models import BackendModel, DecryptedString, ProjectModel
-from dstack._internal.server.services.backends.configurators.base import (
-    Configurator,
-    raise_invalid_credentials_error,
-)
 
 # where dstack images are published
 SUPPORTED_REGIONS = frozenset(
@@ -48,56 +44,23 @@ SUPPORTED_REGIONS = frozenset(
 
 
 class OCIConfigurator(Configurator):
-    TYPE: BackendType = BackendType.OCI
-
-    def get_default_configs(self) -> List[OCIConfigInfoWithCreds]:
-        creds = OCIDefaultCreds()
-        try:
-            subscribed_regions = get_subscribed_regions(creds).names
-        except any_oci_exception:
-            return []
-        return [
-            OCIConfigInfoWithCreds(
-                regions=list(subscribed_regions & SUPPORTED_REGIONS),
-                creds=creds,
-            )
-        ]
+    TYPE = BackendType.OCI
+    BACKEND_CLASS = OCIBackend
 
-    def get_config_values(self, config: OCIConfigInfoWithCredsPartial) -> OCIConfigValues:
-        config_values = OCIConfigValues(regions=None)
-        config_values.default_creds = (
-            settings.DEFAULT_CREDS_ENABLED and auth.default_creds_available()
-        )
-        if config.creds is None:
-            return config_values
-        if (
-            is_core_model_instance(config.creds, OCIDefaultCreds)
-            and not settings.DEFAULT_CREDS_ENABLED
-        ):
+    def validate_config(self, config: OCIBackendConfigWithCreds, default_creds_enabled: bool):
+        if is_core_model_instance(config.creds, OCIDefaultCreds) and not default_creds_enabled:
             raise_invalid_credentials_error(
                 fields=[["creds"]],
                 details="Default credentials are forbidden by dstack settings",
             )
-
         try:
-            available_regions = get_subscribed_regions(config.creds).names & SUPPORTED_REGIONS
+            get_subscribed_regions(config.creds).names
         except any_oci_exception as e:
             raise_invalid_credentials_error(fields=[["creds"]], details=e)
 
-        if config.regions:
-            selected_regions = [r for r in config.regions if r in available_regions]
-        else:
-            selected_regions = list(available_regions)
-
-        config_values.regions = self._get_regions_element(
-            available=available_regions,
-            selected=selected_regions,
-        )
-        return config_values
-
     def create_backend(
-        self, project: ProjectModel, config: OCIConfigInfoWithCreds
-    ) -> BackendModel:
+        self, project_name: str, config: OCIBackendConfigWithCreds
+    ) -> BackendRecord:
         try:
             subscribed_regions = get_subscribed_regions(config.creds)
         except any_oci_exception as e:
@@ -109,44 +72,36 @@ class OCIConfigurator(Configurator):
             _raise_if_regions_unavailable(config.regions, subscribed_regions.names)
 
         compartment_id, subnet_ids_per_region = _create_resources(
-            project, config, subscribed_regions.home_region_name
+            project_name, config, subscribed_regions.home_region_name
         )
         config.compartment_id = compartment_id
         stored_config = OCIStoredConfig.__response__(
             **config.dict(), subnet_ids_per_region=subnet_ids_per_region
         )
 
-        return BackendModel(
-            project_id=project.id,
-            type=self.TYPE.value,
+        return BackendRecord(
             config=stored_config.json(),
-            auth=DecryptedString(plaintext=OCICreds.parse_obj(config.creds).json()),
+            auth=OCICreds.parse_obj(config.creds).json(),
         )
 
-    def get_config_info(self, model: BackendModel, include_creds: bool) -> AnyOCIConfigInfo:
-        config = self._get_backend_config(model)
+    def get_backend_config(
+        self, record: BackendRecord, include_creds: bool
+    ) -> AnyOCIBackendConfig:
+        config = self._get_config(record)
         if include_creds:
-            return OCIConfigInfoWithCreds.__response__.parse_obj(config)
-        return OCIConfigInfo.__response__.parse_obj(config)
+            return OCIBackendConfigWithCreds.__response__.parse_obj(config)
+        return OCIBackendConfig.__response__.parse_obj(config)
 
-    def get_backend(self, model: BackendModel) -> OCIBackend:
-        config = self._get_backend_config(model)
+    def get_backend(self, record: BackendRecord) -> OCIBackend:
+        config = self._get_config(record)
         return OCIBackend(config=config)
 
-    def _get_backend_config(self, model: BackendModel) -> OCIConfig:
+    def _get_config(self, record: BackendRecord) -> OCIConfig:
         return OCIConfig.__response__(
-            **json.loads(model.config),
-            creds=OCICreds.parse_raw(model.auth.get_plaintext_or_error()).__root__,
+            **json.loads(record.config),
+            creds=OCICreds.parse_raw(record.auth).__root__,
         )
 
-    def _get_regions_element(
-        self, available: Iterable[str], selected: List[str]
-    ) -> ConfigMultiElement:
-        element = ConfigMultiElement(selected=selected)
-        for region in available:
-            element.values.append(ConfigElementValue(value=region, label=region))
-        return element
-
 
 def _filter_supported_regions(subscribed_region_names: Set[str]) -> List[str]:
     available_regions = subscribed_region_names & SUPPORTED_REGIONS
@@ -178,13 +133,13 @@ def _raise_if_regions_unavailable(
 
 
 def _create_resources(
-    project: ProjectModel, config: OCIConfigInfoWithCreds, home_region: str
+    project_name: str, config: OCIBackendConfigWithCreds, home_region: str
 ) -> Tuple[str, Dict[str, str]]:
     compartment_id = config.compartment_id
     if not compartment_id:
         home_region_client = make_region_client(home_region, config.creds)
         compartment_id = resources.get_or_create_compartment(
-            f"dstack-{project.name}",
+            f"dstack-{project_name}",
             home_region_client.client_config["tenancy"],
             home_region_client.identity_client,
         ).id
@@ -192,7 +147,7 @@ def _create_resources(
     region_clients = make_region_clients_map(config.regions, config.creds)
     resources.wait_until_compartment_active(compartment_id, region_clients)
     subnets_per_region = resources.set_up_network_resources(
-        compartment_id, project.name, region_clients
+        compartment_id, project_name, region_clients
     )
 
     return compartment_id, subnets_per_region