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

dstack/_internal/cli/commands/init.py

@@ -6,12 +6,12 @@ from typing import Optional
 from dstack._internal.cli.commands import BaseCommand
 from dstack._internal.cli.services.repos import (
     get_repo_from_dir,
-    get_repo_from_url,
     is_git_repo_url,
     register_init_repo_args,
 )
 from dstack._internal.cli.utils.common import configure_logging, confirm_ask, console, warn
 from dstack._internal.core.errors import ConfigurationError
+from dstack._internal.core.models.repos.remote import RemoteRepo
 from dstack._internal.core.services.configs import ConfigManager
 from dstack.api import Client
 
@@ -101,7 +101,7 @@ class InitCommand(BaseCommand):
         if repo_url is not None:
             # Dummy repo branch to avoid autodetection that fails on private repos.
             # We don't need branch/hash for repo_id anyway.
-            repo = get_repo_from_url(repo_url, repo_branch="master")
+            repo = RemoteRepo.from_url(repo_url, repo_branch="master")
         elif repo_path is not None:
             repo = get_repo_from_dir(repo_path, local=local)
         else:

dstack/_internal/cli/services/configurators/run.py

@@ -2,7 +2,7 @@ import argparse
 import subprocess
 import sys
 import time
-from pathlib import Path
+from pathlib import Path, PurePosixPath
 from typing import Dict, List, Optional, Set, TypeVar
 
 import gpuhunt
@@ -17,7 +17,6 @@ from dstack._internal.cli.services.configurators.base import (
 from dstack._internal.cli.services.profile import apply_profile_args, register_profile_args
 from dstack._internal.cli.services.repos import (
     get_repo_from_dir,
-    get_repo_from_url,
     init_default_virtual_repo,
     is_git_repo_url,
     register_init_repo_args,
@@ -33,6 +32,7 @@ from dstack._internal.core.errors import (
 )
 from dstack._internal.core.models.common import ApplyAction, RegistryAuth
 from dstack._internal.core.models.configurations import (
+    LEGACY_REPO_DIR,
     AnyRunConfiguration,
     ApplyConfigurationType,
     ConfigurationWithPortsParams,
@@ -42,19 +42,27 @@ from dstack._internal.core.models.configurations import (
     ServiceConfiguration,
     TaskConfiguration,
 )
+from dstack._internal.core.models.repos import RepoHeadWithCreds
 from dstack._internal.core.models.repos.base import Repo
 from dstack._internal.core.models.repos.local import LocalRepo
+from dstack._internal.core.models.repos.remote import RemoteRepo, RemoteRepoCreds
 from dstack._internal.core.models.resources import CPUSpec
 from dstack._internal.core.models.runs import JobStatus, JobSubmission, RunSpec, RunStatus
 from dstack._internal.core.services.configs import ConfigManager
 from dstack._internal.core.services.diff import diff_models
-from dstack._internal.core.services.repos import load_repo
+from dstack._internal.core.services.repos import (
+    InvalidRepoCredentialsError,
+    get_repo_creds_and_default_branch,
+    load_repo,
+)
 from dstack._internal.utils.common import local_time
 from dstack._internal.utils.interpolator import InterpolatorError, VariablesInterpolator
 from dstack._internal.utils.logging import get_logger
 from dstack._internal.utils.nested_list import NestedList, NestedListItem
+from dstack._internal.utils.path import is_absolute_posix_path
 from dstack.api._public.repos import get_ssh_keypair
 from dstack.api._public.runs import Run
+from dstack.api.server import APIClient
 from dstack.api.utils import load_profile
 
 _KNOWN_AMD_GPUS = {gpu.name.lower() for gpu in gpuhunt.KNOWN_AMD_GPUS}
@@ -89,6 +97,43 @@ class BaseRunConfigurator(
         self.validate_gpu_vendor_and_image(conf)
         self.validate_cpu_arch_and_image(conf)
 
+        working_dir = conf.working_dir
+        if working_dir is None:
+            # Use the default working dir for the image for tasks and services if `commands`
+            # is not set (emulate pre-0.19.27 JobConfigutor logic), otherwise fall back to
+            # `/workflow`.
+            if isinstance(conf, DevEnvironmentConfiguration) or conf.commands:
+                # relative path for compatibility with pre-0.19.27 servers
+                conf.working_dir = "."
+                warn(
+                    f'The [code]working_dir[/code] is not set — using legacy default [code]"{LEGACY_REPO_DIR}"[/code].'
+                    " Future versions will default to the [code]image[/code]'s working directory."
+                )
+        elif not is_absolute_posix_path(working_dir):
+            legacy_working_dir = PurePosixPath(LEGACY_REPO_DIR) / working_dir
+            warn(
+                "[code]working_dir[/code] is relative."
+                f" Using legacy working directory [code]{legacy_working_dir}[/code]\n\n"
+                "Future versions will require absolute path\n"
+                f"To keep using legacy working directory, set"
+                f" [code]working_dir[/code] to [code]{legacy_working_dir}[/code]\n"
+            )
+        else:
+            # relative path for compatibility with pre-0.19.27 servers
+            try:
+                conf.working_dir = str(PurePosixPath(working_dir).relative_to(LEGACY_REPO_DIR))
+            except ValueError:
+                pass
+
+        if conf.repos and conf.repos[0].path is None:
+            warn(
+                "[code]repos[0].path[/code] is not set,"
+                f" using legacy repo path [code]{LEGACY_REPO_DIR}[/code]\n\n"
+                "In a future version the default value will be changed."
+                f" To keep using [code]{LEGACY_REPO_DIR}[/code], explicitly set"
+                f" [code]repos[0].path[/code] to [code]{LEGACY_REPO_DIR}[/code]\n"
+            )
+
         config_manager = ConfigManager()
         repo = self.get_repo(conf, configuration_path, configurator_args, config_manager)
         self.api.ssh_identity_file = get_ssh_keypair(
@@ -184,6 +229,9 @@ class BaseRunConfigurator(
                         format_date=local_time,
                     )
                 )
+
+                _warn_fleet_autocreated(self.api.client, run)
+
                 console.print(
                     f"\n[code]{run.name}[/] provisioning completed [secondary]({run.status.value})[/]"
                 )
@@ -486,15 +534,17 @@ class BaseRunConfigurator(
             return init_default_virtual_repo(api=self.api)
 
         repo: Optional[Repo] = None
+        repo_head: Optional[RepoHeadWithCreds] = None
         repo_branch: Optional[str] = configurator_args.repo_branch
         repo_hash: Optional[str] = configurator_args.repo_hash
+        repo_creds: Optional[RemoteRepoCreds] = None
+        git_identity_file: Optional[str] = configurator_args.git_identity_file
+        git_private_key: Optional[str] = None
+        oauth_token: Optional[str] = configurator_args.gh_token
         # Should we (re)initialize the repo?
         # If any Git credentials provided, we reinitialize the repo, as the user may have provided
         # updated credentials.
-        init = (
-            configurator_args.git_identity_file is not None
-            or configurator_args.gh_token is not None
-        )
+        init = git_identity_file is not None or oauth_token is not None
 
         url: Optional[str] = None
         local_path: Optional[Path] = None
@@ -527,15 +577,15 @@ class BaseRunConfigurator(
             local_path = Path.cwd()
             legacy_local_path = True
         if url:
-            repo = get_repo_from_url(repo_url=url, repo_branch=repo_branch, repo_hash=repo_hash)
-            if not self.api.repos.is_initialized(repo, by_user=True):
-                init = True
+            # "master" is a dummy value, we'll fetch the actual default branch later
+            repo = RemoteRepo.from_url(repo_url=url, repo_branch="master")
+            repo_head = self.api.repos.get(repo_id=repo.repo_id, with_creds=True)
         elif local_path:
             if legacy_local_path:
                 if repo_config := config_manager.get_repo_config(local_path):
                     repo = load_repo(repo_config)
-                    # allow users with legacy configurations use shared repo creds
-                    if self.api.repos.is_initialized(repo, by_user=False):
+                    repo_head = self.api.repos.get(repo_id=repo.repo_id, with_creds=True)
+                    if repo_head is not None:
                         warn(
                             "The repo is not specified but found and will be used in the run\n"
                             "Future versions will not load repos automatically\n"
@@ -562,20 +612,55 @@ class BaseRunConfigurator(
                     )
                 local: bool = configurator_args.local
                 repo = get_repo_from_dir(local_path, local=local)
-                if not self.api.repos.is_initialized(repo, by_user=True):
-                    init = True
+                repo_head = self.api.repos.get(repo_id=repo.repo_id, with_creds=True)
+            if isinstance(repo, RemoteRepo):
+                repo_branch = repo.run_repo_data.repo_branch
+                repo_hash = repo.run_repo_data.repo_hash
         else:
             assert False, "should not reach here"
 
         if repo is None:
             return init_default_virtual_repo(api=self.api)
 
+        if isinstance(repo, RemoteRepo):
+            assert repo.repo_url is not None
+
+            if repo_head is not None and repo_head.repo_creds is not None:
+                if git_identity_file is None and oauth_token is None:
+                    git_private_key = repo_head.repo_creds.private_key
+                    oauth_token = repo_head.repo_creds.oauth_token
+            else:
+                init = True
+
+            try:
+                repo_creds, default_repo_branch = get_repo_creds_and_default_branch(
+                    repo_url=repo.repo_url,
+                    identity_file=git_identity_file,
+                    private_key=git_private_key,
+                    oauth_token=oauth_token,
+                )
+            except InvalidRepoCredentialsError as e:
+                raise CLIError(*e.args) from e
+
+            if repo_branch is None and repo_hash is None:
+                repo_branch = default_repo_branch
+                if repo_branch is None:
+                    raise CLIError(
+                        "Failed to automatically detect remote repo branch."
+                        " Specify branch or hash."
+                    )
+            repo = RemoteRepo.from_url(
+                repo_url=repo.repo_url, repo_branch=repo_branch, repo_hash=repo_hash
+            )
+
         if init:
             self.api.repos.init(
                 repo=repo,
-                git_identity_file=configurator_args.git_identity_file,
-                oauth_token=configurator_args.gh_token,
+                git_identity_file=git_identity_file,
+                oauth_token=oauth_token,
+                creds=repo_creds,
             )
+
         if isinstance(repo, LocalRepo):
             warn(
                 f"{repo.repo_dir} is a local repo\n"
@@ -827,3 +912,16 @@ def render_run_spec_diff(old_spec: RunSpec, new_spec: RunSpec) -> Optional[str]:
             item = NestedListItem(spec_field.replace("_", " ").capitalize())
         nested_list.children.append(item)
     return nested_list.render()
+
+
+def _warn_fleet_autocreated(api: APIClient, run: Run):
+    if run._run.fleet is None:
+        return
+    fleet = api.fleets.get(project_name=run._project, name=run._run.fleet.name)
+    if not fleet.spec.autocreated:
+        return
+    warn(
+        f"\nNo existing fleet matched, so the run created a new fleet [code]{fleet.name}[/code].\n"
+        "Future dstack versions won't create fleets automatically.\n"
+        "Create a fleet explicitly: https://dstack.ai/docs/concepts/fleets/"
+    )

dstack/_internal/cli/services/repos.py

@@ -1,5 +1,5 @@
 import argparse
-from typing import Literal, Optional, Union, overload
+from typing import Literal, Union, overload
 
 import git
 
@@ -8,7 +8,6 @@ from dstack._internal.core.errors import CLIError
 from dstack._internal.core.models.repos.local import LocalRepo
 from dstack._internal.core.models.repos.remote import GitRepoURL, RemoteRepo, RepoError
 from dstack._internal.core.models.repos.virtual import VirtualRepo
-from dstack._internal.core.services.repos import get_default_branch
 from dstack._internal.utils.path import PathLike
 from dstack.api._public import Client
 
@@ -43,22 +42,6 @@ def init_default_virtual_repo(api: Client) -> VirtualRepo:
     return repo
 
 
-def get_repo_from_url(
-    repo_url: str, repo_branch: Optional[str] = None, repo_hash: Optional[str] = None
-) -> RemoteRepo:
-    if repo_branch is None and repo_hash is None:
-        repo_branch = get_default_branch(repo_url)
-        if repo_branch is None:
-            raise CLIError(
-                "Failed to automatically detect remote repo branch. Specify branch or hash."
-            )
-    return RemoteRepo.from_url(
-        repo_url=repo_url,
-        repo_branch=repo_branch,
-        repo_hash=repo_hash,
-    )
-
-
 @overload
 def get_repo_from_dir(repo_dir: PathLike, local: Literal[False] = False) -> RemoteRepo: ...
 

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

@@ -0,0 +1 @@
+# This package contains the implementation for the AMDDevCloud backend.

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

@@ -0,0 +1,16 @@
+from dstack._internal.core.backends.amddevcloud.compute import AMDDevCloudCompute
+from dstack._internal.core.backends.digitalocean_base.backend import BaseDigitalOceanBackend
+from dstack._internal.core.backends.digitalocean_base.models import BaseDigitalOceanConfig
+from dstack._internal.core.models.backends.base import BackendType
+
+
+class AMDDevCloudBackend(BaseDigitalOceanBackend):
+    TYPE = BackendType.AMDDEVCLOUD
+    COMPUTE_CLASS = AMDDevCloudCompute
+
+    def __init__(self, config: BaseDigitalOceanConfig, api_url: str):
+        self.config = config
+        self._compute = AMDDevCloudCompute(self.config, api_url=api_url, type=self.TYPE)
+
+    def compute(self) -> AMDDevCloudCompute:
+        return self._compute

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

@@ -0,0 +1,5 @@
+from dstack._internal.core.backends.digitalocean_base.compute import BaseDigitalOceanCompute
+
+
+class AMDDevCloudCompute(BaseDigitalOceanCompute):
+    pass

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

@@ -0,0 +1,29 @@
+from typing import Optional
+
+from dstack._internal.core.backends.amddevcloud.backend import AMDDevCloudBackend
+from dstack._internal.core.backends.base.configurator import BackendRecord
+from dstack._internal.core.backends.digitalocean_base.api_client import DigitalOceanAPIClient
+from dstack._internal.core.backends.digitalocean_base.backend import BaseDigitalOceanBackend
+from dstack._internal.core.backends.digitalocean_base.configurator import (
+    BaseDigitalOceanConfigurator,
+)
+from dstack._internal.core.backends.digitalocean_base.models import AnyBaseDigitalOceanCreds
+from dstack._internal.core.models.backends.base import (
+    BackendType,
+)
+
+
+class AMDDevCloudConfigurator(BaseDigitalOceanConfigurator):
+    TYPE = BackendType.AMDDEVCLOUD
+    BACKEND_CLASS = AMDDevCloudBackend
+    API_URL = "https://api-amd.digitalocean.com"
+
+    def get_backend(self, record: BackendRecord) -> BaseDigitalOceanBackend:
+        config = self._get_config(record)
+        return AMDDevCloudBackend(config=config, api_url=self.API_URL)
+
+    def _validate_creds(self, creds: AnyBaseDigitalOceanCreds, project_name: Optional[str] = None):
+        api_client = DigitalOceanAPIClient(creds.api_key, self.API_URL)
+        api_client.validate_api_key()
+        if project_name:
+            api_client.validate_project_name(project_name)

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

@@ -292,7 +292,12 @@ class AWSCompute(
                         image_id=image_id,
                         instance_type=instance_offer.instance.name,
                         iam_instance_profile=self.config.iam_instance_profile,
-                        user_data=get_user_data(authorized_keys=instance_config.get_public_keys()),
+                        user_data=get_user_data(
+                            authorized_keys=instance_config.get_public_keys(),
+                            # Custom OS images may lack ufw, so don't attempt to set up the firewall.
+                            # Rely on security groups and the image's built-in firewall rules instead.
+                            skip_firewall_setup=self.config.os_images is not None,
+                        ),
                         tags=aws_resources.make_tags(tags),
                         security_group_id=security_group_id,
                         spot=instance_offer.instance.resources.spot,

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

@@ -4,6 +4,7 @@ import re
 import string
 import threading
 from abc import ABC, abstractmethod
+from collections.abc import Iterable
 from functools import lru_cache
 from pathlib import Path
 from typing import Dict, List, Literal, Optional
@@ -19,7 +20,7 @@ from dstack._internal.core.consts import (
     DSTACK_RUNNER_SSH_PORT,
     DSTACK_SHIM_HTTP_PORT,
 )
-from dstack._internal.core.models.configurations import DEFAULT_REPO_DIR
+from dstack._internal.core.models.configurations import LEGACY_REPO_DIR
 from dstack._internal.core.models.gateways import (
     GatewayComputeConfiguration,
     GatewayProvisioningData,
@@ -45,6 +46,7 @@ logger = get_logger(__name__)
 
 DSTACK_SHIM_BINARY_NAME = "dstack-shim"
 DSTACK_RUNNER_BINARY_NAME = "dstack-runner"
+DEFAULT_PRIVATE_SUBNETS = ("10.0.0.0/8", "172.16.0.0/12", "192.168.0.0/16")
 
 GoArchType = Literal["amd64", "arm64"]
 
@@ -507,12 +509,16 @@ def get_user_data(
     base_path: Optional[PathLike] = None,
     bin_path: Optional[PathLike] = None,
     backend_shim_env: Optional[Dict[str, str]] = None,
+    skip_firewall_setup: bool = False,
+    firewall_allow_from_subnets: Iterable[str] = DEFAULT_PRIVATE_SUBNETS,
 ) -> str:
     shim_commands = get_shim_commands(
         authorized_keys=authorized_keys,
         base_path=base_path,
         bin_path=bin_path,
         backend_shim_env=backend_shim_env,
+        skip_firewall_setup=skip_firewall_setup,
+        firewall_allow_from_subnets=firewall_allow_from_subnets,
     )
     commands = (backend_specific_commands or []) + shim_commands
     return get_cloud_config(
@@ -554,8 +560,13 @@ def get_shim_commands(
     bin_path: Optional[PathLike] = None,
     backend_shim_env: Optional[Dict[str, str]] = None,
     arch: Optional[str] = None,
+    skip_firewall_setup: bool = False,
+    firewall_allow_from_subnets: Iterable[str] = DEFAULT_PRIVATE_SUBNETS,
 ) -> List[str]:
-    commands = get_setup_cloud_instance_commands()
+    commands = get_setup_cloud_instance_commands(
+        skip_firewall_setup=skip_firewall_setup,
+        firewall_allow_from_subnets=firewall_allow_from_subnets,
+    )
     commands += get_shim_pre_start_commands(
         base_path=base_path,
         bin_path=bin_path,
@@ -638,8 +649,11 @@ def get_dstack_shim_download_url(arch: Optional[str] = None) -> str:
     return url_template.format(version=version, arch=arch)
 
 
-def get_setup_cloud_instance_commands() -> list[str]:
-    return [
+def get_setup_cloud_instance_commands(
+    skip_firewall_setup: bool,
+    firewall_allow_from_subnets: Iterable[str],
+) -> list[str]:
+    commands = [
         # Workaround for https://github.com/NVIDIA/nvidia-container-toolkit/issues/48
         # Attempts to patch /etc/docker/daemon.json while keeping any custom settings it may have.
         (
@@ -653,6 +667,19 @@ def get_setup_cloud_instance_commands() -> list[str]:
             "'"
         ),
     ]
+    if not skip_firewall_setup:
+        commands += [
+            "ufw --force reset",  # Some OS images have default rules like `allow 80`. Delete them
+            "ufw default deny incoming",
+            "ufw default allow outgoing",
+            "ufw allow ssh",
+        ]
+        for subnet in firewall_allow_from_subnets:
+            commands.append(f"ufw allow from {subnet}")
+        commands += [
+            "ufw --force enable",
+        ]
+    return commands
 
 
 def get_shim_pre_start_commands(
@@ -773,7 +800,8 @@ def get_docker_commands(
             f" --ssh-port {DSTACK_RUNNER_SSH_PORT}"
             " --temp-dir /tmp/runner"
             " --home-dir /root"
-            f" --working-dir {DEFAULT_REPO_DIR}"
+            # TODO: Not used, left for compatibility with old runners. Remove eventually.
+            f" --working-dir {LEGACY_REPO_DIR}"
         ),
     ]
 

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

@@ -34,6 +34,8 @@ def get_catalog_offers(
     provider = backend.value
     if backend == BackendType.LAMBDA:
         provider = "lambdalabs"
+    if backend == BackendType.AMDDEVCLOUD:
+        provider = "digitalocean"
     q = requirements_to_query_filter(requirements)
     q.provider = [provider]
     offers = []

dstack/_internal/core/backends/configurators.py

@@ -5,6 +5,12 @@ from dstack._internal.core.models.backends.base import BackendType
 
 _CONFIGURATOR_CLASSES: List[Type[Configurator]] = []
 
+try:
+    from dstack._internal.core.backends.amddevcloud.configurator import AMDDevCloudConfigurator
+
+    _CONFIGURATOR_CLASSES.append(AMDDevCloudConfigurator)
+except ImportError:
+    pass
 
 try:
     from dstack._internal.core.backends.aws.configurator import AWSConfigurator
@@ -47,6 +53,15 @@ try:
 except ImportError:
     pass
 
+try:
+    from dstack._internal.core.backends.digitalocean.configurator import (
+        DigitalOceanConfigurator,
+    )
+
+    _CONFIGURATOR_CLASSES.append(DigitalOceanConfigurator)
+except ImportError:
+    pass
+
 try:
     from dstack._internal.core.backends.gcp.configurator import GCPConfigurator
 

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

@@ -0,0 +1 @@
+# DigitalOcean backend for dstack

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

@@ -0,0 +1,16 @@
+from dstack._internal.core.backends.digitalocean.compute import DigitalOceanCompute
+from dstack._internal.core.backends.digitalocean_base.backend import BaseDigitalOceanBackend
+from dstack._internal.core.backends.digitalocean_base.models import BaseDigitalOceanConfig
+from dstack._internal.core.models.backends.base import BackendType
+
+
+class DigitalOceanBackend(BaseDigitalOceanBackend):
+    TYPE = BackendType.DIGITALOCEAN
+    COMPUTE_CLASS = DigitalOceanCompute
+
+    def __init__(self, config: BaseDigitalOceanConfig, api_url: str):
+        self.config = config
+        self._compute = DigitalOceanCompute(self.config, api_url=api_url, type=self.TYPE)
+
+    def compute(self) -> DigitalOceanCompute:
+        return self._compute

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

@@ -0,0 +1,5 @@
+from ..digitalocean_base.compute import BaseDigitalOceanCompute
+
+
+class DigitalOceanCompute(BaseDigitalOceanCompute):
+    pass

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

@@ -0,0 +1,31 @@
+from typing import Optional
+
+from dstack._internal.core.backends.base.configurator import BackendRecord
+from dstack._internal.core.backends.digitalocean.backend import DigitalOceanBackend
+from dstack._internal.core.backends.digitalocean_base.api_client import DigitalOceanAPIClient
+from dstack._internal.core.backends.digitalocean_base.backend import BaseDigitalOceanBackend
+from dstack._internal.core.backends.digitalocean_base.configurator import (
+    BaseDigitalOceanConfigurator,
+)
+from dstack._internal.core.backends.digitalocean_base.models import (
+    AnyBaseDigitalOceanCreds,
+)
+from dstack._internal.core.models.backends.base import (
+    BackendType,
+)
+
+
+class DigitalOceanConfigurator(BaseDigitalOceanConfigurator):
+    TYPE = BackendType.DIGITALOCEAN
+    BACKEND_CLASS = DigitalOceanBackend
+    API_URL = "https://api.digitalocean.com"
+
+    def get_backend(self, record: BackendRecord) -> BaseDigitalOceanBackend:
+        config = self._get_config(record)
+        return DigitalOceanBackend(config=config, api_url=self.API_URL)
+
+    def _validate_creds(self, creds: AnyBaseDigitalOceanCreds, project_name: Optional[str] = None):
+        api_client = DigitalOceanAPIClient(creds.api_key, self.API_URL)
+        api_client.validate_api_key()
+        if project_name:
+            api_client.validate_project_name(project_name)

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

@@ -0,0 +1 @@
+# This package contains the base classes for DigitalOcean and AMDDevCloud backends.

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

@@ -0,0 +1,104 @@
+from typing import Any, Dict, List, Optional
+
+import requests
+
+from dstack._internal.core.backends.base.configurator import raise_invalid_credentials_error
+from dstack._internal.core.errors import NoCapacityError
+from dstack._internal.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+class DigitalOceanAPIClient:
+    def __init__(self, api_key: str, api_url: str):
+        self.api_key = api_key
+        self.base_url = api_url
+
+    def validate_api_key(self) -> bool:
+        try:
+            response = self._make_request("GET", "/v2/account")
+            response.raise_for_status()
+            return True
+        except requests.HTTPError as e:
+            status = e.response.status_code
+            if status == 401:
+                raise_invalid_credentials_error(
+                    fields=[["creds", "api_key"]], details="Invaild API key"
+                )
+            raise e
+
+    def validate_project_name(self, project_name: str) -> bool:
+        if self.get_project_id(project_name) is None:
+            raise_invalid_credentials_error(
+                fields=[["project_name"]],
+                details=f"Project with name '{project_name}' does not exist",
+            )
+        return True
+
+    def list_ssh_keys(self) -> List[Dict[str, Any]]:
+        response = self._make_request("GET", "/v2/account/keys")
+        response.raise_for_status()
+        return response.json()["ssh_keys"]
+
+    def list_projects(self) -> List[Dict[str, Any]]:
+        response = self._make_request("GET", "/v2/projects")
+        response.raise_for_status()
+        return response.json()["projects"]
+
+    def get_project_id(self, project_name: str) -> Optional[str]:
+        projects = self.list_projects()
+        for project in projects:
+            if project["name"] == project_name:
+                return project["id"]
+        return None
+
+    def create_ssh_key(self, name: str, public_key: str) -> Dict[str, Any]:
+        payload = {"name": name, "public_key": public_key}
+        response = self._make_request("POST", "/v2/account/keys", json=payload)
+        response.raise_for_status()
+        return response.json()["ssh_key"]
+
+    def get_or_create_ssh_key(self, name: str, public_key: str) -> int:
+        ssh_keys = self.list_ssh_keys()
+        for ssh_key in ssh_keys:
+            if ssh_key["public_key"].strip() == public_key.strip():
+                return ssh_key["id"]
+
+        ssh_key = self.create_ssh_key(name, public_key)
+        return ssh_key["id"]
+
+    def create_droplet(self, droplet_config: Dict[str, Any]) -> Dict[str, Any]:
+        response = self._make_request("POST", "/v2/droplets", json=droplet_config)
+        if response.status_code == 422:
+            raise NoCapacityError(response.json()["message"])
+        response.raise_for_status()
+        return response.json()["droplet"]
+
+    def get_droplet(self, droplet_id: str) -> Dict[str, Any]:
+        response = self._make_request("GET", f"/v2/droplets/{droplet_id}")
+        response.raise_for_status()
+        return response.json()["droplet"]
+
+    def delete_droplet(self, droplet_id: str) -> None:
+        response = self._make_request("DELETE", f"/v2/droplets/{droplet_id}")
+        if response.status_code == 404:
+            logger.debug("DigitalOcean droplet %s not found", droplet_id)
+            return
+        response.raise_for_status()
+
+    def _make_request(
+        self, method: str, endpoint: str, json: Optional[Dict[str, Any]] = None, timeout: int = 30
+    ) -> requests.Response:
+        url = f"{self.base_url}{endpoint}"
+        headers = {
+            "Authorization": f"Bearer {self.api_key}",
+        }
+
+        response = requests.request(
+            method=method,
+            url=url,
+            headers=headers,
+            json=json,
+            timeout=timeout,
+        )
+        return response

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

@@ -0,0 +1,5 @@
+from dstack._internal.core.backends.base.backend import Backend
+
+
+class BaseDigitalOceanBackend(Backend):
+    pass