dstack 0.19.26__py3-none-any.whl → 0.19.28__py3-none-any.whl

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

@@ -0,0 +1,57 @@
+import json
+from typing import Optional
+
+from dstack._internal.core.backends.base.configurator import (
+    BackendRecord,
+    Configurator,
+)
+from dstack._internal.core.backends.digitalocean_base.backend import BaseDigitalOceanBackend
+from dstack._internal.core.backends.digitalocean_base.models import (
+    AnyBaseDigitalOceanCreds,
+    BaseDigitalOceanBackendConfig,
+    BaseDigitalOceanBackendConfigWithCreds,
+    BaseDigitalOceanConfig,
+    BaseDigitalOceanCreds,
+    BaseDigitalOceanStoredConfig,
+)
+
+
+class BaseDigitalOceanConfigurator(Configurator):
+    def validate_config(
+        self, config: BaseDigitalOceanBackendConfigWithCreds, default_creds_enabled: bool
+    ):
+        self._validate_creds(config.creds, config.project_name)
+
+    def create_backend(
+        self, project_name: str, config: BaseDigitalOceanBackendConfigWithCreds
+    ) -> BackendRecord:
+        return BackendRecord(
+            config=BaseDigitalOceanStoredConfig(
+                **BaseDigitalOceanBackendConfig.__response__.parse_obj(config).dict()
+            ).json(),
+            auth=BaseDigitalOceanCreds.parse_obj(config.creds).json(),
+        )
+
+    def get_backend_config_with_creds(
+        self, record: BackendRecord
+    ) -> BaseDigitalOceanBackendConfigWithCreds:
+        config = self._get_config(record)
+        return BaseDigitalOceanBackendConfigWithCreds.__response__.parse_obj(config)
+
+    def get_backend_config_without_creds(
+        self, record: BackendRecord
+    ) -> BaseDigitalOceanBackendConfig:
+        config = self._get_config(record)
+        return BaseDigitalOceanBackendConfig.__response__.parse_obj(config)
+
+    def get_backend(self, record: BackendRecord) -> BaseDigitalOceanBackend:
+        raise NotImplementedError("Subclasses must implement get_backend")
+
+    def _get_config(self, record: BackendRecord) -> BaseDigitalOceanConfig:
+        return BaseDigitalOceanConfig.__response__(
+            **json.loads(record.config),
+            creds=BaseDigitalOceanCreds.parse_raw(record.auth),
+        )
+
+    def _validate_creds(self, creds: AnyBaseDigitalOceanCreds, project_name: Optional[str] = None):
+        pass

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

@@ -0,0 +1,43 @@
+from typing import Annotated, List, Literal, Optional, Union
+
+from pydantic import Field
+
+from dstack._internal.core.models.common import CoreModel
+
+
+class BaseDigitalOceanAPIKeyCreds(CoreModel):
+    type: Annotated[Literal["api_key"], Field(description="The type of credentials")] = "api_key"
+    api_key: Annotated[str, Field(description="The API key")]
+
+
+AnyBaseDigitalOceanCreds = BaseDigitalOceanAPIKeyCreds
+BaseDigitalOceanCreds = AnyBaseDigitalOceanCreds
+
+
+class BaseDigitalOceanBackendConfig(CoreModel):
+    type: Annotated[
+        Literal["amddevcloud", "digitalocean"],
+        Field(description="The type of backend"),
+    ]
+    project_name: Annotated[Optional[str], Field(description="The name of the project")] = None
+    regions: Annotated[
+        Optional[List[str]],
+        Field(description="The list of regions. Omit to use all regions"),
+    ] = None
+
+
+class BaseDigitalOceanBackendConfigWithCreds(BaseDigitalOceanBackendConfig):
+    creds: Annotated[AnyBaseDigitalOceanCreds, Field(description="The credentials")]
+
+
+AnyBaseDigitalOceanBackendConfig = Union[
+    BaseDigitalOceanBackendConfig, BaseDigitalOceanBackendConfigWithCreds
+]
+
+
+class BaseDigitalOceanStoredConfig(BaseDigitalOceanBackendConfig):
+    pass
+
+
+class BaseDigitalOceanConfig(BaseDigitalOceanStoredConfig):
+    creds: AnyBaseDigitalOceanCreds

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

@@ -2,6 +2,7 @@ import concurrent.futures
 import json
 import threading
 from collections import defaultdict
+from dataclasses import dataclass
 from typing import Callable, Dict, List, Literal, Optional, Tuple
 
 import google.api_core.exceptions
@@ -285,16 +286,18 @@ class GCPCompute(
                 )
             raise NoCapacityError()
 
+        image = _get_image(
+            instance_type_name=instance_offer.instance.name,
+            cuda=len(instance_offer.instance.resources.gpus) > 0,
+        )
+
         for zone in zones:
             request = compute_v1.InsertInstanceRequest()
             request.zone = zone
             request.project = self.config.project_id
             request.instance_resource = gcp_resources.create_instance_struct(
                 disk_size=disk_size,
-                image_id=_get_image_id(
-                    instance_type_name=instance_offer.instance.name,
-                    cuda=len(instance_offer.instance.resources.gpus) > 0,
-                ),
+                image_id=image.id,
                 machine_type=instance_offer.instance.name,
                 accelerators=gcp_resources.get_accelerators(
                     project_id=self.config.project_id,
@@ -305,6 +308,7 @@ class GCPCompute(
                 user_data=_get_user_data(
                     authorized_keys=authorized_keys,
                     instance_type_name=instance_offer.instance.name,
+                    is_ufw_installed=image.is_ufw_installed,
                 ),
                 authorized_keys=authorized_keys,
                 labels=labels,
@@ -889,24 +893,41 @@ def _get_vpc_subnet(
     )
 
 
-def _get_image_id(instance_type_name: str, cuda: bool) -> str:
+@dataclass
+class GCPImage:
+    id: str
+    is_ufw_installed: bool
+
+
+def _get_image(instance_type_name: str, cuda: bool) -> GCPImage:
     if instance_type_name == "a3-megagpu-8g":
         image_name = "dstack-a3mega-5"
+        is_ufw_installed = False
     elif instance_type_name in ["a3-edgegpu-8g", "a3-highgpu-8g"]:
-        return "projects/cos-cloud/global/images/cos-105-17412-535-78"
+        return GCPImage(
+            id="projects/cos-cloud/global/images/cos-105-17412-535-78",
+            is_ufw_installed=False,
+        )
     elif cuda:
         image_name = f"dstack-cuda-{version.base_image}"
+        is_ufw_installed = True
     else:
         image_name = f"dstack-{version.base_image}"
+        is_ufw_installed = True
     image_name = image_name.replace(".", "-")
-    return f"projects/dstack/global/images/{image_name}"
+    return GCPImage(
+        id=f"projects/dstack/global/images/{image_name}",
+        is_ufw_installed=is_ufw_installed,
+    )
 
 
 def _get_gateway_image_id() -> str:
     return "projects/ubuntu-os-cloud/global/images/ubuntu-2204-jammy-v20230714"
 
 
-def _get_user_data(authorized_keys: List[str], instance_type_name: str) -> str:
+def _get_user_data(
+    authorized_keys: List[str], instance_type_name: str, is_ufw_installed: bool
+) -> str:
     base_path = None
     bin_path = None
     backend_shim_env = None
@@ -929,6 +950,9 @@ def _get_user_data(authorized_keys: List[str], instance_type_name: str) -> str:
         base_path=base_path,
         bin_path=bin_path,
         backend_shim_env=backend_shim_env,
+        # Instance-level firewall is optional on GCP. The main protection comes from GCP firewalls.
+        # So only set up instance-level firewall as an additional measure if ufw is available.
+        skip_firewall_setup=not is_ufw_installed,
     )
 
 

dstack/_internal/core/backends/hotaisle/api_client.py

@@ -16,46 +16,38 @@ class HotAisleAPIClient:
         self.team_handle = team_handle
 
     def validate_api_key(self) -> bool:
+        url = f"{API_URL}/user/"
         try:
-            self._validate_user_and_team()
-            return True
+            response = self._make_request("GET", url)
+            response.raise_for_status()
         except requests.HTTPError as e:
-            if e.response.status_code == 401:
-                raise_invalid_credentials_error(
-                    fields=[["creds", "api_key"]], details="Invalid API key"
-                )
-            elif e.response.status_code == 403:
-                raise_invalid_credentials_error(
-                    fields=[["creds", "api_key"]],
-                    details="Authenticated user does note have required permissions",
-                )
-            raise e
-        except ValueError as e:
-            error_message = str(e)
-            if "No Hot Aisle teams found" in error_message:
-                raise_invalid_credentials_error(
-                    fields=[["creds", "api_key"]],
-                    details="Valid API key but no teams found for this user",
-                )
-            elif "not found" in error_message:
-                raise_invalid_credentials_error(
-                    fields=[["team_handle"]], details=f"Team handle '{self.team_handle}' not found"
-                )
-            raise e
-
-    def _validate_user_and_team(self) -> None:
-        url = f"{API_URL}/user/"
-        response = self._make_request("GET", url)
-        response.raise_for_status()
-        user_data = response.json()
+            if e.response is not None:
+                if e.response.status_code == 401:
+                    raise_invalid_credentials_error(
+                        fields=[["creds", "api_key"]], details="Invalid API key"
+                    )
+                if e.response.status_code == 403:
+                    raise_invalid_credentials_error(
+                        fields=[["creds", "api_key"]],
+                        details="Authenticated user does not have required permissions",
+                    )
+            raise
 
-        teams = user_data.get("teams", [])
+        user_data = response.json()
+        teams = user_data["teams"]
         if not teams:
-            raise ValueError("No Hot Aisle teams found for this user")
+            raise_invalid_credentials_error(
+                fields=[["creds", "api_key"]],
+                details="Valid API key but no teams found for this user",
+            )
 
         available_teams = [team["handle"] for team in teams]
         if self.team_handle not in available_teams:
-            raise ValueError(f"Hot Aisle team '{self.team_handle}' not found.")
+            raise_invalid_credentials_error(
+                fields=[["team_handle"]],
+                details=f"Team handle '{self.team_handle}' not found",
+            )
+        return True
 
     def upload_ssh_key(self, public_key: str) -> bool:
         url = f"{API_URL}/user/ssh_keys/"

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

@@ -28,8 +28,6 @@ from dstack._internal.utils.logging import get_logger
 
 logger = get_logger(__name__)
 
-MAX_INSTANCE_NAME_LEN = 60
-
 
 INSTANCE_TYPE_SPECS = {
     "1x MI300X 8x Xeon Platinum 8462Y+": {
@@ -130,9 +128,7 @@ class HotAisleCompute(
             ssh_port=22,
             dockerized=True,
             ssh_proxy=None,
-            backend_data=HotAisleInstanceBackendData(
-                ip_address=vm_data["ip_address"], vm_id=vm_data["name"]
-            ).json(),
+            backend_data=HotAisleInstanceBackendData(ip_address=vm_data["ip_address"]).json(),
         )
 
     def update_provisioning_data(
@@ -217,7 +213,6 @@ def _run_ssh_command(hostname: str, ssh_private_key: str, command: str):
 
 class HotAisleInstanceBackendData(CoreModel):
     ip_address: str
-    vm_id: Optional[str] = None
 
     @classmethod
     def load(cls, raw: Optional[str]) -> "HotAisleInstanceBackendData":

dstack/_internal/core/backends/models.py

@@ -20,6 +20,10 @@ from dstack._internal.core.backends.datacrunch.models import (
     DataCrunchBackendConfig,
     DataCrunchBackendConfigWithCreds,
 )
+from dstack._internal.core.backends.digitalocean_base.models import (
+    BaseDigitalOceanBackendConfig,
+    BaseDigitalOceanBackendConfigWithCreds,
+)
 from dstack._internal.core.backends.dstack.models import (
     DstackBackendConfig,
     DstackBaseBackendConfig,
@@ -77,6 +81,7 @@ AnyBackendConfigWithoutCreds = Union[
     CloudRiftBackendConfig,
     CudoBackendConfig,
     DataCrunchBackendConfig,
+    BaseDigitalOceanBackendConfig,
     GCPBackendConfig,
     HotAisleBackendConfig,
     KubernetesBackendConfig,
@@ -100,6 +105,7 @@ AnyBackendConfigWithCreds = Union[
     CloudRiftBackendConfigWithCreds,
     CudoBackendConfigWithCreds,
     DataCrunchBackendConfigWithCreds,
+    BaseDigitalOceanBackendConfigWithCreds,
     GCPBackendConfigWithCreds,
     HotAisleBackendConfigWithCreds,
     KubernetesBackendConfigWithCreds,
@@ -122,6 +128,7 @@ AnyBackendFileConfigWithCreds = Union[
     CloudRiftBackendConfigWithCreds,
     CudoBackendConfigWithCreds,
     DataCrunchBackendConfigWithCreds,
+    BaseDigitalOceanBackendConfigWithCreds,
     GCPBackendFileConfigWithCreds,
     HotAisleBackendFileConfigWithCreds,
     KubernetesBackendFileConfigWithCreds,

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

@@ -59,13 +59,6 @@ DOCKER_DAEMON_CONFIG = {
     "exec-opts": ["native.cgroupdriver=cgroupfs"],
 }
 SETUP_COMMANDS = [
-    "ufw allow ssh",
-    "ufw allow from 10.0.0.0/8",
-    "ufw allow from 172.16.0.0/12",
-    "ufw allow from 192.168.0.0/16",
-    "ufw default deny incoming",
-    "ufw default allow outgoing",
-    "ufw enable",
     'sed -i "s/.*AllowTcpForwarding.*/AllowTcpForwarding yes/g" /etc/ssh/sshd_config',
     "service ssh restart",
     f"echo {shlex.quote(json.dumps(DOCKER_DAEMON_CONFIG))} > /etc/docker/daemon.json",

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

@@ -135,11 +135,10 @@ class OCICompute(
             security_group.id, region.virtual_network_client
         )
 
-        setup_commands = [
-            f"sudo iptables -I INPUT -s {resources.VCN_CIDR} -j ACCEPT",
-            "sudo netfilter-persistent save",
-        ]
-        cloud_init_user_data = get_user_data(instance_config.get_public_keys(), setup_commands)
+        cloud_init_user_data = get_user_data(
+            authorized_keys=instance_config.get_public_keys(),
+            firewall_allow_from_subnets=[resources.VCN_CIDR],
+        )
 
         display_name = generate_unique_instance_name(instance_config)
         try:

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

@@ -75,17 +75,13 @@ class VultrCompute(
         subnet = vpc["v4_subnet"]
         subnet_mask = vpc["v4_subnet_mask"]
 
-        setup_commands = [
-            f"sudo ufw allow from {subnet}/{subnet_mask}",
-            "sudo ufw reload",
-        ]
         instance_id = self.api_client.launch_instance(
             region=instance_offer.region,
             label=instance_name,
             plan=instance_offer.instance.name,
             user_data=get_user_data(
                 authorized_keys=instance_config.get_public_keys(),
-                backend_specific_commands=setup_commands,
+                firewall_allow_from_subnets=[f"{subnet}/{subnet_mask}"],
             ),
             vpc_id=vpc["id"],
         )

dstack/_internal/core/compatibility/fleets.py

@@ -59,6 +59,11 @@ def get_fleet_spec_excludes(fleet_spec: FleetSpec) -> Optional[IncludeExcludeDic
         profile_excludes.add("stop_criteria")
     if profile.schedule is None:
         profile_excludes.add("schedule")
+    if (
+        fleet_spec.configuration.nodes
+        and fleet_spec.configuration.nodes.min == fleet_spec.configuration.nodes.target
+    ):
+        configuration_excludes["nodes"] = {"target"}
     if configuration_excludes:
         spec_excludes["configuration"] = configuration_excludes
     if profile_excludes:

dstack/_internal/core/compatibility/runs.py

@@ -1,7 +1,7 @@
 from typing import Optional
 
 from dstack._internal.core.models.common import IncludeExcludeDictType, IncludeExcludeSetType
-from dstack._internal.core.models.configurations import ServiceConfiguration
+from dstack._internal.core.models.configurations import LEGACY_REPO_DIR, ServiceConfiguration
 from dstack._internal.core.models.runs import ApplyRunPlanInput, JobSpec, JobSubmission, RunSpec
 from dstack._internal.server.schemas.runs import GetRunPlanRequest, ListRunsRequest
 
@@ -31,6 +31,8 @@ def get_apply_plan_excludes(plan: ApplyRunPlanInput) -> Optional[IncludeExcludeD
         current_resource_excludes["status_message"] = True
         if current_resource.deployment_num == 0:
             current_resource_excludes["deployment_num"] = True
+        if current_resource.fleet is None:
+            current_resource_excludes["fleet"] = True
         apply_plan_excludes["current_resource"] = current_resource_excludes
         current_resource_excludes["run_spec"] = get_run_spec_excludes(current_resource.run_spec)
         job_submissions_excludes: IncludeExcludeDictType = {}
@@ -102,6 +104,11 @@ def get_run_spec_excludes(run_spec: RunSpec) -> IncludeExcludeDictType:
     configuration = run_spec.configuration
     profile = run_spec.profile
 
+    if run_spec.repo_dir in [None, LEGACY_REPO_DIR]:
+        spec_excludes["repo_dir"] = True
+    elif run_spec.repo_dir == "." and configuration.working_dir in [None, LEGACY_REPO_DIR, "."]:
+        spec_excludes["repo_dir"] = True
+
     if configuration.fleets is None:
         configuration_excludes["fleets"] = True
     if profile is not None and profile.fleets is None:
@@ -163,6 +170,8 @@ def get_job_spec_excludes(job_specs: list[JobSpec]) -> IncludeExcludeDictType:
         spec_excludes["service_port"] = True
     if all(not s.probes for s in job_specs):
         spec_excludes["probes"] = True
+    if all(s.repo_dir in [None, LEGACY_REPO_DIR] for s in job_specs):
+        spec_excludes["repo_dir"] = True
 
     return spec_excludes
 

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

@@ -4,13 +4,15 @@ import enum
 class BackendType(str, enum.Enum):
     """
     Attributes:
+        AMDDEVCLOUD (BackendType): AMD Developer Cloud
         AWS (BackendType): Amazon Web Services
         AZURE (BackendType): Microsoft Azure
         CLOUDRIFT (BackendType): CloudRift
         CUDO (BackendType): Cudo
+        DATACRUNCH (BackendType): DataCrunch
+        DIGITALOCEAN (BackendType): DigitalOcean
         DSTACK (BackendType): dstack Sky
         GCP (BackendType): Google Cloud Platform
-        DATACRUNCH (BackendType): DataCrunch
         HOTAISLE (BackendType): Hot Aisle
         KUBERNETES (BackendType): Kubernetes
         LAMBDA (BackendType): Lambda Cloud
@@ -22,11 +24,13 @@ class BackendType(str, enum.Enum):
         VULTR (BackendType): Vultr
     """
 
+    AMDDEVCLOUD = "amddevcloud"
     AWS = "aws"
     AZURE = "azure"
     CLOUDRIFT = "cloudrift"
     CUDO = "cudo"
     DATACRUNCH = "datacrunch"
+    DIGITALOCEAN = "digitalocean"
     DSTACK = "dstack"
     GCP = "gcp"
     HOTAISLE = "hotaisle"

dstack/_internal/core/models/common.py

@@ -1,10 +1,10 @@
 import re
 from enum import Enum
-from typing import Any, Callable, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable, Mapping, Optional, Union
 
 import orjson
 from pydantic import Field
-from pydantic_duality import DualBaseModel
+from pydantic_duality import generate_dual_base_model
 from typing_extensions import Annotated
 
 from dstack._internal.utils.json_utils import pydantic_orjson_dumps
@@ -17,46 +17,73 @@ IncludeExcludeDictType = dict[
 IncludeExcludeType = Union[IncludeExcludeSetType, IncludeExcludeDictType]
 
 
+class CoreConfig:
+    json_loads = orjson.loads
+    json_dumps = pydantic_orjson_dumps
+
+
+# All dstack models inherit from pydantic-duality's DualBaseModel.
 # DualBaseModel creates two classes for the model:
 # one with extra = "forbid" (CoreModel/CoreModel.__request__),
 # and another with extra = "ignore" (CoreModel.__response__).
-# This allows to use the same model both for a strict parsing of the user input and
-# for a permissive parsing of the server responses.
-class CoreModel(DualBaseModel):
-    class Config:
-        json_loads = orjson.loads
-        json_dumps = pydantic_orjson_dumps
-
-    def json(
-        self,
-        *,
-        include: Optional[IncludeExcludeType] = None,
-        exclude: Optional[IncludeExcludeType] = None,
-        by_alias: bool = False,
-        skip_defaults: Optional[bool] = None,  # ignore as it's deprecated
-        exclude_unset: bool = False,
-        exclude_defaults: bool = False,
-        exclude_none: bool = False,
-        encoder: Optional[Callable[[Any], Any]] = None,
-        models_as_dict: bool = True,  # does not seems to be needed by dstack or dependencies
-        **dumps_kwargs: Any,
-    ) -> str:
-        """
-        Override `json()` method so that it calls `dict()`.
-        Allows changing how models are serialized by overriding `dict()` only.
-        By default, `json()` won't call `dict()`, so changes applied in `dict()` won't take place.
-        """
-        data = self.dict(
-            by_alias=by_alias,
-            include=include,
-            exclude=exclude,
-            exclude_unset=exclude_unset,
-            exclude_defaults=exclude_defaults,
-            exclude_none=exclude_none,
-        )
-        if self.__custom_root_type__:
-            data = data["__root__"]
-        return self.__config__.json_dumps(data, default=encoder, **dumps_kwargs)
+# This allows to use the same model both for strict parsing of the user input and
+# for permissive parsing of the server responses.
+#
+# We define a func to generate CoreModel dynamically that can be used
+# to define custom Config for both __request__ and __response__ models.
+# Note: Defining config in the model class directly overrides
+# pydantic-duality's base config, breaking __response__.
+def generate_dual_core_model(
+    custom_config: Union[type, Mapping],
+) -> "type[CoreModel]":
+    class CoreModel(generate_dual_base_model(custom_config)):
+        def json(
+            self,
+            *,
+            include: Optional[IncludeExcludeType] = None,
+            exclude: Optional[IncludeExcludeType] = None,
+            by_alias: bool = False,
+            skip_defaults: Optional[bool] = None,  # ignore as it's deprecated
+            exclude_unset: bool = False,
+            exclude_defaults: bool = False,
+            exclude_none: bool = False,
+            encoder: Optional[Callable[[Any], Any]] = None,
+            models_as_dict: bool = True,  # does not seems to be needed by dstack or dependencies
+            **dumps_kwargs: Any,
+        ) -> str:
+            """
+            Override `json()` method so that it calls `dict()`.
+            Allows changing how models are serialized by overriding `dict()` only.
+            By default, `json()` won't call `dict()`, so changes applied in `dict()` won't take place.
+            """
+            data = self.dict(
+                by_alias=by_alias,
+                include=include,
+                exclude=exclude,
+                exclude_unset=exclude_unset,
+                exclude_defaults=exclude_defaults,
+                exclude_none=exclude_none,
+            )
+            if self.__custom_root_type__:
+                data = data["__root__"]
+            return self.__config__.json_dumps(data, default=encoder, **dumps_kwargs)
+
+    return CoreModel
+
+
+if TYPE_CHECKING:
+
+    class CoreModel(generate_dual_base_model(CoreConfig)):
+        pass
+else:
+    CoreModel = generate_dual_core_model(CoreConfig)
+
+
+class FrozenConfig(CoreConfig):
+    frozen = True
+
+
+FrozenCoreModel = generate_dual_core_model(FrozenConfig)
 
 
 class Duration(int):
@@ -93,7 +120,7 @@ class Duration(int):
         raise ValueError(f"Cannot parse the duration {v}")
 
 
-class RegistryAuth(CoreModel):
+class RegistryAuth(FrozenCoreModel):
     """
     Credentials for pulling a private Docker image.
 
@@ -105,9 +132,6 @@ class RegistryAuth(CoreModel):
     username: Annotated[str, Field(description="The username")]
     password: Annotated[str, Field(description="The password or access token")]
 
-    class Config(CoreModel.Config):
-        frozen = True
-
 
 class ApplyAction(str, Enum):
     CREATE = "create"  # resource is to be created or overridden