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

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

@@ -8,20 +8,25 @@ from pydantic import ValidationError
 
 import dstack._internal.core.backends.aws.resources as aws_resources
 from dstack._internal import settings
-from dstack._internal.core.backends.aws.config import AWSConfig
+from dstack._internal.core.backends.aws.models import AWSAccessKeyCreds, AWSConfig
 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,
     get_gateway_user_data,
-    get_job_instance_name,
     get_user_data,
     merge_tags,
 )
 from dstack._internal.core.backends.base.offers import get_catalog_offers
 from dstack._internal.core.errors import ComputeError, NoCapacityError, PlacementGroupInUseError
-from dstack._internal.core.models.backends.aws import AWSAccessKeyCreds
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.common import CoreModel, is_core_model_instance
 from dstack._internal.core.models.gateways import (
@@ -33,11 +38,10 @@ from dstack._internal.core.models.instances import (
     InstanceConfiguration,
     InstanceOffer,
     InstanceOfferWithAvailability,
-    SSHKey,
 )
 from dstack._internal.core.models.placement import PlacementGroup, PlacementGroupProvisioningData
 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,
@@ -62,7 +66,16 @@ class AWSVolumeBackendData(CoreModel):
     iops: int
 
 
-class AWSCompute(Compute):
+class AWSCompute(
+    ComputeWithCreateInstanceSupport,
+    ComputeWithMultinodeSupport,
+    ComputeWithReservationSupport,
+    ComputeWithPlacementGroupSupport,
+    ComputeWithGatewaySupport,
+    ComputeWithPrivateGatewaySupport,
+    ComputeWithVolumeSupport,
+    Compute,
+):
     def __init__(self, config: AWSConfig):
         super().__init__()
         self.config = config
@@ -270,44 +283,6 @@ class AWSCompute(Compute):
                 continue
         raise NoCapacityError()
 
-    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_placement_group(
         self,
         placement_group: PlacementGroup,
@@ -382,7 +357,7 @@ class AWSCompute(Compute):
             **aws_resources.create_instances_struct(
                 disk_size=10,
                 image_id=aws_resources.get_gateway_image_id(ec2_client),
-                instance_type="t2.micro",
+                instance_type="t3.micro",
                 iam_instance_profile=None,
                 user_data=get_gateway_user_data(configuration.ssh_key_pub),
                 tags=tags,

dstack/_internal/{server/services/backends/configurators/aws.py → core/backends/aws/configurator.py}

@@ -1,41 +1,35 @@
 import concurrent.futures
 import json
-from typing import List
 
 import botocore.exceptions
 from boto3.session import Session
 
-from dstack._internal.core.backends.aws import AWSBackend, auth, compute, resources
-from dstack._internal.core.backends.aws.config import AWSConfig
-from dstack._internal.core.errors import (
-    BackendAuthError,
-    BackendError,
-    ServerClientError,
-)
-from dstack._internal.core.models.backends.aws import (
-    AnyAWSConfigInfo,
+from dstack._internal.core.backends.aws import auth, compute, resources
+from dstack._internal.core.backends.aws.backend import AWSBackend
+from dstack._internal.core.backends.aws.models import (
+    AnyAWSBackendConfig,
     AWSAccessKeyCreds,
-    AWSConfigInfo,
-    AWSConfigInfoWithCreds,
-    AWSConfigInfoWithCredsPartial,
-    AWSConfigValues,
+    AWSBackendConfig,
+    AWSBackendConfigWithCreds,
+    AWSConfig,
     AWSCreds,
     AWSDefaultCreds,
     AWSStoredConfig,
 )
-from dstack._internal.core.models.backends.base import (
-    BackendType,
-    ConfigElementValue,
-    ConfigMultiElement,
-)
-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 (
+from dstack._internal.core.backends.base.configurator import (
     TAGS_MAX_NUM,
+    BackendRecord,
     Configurator,
     raise_invalid_credentials_error,
 )
+from dstack._internal.core.errors import (
+    BackendError,
+    ServerClientError,
+)
+from dstack._internal.core.models.backends.base import (
+    BackendType,
+)
+from dstack._internal.core.models.common import is_core_model_instance
 from dstack._internal.utils.logging import get_logger
 
 logger = get_logger(__name__)
@@ -59,33 +53,11 @@ MAIN_REGION = "us-east-1"
 
 
 class AWSConfigurator(Configurator):
-    TYPE: BackendType = BackendType.AWS
+    TYPE = BackendType.AWS
+    BACKEND_CLASS = AWSBackend
 
-    def get_default_configs(self) -> List[AWSConfigInfoWithCreds]:
-        if not auth.default_creds_available():
-            return []
-        try:
-            auth.authenticate(creds=AWSDefaultCreds(), region=MAIN_REGION)
-        except BackendAuthError:
-            return []
-        return [
-            AWSConfigInfoWithCreds(
-                regions=DEFAULT_REGIONS,
-                creds=AWSDefaultCreds(),
-            )
-        ]
-
-    def get_config_values(self, config: AWSConfigInfoWithCredsPartial) -> AWSConfigValues:
-        config_values = AWSConfigValues(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, AWSDefaultCreds)
-            and not settings.DEFAULT_CREDS_ENABLED
-        ):
+    def validate_config(self, config: AWSBackendConfigWithCreds, default_creds_enabled: bool):
+        if is_core_model_instance(config.creds, AWSDefaultCreds) and not default_creds_enabled:
             raise_invalid_credentials_error(fields=[["creds"]])
         try:
             session = auth.authenticate(creds=config.creds, region=MAIN_REGION)
@@ -99,52 +71,41 @@ class AWSConfigurator(Configurator):
                 )
             else:
                 raise_invalid_credentials_error(fields=[["creds"]])
-        config_values.regions = self._get_regions_element(
-            selected=config.regions or DEFAULT_REGIONS
-        )
-        self._check_config(session=session, config=config)
-        return config_values
+        self._check_config_tags(config)
+        self._check_config_iam_instance_profile(session, config)
+        self._check_config_vpc(session, config)
 
     def create_backend(
-        self, project: ProjectModel, config: AWSConfigInfoWithCreds
-    ) -> BackendModel:
+        self, project_name: str, config: AWSBackendConfigWithCreds
+    ) -> BackendRecord:
         if config.regions is None:
             config.regions = DEFAULT_REGIONS
-        return BackendModel(
-            project_id=project.id,
-            type=self.TYPE.value,
-            config=AWSStoredConfig(**AWSConfigInfo.__response__.parse_obj(config).dict()).json(),
-            auth=DecryptedString(plaintext=AWSCreds.parse_obj(config.creds).json()),
+        return BackendRecord(
+            config=AWSStoredConfig(
+                **AWSBackendConfig.__response__.parse_obj(config).dict()
+            ).json(),
+            auth=AWSCreds.parse_obj(config.creds).json(),
         )
 
-    def get_config_info(self, model: BackendModel, include_creds: bool) -> AnyAWSConfigInfo:
-        config = self._get_backend_config(model)
+    def get_backend_config(
+        self, record: BackendRecord, include_creds: bool
+    ) -> AnyAWSBackendConfig:
+        config = self._get_config(record)
         if include_creds:
-            return AWSConfigInfoWithCreds.__response__.parse_obj(config)
-        return AWSConfigInfo.__response__.parse_obj(config)
+            return AWSBackendConfigWithCreds.__response__.parse_obj(config)
+        return AWSBackendConfig.__response__.parse_obj(config)
 
-    def get_backend(self, model: BackendModel) -> AWSBackend:
-        config = self._get_backend_config(model)
+    def get_backend(self, record: BackendRecord) -> AWSBackend:
+        config = self._get_config(record)
         return AWSBackend(config=config)
 
-    def _get_backend_config(self, model: BackendModel) -> AWSConfig:
+    def _get_config(self, record: BackendRecord) -> AWSConfig:
         return AWSConfig.__response__(
-            **json.loads(model.config),
-            creds=AWSCreds.parse_raw(model.auth.get_plaintext_or_error()).__root__,
+            **json.loads(record.config),
+            creds=AWSCreds.parse_raw(record.auth).__root__,
         )
 
-    def _get_regions_element(self, selected: List[str]) -> ConfigMultiElement:
-        element = ConfigMultiElement(selected=selected)
-        for r in REGION_VALUES:
-            element.values.append(ConfigElementValue(value=r, label=r))
-        return element
-
-    def _check_config(self, session: Session, config: AWSConfigInfoWithCredsPartial):
-        self._check_tags_config(config)
-        self._check_iam_instance_profile_config(session, config)
-        self._check_vpc_config(session, config)
-
-    def _check_tags_config(self, config: AWSConfigInfoWithCredsPartial):
+    def _check_config_tags(self, config: AWSBackendConfigWithCreds):
         if not config.tags:
             return
         if len(config.tags) > TAGS_MAX_NUM:
@@ -156,8 +117,8 @@ class AWSConfigurator(Configurator):
         except BackendError as e:
             raise ServerClientError(e.args[0])
 
-    def _check_iam_instance_profile_config(
-        self, session: Session, config: AWSConfigInfoWithCredsPartial
+    def _check_config_iam_instance_profile(
+        self, session: Session, config: AWSBackendConfigWithCreds
     ):
         if config.iam_instance_profile is None:
             return
@@ -181,7 +142,7 @@ class AWSConfigurator(Configurator):
                 f"Failed to check IAM instance profile {config.iam_instance_profile}"
             )
 
-    def _check_vpc_config(self, session: Session, config: AWSConfigInfoWithCredsPartial):
+    def _check_config_vpc(self, session: Session, config: AWSBackendConfigWithCreds):
         allocate_public_ip = config.public_ips if config.public_ips is not None else True
         use_default_vpcs = config.default_vpcs if config.default_vpcs is not None else True
         if config.vpc_name is not None and config.vpc_ids is not None:

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

@@ -0,0 +1,135 @@
+from typing import Annotated, Dict, List, Literal, Optional, Union
+
+from pydantic import Field
+
+from dstack._internal.core.models.common import CoreModel
+
+
+class AWSOSImage(CoreModel):
+    name: Annotated[str, Field(description="The AMI name")]
+    owner: Annotated[
+        str,
+        Field(regex=r"^(\d{12}|self)$", description="The AMI owner, account ID or `self`"),
+    ] = "self"
+    user: Annotated[str, Field(description="The OS user for provisioning")]
+
+
+class AWSOSImageConfig(CoreModel):
+    cpu: Annotated[Optional[AWSOSImage], Field(description="The AMI used for CPU instances")] = (
+        None
+    )
+    nvidia: Annotated[
+        Optional[AWSOSImage], Field(description="The AMI used for NVIDIA GPU instances")
+    ] = None
+
+
+class AWSAccessKeyCreds(CoreModel):
+    type: Annotated[Literal["access_key"], Field(description="The type of credentials")] = (
+        "access_key"
+    )
+    access_key: Annotated[str, Field(description="The access key")]
+    secret_key: Annotated[str, Field(description="The secret key")]
+
+
+class AWSDefaultCreds(CoreModel):
+    type: Annotated[Literal["default"], Field(description="The type of credentials")] = "default"
+
+
+AnyAWSCreds = Union[AWSAccessKeyCreds, AWSDefaultCreds]
+
+
+class AWSCreds(CoreModel):
+    __root__: AnyAWSCreds = Field(..., discriminator="type")
+
+
+class AWSBackendConfig(CoreModel):
+    type: Annotated[Literal["aws"], Field(description="The type of the backend")] = "aws"
+    regions: Annotated[
+        Optional[List[str]], Field(description="The list of AWS regions. Omit to use all regions")
+    ] = None
+    vpc_name: Annotated[
+        Optional[str],
+        Field(
+            description=(
+                "The name of custom VPCs. All configured regions must have a VPC with this name."
+                " If your custom VPCs don't have names or have different names in different regions, use `vpc_ids` instead."
+            )
+        ),
+    ] = None
+    vpc_ids: Annotated[
+        Optional[Dict[str, str]],
+        Field(
+            description=(
+                "The mapping from AWS regions to VPC IDs."
+                " If `default_vpcs: true`, omitted regions will use default VPCs"
+            )
+        ),
+    ] = None
+    default_vpcs: Annotated[
+        Optional[bool],
+        Field(
+            description=(
+                "A flag to enable/disable using default VPCs in regions not configured by `vpc_ids`."
+                " Set to `false` if default VPCs should never be used."
+                " Defaults to `true`"
+            )
+        ),
+    ] = None
+    public_ips: Annotated[
+        Optional[bool],
+        Field(
+            description=(
+                "A flag to enable/disable public IP assigning on instances."
+                " `public_ips: false` requires at least one private subnet with outbound internet connectivity"
+                " provided by a NAT Gateway or a Transit Gateway."
+                " Defaults to `true`"
+            )
+        ),
+    ] = None
+    iam_instance_profile: Annotated[
+        Optional[str],
+        Field(
+            description=(
+                "The name of the IAM instance profile to associate with EC2 instances."
+                " You can also specify the IAM role name for roles created via the AWS console."
+                " AWS automatically creates an instance profile and gives it the same name as the role"
+            )
+        ),
+    ] = None
+    tags: Annotated[
+        Optional[Dict[str, str]],
+        Field(description="The tags that will be assigned to resources created by `dstack`"),
+    ] = None
+    os_images: Annotated[
+        Optional[AWSOSImageConfig],
+        Field(
+            description="The mapping of instance categories (CPU, NVIDIA GPU) to AMI configurations"
+        ),
+    ] = None
+
+
+class AWSBackendConfigWithCreds(AWSBackendConfig):
+    creds: AnyAWSCreds = Field(..., description="The credentials", discriminator="type")
+
+
+AnyAWSBackendConfig = Union[AWSBackendConfig, AWSBackendConfigWithCreds]
+
+
+class AWSStoredConfig(AWSBackendConfig):
+    pass
+
+
+class AWSConfig(AWSStoredConfig):
+    creds: AnyAWSCreds
+
+    @property
+    def allocate_public_ips(self) -> bool:
+        if self.public_ips is not None:
+            return self.public_ips
+        return True
+
+    @property
+    def use_default_vpcs(self) -> bool:
+        if self.default_vpcs is not None:
+            return self.default_vpcs
+        return True

dstack/_internal/core/backends/aws/resources.py

@@ -5,8 +5,8 @@ import botocore.client
 import botocore.exceptions
 
 import dstack.version as version
+from dstack._internal.core.backends.aws.models import AWSOSImageConfig
 from dstack._internal.core.errors import BackendError, ComputeError, ComputeResourceNotFoundError
-from dstack._internal.core.models.backends.aws import AWSOSImageConfig
 from dstack._internal.utils.logging import get_logger
 
 logger = get_logger(__name__)

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

@@ -1,20 +0,0 @@
-from dstack._internal.core.backends.azure import auth
-from dstack._internal.core.backends.azure.compute import AzureCompute
-from dstack._internal.core.backends.azure.config import AzureConfig
-from dstack._internal.core.backends.base import Backend
-from dstack._internal.core.models.backends.base import BackendType
-
-
-class AzureBackend(Backend):
-    TYPE: BackendType = BackendType.AZURE
-
-    def __init__(self, config: AzureConfig):
-        self.config = config
-        self.credential, _ = auth.authenticate(self.config.creds)
-        self._compute = AzureCompute(
-            config=self.config,
-            credential=self.credential,
-        )
-
-    def compute(self) -> AzureCompute:
-        return self._compute

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

@@ -4,12 +4,11 @@ from azure.core.exceptions import ClientAuthenticationError
 from azure.identity import ClientSecretCredential, DefaultAzureCredential
 from azure.mgmt.subscription import SubscriptionClient
 
-from dstack._internal.core.errors import BackendAuthError
-from dstack._internal.core.models.backends.azure import (
+from dstack._internal.core.backends.azure.models import (
     AnyAzureCreds,
     AzureClientCreds,
-    AzureDefaultCreds,
 )
+from dstack._internal.core.errors import BackendAuthError
 from dstack._internal.core.models.common import is_core_model_instance
 
 AzureCredential = Union[ClientSecretCredential, DefaultAzureCredential]
@@ -39,11 +38,3 @@ def check_credential(credential: AzureCredential):
         list(client.subscriptions.list())
     except ClientAuthenticationError:
         raise BackendAuthError()
-
-
-def default_creds_available() -> bool:
-    try:
-        authenticate(AzureDefaultCreds())
-    except BackendAuthError:
-        return False
-    return True

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

@@ -0,0 +1,21 @@
+from dstack._internal.core.backends.azure import auth
+from dstack._internal.core.backends.azure.compute import AzureCompute
+from dstack._internal.core.backends.azure.models import AzureConfig
+from dstack._internal.core.backends.base.backend import Backend
+from dstack._internal.core.models.backends.base import BackendType
+
+
+class AzureBackend(Backend):
+    TYPE = BackendType.AZURE
+    COMPUTE_CLASS = AzureCompute
+
+    def __init__(self, config: AzureConfig):
+        self.config = config
+        self.credential, _ = auth.authenticate(self.config.creds)
+        self._compute = AzureCompute(
+            config=self.config,
+            credential=self.credential,
+        )
+
+    def compute(self) -> AzureCompute:
+        return self._compute

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

@@ -36,13 +36,15 @@ from dstack import version
 from dstack._internal import settings
 from dstack._internal.core.backends.azure import resources as azure_resources
 from dstack._internal.core.backends.azure import utils as azure_utils
-from dstack._internal.core.backends.azure.config import AzureConfig
+from dstack._internal.core.backends.azure.models 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,
-    get_job_instance_name,
     get_user_data,
     merge_tags,
 )
@@ -59,11 +61,9 @@ from dstack._internal.core.models.instances import (
     InstanceOffer,
     InstanceOfferWithAvailability,
     InstanceType,
-    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(__name__)
@@ -71,7 +71,12 @@ logger = get_logger(__name__)
 CONFIGURABLE_DISK_SIZE = Range[Memory](min=Memory.parse("30GB"), max=Memory.parse("4095GB"))
 
 
-class AzureCompute(Compute):
+class AzureCompute(
+    ComputeWithCreateInstanceSupport,
+    ComputeWithMultinodeSupport,
+    ComputeWithGatewaySupport,
+    Compute,
+):
     def __init__(self, config: AzureConfig, credential: TokenCredential):
         super().__init__()
         self.config = config
@@ -88,14 +93,14 @@ class AzureCompute(Compute):
     ) -> List[InstanceOfferWithAvailability]:
         offers = get_catalog_offers(
             backend=BackendType.AZURE,
-            locations=self.config.locations,
+            locations=self.config.regions,
             requirements=requirements,
             configurable_disk_size=CONFIGURABLE_DISK_SIZE,
             extra_filter=_supported_instances,
         )
         offers_with_availability = _get_offers_with_availability(
             compute_client=self._compute_client,
-            config_locations=self.config.locations,
+            config_locations=self.config.regions,
             offers=offers,
         )
         return offers_with_availability
@@ -190,25 +195,6 @@ class AzureCompute(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=get_job_instance_name(run, job),  # 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
     ):