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

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

@@ -19,39 +19,34 @@ from azure.mgmt.network.models import (
 )
 from azure.mgmt.resource.resources.models import ResourceGroup
 
-from dstack._internal.core.backends.azure import AzureBackend, auth, compute, resources
+from dstack._internal.core.backends.azure import auth, compute, 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.errors import (
-    BackendAuthError,
-    BackendError,
-    ServerClientError,
-)
-from dstack._internal.core.models.backends.azure import (
-    AnyAzureConfigInfo,
+from dstack._internal.core.backends.azure.backend import AzureBackend
+from dstack._internal.core.backends.azure.models import (
+    AnyAzureBackendConfig,
+    AzureBackendConfig,
+    AzureBackendConfigWithCreds,
     AzureClientCreds,
-    AzureConfigInfo,
-    AzureConfigInfoWithCreds,
-    AzureConfigInfoWithCredsPartial,
-    AzureConfigValues,
+    AzureConfig,
     AzureCreds,
     AzureDefaultCreds,
     AzureStoredConfig,
 )
-from dstack._internal.core.models.backends.base import (
-    BackendType,
-    ConfigElement,
-    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 (
+    BackendAuthError,
+    BackendError,
+    ServerClientError,
+)
+from dstack._internal.core.models.backends.base import (
+    BackendType,
+)
+from dstack._internal.core.models.common import is_core_model_instance
 
 LOCATIONS = [
     ("(US) Central US", "centralus"),
@@ -77,47 +72,16 @@ MAIN_LOCATION = "eastus"
 
 
 class AzureConfigurator(Configurator):
-    TYPE: BackendType = BackendType.AZURE
-
-    def get_default_configs(self) -> List[AzureConfigInfoWithCreds]:
-        if not auth.default_creds_available():
-            return []
-        try:
-            credential, _ = auth.authenticate(AzureDefaultCreds())
-        except BackendAuthError:
-            return []
-        tenant_id_element = self._get_tenant_id_element(credential=credential)
-        tenant_ids = [v.value for v in tenant_id_element.values]
-        subscription_id_element = self._get_subscription_id_element(credential=credential)
-        subscription_ids = [v.value for v in subscription_id_element.values]
-        configs = []
-        for tenant_id in tenant_ids:
-            for subscription_id in subscription_ids:
-                config = AzureConfigInfoWithCreds(
-                    tenant_id=tenant_id,
-                    subscription_id=subscription_id,
-                    locations=DEFAULT_LOCATIONS,
-                    creds=AzureDefaultCreds(),
-                )
-                configs.append(config)
-        return configs
+    TYPE = BackendType.AZURE
+    BACKEND_CLASS = AzureBackend
 
-    def get_config_values(self, config: AzureConfigInfoWithCredsPartial) -> AzureConfigValues:
-        config_values = AzureConfigValues()
-        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, AzureDefaultCreds)
-            and not settings.DEFAULT_CREDS_ENABLED
-        ):
+    def validate_config(self, config: AzureBackendConfigWithCreds, default_creds_enabled: bool):
+        if is_core_model_instance(config.creds, AzureDefaultCreds) and not default_creds_enabled:
             raise_invalid_credentials_error(fields=[["creds"]])
         if is_core_model_instance(config.creds, AzureClientCreds):
             self._set_client_creds_tenant_id(config.creds, config.tenant_id)
         try:
-            credential, creds_tenant_id = auth.authenticate(config.creds)
+            credential, _ = auth.authenticate(config.creds)
         except BackendAuthError:
             if is_core_model_instance(config.creds, AzureClientCreds):
                 raise_invalid_credentials_error(
@@ -129,29 +93,18 @@ class AzureConfigurator(Configurator):
                 )
             else:
                 raise_invalid_credentials_error(fields=[["creds"]])
-        config_values.tenant_id = self._get_tenant_id_element(
-            credential=credential,
-            selected=config.tenant_id or creds_tenant_id,
-        )
-        if config_values.tenant_id.selected is None:
-            return config_values
-        config_values.subscription_id = self._get_subscription_id_element(
-            credential=credential,
-            selected=config.subscription_id,
-        )
-        if config_values.subscription_id.selected is None:
-            return config_values
-        config_values.locations = self._get_locations_element(
-            selected=config.locations or DEFAULT_LOCATIONS
-        )
-        self._check_config(config=config, credential=credential)
-        return config_values
+        self._check_config_tenant_id(config=config, credential=credential)
+        self._check_config_subscription_id(config=config, credential=credential)
+        self._check_config_locations(config)
+        self._check_config_tags(config)
+        self._check_config_resource_group(config=config, credential=credential)
+        self._check_config_vpc(config=config, credential=credential)
 
     def create_backend(
-        self, project: ProjectModel, config: AzureConfigInfoWithCreds
-    ) -> BackendModel:
-        if config.locations is None:
-            config.locations = DEFAULT_LOCATIONS
+        self, project_name: str, config: AzureBackendConfigWithCreds
+    ) -> BackendRecord:
+        if config.regions is None:
+            config.regions = DEFAULT_LOCATIONS
         if is_core_model_instance(config.creds, AzureClientCreds):
             self._set_client_creds_tenant_id(config.creds, config.tenant_id)
         credential, _ = auth.authenticate(config.creds)
@@ -160,173 +113,85 @@ class AzureConfigurator(Configurator):
                 credential=credential,
                 subscription_id=config.subscription_id,
                 location=MAIN_LOCATION,
-                project_name=project.name,
+                project_name=project_name,
             )
         self._create_network_resources(
             credential=credential,
             subscription_id=config.subscription_id,
             resource_group=config.resource_group,
-            locations=config.locations,
+            locations=config.regions,
             create_default_network=config.vpc_ids is None,
         )
-        return BackendModel(
-            project_id=project.id,
-            type=self.TYPE.value,
+        return BackendRecord(
             config=AzureStoredConfig(
-                **AzureConfigInfo.__response__.parse_obj(config).dict(),
+                **AzureBackendConfig.__response__.parse_obj(config).dict()
             ).json(),
-            auth=DecryptedString(plaintext=AzureCreds.parse_obj(config.creds).__root__.json()),
+            auth=AzureCreds.parse_obj(config.creds).__root__.json(),
         )
 
-    def get_config_info(self, model: BackendModel, include_creds: bool) -> AnyAzureConfigInfo:
-        config = self._get_backend_config(model)
+    def get_backend_config(
+        self, record: BackendRecord, include_creds: bool
+    ) -> AnyAzureBackendConfig:
+        config = self._get_config(record)
         if include_creds:
-            return AzureConfigInfoWithCreds.__response__.parse_obj(config)
-        return AzureConfigInfo.__response__.parse_obj(config)
+            return AzureBackendConfigWithCreds.__response__.parse_obj(config)
+        return AzureBackendConfig.__response__.parse_obj(config)
 
-    def get_backend(self, model: BackendModel) -> AzureBackend:
-        config = self._get_backend_config(model)
+    def get_backend(self, record: BackendRecord) -> AzureBackend:
+        config = self._get_config(record)
         return AzureBackend(config=config)
 
-    def _get_backend_config(self, model: BackendModel) -> AzureConfig:
+    def _get_config(self, record: BackendRecord) -> AzureConfig:
+        config_dict = json.loads(record.config)
+        regions = config_dict.pop("regions", None)
+        if regions is None:
+            # Legacy config stores regions as locations
+            regions = config_dict.pop("locations")
         return AzureConfig.__response__(
-            **json.loads(model.config),
-            creds=AzureCreds.parse_raw(model.auth.get_plaintext_or_error()).__root__,
+            **config_dict,
+            regions=regions,
+            creds=AzureCreds.parse_raw(record.auth).__root__,
         )
 
-    def _set_client_creds_tenant_id(
-        self,
-        creds: AzureClientCreds,
-        tenant_id: Optional[str],
+    def _check_config_tenant_id(
+        self, config: AzureBackendConfigWithCreds, credential: auth.AzureCredential
     ):
-        if creds.tenant_id is not None:
-            return
-        if tenant_id is None:
-            raise_invalid_credentials_error(
-                fields=[
-                    ["creds", "tenant_id"],
-                    ["tenant_id"],
-                ]
-            )
-        creds.tenant_id = tenant_id
-
-    def _get_tenant_id_element(
-        self,
-        credential: auth.AzureCredential,
-        selected: Optional[str] = None,
-    ) -> ConfigElement:
         subscription_client = subscription_mgmt.SubscriptionClient(credential=credential)
-        element = ConfigElement(selected=selected)
         tenant_ids = []
         for tenant in subscription_client.tenants.list():
             tenant_ids.append(tenant.tenant_id)
-            element.values.append(
-                ConfigElementValue(value=tenant.tenant_id, label=tenant.tenant_id)
-            )
-        if selected is not None and selected not in tenant_ids:
+        if config.tenant_id not in tenant_ids:
             raise ServerClientError(
                 "Invalid tenant_id",
                 fields=[["tenant_id"]],
             )
-        if len(tenant_ids) == 1:
-            element.selected = tenant_ids[0]
-        return element
 
-    def _get_subscription_id_element(
-        self,
-        credential: auth.AzureCredential,
-        selected: Optional[str] = None,
-    ) -> ConfigElement:
+    def _check_config_subscription_id(
+        self, config: AzureBackendConfigWithCreds, credential: auth.AzureCredential
+    ):
         subscription_client = subscription_mgmt.SubscriptionClient(credential=credential)
-        element = ConfigElement(selected=selected)
         subscription_ids = []
         for subscription in subscription_client.subscriptions.list():
             subscription_ids.append(subscription.subscription_id)
-            element.values.append(
-                ConfigElementValue(
-                    value=subscription.subscription_id,
-                    label=f"{subscription.display_name} ({subscription.subscription_id})",
-                )
-            )
-        if selected is not None and selected not in subscription_ids:
+        if config.subscription_id not in subscription_ids:
             raise ServerClientError(
                 "Invalid subscription_id",
                 fields=[["subscription_id"]],
             )
-        if len(subscription_ids) == 1:
-            element.selected = subscription_ids[0]
         if len(subscription_ids) == 0:
             # Credentials without granted roles don't see any subscriptions
             raise ServerClientError(
                 msg="No Azure subscriptions found for provided credentials. Ensure the account has enough permissions.",
             )
-        return element
-
-    def _get_locations_element(self, selected: List[str]) -> ConfigMultiElement:
-        element = ConfigMultiElement()
-        for loc in LOCATION_VALUES:
-            element.values.append(ConfigElementValue(value=loc, label=loc))
-        element.selected = selected
-        return element
-
-    def _create_resource_group(
-        self,
-        credential: auth.AzureCredential,
-        subscription_id: str,
-        location: str,
-        project_name: str,
-    ) -> str:
-        resource_manager = ResourceManager(
-            credential=credential,
-            subscription_id=subscription_id,
-        )
-        return resource_manager.create_resource_group(
-            name=_get_resource_group_name(project_name),
-            location=location,
-        )
-
-    def _create_network_resources(
-        self,
-        credential: auth.AzureCredential,
-        subscription_id: str,
-        resource_group: str,
-        locations: List[str],
-        create_default_network: bool,
-    ):
-        def func(location: str):
-            network_manager = NetworkManager(
-                credential=credential, subscription_id=subscription_id
-            )
-            if create_default_network:
-                network_manager.create_virtual_network(
-                    resource_group=resource_group,
-                    location=location,
-                    name=azure_utils.get_default_network_name(resource_group, location),
-                    subnet_name=azure_utils.get_default_subnet_name(resource_group, location),
-                )
-            network_manager.create_network_security_group(
-                resource_group=resource_group,
-                location=location,
-                name=azure_utils.get_default_network_security_group_name(resource_group, location),
-            )
-            network_manager.create_gateway_network_security_group(
-                resource_group=resource_group,
-                location=location,
-                name=azure_utils.get_gateway_network_security_group_name(resource_group, location),
-            )
-
-        with ThreadPoolExecutor(max_workers=8) as executor:
-            for location in locations:
-                executor.submit(func, location)
 
-    def _check_config(
-        self, config: AzureConfigInfoWithCredsPartial, credential: auth.AzureCredential
-    ):
-        self._check_tags_config(config)
-        self._check_resource_group(config=config, credential=credential)
-        self._check_vpc_config(config=config, credential=credential)
+    def _check_config_locations(self, config: AzureBackendConfigWithCreds):
+        if config.regions is None:
+            return
+        for location in config.regions:
+            if location not in LOCATION_VALUES:
+                raise ServerClientError(f"Unknown Azure location {location}")
 
-    def _check_tags_config(self, config: AzureConfigInfoWithCredsPartial):
+    def _check_config_tags(self, config: AzureBackendConfigWithCreds):
         if not config.tags:
             return
         if len(config.tags) > TAGS_MAX_NUM:
@@ -338,8 +203,8 @@ class AzureConfigurator(Configurator):
         except BackendError as e:
             raise ServerClientError(e.args[0])
 
-    def _check_resource_group(
-        self, config: AzureConfigInfoWithCredsPartial, credential: auth.AzureCredential
+    def _check_config_resource_group(
+        self, config: AzureBackendConfigWithCreds, credential: auth.AzureCredential
     ):
         if config.resource_group is None:
             return
@@ -350,22 +215,22 @@ class AzureConfigurator(Configurator):
         if not resource_manager.resource_group_exists(config.resource_group):
             raise ServerClientError(f"Resource group {config.resource_group} not found")
 
-    def _check_vpc_config(
-        self, config: AzureConfigInfoWithCredsPartial, credential: auth.AzureCredential
+    def _check_config_vpc(
+        self, config: AzureBackendConfigWithCreds, credential: auth.AzureCredential
     ):
         if config.subscription_id is None:
             return None
         allocate_public_ip = config.public_ips if config.public_ips is not None else True
         if config.public_ips is False and config.vpc_ids is None:
             raise ServerClientError(msg="`vpc_ids` must be specified if `public_ips: false`.")
-        locations = config.locations
+        locations = config.regions
         if locations is None:
             locations = DEFAULT_LOCATIONS
         if config.vpc_ids is not None:
             vpc_ids_locations = list(config.vpc_ids.keys())
             not_configured_locations = [loc for loc in locations if loc not in vpc_ids_locations]
             if len(not_configured_locations) > 0:
-                if config.locations is None:
+                if config.regions is None:
                     raise ServerClientError(
                         f"`vpc_ids` not configured for regions {not_configured_locations}. "
                         "Configure `vpc_ids` for all regions or specify `regions`."
@@ -396,6 +261,72 @@ class AzureConfigurator(Configurator):
                     except BackendError as e:
                         raise ServerClientError(e.args[0])
 
+    def _set_client_creds_tenant_id(
+        self,
+        creds: AzureClientCreds,
+        tenant_id: Optional[str],
+    ):
+        if creds.tenant_id is not None:
+            return
+        if tenant_id is None:
+            raise_invalid_credentials_error(
+                fields=[
+                    ["creds", "tenant_id"],
+                    ["tenant_id"],
+                ]
+            )
+        creds.tenant_id = tenant_id
+
+    def _create_resource_group(
+        self,
+        credential: auth.AzureCredential,
+        subscription_id: str,
+        location: str,
+        project_name: str,
+    ) -> str:
+        resource_manager = ResourceManager(
+            credential=credential,
+            subscription_id=subscription_id,
+        )
+        return resource_manager.create_resource_group(
+            name=_get_resource_group_name(project_name),
+            location=location,
+        )
+
+    def _create_network_resources(
+        self,
+        credential: auth.AzureCredential,
+        subscription_id: str,
+        resource_group: str,
+        locations: List[str],
+        create_default_network: bool,
+    ):
+        def func(location: str):
+            network_manager = NetworkManager(
+                credential=credential, subscription_id=subscription_id
+            )
+            if create_default_network:
+                network_manager.create_virtual_network(
+                    resource_group=resource_group,
+                    location=location,
+                    name=azure_utils.get_default_network_name(resource_group, location),
+                    subnet_name=azure_utils.get_default_subnet_name(resource_group, location),
+                )
+            network_manager.create_network_security_group(
+                resource_group=resource_group,
+                location=location,
+                name=azure_utils.get_default_network_security_group_name(resource_group, location),
+            )
+            network_manager.create_gateway_network_security_group(
+                resource_group=resource_group,
+                location=location,
+                name=azure_utils.get_gateway_network_security_group_name(resource_group, location),
+            )
+
+        with ThreadPoolExecutor(max_workers=8) as executor:
+            for location in locations:
+                executor.submit(func, location)
+
 
 def _get_resource_group_name(project_name: str) -> str:
     return f"dstack-{project_name}"

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

@@ -0,0 +1,89 @@
+from typing import Annotated, Dict, List, Literal, Optional, Union
+
+from pydantic import Field
+
+from dstack._internal.core.models.common import CoreModel
+
+
+class AzureClientCreds(CoreModel):
+    type: Annotated[Literal["client"], Field(description="The type of credentials")] = "client"
+    client_id: Annotated[str, Field(description="The client ID")]
+    client_secret: Annotated[str, Field(description="The client secret")]
+    # if tenant_id is missing, it will be populated from config info
+    tenant_id: Optional[str]
+
+
+class AzureDefaultCreds(CoreModel):
+    type: Annotated[Literal["default"], Field(description="The type of credentials")] = "default"
+
+
+AnyAzureCreds = Union[AzureClientCreds, AzureDefaultCreds]
+
+
+class AzureCreds(CoreModel):
+    __root__: AnyAzureCreds = Field(..., discriminator="type")
+
+
+class AzureBackendConfig(CoreModel):
+    type: Annotated[Literal["azure"], Field(description="The type of the backend")] = "azure"
+    tenant_id: Annotated[str, Field(description="The tenant ID")]
+    subscription_id: Annotated[str, Field(description="The subscription ID")]
+    resource_group: Annotated[
+        Optional[str],
+        Field(
+            description=(
+                "The resource group for resources created by `dstack`."
+                " If not specified, `dstack` will create a new resource group"
+            )
+        ),
+    ] = None
+    regions: Annotated[
+        Optional[List[str]],
+        Field(description="The list of Azure regions (locations). Omit to use all regions"),
+    ] = None
+    vpc_ids: Annotated[
+        Optional[Dict[str, str]],
+        Field(
+            description=(
+                "The mapping from configured Azure locations to network IDs."
+                " A network ID must have a format `networkResourceGroup/networkName`"
+                " If not specified, `dstack` will create a new network for every configured region"
+            )
+        ),
+    ] = None
+    public_ips: Annotated[
+        Optional[bool],
+        Field(
+            description=(
+                "A flag to enable/disable public IP assigning on instances."
+                " `public_ips: false` requires `vpc_ids` that specifies custom networks with outbound internet connectivity"
+                " provided by NAT Gateway or other mechanism."
+                " Defaults to `true`"
+            )
+        ),
+    ] = None
+    tags: Annotated[
+        Optional[Dict[str, str]],
+        Field(description="The tags that will be assigned to resources created by `dstack`"),
+    ] = None
+
+
+class AzureBackendConfigWithCreds(AzureBackendConfig):
+    creds: AnyAzureCreds = Field(..., description="The credentials", discriminator="type")
+
+
+AnyAzureBackendConfig = Union[AzureBackendConfig, AzureBackendConfigWithCreds]
+
+
+class AzureStoredConfig(AzureBackendConfig):
+    resource_group: str = ""
+
+
+class AzureConfig(AzureStoredConfig):
+    creds: AnyAzureCreds
+
+    @property
+    def allocate_public_ips(self) -> bool:
+        if self.public_ips is not None:
+            return self.public_ips
+        return True

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

@@ -1,12 +0,0 @@
-from abc import ABC, abstractmethod
-
-from dstack._internal.core.backends.base.compute import Compute
-from dstack._internal.core.models.backends.base import BackendType
-
-
-class Backend(ABC):
-    TYPE: BackendType
-
-    @abstractmethod
-    def compute(self) -> Compute:
-        pass

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

@@ -0,0 +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: 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