diff --git a/frontend/src/types/backend.d.ts b/frontend/src/types/backend.d.ts index f09bf1b208..d42a26ba17 100644 --- a/frontend/src/types/backend.d.ts +++ b/frontend/src/types/backend.d.ts @@ -4,6 +4,7 @@ declare type TBackendType = | 'crusoe' | 'cudo' | 'datacrunch' + | 'daytona' | 'dstack' | 'gcp' | 'kubernetes' diff --git a/frontend/src/types/volume.d.ts b/frontend/src/types/volume.d.ts index e704e02b07..c72878f3d3 100644 --- a/frontend/src/types/volume.d.ts +++ b/frontend/src/types/volume.d.ts @@ -11,7 +11,7 @@ declare interface IVolumeConfiguration { type: "volume", name?: string, backend: TBackendType, - region: string, + region?: string, size?: number, volume_id?: string } @@ -19,7 +19,7 @@ declare interface IVolumeConfiguration { declare interface IVolumeProvisioningData { backend?: TBackendType, volume_id: string, - size_gb: number, + size_gb: number | null, availability_zone?: string price?:number attachable: boolean diff --git a/mkdocs/docs/concepts/backends.md b/mkdocs/docs/concepts/backends.md index 110c1b78a5..fd23b1ad5f 100644 --- a/mkdocs/docs/concepts/backends.md +++ b/mkdocs/docs/concepts/backends.md @@ -1473,6 +1473,34 @@ projects: - **Runs always execute as `root`.** `enroot` is a single-user container runtime: it can only map the invoking host user inside the container to either the same UID/GID or to `root` (`0:0`). `dstack` uses the latter, so runs always execute as `root`, and both the image's default user and the run configuration's `user` property are ignored. The single-user model is a fundamental `enroot` limitation, not specific to the `dstack` integration. - **Private registries via `registry_auth` are not supported.** The `registry_auth` run configuration property is rejected. You can still pull images from a private registry by preconfiguring `enroot` credentials on the cluster's compute nodes — see the [enroot import documentation](https://github.com/NVIDIA/enroot/blob/main/doc/cmd/import.md#description). +### Daytona + +Create an API key in the [Daytona dashboard](https://app.daytona.io/) for the organization +that will own the sandboxes. Then configure the backend: + +
+ +```yaml +projects: + - name: main + backends: + - type: daytona + creds: + type: api_key + api_key: YOUR_DAYTONA_API_KEY +``` + +
+ +??? info "Required permissions" + The API key requires `write:sandboxes`, `delete:sandboxes`, and `read:limits`. + + Additional permissions depend on the features used: + + * Private images: `write:registries` and `delete:registries` + * Creating volumes: `read:volumes`, `write:volumes`, and `delete:volumes` + * Registering existing volumes: `read:volumes` + ### Runpod Log into your [Runpod](https://www.runpod.io/console/) console, click Settings in the sidebar, expand the `API Keys` section, and click diff --git a/mkdocs/docs/concepts/volumes.md b/mkdocs/docs/concepts/volumes.md index 19be81952d..6aa26ead5a 100644 --- a/mkdocs/docs/concepts/volumes.md +++ b/mkdocs/docs/concepts/volumes.md @@ -16,7 +16,7 @@ Useful as a cache for cloud fleets or for persistent storage with SSH fleets. ## Network volumes -> Network volumes are currently supported for the `aws`, `gcp`, `runpod`, and `kubernetes` backends. +> Network volumes are currently supported for the `aws`, `gcp`, `runpod`, `kubernetes`, and `daytona` backends. ### Apply a configuration @@ -42,6 +42,16 @@ size: 100GB If you use this configuration, `dstack` will create a new volume based on the specified options. +??? info "Daytona" + Daytona volumes use [S3-backed FUSE storage](https://www.daytona.io/docs/en/volumes/). + Storage grows with usage, and the same volume can be mounted in any Daytona region. + + ```yaml + type: volume + backend: daytona + name: my-volume + ``` + ??? info "Kubernetes" Set `region` to a kubeconfig context name enabled in the [backend configuration](backends.md#kubernetes). Omit it only for legacy configurations without `contexts`. @@ -86,7 +96,7 @@ Volume my-volume does not exist yet. Create the volume? [y/n]: y Once created, the volume can be attached to dev environments, tasks, and services. -> When creating a new network volume, `dstack` automatically creates an `ext4` filesystem on it. +> When creating a new network volume in AWS or GCP, `dstack` automatically creates an `ext4` filesystem on it. #### Register existing volumes @@ -110,7 +120,7 @@ volume_id: vol1235 -If you register an existing volume, you must ensure the volume already has a filesystem. +If you register an existing volume in AWS or GCP, you must ensure the volume already has a filesystem. ??? info "Kubernetes" @@ -254,7 +264,7 @@ If you've registered an existing volume, it will be de-registered with `dstack` ??? info "Can I use network volumes across backends?" - Since volumes are backed up by cloud network disks, you can only use them within the same cloud. If you need to access + Volumes can only be used with their backend. If you need to access data across different backends, you should either use object storage or replicate the data across multiple volumes. ??? info "Can I use network volumes across regions?" @@ -262,19 +272,22 @@ If you've registered an existing volume, it will be de-registered with `dstack` Typically, network volumes are associated with specific regions, so you can't use them in other regions. Often, volumes are also linked to availability zones, but some providers support volumes that can be used across different availability zones within the same region. - + If you don't want to limit a run to one particular region, you can create different volumes for different regions and specify them for the same mount point as [documented above](#attach-network-volume). + Daytona volumes can be mounted in any Daytona region. + ??? info "Can I attach network volumes to multiple runs or instances?" - You can mount a volume in multiple runs. This feature is currently supported only by the `runpod` backend. + The `runpod` and `daytona` backends support mounting the same volume in multiple runs at the same time. + For `kubernetes`, the storage class and access modes must support sharing, such as `ReadWriteMany`. ## Instance volumes Instance volumes allow mapping any directory on the instance where the run is executed to any path inside the container. This means that the data in instance volumes is persisted only if the run is executed on the same instance. -> Instance volumes are currently supported for all backends except `runpod` and `vastai`, and can also be used with [SSH fleets](fleets.md#ssh-fleets). +> Instance volumes are currently supported for all backends except `daytona`, `runpod`, and `vastai`, and can also be used with [SSH fleets](fleets.md#ssh-fleets). ### Attach a volume diff --git a/mkdocs/docs/reference/dstack.yml/volume.md b/mkdocs/docs/reference/dstack.yml/volume.md index 23eccccd11..58883b4744 100644 --- a/mkdocs/docs/reference/dstack.yml/volume.md +++ b/mkdocs/docs/reference/dstack.yml/volume.md @@ -26,6 +26,14 @@ The `volume` configuration type allows creating, registering, and updating [volu backend: required: true +=== "Daytona" + + #SCHEMA# dstack._internal.core.models.volumes.DaytonaVolumeConfiguration + overrides: + show_root_heading: false + backend: + required: true + === "Kubernetes" Kubernetes backend volumes are mapped to [`PersistentVolumeClaim`](https://kubernetes.io/docs/concepts/storage/persistent-volumes/#persistentvolumeclaims) objects. diff --git a/mkdocs/docs/reference/server/config.yml.md b/mkdocs/docs/reference/server/config.yml.md index e6b5359084..538f703a89 100644 --- a/mkdocs/docs/reference/server/config.yml.md +++ b/mkdocs/docs/reference/server/config.yml.md @@ -92,6 +92,23 @@ to configure [backends](../../concepts/backends.md) and other [server-level sett type: required: true +##### `projects[n].backends[type=daytona]` { #daytona data-toc-label="daytona" } + +#SCHEMA# dstack._internal.core.backends.daytona.models.DaytonaBackendConfigWithCreds + overrides: + show_root_heading: false + type: + required: true + item_id_prefix: daytona- + +###### `projects[n].backends[type=daytona].creds` { #daytona-creds data-toc-label="creds" } + +#SCHEMA# dstack._internal.core.backends.daytona.models.DaytonaAPIKeyCreds + overrides: + show_root_heading: false + type: + required: true + ##### `projects[n].backends[type=gcp]` { #gcp data-toc-label="gcp" } #SCHEMA# dstack._internal.core.backends.gcp.models.GCPBackendConfigWithCreds diff --git a/pyproject.toml b/pyproject.toml index 4c623306c4..5c51736acb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,7 +33,7 @@ dependencies = [ "python-multipart>=0.0.16", "filelock", "psutil", - "gpuhunt==0.1.30", + "gpuhunt==0.1.31", "argcomplete>=3.5.0", "ignore-python>=0.2.0", "apscheduler<4", @@ -110,6 +110,7 @@ include = [ "src/dstack/_internal/server", "src/dstack/_internal/core/services", "src/dstack/_internal/core/backends/aws", + "src/dstack/_internal/core/backends/daytona", "src/dstack/_internal/core/backends/kubernetes", "src/dstack/_internal/core/backends/runpod", "src/dstack/_internal/core/backends/slurm", diff --git a/src/dstack/_internal/cli/services/configurators/volume.py b/src/dstack/_internal/cli/services/configurators/volume.py index 2449de0da7..f8e2dea334 100644 --- a/src/dstack/_internal/cli/services/configurators/volume.py +++ b/src/dstack/_internal/cli/services/configurators/volume.py @@ -17,6 +17,7 @@ AnyVolumeConfiguration, Volume, VolumeConfigurationWithRegion, + VolumeConfigurationWithSize, VolumePlan, VolumeSpec, VolumeStatus, @@ -204,9 +205,6 @@ def th(s: str) -> str: configuration_table.add_row(th("Type"), plan.spec.configuration.type) volume_type = "managed" - size = "-" - if plan.spec.configuration.size is not None: - size = str(plan.spec.configuration.size) if plan.spec.configuration.is_external: volume_type = "external" @@ -214,7 +212,11 @@ def th(s: str) -> str: configuration_table.add_row(th("Backend"), plan.spec.configuration.backend.value) if isinstance(plan.spec.configuration, VolumeConfigurationWithRegion): configuration_table.add_row(th("Region"), plan.spec.configuration.region) - configuration_table.add_row(th("Size"), size) + if isinstance(plan.spec.configuration, VolumeConfigurationWithSize): + size = "-" + if plan.spec.configuration.size is not None: + size = str(plan.spec.configuration.size) + configuration_table.add_row(th("Size"), size) console.print(configuration_table) console.print() diff --git a/src/dstack/_internal/core/backends/configurators.py b/src/dstack/_internal/core/backends/configurators.py index 2fa3267be4..7bad991340 100644 --- a/src/dstack/_internal/core/backends/configurators.py +++ b/src/dstack/_internal/core/backends/configurators.py @@ -53,6 +53,13 @@ except ImportError: pass +try: + from dstack._internal.core.backends.daytona.configurator import DaytonaConfigurator + + _CONFIGURATOR_CLASSES.append(DaytonaConfigurator) +except ImportError: + pass + try: from dstack._internal.core.backends.digitalocean.configurator import ( DigitalOceanConfigurator, diff --git a/src/dstack/_internal/core/backends/daytona/__init__.py b/src/dstack/_internal/core/backends/daytona/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/dstack/_internal/core/backends/daytona/api_client.py b/src/dstack/_internal/core/backends/daytona/api_client.py new file mode 100644 index 0000000000..c1a9981059 --- /dev/null +++ b/src/dstack/_internal/core/backends/daytona/api_client.py @@ -0,0 +1,267 @@ +import time +from typing import Any, Dict, List, Optional +from urllib.parse import quote, urlsplit + +import requests +from gpuhunt.providers.daytona import API_URL + +from dstack._internal.core.errors import BackendError, BackendInvalidCredentialsError + +REQUEST_TIMEOUT_SECONDS = 30 + + +class DaytonaAPIError(BackendError): + def __init__(self, message: str, status_code: Optional[int] = None): + super().__init__(message) + self.status_code = status_code + + +class DaytonaNotFoundError(DaytonaAPIError): + pass + + +class DaytonaAPIClient: + def __init__(self, api_key: str): + self.api_key = api_key + + def get_current_api_key(self) -> Dict[str, Any]: + return _object(self._make_request("GET", "/api-keys/current")) + + def get_organization_usage(self, organization_id: str) -> Dict[str, Any]: + return _object( + self._make_request("GET", f"/organizations/{_path_id(organization_id)}/usage") + ) + + def get_available_sandbox_classes(self, organization_id: str) -> List[Dict[str, Any]]: + return _objects( + self._make_request( + "GET", f"/organizations/{_path_id(organization_id)}/available-sandbox-classes" + ) + ) + + def get_shared_regions(self) -> List[Dict[str, Any]]: + return _objects(self._make_request("GET", "/shared-regions")) + + def create_sandbox(self, payload: Dict[str, Any]) -> Dict[str, Any]: + # Do not retry creates: a lost response does not mean the sandbox was not created. + return _object(self._make_request("POST", "/sandbox", json=payload)) + + def get_sandbox(self, sandbox_id_or_name: str) -> Optional[Dict[str, Any]]: + try: + return _object(self._make_request("GET", f"/sandbox/{_path_id(sandbox_id_or_name)}")) + except DaytonaNotFoundError: + return None + + def delete_sandbox(self, sandbox_id_or_name: str) -> None: + self._delete(f"/sandbox/{_path_id(sandbox_id_or_name)}") + + def create_registry(self, payload: Dict[str, Any]) -> Dict[str, Any]: + return _object(self._make_request("POST", "/docker-registry", json=payload)) + + def get_registries(self) -> List[Dict[str, Any]]: + return _objects(self._make_request("GET", "/docker-registry")) + + def delete_registry(self, registry_id: str) -> None: + self._delete(f"/docker-registry/{_path_id(registry_id)}") + + def create_volume(self, name: str) -> Dict[str, Any]: + return _object(self._make_request("POST", "/volumes", json={"name": name})) + + def get_volume(self, volume_id: str) -> Optional[Dict[str, Any]]: + try: + return _object(self._make_request("GET", f"/volumes/{_path_id(volume_id)}")) + except DaytonaNotFoundError: + return None + + def get_volume_by_name(self, name: str) -> Optional[Dict[str, Any]]: + try: + return _object(self._make_request("GET", f"/volumes/by-name/{_path_id(name)}")) + except DaytonaNotFoundError: + return None + + def delete_volume(self, volume_id: str) -> None: + self._delete(f"/volumes/{_path_id(volume_id)}") + + def get_toolbox_url(self, sandbox_id: str) -> str: + data = _object( + self._make_request("GET", f"/sandbox/{_path_id(sandbox_id)}/toolbox-proxy-url") + ) + proxy_url = data.get("url") + if not isinstance(proxy_url, str): + raise DaytonaAPIError("Daytona returned an invalid toolbox proxy URL") + _validate_toolbox_url(proxy_url) + # The API returns the proxy base; Toolbox routes also need the sandbox ID. + return f"{proxy_url.rstrip('/')}/{_path_id(sandbox_id)}" + + def get_session(self, toolbox_url: str, session_id: str) -> Optional[Dict[str, Any]]: + try: + return _object( + self._make_toolbox_request( + "GET", toolbox_url, f"/process/session/{_path_id(session_id)}" + ) + ) + except DaytonaNotFoundError: + return None + + def create_session(self, toolbox_url: str, session_id: str) -> None: + self._make_toolbox_request( + "POST", toolbox_url, "/process/session", json={"sessionId": session_id} + ) + + def execute_session_command( + self, + toolbox_url: str, + session_id: str, + command: str, + *, + run_async: bool = True, + ) -> Dict[str, Any]: + return _object( + self._make_toolbox_request( + "POST", + toolbox_url, + f"/process/session/{_path_id(session_id)}/exec", + json={"command": command, "runAsync": run_async}, + ) + ) + + def create_ssh_access( + self, sandbox_id_or_name: str, *, expires_in_minutes: int + ) -> Dict[str, Any]: + if expires_in_minutes <= 0: + raise ValueError("SSH access lifetime must be positive") + return _object( + self._make_request( + "POST", + f"/sandbox/{_path_id(sandbox_id_or_name)}/ssh-access", + params={"expiresInMinutes": expires_in_minutes}, + ) + ) + + def _delete(self, path: str) -> None: + # Deletes are idempotent. Retry transient failures, including during rollback + # when the server has not yet persisted the resource's provisioning data. + for attempt in range(3): + try: + self._make_request("DELETE", path) + return + except DaytonaNotFoundError: + return + except DaytonaAPIError as e: + if attempt == 2 or ( + e.status_code is not None + and e.status_code not in (408, 429) + and e.status_code < 500 + ): + raise + time.sleep(attempt + 1) + + def _make_request( + self, + method: str, + path: str, + *, + json: Optional[Dict[str, Any]] = None, + params: Optional[Dict[str, Any]] = None, + ) -> Any: + return self._request(method, API_URL + path, json=json, params=params) + + def _make_toolbox_request( + self, + method: str, + toolbox_url: str, + path: str, + *, + json: Optional[Dict[str, Any]] = None, + ) -> Any: + _validate_toolbox_url(toolbox_url) + return self._request(method, toolbox_url.rstrip("/") + path, json=json) + + def _request( + self, + method: str, + url: str, + *, + json: Optional[Dict[str, Any]] = None, + params: Optional[Dict[str, Any]] = None, + ) -> Any: + try: + response = requests.request( + method=method, + url=url, + headers={"Authorization": f"Bearer {self.api_key}"}, + json=json, + params=params, + timeout=REQUEST_TIMEOUT_SECONDS, + allow_redirects=False, + ) + except requests.Timeout as e: + raise DaytonaAPIError("Daytona API request timed out") from e + except requests.RequestException as e: + raise DaytonaAPIError("Daytona API request failed") from e + if response.status_code == 401: + raise BackendInvalidCredentialsError(fields=[["creds", "api_key"]]) + if not 200 <= response.status_code < 300: + message = _response_message(response) + if self.api_key: + message = message.replace(self.api_key, "[redacted]") + error = DaytonaNotFoundError if response.status_code == 404 else DaytonaAPIError + raise error(message, status_code=response.status_code) + if not response.content: + return None + try: + return response.json() + except ValueError as e: + raise DaytonaAPIError("Daytona API returned a non-JSON response") from e + + +def _object(data: Any) -> Dict[str, Any]: + if not isinstance(data, dict): + raise DaytonaAPIError("Daytona API returned an invalid object response") + return data + + +def _objects(data: Any) -> List[Dict[str, Any]]: + if not isinstance(data, list) or not all(isinstance(item, dict) for item in data): + raise DaytonaAPIError("Daytona API returned an invalid list response") + return data + + +def _path_id(value: str) -> str: + return quote(value, safe="") + + +def _validate_toolbox_url(url: str) -> None: + try: + parsed = urlsplit(url) + valid = ( + parsed.scheme == "https" + and parsed.hostname is not None + and parsed.hostname.endswith((".daytona.io", ".daytona.work")) + and parsed.port in (None, 443) + and parsed.username is None + and parsed.password is None + and not parsed.query + and not parsed.fragment + ) + except ValueError: + valid = False + if not valid: + raise DaytonaAPIError("Daytona returned an invalid toolbox proxy URL") + + +def _response_message(response: requests.Response) -> str: + message = f"Daytona API returned HTTP {response.status_code}" + try: + data = response.json() + except ValueError: + data = None + if isinstance(data, dict): + for field in ("message", "error", "detail"): + value = data.get(field) + if isinstance(value, str) and value: + return f"{message}: {value}" + if isinstance(value, list) and all(isinstance(item, str) for item in value): + return f"{message}: {'; '.join(value)}" + # Do not include arbitrary HTML/text responses, which can reflect request credentials. + return message diff --git a/src/dstack/_internal/core/backends/daytona/backend.py b/src/dstack/_internal/core/backends/daytona/backend.py new file mode 100644 index 0000000000..e89594c514 --- /dev/null +++ b/src/dstack/_internal/core/backends/daytona/backend.py @@ -0,0 +1,16 @@ +from dstack._internal.core.backends.base.backend import Backend +from dstack._internal.core.backends.daytona.compute import DaytonaCompute +from dstack._internal.core.backends.daytona.models import DaytonaConfig +from dstack._internal.core.models.backends.base import BackendType + + +class DaytonaBackend(Backend): + TYPE = BackendType.DAYTONA + COMPUTE_CLASS = DaytonaCompute + + def __init__(self, config: DaytonaConfig): + self.config = config + self._compute = DaytonaCompute(self.config) + + def compute(self) -> DaytonaCompute: + return self._compute diff --git a/src/dstack/_internal/core/backends/daytona/compute.py b/src/dstack/_internal/core/backends/daytona/compute.py new file mode 100644 index 0000000000..c30a4a661a --- /dev/null +++ b/src/dstack/_internal/core/backends/daytona/compute.py @@ -0,0 +1,449 @@ +import re +import shlex +import time +from typing import Any, List, Optional + +import gpuhunt +from gpuhunt.providers.daytona import DaytonaProvider + +from dstack._internal.core.backends.base.authorized_keys import build_authorized_keys +from dstack._internal.core.backends.base.compute import ( + Compute, + ComputeWithFilteredOffersCached, + ComputeWithVolumeSupport, + generate_unique_instance_name_for_job, + generate_unique_volume_name, + get_docker_commands, +) +from dstack._internal.core.backends.base.offers import get_catalog_offers +from dstack._internal.core.backends.daytona.api_client import DaytonaAPIClient, DaytonaAPIError +from dstack._internal.core.backends.daytona.models import DaytonaConfig +from dstack._internal.core.consts import DSTACK_RUNNER_SSH_PORT +from dstack._internal.core.errors import ComputeError, NotYetTerminated, ProvisioningError +from dstack._internal.core.models.backends.base import BackendType +from dstack._internal.core.models.common import CoreModel, RegistryAuth, validate_json_extra_ignore +from dstack._internal.core.models.instances import ( + InstanceAvailability, + InstanceOffer, + InstanceOfferWithAvailability, + InstanceRuntime, + SSHConnectionParams, +) +from dstack._internal.core.models.placement import PlacementGroup +from dstack._internal.core.models.runs import Job, JobProvisioningData, Requirements, Run +from dstack._internal.core.models.volumes import ( + DaytonaVolumeConfiguration, + Volume, + VolumeMountPoint, + VolumeProvisioningData, +) +from dstack._internal.utils.common import get_or_error +from dstack._internal.utils.docker import DEFAULT_REGISTRY, is_default_registry, parse_image_name +from dstack._internal.utils.logging import get_logger + +logger = get_logger(__name__) + +MAX_INSTANCE_NAME_LEN = 63 +RUNNER_SESSION = "dstack-runner" +# A replacement token revokes the previous one for new connections. Request one +# ten-year credential at provisioning and reuse it for the sandbox's lifetime. +SSH_ACCESS_MINUTES = 3650 * 24 * 60 +VOLUME_READY_TIMEOUT_SECONDS = 60 + + +class DaytonaCompute(ComputeWithFilteredOffersCached, ComputeWithVolumeSupport, Compute): + def __init__(self, config: DaytonaConfig): + super().__init__() + self.config = config + self.api_client = DaytonaAPIClient(config.creds.api_key) + self._catalog = gpuhunt.Catalog(balance_resources=False, auto_reload=False) + self._catalog.add_provider(DaytonaProvider(api_key=config.creds.api_key)) + self._organization_id: Optional[str] = None + + def get_offers_by_requirements( + self, requirements: Requirements, full_offers: bool, unallocated_resources: bool + ) -> List[InstanceOfferWithAvailability]: + offers = get_catalog_offers( + backend=BackendType.DAYTONA, + locations=self.config.regions or None, + requirements=requirements, + catalog=self._catalog, + ) + if not offers: + return [] + organization_id = self._organization_id + if organization_id is None: + organization_id = self.api_client.get_current_api_key().get("organizationId") + if not isinstance(organization_id, str) or not organization_id: + raise DaytonaAPIError("Daytona API key has no organization") + self._organization_id = organization_id + usage = self.api_client.get_organization_usage(organization_id) + classes = self.api_client.get_available_sandbox_classes(organization_id) + region_usage = { + row["regionId"]: row + for row in usage["regionUsage"] + if row["sandboxClass"] == "container" + } + region_classes = { + row["regionId"]: row for row in classes if row["sandboxClass"] == "container" + } + return [ + offer.with_availability( + availability=( + InstanceAvailability.AVAILABLE + if _has_quota( + offer, + region_usage.get(offer.region), + region_classes.get(offer.region), + ) + else InstanceAvailability.NO_QUOTA + ), + instance_runtime=InstanceRuntime.RUNNER, + ) + for offer in offers + ] + + def run_job( + self, + run: Run, + job: Job, + instance_offer: InstanceOfferWithAvailability, + project_ssh_public_key: str, + project_ssh_private_key: str, + volumes: List[Volume], + placement_group: Optional[PlacementGroup], + requirements: Requirements, + extra_authorized_keys: list[str], + ) -> JobProvisioningData: + image_name = job.job_spec.image_name + # Whitespace could inject additional Dockerfile instructions. + if not image_name or re.search(r"\s", image_name): + raise ComputeError("Invalid Daytona container image name") + instance_name = generate_unique_instance_name_for_job( + run, job, max_length=MAX_INSTANCE_NAME_LEN + ) + commands = get_docker_commands( + build_authorized_keys(project_ssh_public_key, extra_authorized_keys) + ) + backend_data = DaytonaInstanceBackendData( + startup_command=shlex.join(["sh", "-c", " && ".join(commands)]) + ) + registry: Optional[str] = None + if job.job_spec.registry_auth is not None: + image = parse_image_name(image_name) + registry = image.registry or DEFAULT_REGISTRY + if is_default_registry(registry): + registry = DEFAULT_REGISTRY + # Daytona's Dockerfile builder includes organization credentials only + # for FROM images with a registry prefix. + suffix = f"@{image.digest}" if image.digest else f":{image.tag}" + image_name = f"{registry}/{image.repo}{suffix}" + resources = instance_offer.instance.resources + payload = { + "name": instance_name, + "buildInfo": { + # The sandbox's user field does not override the image USER. + # Daytona caches builds, so updating a mutable tag such as latest may + # still reuse an older image. A digest-pinned image avoids this. + "dockerfileContent": ( + f'FROM {image_name}\nUSER root\nENTRYPOINT []\nCMD ["sleep", "infinity"]\n' + ) + }, + "user": "root", + "target": instance_offer.region, + "cpu": resources.cpus, + "memory": round(resources.memory_mib / 1024), + "disk": round(resources.disk.size_mib / 1024), + "gpu": len(resources.gpus), + "spot": resources.spot, + "public": False, + "autoStopInterval": 0, + "autoPauseInterval": 0, + "autoDeleteInterval": 0, + "ttlMinutes": 0, + } + if resources.gpus: + payload["gpuType"] = [_get_gpu_type(instance_offer)] + if volumes: + mount_points = job.job_spec.volumes + if mount_points is None: + mount_points = run.run_spec.configuration.volumes + paths = {} + for mount in mount_points: + if isinstance(mount, VolumeMountPoint): + names = [mount.name] if isinstance(mount.name, str) else mount.name + paths.update((name, mount.path) for name in names) + payload["volumes"] = [ + {"volumeId": get_or_error(volume.volume_id), "mountPath": paths[volume.name]} + for volume in volumes + ] + if job.job_spec.registry_auth is not None: + # Daytona selects registry credentials at organization scope, without a + # per-sandbox selector. Another record for the same registry can affect this pull. + backend_data.registry_id = self._create_registry( + f"{instance_name}-registry", get_or_error(registry), job.job_spec.registry_auth + ) + instance_id = instance_name + try: + sandbox = self.api_client.create_sandbox(payload) + if isinstance(sandbox.get("id"), str) and sandbox["id"]: + instance_id = sandbox["id"] + except DaytonaAPIError as e: + if e.status_code is not None and e.status_code < 500: + self._cleanup_failed_registry(backend_data.registry_id) + raise + # A timed-out create may have succeeded. Persist its unique name so + # provisioning can find it and termination can clean it up; never retry POST. + logger.warning( + "Daytona sandbox %s creation response was lost; checking by name", instance_name + ) + except Exception: + self._cleanup_failed_registry(backend_data.registry_id) + raise + return JobProvisioningData( + backend=BackendType.DAYTONA, + instance_type=instance_offer.instance, + instance_id=instance_id, + region=instance_offer.region, + price=instance_offer.price, + username="root", + dockerized=False, + backend_data=backend_data.model_dump_json(), + ) + + def _create_registry(self, name: str, url: str, auth: RegistryAuth) -> str: + try: + registry = self.api_client.create_registry( + {"name": name, "url": url, "username": auth.username, "password": auth.password} + ) + return registry["id"] + except DaytonaAPIError as e: + if e.status_code is not None and e.status_code < 500: + raise + # Recover the exact per-job record if its create response was lost. + for registry in self.api_client.get_registries(): + if registry["name"] == name: + return registry["id"] + raise + + def _cleanup_failed_registry(self, registry_id: Optional[str]) -> None: + if registry_id is None: + return + try: + self.api_client.delete_registry(registry_id) + except Exception: + logger.exception( + "Failed to delete Daytona registry %s after provisioning failed. " + "Delete it manually in Daytona.", + registry_id, + ) + + def update_provisioning_data( + self, + provisioning_data: JobProvisioningData, + project_ssh_public_key: str, + project_ssh_private_key: str, + ): + if provisioning_data.hostname is not None and provisioning_data.ssh_port is not None: + return + sandbox = self.api_client.get_sandbox(provisioning_data.instance_id) + if sandbox is None: + # The create response may have been lost before the sandbox became visible. + # The server's provisioning timeout still schedules deletion by name. + return + state = sandbox.get("state") + if state in { + "error", + "build_failed", + "destroyed", + "destroying", + "stopped", + "stopping", + "archived", + "archiving", + "paused", + "pausing", + }: + reason = sandbox.get("errorReason") or state + raise ProvisioningError(f"Daytona sandbox failed to start: {reason}") + if state != "started": + return + backend_data = validate_json_extra_ignore( + DaytonaInstanceBackendData, provisioning_data.backend_data or "{}" + ) + toolbox_url = self.api_client.get_toolbox_url(sandbox["id"]) + self._start_runner(toolbox_url, get_or_error(backend_data.startup_command)) + access = self.api_client.create_ssh_access( + sandbox["id"], expires_in_minutes=SSH_ACCESS_MINUTES + ) + provisioning_data.ssh_proxy = _get_ssh_proxy(access) + provisioning_data.hostname = "localhost" + provisioning_data.ssh_port = DSTACK_RUNNER_SSH_PORT + # The server checks runner readiness through the normal SSH connection. + backend_data.startup_command = None + provisioning_data.backend_data = ( + backend_data.model_dump_json() if backend_data.registry_id is not None else None + ) + + def _start_runner(self, toolbox_url: str, startup_command: str) -> None: + session = self.api_client.get_session(toolbox_url, RUNNER_SESSION) + if session is None: + self.api_client.create_session(toolbox_url, RUNNER_SESSION) + session = {"commands": []} + commands = session["commands"] + if not commands: + self.api_client.execute_session_command( + toolbox_url, RUNNER_SESSION, startup_command, run_async=True + ) + return + for command in commands: + if command.get("exitCode") is not None: + raise ProvisioningError( + "Daytona runner setup exited with code " + f"{command['exitCode']}. Check the dstack-runner session in Daytona." + ) + + def terminate_instance( + self, instance_id: str, region: str, backend_data: Optional[str] = None + ): + sandbox = self.api_client.get_sandbox(instance_id) + if sandbox is not None and sandbox.get("state") != "destroyed": + if sandbox.get("state") != "destroying": + self.api_client.delete_sandbox(instance_id) + raise NotYetTerminated("Waiting for Daytona sandbox deletion") + data = validate_json_extra_ignore(DaytonaInstanceBackendData, backend_data or "{}") + if data.registry_id is not None: + self.api_client.delete_registry(data.registry_id) + + def register_volume(self, volume: Volume) -> VolumeProvisioningData: + assert isinstance(volume.configuration, DaytonaVolumeConfiguration) + volume_id = get_or_error(volume.configuration.volume_id) + self._wait_for_volume(volume_id) + return _volume_provisioning_data(volume_id) + + def create_volume(self, volume: Volume) -> VolumeProvisioningData: + name = generate_unique_volume_name(volume) + try: + data = self.api_client.create_volume(name) + except DaytonaAPIError as e: + if e.status_code is not None and e.status_code < 500: + raise + data = self.api_client.get_volume_by_name(name) + if data is None: + raise + volume_id = data["id"] + try: + self._wait_for_volume(volume_id) + except Exception: + try: + self.api_client.delete_volume(volume_id) + except Exception: + logger.exception( + "Failed to delete Daytona volume %s after provisioning failed. " + "Delete it manually in Daytona.", + volume_id, + ) + raise + return _volume_provisioning_data(volume_id) + + def _wait_for_volume(self, volume_id: str) -> None: + deadline = time.monotonic() + VOLUME_READY_TIMEOUT_SECONDS + while True: + data = self.api_client.get_volume(volume_id) + if data is None: + raise ComputeError(f"Daytona volume {volume_id} not found") + state = data["state"] + if state == "ready": + return + if state not in {"pending_create", "creating"}: + raise ComputeError( + f"Daytona volume {volume_id} is {state}: {data.get('errorReason') or state}" + ) + if time.monotonic() >= deadline: + raise ComputeError(f"Timed out waiting for Daytona volume {volume_id}") + time.sleep(2) + + def delete_volume(self, volume: Volume) -> None: + if volume.volume_id is None: + return + data = self.api_client.get_volume(volume.volume_id) + if data is None or data["state"] in {"pending_delete", "deleting", "deleted"}: + return + self.api_client.delete_volume(volume.volume_id) + + +class DaytonaInstanceBackendData(CoreModel): + startup_command: Optional[str] = None + registry_id: Optional[str] = None + + +def _volume_provisioning_data(volume_id: str) -> VolumeProvisioningData: + return VolumeProvisioningData( + backend=BackendType.DAYTONA, + volume_id=volume_id, + size_gb=None, + price=0, + attachable=False, + detachable=False, + ) + + +def _has_quota( + offer: InstanceOffer, + usage: Optional[dict[str, Any]], + sandbox_class: Optional[dict[str, Any]], +) -> bool: + if usage is None or sandbox_class is None: + return False + resources = offer.instance.resources + sizes = { + "Cpu": resources.cpus, + "Memory": resources.memory_mib / 1024, + "Disk": resources.disk.size_mib / 1024, + } + gpu_count = len(resources.gpus) + if gpu_count: + if not sandbox_class.get("gpuAvailable"): + return False + if not resources.spot and gpu_count > usage["totalGpuQuota"] - usage["currentGpuUsage"]: + return False + for limits in (usage, sandbox_class): + allowed_types = limits.get("allowedGpuTypes") + if allowed_types is not None and _get_gpu_type(offer) not in allowed_types: + return False + # GPU sandboxes use separate quotas: CPU/RAM/disk aggregate totals may be zero. + # Check per-GPU limits, falling back to per-sandbox limits. + for name, size in sizes.items(): + limit = usage.get(f"max{name}PerGpu") + if limit is not None: + limit *= gpu_count + else: + limit = usage.get(f"max{name}PerSandbox") + if limit is not None and size > limit: + return False + return True + for name, size in sizes.items(): + if size > usage[f"total{name}Quota"] - usage[f"current{name}Usage"]: + return False + limit = usage.get(f"max{name}PerSandbox") + if limit is not None and size > limit: + return False + return True + + +def _get_gpu_type(offer: InstanceOffer) -> str: + return offer.backend_data.get("gpu_type") or offer.instance.resources.gpus[0].name + + +def _get_ssh_proxy(access: dict[str, Any]) -> SSHConnectionParams: + # Parse only the endpoint; never execute the provider's command or log its token. + try: + parts = shlex.split(access["sshCommand"]) + target = next(part for part in parts[1:] if "@" in part) + _, hostname = target.rsplit("@", 1) + port = int(parts[parts.index("-p") + 1]) if "-p" in parts else 22 + if parts[0] != "ssh" or not hostname or not access["token"] or not 1 <= port <= 65535: + raise ValueError + except (KeyError, IndexError, StopIteration, TypeError, ValueError) as e: + raise ProvisioningError("Daytona returned invalid SSH connection details") from e + return SSHConnectionParams(hostname=hostname, port=port, username=access["token"]) diff --git a/src/dstack/_internal/core/backends/daytona/configurator.py b/src/dstack/_internal/core/backends/daytona/configurator.py new file mode 100644 index 0000000000..184e51d4ac --- /dev/null +++ b/src/dstack/_internal/core/backends/daytona/configurator.py @@ -0,0 +1,75 @@ +import json + +from gpuhunt.providers.daytona import GPU_REGION + +from dstack._internal.core.backends.base.configurator import BackendRecord, Configurator +from dstack._internal.core.backends.daytona.api_client import DaytonaAPIClient +from dstack._internal.core.backends.daytona.backend import DaytonaBackend +from dstack._internal.core.backends.daytona.models import ( + DaytonaBackendConfig, + DaytonaBackendConfigWithCreds, + DaytonaConfig, + DaytonaCreds, + DaytonaStoredConfig, +) +from dstack._internal.core.errors import BackendInvalidCredentialsError, ServerClientError +from dstack._internal.core.models.backends.base import BackendType +from dstack._internal.core.models.common import validate_extra_ignore, validate_json_extra_ignore + + +class DaytonaConfigurator(Configurator[DaytonaBackendConfig, DaytonaBackendConfigWithCreds]): + TYPE = BackendType.DAYTONA + BACKEND_CLASS = DaytonaBackend + + def validate_config(self, config: DaytonaBackendConfigWithCreds, default_creds_enabled: bool): + client = DaytonaAPIClient(api_key=config.creds.api_key) + api_key = client.get_current_api_key() + required_permissions = {"write:sandboxes", "delete:sandboxes", "read:limits"} + missing_permissions = sorted(required_permissions - set(api_key["permissions"])) + if missing_permissions: + raise BackendInvalidCredentialsError( + msg=f"Daytona API key is missing permissions: {', '.join(missing_permissions)}", + fields=[["creds", "api_key"]], + ) + if not config.regions: + return + regions = {GPU_REGION, *(region["id"] for region in client.get_shared_regions())} + invalid_regions = sorted(set(config.regions) - regions) + if invalid_regions: + raise ServerClientError( + msg=( + f"Unsupported Daytona regions: {invalid_regions}. " + f"Supported regions: {sorted(regions)}." + ), + fields=[["regions"]], + ) + + def create_backend( + self, project_name: str, config: DaytonaBackendConfigWithCreds + ) -> BackendRecord: + return BackendRecord( + config=DaytonaStoredConfig( + **validate_extra_ignore(DaytonaBackendConfig, config).model_dump() + ).model_dump_json(), + auth=DaytonaCreds.model_validate(config.creds).model_dump_json(), + ) + + def get_backend_config_with_creds( + self, record: BackendRecord + ) -> DaytonaBackendConfigWithCreds: + return validate_extra_ignore(DaytonaBackendConfigWithCreds, self._get_config(record)) + + def get_backend_config_without_creds(self, record: BackendRecord) -> DaytonaBackendConfig: + return validate_extra_ignore(DaytonaBackendConfig, self._get_config(record)) + + def get_backend(self, record: BackendRecord) -> DaytonaBackend: + return DaytonaBackend(config=self._get_config(record)) + + def _get_config(self, record: BackendRecord) -> DaytonaConfig: + return validate_extra_ignore( + DaytonaConfig, + { + **json.loads(record.config), + "creds": validate_json_extra_ignore(DaytonaCreds, record.auth), + }, + ) diff --git a/src/dstack/_internal/core/backends/daytona/models.py b/src/dstack/_internal/core/backends/daytona/models.py new file mode 100644 index 0000000000..3ee230d33d --- /dev/null +++ b/src/dstack/_internal/core/backends/daytona/models.py @@ -0,0 +1,42 @@ +from typing import Annotated, Literal, Optional, Union + +from pydantic import Field + +from dstack._internal.core.models.common import CoreModel + + +class DaytonaAPIKeyCreds(CoreModel): + type: Annotated[Literal["api_key"], Field(description="The type of credentials")] = "api_key" + api_key: Annotated[str, Field(description="The Daytona API key", min_length=1)] + + +AnyDaytonaCreds = DaytonaAPIKeyCreds +DaytonaCreds = AnyDaytonaCreds + + +class DaytonaBackendConfig(CoreModel): + type: Annotated[Literal["daytona"], Field(description="The type of backend")] = "daytona" + regions: Annotated[ + Optional[list[str]], + Field( + description=( + "The list of Daytona regions. GPU sandboxes use `earth`; CPU sandboxes use " + "shared regions such as `us` and `eu`. Omit to use all regions" + ) + ), + ] = None + + +class DaytonaBackendConfigWithCreds(DaytonaBackendConfig): + creds: Annotated[AnyDaytonaCreds, Field(description="The credentials")] + + +AnyDaytonaBackendConfig = Union[DaytonaBackendConfig, DaytonaBackendConfigWithCreds] + + +class DaytonaStoredConfig(DaytonaBackendConfig): + pass + + +class DaytonaConfig(DaytonaStoredConfig): + creds: AnyDaytonaCreds diff --git a/src/dstack/_internal/core/backends/models.py b/src/dstack/_internal/core/backends/models.py index 071d2b1b9f..94d9dcd094 100644 --- a/src/dstack/_internal/core/backends/models.py +++ b/src/dstack/_internal/core/backends/models.py @@ -23,6 +23,10 @@ CudoBackendConfig, CudoBackendConfigWithCreds, ) +from dstack._internal.core.backends.daytona.models import ( + DaytonaBackendConfig, + DaytonaBackendConfigWithCreds, +) from dstack._internal.core.backends.digitalocean_base.models import ( BaseDigitalOceanBackendConfig, BaseDigitalOceanBackendConfigWithCreds, @@ -102,6 +106,7 @@ CloudRiftBackendConfig, CrusoeBackendConfig, CudoBackendConfig, + DaytonaBackendConfig, BaseDigitalOceanBackendConfig, GCPBackendConfig, HotAisleBackendConfig, @@ -131,6 +136,7 @@ CrusoeBackendConfigWithCreds, CudoBackendConfigWithCreds, VerdaBackendConfigWithCreds, + DaytonaBackendConfigWithCreds, BaseDigitalOceanBackendConfigWithCreds, GCPBackendConfigWithCreds, HotAisleBackendConfigWithCreds, @@ -175,6 +181,7 @@ class BackendConfigWithCreds(RootModel[AnyBackendConfigWithCredsTagged]): CrusoeBackendFileConfigWithCreds, CudoBackendConfigWithCreds, VerdaBackendConfigWithCreds, + DaytonaBackendConfigWithCreds, BaseDigitalOceanBackendConfigWithCreds, GCPBackendFileConfigWithCreds, HotAisleBackendFileConfigWithCreds, diff --git a/src/dstack/_internal/core/models/backends/base.py b/src/dstack/_internal/core/models/backends/base.py index b42b7bf575..6ee2931b18 100644 --- a/src/dstack/_internal/core/models/backends/base.py +++ b/src/dstack/_internal/core/models/backends/base.py @@ -11,6 +11,7 @@ class BackendType(str, enum.Enum): CRUSOE (BackendType): Crusoe CUDO (BackendType): Cudo DATACRUNCH (BackendType): DataCrunch (for backward compatibility) + DAYTONA (BackendType): Daytona DIGITALOCEAN (BackendType): DigitalOcean DSTACK (BackendType): dstack Sky GCP (BackendType): Google Cloud Platform @@ -37,6 +38,7 @@ class BackendType(str, enum.Enum): CUDO = "cudo" DATACRUNCH = "datacrunch" """`DATACRUNCH` is kept as a `BackendType` for backward compatibility.""" + DAYTONA = "daytona" DIGITALOCEAN = "digitalocean" DSTACK = "dstack" GCP = "gcp" diff --git a/src/dstack/_internal/core/models/volumes.py b/src/dstack/_internal/core/models/volumes.py index d651e717a4..92b3a650ee 100644 --- a/src/dstack/_internal/core/models/volumes.py +++ b/src/dstack/_internal/core/models/volumes.py @@ -42,10 +42,6 @@ class BaseVolumeConfiguration(CoreModel): > Variable is mutable so its type is invariant """ name: Annotated[Optional[str], Field(description="The volume name")] = None - size: Annotated[ - Optional[Memory], - Field(description="The volume size. Must be specified when creating new volumes"), - ] = None auto_cleanup_duration: Annotated[ OptionalIdleDuration, Field( @@ -81,6 +77,13 @@ def external_volume_id(self) -> Optional[str]: def is_external(self) -> bool: return self.external_volume_id is not None + +class VolumeConfigurationWithSize(BaseVolumeConfiguration): + size: Annotated[ + Optional[Memory], + Field(description="The volume size. Must be specified when creating new volumes"), + ] = None + @property def size_gb(self) -> int: return int(get_or_error(self.size)) @@ -108,7 +111,9 @@ def external_volume_id(self) -> Optional[str]: class AWSVolumeConfiguration( - VolumeConfigurationWithAvailibilityZone, VolumeConfigurationWithVolumeID + VolumeConfigurationWithAvailibilityZone, + VolumeConfigurationWithVolumeID, + VolumeConfigurationWithSize, ): backend: Annotated[Literal[BackendType.AWS], Field(description="The volume backend")] = ( BackendType.AWS @@ -116,14 +121,18 @@ class AWSVolumeConfiguration( class GCPVolumeConfiguration( - VolumeConfigurationWithAvailibilityZone, VolumeConfigurationWithVolumeID + VolumeConfigurationWithAvailibilityZone, + VolumeConfigurationWithVolumeID, + VolumeConfigurationWithSize, ): backend: Annotated[Literal[BackendType.GCP], Field(description="The volume backend")] = ( BackendType.GCP ) -class RunpodVolumeConfiguration(VolumeConfigurationWithRegion, VolumeConfigurationWithVolumeID): +class RunpodVolumeConfiguration( + VolumeConfigurationWithRegion, VolumeConfigurationWithVolumeID, VolumeConfigurationWithSize +): backend: Annotated[Literal[BackendType.RUNPOD], Field(description="The volume backend")] = ( BackendType.RUNPOD ) @@ -131,7 +140,7 @@ class RunpodVolumeConfiguration(VolumeConfigurationWithRegion, VolumeConfigurati """Runpod doesn't have AZs but we accept this field for compatibility with older clients.""" -class KubernetesVolumeConfiguration(VolumeConfigurationWithRegion): +class KubernetesVolumeConfiguration(VolumeConfigurationWithRegion, VolumeConfigurationWithSize): backend: Annotated[ Literal[BackendType.KUBERNETES], Field(description="The volume backend") ] = BackendType.KUBERNETES @@ -173,11 +182,18 @@ def external_volume_id(self) -> Optional[str]: return self.claim_name +class DaytonaVolumeConfiguration(VolumeConfigurationWithVolumeID): + backend: Annotated[Literal[BackendType.DAYTONA], Field(description="The volume backend")] = ( + BackendType.DAYTONA + ) + + AnyVolumeConfiguration = Union[ AWSVolumeConfiguration, GCPVolumeConfiguration, RunpodVolumeConfiguration, KubernetesVolumeConfiguration, + DaytonaVolumeConfiguration, ] @@ -202,7 +218,8 @@ class VolumeSpec(CoreModel): class VolumeProvisioningData(CoreModel): backend: Optional[BackendType] = None volume_id: str - size_gb: int + size_gb: Optional[int] + """The provisioned capacity, or `None` for volumes without a fixed capacity.""" availability_zone: Optional[str] = None price: Optional[float] = None """`price` stores the monthly price.""" @@ -263,6 +280,12 @@ def get_attachment_data_for_instance(self, instance_id: str) -> Optional[VolumeA def get_backend(self) -> BackendType: return self.configuration.backend + def matches_location(self, backend: BackendType, region: str) -> bool: + return self.get_backend() == backend and ( + not isinstance(self.configuration, VolumeConfigurationWithRegion) + or self.configuration.region.lower() == region.lower() + ) + def get_region(self) -> str: """ Returns the volume region or an empty string if the volume (that is, its backend) diff --git a/src/dstack/_internal/server/background/pipeline_tasks/common.py b/src/dstack/_internal/server/background/pipeline_tasks/common.py index 3d0327075c..d6f9564452 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/common.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/common.py @@ -12,6 +12,8 @@ def get_provisioning_timeout(backend_type: BackendType, instance_type_name: str) return timedelta(minutes=30) if backend_type == BackendType.RUNPOD: return timedelta(minutes=20) + if backend_type == BackendType.DAYTONA: + return timedelta(minutes=20) if backend_type == BackendType.SEEWEB: # Seeweb GPU servers install drivers/Docker on first boot, which can exceed 10 minutes. return timedelta(minutes=20) diff --git a/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py b/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py index 4a1796847d..b0f592b0e6 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py @@ -2724,10 +2724,7 @@ def _get_offer_mount_point_volume( offer: InstanceOfferWithAvailability, ) -> Volume: for volume in volumes: - if ( - volume.get_backend() != offer.backend - or volume.get_region().lower() != offer.region.lower() - ): + if not volume.matches_location(offer.backend, offer.region): continue return volume raise ServerClientError("Failed to find an eligible volume for the mount point") diff --git a/src/dstack/_internal/server/services/instances.py b/src/dstack/_internal/server/services/instances.py index 2b763c93ad..cdcab7ce34 100644 --- a/src/dstack/_internal/server/services/instances.py +++ b/src/dstack/_internal/server/services/instances.py @@ -43,7 +43,7 @@ TerminationPolicy, ) from dstack._internal.core.models.runs import JobProvisioningData, Requirements -from dstack._internal.core.models.volumes import Volume +from dstack._internal.core.models.volumes import Volume, VolumeConfigurationWithRegion from dstack._internal.core.services.profiles import get_termination from dstack._internal.server import settings as server_settings from dstack._internal.server.models import ( @@ -610,14 +610,17 @@ def filter_instances( backend_types: Optional[list[BackendType]] = profile.backends regions: Optional[list[str]] = profile.regions zones: Optional[list[str]] = profile.availability_zones - # (BackendType, region.lower() | "", availability_zone.lower() | None) - volumes_locations: Optional[set[tuple[BackendType, str, Optional[str]]]] = None + # (BackendType, region.lower() | None, availability_zone.lower() | None). + # A None region matches any region; an empty string remains an exact region. + volumes_locations: Optional[set[tuple[BackendType, Optional[str], Optional[str]]]] = None if volumes: volumes_locations = set() for volume in volumes[0]: volume_backend = volume.get_backend() - volume_region = volume.get_region().lower() + volume_region = None + if isinstance(volume.configuration, VolumeConfigurationWithRegion): + volume_region = volume.get_region().lower() # If the volume has an AZ, it's added twice -- with and without an AZ. # When the instance location is checked against the available volumes locations (see # below) the instance with an AZ matches only the volume with the same AZ, while @@ -677,7 +680,12 @@ def filter_instances( instance_zone = jpd.availability_zone if instance_zone is not None: instance_zone = instance_zone.lower() - if (instance_backend, instance_region, instance_zone) not in volumes_locations: + instance_location = (instance_backend, instance_region, instance_zone) + regionless_location = (instance_backend, None, instance_zone) + if ( + instance_location not in volumes_locations + and regionless_location not in volumes_locations + ): continue filtered_instances.append(instance) return filtered_instances diff --git a/src/dstack/_internal/server/services/jobs/__init__.py b/src/dstack/_internal/server/services/jobs/__init__.py index ffd3078c21..0326be9aca 100644 --- a/src/dstack/_internal/server/services/jobs/__init__.py +++ b/src/dstack/_internal/server/services/jobs/__init__.py @@ -597,9 +597,8 @@ def _get_job_mount_point_attached_volume( Returns the volume attached to the job among the list of possible mount point volumes. """ for volume in volumes: - if ( - volume.get_backend() != job_provisioning_data.get_base_backend() - or volume.get_region().lower() != job_provisioning_data.region.lower() + if not volume.matches_location( + job_provisioning_data.get_base_backend(), job_provisioning_data.region ): continue if ( diff --git a/src/dstack/_internal/server/services/offers.py b/src/dstack/_internal/server/services/offers.py index 8193ab0863..7fbe24f5c3 100644 --- a/src/dstack/_internal/server/services/offers.py +++ b/src/dstack/_internal/server/services/offers.py @@ -22,7 +22,7 @@ from dstack._internal.core.models.placement import PlacementGroup from dstack._internal.core.models.profiles import Profile from dstack._internal.core.models.runs import JobProvisioningData, Requirements -from dstack._internal.core.models.volumes import Volume +from dstack._internal.core.models.volumes import Volume, VolumeConfigurationWithRegion from dstack._internal.server.models import ProjectModel from dstack._internal.server.services import backends as backends_services @@ -49,11 +49,19 @@ async def get_offers_by_requirements( regions: Optional[list[str]] = profile.regions availability_zones: Optional[list[str]] = profile.availability_zones instance_types: Optional[list[str]] = profile.instance_types - # (BackendType, region.lower() | "") - volumes_locations: Optional[set[tuple[BackendType, str]]] = None + # (BackendType, region.lower() | None). None means any region. + volumes_locations: Optional[set[tuple[BackendType, Optional[str]]]] = None if volumes: - volumes_locations = {(v.get_backend(), v.get_region().lower()) for v in volumes[0]} + volumes_locations = { + ( + v.get_backend(), + v.get_region().lower() + if isinstance(v.configuration, VolumeConfigurationWithRegion) + else None, + ) + for v in volumes[0] + } if multinode: if backend_types is None: @@ -206,7 +214,7 @@ def _filter_offers( availability_zones: Optional[List[str]] = None, instance_types: Optional[List[str]] = None, placement_group: Optional[PlacementGroup] = None, - volumes_locations: Optional[Container[tuple[BackendType, str]]] = None, + volumes_locations: Optional[Container[tuple[BackendType, Optional[str]]]] = None, ) -> Iterator[Tuple[Backend, InstanceOfferWithAvailability]]: """ Yields filtered offers. May return modified offers to match the filters. @@ -246,6 +254,7 @@ def _filter_offers( if ( volumes_locations is not None and (offer.backend, offer.region.lower()) not in volumes_locations + and (offer.backend, None) not in volumes_locations ): continue yield (b, offer) diff --git a/src/dstack/_internal/server/services/volumes.py b/src/dstack/_internal/server/services/volumes.py index 374fe76a4c..9d648eee46 100644 --- a/src/dstack/_internal/server/services/volumes.py +++ b/src/dstack/_internal/server/services/volumes.py @@ -22,6 +22,7 @@ VolumeAttachment, VolumeAttachmentData, VolumeConfiguration, + VolumeConfigurationWithSize, VolumeInstance, VolumeProvisioningData, VolumeSpec, @@ -474,7 +475,11 @@ async def generate_volume_name(session: AsyncSession, project: ProjectModel) -> def _validate_volume_configuration(configuration: AnyVolumeConfiguration): - if configuration.external_volume_id is None and configuration.size is None: + if ( + isinstance(configuration, VolumeConfigurationWithSize) + and configuration.external_volume_id is None + and configuration.size is None + ): raise ServerClientError("Volume must specify either existing identifier or size") backends_services.check_backend_type_available(configuration.backend) if configuration.backend not in BACKENDS_WITH_VOLUMES_SUPPORT: diff --git a/src/tests/_internal/core/backends/daytona/__init__.py b/src/tests/_internal/core/backends/daytona/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/tests/_internal/core/backends/daytona/test_api_client.py b/src/tests/_internal/core/backends/daytona/test_api_client.py new file mode 100644 index 0000000000..a3198c6a81 --- /dev/null +++ b/src/tests/_internal/core/backends/daytona/test_api_client.py @@ -0,0 +1,242 @@ +import pytest +import requests + +from dstack._internal.core.backends.daytona.api_client import ( + API_URL, + DaytonaAPIClient, + DaytonaAPIError, +) +from dstack._internal.core.errors import BackendInvalidCredentialsError, NoCapacityError + +TOOLBOX_URL = "https://proxy.app.daytona.io/toolbox/sandbox-id" + + +class TestDaytonaAPIClient: + def test_identity_uses_bearer_auth(self, requests_mock): + requests_mock.get( + f"{API_URL}/api-keys/current", json={"organizationId": "organization-id"} + ) + + identity = DaytonaAPIClient("test-api-key").get_current_api_key() + + assert identity["organizationId"] == "organization-id" + assert requests_mock.last_request.headers["Authorization"] == "Bearer test-api-key" + assert requests_mock.last_request.timeout == 30 + + def test_unauthorized_maps_to_invalid_credentials(self, requests_mock): + requests_mock.get(f"{API_URL}/api-keys/current", status_code=401) + + with pytest.raises(BackendInvalidCredentialsError) as exc: + DaytonaAPIClient("bad-key").get_current_api_key() + + assert exc.value.fields == [["creds", "api_key"]] + + @pytest.mark.parametrize( + ["status", "message"], + [ + (400, "No capacity available"), + (429, "Too many requests"), + (503, "Service unavailable"), + ], + ) + def test_http_errors_preserve_status_and_reason(self, requests_mock, status, message): + requests_mock.post(f"{API_URL}/sandbox", status_code=status, json={"message": message}) + + with pytest.raises(DaytonaAPIError, match=message) as exc: + DaytonaAPIClient("test-api-key").create_sandbox({"name": "test"}) + + assert exc.value.status_code == status + assert not isinstance(exc.value, NoCapacityError) + assert requests_mock.call_count == 1 + + @pytest.mark.parametrize("error", [requests.ReadTimeout, requests.ConnectionError]) + def test_network_errors_preserve_cause_without_retrying(self, requests_mock, error): + cause = error("Connection failed") + requests_mock.post(f"{API_URL}/sandbox", exc=cause) + client = DaytonaAPIClient("test-api-key") + + with pytest.raises(DaytonaAPIError) as exc: + client.create_sandbox({"name": "test"}) + + assert exc.value.__cause__ is cause + assert exc.value.status_code is None + assert requests_mock.call_count == 1 + + def test_error_response_redacts_api_key(self, requests_mock): + requests_mock.get( + f"{API_URL}/api-keys/current", + status_code=403, + json={"message": "Denied for Bearer secret-api-key"}, + ) + + with pytest.raises(DaytonaAPIError, match="Denied for Bearer \\[redacted\\]"): + DaytonaAPIClient("secret-api-key").get_current_api_key() + + def test_non_json_error_omits_response_body(self, requests_mock): + requests_mock.get( + f"{API_URL}/api-keys/current", status_code=502, text="gateway HTML with credentials" + ) + + with pytest.raises(DaytonaAPIError) as exc: + DaytonaAPIClient("test-api-key").get_current_api_key() + + assert str(exc.value) == "Daytona API returned HTTP 502" + + def test_redirect_is_not_followed(self, requests_mock): + requests_mock.get( + f"{API_URL}/api-keys/current", + status_code=302, + headers={"Location": "https://untrusted.example/identity"}, + ) + + with pytest.raises(DaytonaAPIError) as exc: + DaytonaAPIClient("test-api-key").get_current_api_key() + + assert exc.value.status_code == 302 + assert requests_mock.call_count == 1 + + def test_get_sandbox_accepts_id_or_name_and_returns_none_only_for_404(self, requests_mock): + requests_mock.get(f"{API_URL}/sandbox/sandbox-id", json={"id": "sandbox-id"}) + requests_mock.get(f"{API_URL}/sandbox/job-name", json={"id": "sandbox-id"}) + requests_mock.get(f"{API_URL}/sandbox/missing", status_code=404) + requests_mock.get(f"{API_URL}/sandbox/forbidden", status_code=403) + client = DaytonaAPIClient("test-api-key") + + assert client.get_sandbox("sandbox-id") == {"id": "sandbox-id"} + assert client.get_sandbox("job-name") == {"id": "sandbox-id"} + assert client.get_sandbox("missing") is None + with pytest.raises(DaytonaAPIError): + client.get_sandbox("forbidden") + + def test_volume_create_lookup_and_deleted_tombstone(self, requests_mock): + pending = {"id": "volume-id", "state": "pending_create"} + deleted = {"id": "volume-id", "state": "deleted"} + requests_mock.post(f"{API_URL}/volumes", json=pending) + requests_mock.get(f"{API_URL}/volumes/by-name/volume-name", json=pending) + requests_mock.get(f"{API_URL}/volumes/volume-id", json=deleted) + requests_mock.get(f"{API_URL}/volumes/missing", status_code=404) + requests_mock.get(f"{API_URL}/volumes/by-name/missing", status_code=404) + client = DaytonaAPIClient("test-api-key") + + assert client.create_volume("volume-name") == pending + assert requests_mock.last_request.json() == {"name": "volume-name"} + assert client.get_volume_by_name("volume-name") == pending + assert client.get_volume("volume-id") == deleted + assert client.get_volume("missing") is None + assert client.get_volume_by_name("missing") is None + + @pytest.mark.parametrize( + ["method", "path", "status"], + [ + ("delete_sandbox", "/sandbox/resource-id", 200), + ("delete_registry", "/docker-registry/resource-id", 204), + ("delete_volume", "/volumes/resource-id", 404), + ], + ) + def test_delete_endpoints_are_idempotent(self, requests_mock, method, path, status): + requests_mock.delete(API_URL + path, status_code=status) + + getattr(DaytonaAPIClient("test-api-key"), method)("resource-id") + + @pytest.mark.parametrize( + ["method", "path", "failure"], + [ + ("delete_sandbox", "/sandbox/resource-id", {"status_code": 503}), + ("delete_registry", "/docker-registry/resource-id", {"status_code": 429}), + ("delete_volume", "/volumes/resource-id", {"exc": requests.ReadTimeout}), + ], + ) + def test_delete_retries_transient_errors(self, requests_mock, mocker, method, path, failure): + sleep = mocker.patch("dstack._internal.core.backends.daytona.api_client.time.sleep") + requests_mock.delete(API_URL + path, [failure, {"status_code": 404}]) + + getattr(DaytonaAPIClient("test-api-key"), method)("resource-id") + + assert requests_mock.call_count == 2 + sleep.assert_called_once_with(1) + + def test_delete_stops_retrying_after_three_attempts(self, requests_mock, mocker): + sleep = mocker.patch("dstack._internal.core.backends.daytona.api_client.time.sleep") + requests_mock.delete(f"{API_URL}/docker-registry/registry-id", status_code=503) + + with pytest.raises(DaytonaAPIError) as exc: + DaytonaAPIClient("test-api-key").delete_registry("registry-id") + + assert exc.value.status_code == 503 + assert requests_mock.call_count == 3 + assert [call.args[0] for call in sleep.call_args_list] == [1, 2] + + def test_delete_does_not_retry_permission_errors(self, requests_mock, mocker): + sleep = mocker.patch("dstack._internal.core.backends.daytona.api_client.time.sleep") + requests_mock.delete(f"{API_URL}/docker-registry/registry-id", status_code=403) + + with pytest.raises(DaytonaAPIError) as exc: + DaytonaAPIClient("test-api-key").delete_registry("registry-id") + + assert exc.value.status_code == 403 + assert requests_mock.call_count == 1 + sleep.assert_not_called() + + def test_toolbox_url_includes_sandbox_id(self, requests_mock): + requests_mock.get( + f"{API_URL}/sandbox/sandbox-id/toolbox-proxy-url", + json={"url": "https://proxy.app.daytona.io/toolbox/"}, + ) + + assert DaytonaAPIClient("test-api-key").get_toolbox_url("sandbox-id") == TOOLBOX_URL + + @pytest.mark.parametrize( + "url", + [ + "http://proxy.app.daytona.io/toolbox", + "https://daytona.io.attacker.example/toolbox", + "https://attacker.example/toolbox", + "https://user:password@proxy.app.daytona.io/toolbox", + "https://proxy.app.daytona.io/toolbox?token=secret", + "https://proxy.app.daytona.io:8443/toolbox", + ], + ) + def test_toolbox_rejects_untrusted_url_before_sending_credentials(self, requests_mock, url): + with pytest.raises(DaytonaAPIError, match="invalid toolbox proxy URL"): + DaytonaAPIClient("test-api-key").get_session(url, "runner") + + assert requests_mock.call_count == 0 + + def test_session_creation_and_async_command(self, requests_mock): + requests_mock.post(f"{TOOLBOX_URL}/process/session", status_code=201) + requests_mock.post( + f"{TOOLBOX_URL}/process/session/runner/exec", json={"cmdId": "command-id"} + ) + client = DaytonaAPIClient("test-api-key") + + client.create_session(TOOLBOX_URL, "runner") + command = client.execute_session_command(TOOLBOX_URL, "runner", "/tmp/dstack-runner start") + + assert requests_mock.request_history[0].json() == {"sessionId": "runner"} + assert requests_mock.last_request.json() == { + "command": "/tmp/dstack-runner start", + "runAsync": True, + } + assert requests_mock.last_request.headers["Authorization"] == "Bearer test-api-key" + assert command == {"cmdId": "command-id"} + + def test_missing_session(self, requests_mock): + requests_mock.get(f"{TOOLBOX_URL}/process/session/runner", status_code=404) + client = DaytonaAPIClient("test-api-key") + + assert client.get_session(TOOLBOX_URL, "runner") is None + + def test_ssh_access_preserves_explicit_long_lifetime(self, requests_mock): + access = { + "token": "ssh-access-token", + "sshCommand": "ssh -p 2222 ssh-access-token@ssh.app.daytona.io", + "expiresAt": "2027-09-28T00:00:00Z", + } + requests_mock.post(f"{API_URL}/sandbox/sandbox-id/ssh-access", json=access) + + result = DaytonaAPIClient("test-api-key").create_ssh_access( + "sandbox-id", expires_in_minutes=365 * 24 * 60 + ) + + assert result == access + assert requests_mock.last_request.qs == {"expiresinminutes": ["525600"]} diff --git a/src/tests/_internal/core/backends/daytona/test_compute.py b/src/tests/_internal/core/backends/daytona/test_compute.py new file mode 100644 index 0000000000..b4d6566878 --- /dev/null +++ b/src/tests/_internal/core/backends/daytona/test_compute.py @@ -0,0 +1,867 @@ +import json +import shlex +from types import SimpleNamespace +from unittest.mock import patch + +import gpuhunt +import pytest + +from dstack._internal.core.backends.daytona.api_client import ( + API_URL, + DaytonaAPIError, +) +from dstack._internal.core.backends.daytona.compute import DaytonaCompute +from dstack._internal.core.backends.daytona.models import DaytonaConfig, DaytonaCreds +from dstack._internal.core.errors import ComputeError, NotYetTerminated, ProvisioningError +from dstack._internal.core.models.backends.base import BackendType +from dstack._internal.core.models.common import RegistryAuth +from dstack._internal.core.models.instances import ( + InstanceAvailability, + InstanceRuntime, +) +from dstack._internal.core.models.resources import ResourcesSpec +from dstack._internal.core.models.runs import JobProvisioningData, Requirements +from dstack._internal.core.models.volumes import DaytonaVolumeConfiguration, VolumeMountPoint + +COMPUTE_MODULE = "dstack._internal.core.backends.daytona.compute" +PUBLIC_KEY = ( + "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAINOmx0T+hBRaJ6jCi21ZYe2NW3EZS8e0Mdwl+yZJt+kD project" +) +TOOLBOX_URL = "https://proxy.app.daytona.io/sandbox-id" +STARTUP_COMMANDS = ["install runner", "runner --key 'a key with spaces'"] + + +@pytest.fixture +def compute(): + with patch(f"{COMPUTE_MODULE}.DaytonaAPIClient", autospec=True) as client_class: + compute = DaytonaCompute(DaytonaConfig(creds=DaytonaCreds(api_key="test-api-key"))) + client = client_class.return_value + client.get_current_api_key.return_value = {"organizationId": "org-id"} + client.get_organization_usage.return_value = {"regionUsage": [_cpu_usage(), _gpu_usage()]} + client.get_available_sandbox_classes.return_value = [ + {"regionId": "us", "sandboxClass": "container", "gpuAvailable": False}, + {"regionId": "earth", "sandboxClass": "container", "gpuAvailable": True}, + ] + client.create_sandbox.return_value = {"id": "sandbox-id", "state": "creating"} + client.create_registry.return_value = {"id": "registry-id"} + client.get_registries.return_value = [] + client.create_volume.return_value = {"id": "volume-id", "state": "pending_create"} + client.get_volume.return_value = {"id": "volume-id", "state": "ready"} + client.get_sandbox.return_value = {"id": "sandbox-id", "state": "started"} + client.get_toolbox_url.return_value = TOOLBOX_URL + client.get_session.return_value = { + "sessionId": "dstack-runner", + "commands": [{"id": "command-id", "command": "runner"}], + } + client.create_ssh_access.return_value = { + "token": "ssh-token_123", + "sshCommand": "ssh ssh-token_123@ssh.app.daytona.io -p 2222", + } + yield compute + + +def _cpu_usage(**overrides): + return { + "regionId": "us", + "sandboxClass": "container", + "totalCpuQuota": 100, + "totalMemoryQuota": 200, + "totalDiskQuota": 1000, + "currentCpuUsage": 0, + "currentMemoryUsage": 0, + "currentDiskUsage": 0, + "maxCpuPerSandbox": None, + "maxMemoryPerSandbox": None, + "maxDiskPerSandbox": None, + **overrides, + } + + +def _gpu_usage(**overrides): + return { + "regionId": "earth", + "sandboxClass": "container", + # GPU sandbox CPU/RAM/disk are not charged against these aggregate quotas. + "totalCpuQuota": 0, + "totalMemoryQuota": 0, + "totalDiskQuota": 0, + "currentCpuUsage": 0, + "currentMemoryUsage": 0, + "currentDiskUsage": 0, + "totalGpuQuota": 4, + "currentGpuUsage": 0, + "maxCpuPerGpu": 16, + "maxMemoryPerGpu": 192, + "maxDiskPerGpu": 512, + **overrides, + } + + +def _item(*, gpu=False, spot=False, cpu=4, memory=16, disk=101, region=None): + return gpuhunt.CatalogItem( + provider="daytona", + instance_name="test-shape", + location=region or ("earth" if gpu else "us"), + price=0.512345, + cpu=cpu, + memory=memory, + disk_size=disk, + gpu_count=2 if gpu else 0, + gpu_name="RTXPRO6000" if gpu else None, + gpu_vendor=gpuhunt.AcceleratorVendor.NVIDIA if gpu else None, + gpu_memory=96 if gpu else None, + spot=spot, + provider_data={"gpu_type": "RTX-PRO-6000"} if gpu else {}, + ) + + +def _offers(compute, *items): + with patch.object(compute._catalog, "query", return_value=list(items)): + return compute.get_offers_by_requirements( + Requirements(resources=ResourcesSpec()), + full_offers=False, + unallocated_resources=False, + ) + + +def _run_job(compute, offer=None, instance_name="dstack-test-unique", **overrides): + job_spec = SimpleNamespace( + image_name="pytorch/pytorch:2.11.0-cuda12.8-cudnn9-runtime", + registry_auth=None, + volumes=[], + ) + for name in ("image_name", "registry_auth", "mount_points"): + if name in overrides: + setattr(job_spec, "volumes" if name == "mount_points" else name, overrides.pop(name)) + kwargs = { + "run": SimpleNamespace(), + "job": SimpleNamespace(job_spec=job_spec), + "instance_offer": offer or _offers(compute, _item())[0], + "project_ssh_public_key": PUBLIC_KEY, + "project_ssh_private_key": "private-key", + "volumes": [], + "placement_group": None, + "requirements": Requirements(resources=ResourcesSpec()), + "extra_authorized_keys": [], + **overrides, + } + with ( + patch( + f"{COMPUTE_MODULE}.generate_unique_instance_name_for_job", + return_value=instance_name, + ), + patch(f"{COMPUTE_MODULE}.get_docker_commands", return_value=list(STARTUP_COMMANDS)), + ): + return compute.run_job(**kwargs) + + +def _update(compute, data): + compute.update_provisioning_data(data, PUBLIC_KEY, "private-key") + + +class TestRunJob: + @pytest.mark.parametrize("gpu,spot", [(False, False), (True, False), (True, True)]) + def test_preserves_priced_resources_and_native_gpu_type(self, compute, gpu, spot): + offer = _offers(compute, _item(gpu=gpu, spot=spot))[0] + + data = _run_job(compute, offer) + + assert offer.availability == InstanceAvailability.AVAILABLE + assert offer.instance_runtime == InstanceRuntime.RUNNER + payload = compute.api_client.create_sandbox.call_args.args[0] + assert payload["name"] == "dstack-test-unique" + assert payload["target"] == ("earth" if gpu else "us") + assert (payload["cpu"], payload["memory"], payload["disk"]) == (4, 16, 101) + assert payload["gpu"] == (2 if gpu else 0) + assert payload["spot"] is spot + if gpu: + assert offer.backend_data == {"gpu_type": "RTX-PRO-6000"} + assert offer.instance.resources.gpus[0].name == "RTXPRO6000" + assert payload["gpuType"] == ["RTX-PRO-6000"] + else: + assert "gpuType" not in payload + assert data.backend == BackendType.DAYTONA + assert data.instance_id == "sandbox-id" + assert data.price == offer.price == 0.512345 + assert data.instance_type.resources.disk.size_mib == 101 * 1024 + assert data.username == "root" + assert data.dockerized is False + assert data.hostname is None + assert data.ssh_port is None + compute.api_client.create_ssh_access.assert_not_called() + compute.api_client.create_registry.assert_not_called() + + def test_creates_root_image_without_user_entrypoint_or_idle_shutdown(self, compute): + data = _run_job(compute) + + payload = compute.api_client.create_sandbox.call_args.args[0] + assert payload["buildInfo"]["dockerfileContent"] == ( + "FROM pytorch/pytorch:2.11.0-cuda12.8-cudnn9-runtime\n" + 'USER root\nENTRYPOINT []\nCMD ["sleep", "infinity"]\n' + ) + assert payload["user"] == "root" + assert payload["public"] is False + for field in ("autoStopInterval", "autoPauseInterval", "autoDeleteInterval", "ttlMinutes"): + assert payload[field] == 0 + startup = json.loads(data.backend_data)["startup_command"] + assert shlex.split(startup) == [ + "sh", + "-c", + " && ".join(STARTUP_COMMANDS), + ] + + @pytest.mark.parametrize( + "override", + [ + {"image_name": "ubuntu:24.04\nRUN touch /injected"}, + {"image_name": ""}, + ], + ) + def test_rejects_invalid_image_names_before_creating_sandbox(self, compute, override): + with pytest.raises(ComputeError): + _run_job(compute, **override) + + compute.api_client.create_sandbox.assert_not_called() + + @pytest.mark.parametrize("status_code", [None, 500]) + def test_lost_create_response_retains_unique_name_for_recovery_and_cleanup( + self, compute, status_code + ): + compute.api_client.create_sandbox.side_effect = DaytonaAPIError( + "response lost", status_code=status_code + ) + + data = _run_job(compute) + + assert data.instance_id == "dstack-test-unique" + compute.api_client.create_sandbox.assert_called_once() + _update(compute, data) + compute.api_client.get_sandbox.assert_called_with("dstack-test-unique") + compute.api_client.create_ssh_access.assert_called_once_with( + "sandbox-id", expires_in_minutes=3650 * 24 * 60 + ) + with pytest.raises(NotYetTerminated): + compute.terminate_instance(data.instance_id, data.region) + compute.api_client.delete_sandbox.assert_called_once_with("dstack-test-unique") + + @pytest.mark.parametrize("status_code", [400, 429]) + def test_definite_create_failure_is_not_treated_as_success(self, compute, status_code): + error = DaytonaAPIError("creation rejected", status_code=status_code) + compute.api_client.create_sandbox.side_effect = error + + with pytest.raises(DaytonaAPIError) as exc_info: + _run_job(compute) + + assert exc_info.value is error + compute.api_client.create_sandbox.assert_called_once() + + +class TestRegistryAuth: + @pytest.mark.parametrize( + "image,registry,qualified", + [ + ("org/private:v1", "docker.io", "docker.io/org/private:v1"), + ("index.docker.io/org/private:v1", "docker.io", "docker.io/org/private:v1"), + ("ghcr.io/org/image@sha256:abc", "ghcr.io", "ghcr.io/org/image@sha256:abc"), + ( + "registry.example:5000/image:v1", + "registry.example:5000", + "registry.example:5000/image:v1", + ), + ], + ) + def test_creates_job_registry_and_qualifies_private_base_image( + self, compute, image, registry, qualified + ): + auth = RegistryAuth(username="registry-user", password="registry-password") + data = _run_job(compute, image_name=image, registry_auth=auth) + + compute.api_client.create_registry.assert_called_once_with( + { + "name": "dstack-test-unique-registry", + "url": registry, + "username": auth.username, + "password": auth.password, + } + ) + assert compute.api_client.create_sandbox.call_args.args[0]["buildInfo"][ + "dockerfileContent" + ].startswith(f"FROM {qualified}\n") + assert json.loads(data.backend_data)["registry_id"] == "registry-id" + assert auth.password not in data.backend_data + + def test_keeps_separate_ids_for_jobs_with_the_same_registry_credentials(self, compute): + auth = RegistryAuth(username="user", password="password") + compute.api_client.create_registry.side_effect = [ + {"id": "registry-a"}, + {"id": "registry-b"}, + ] + a = _run_job(compute, registry_auth=auth, instance_name="job-a") + b = _run_job(compute, registry_auth=auth, instance_name="job-b") + + assert json.loads(a.backend_data)["registry_id"] == "registry-a" + assert json.loads(b.backend_data)["registry_id"] == "registry-b" + assert [c.args[0]["name"] for c in compute.api_client.create_registry.call_args_list] == [ + "job-a-registry", + "job-b-registry", + ] + compute.api_client.get_sandbox.return_value = None + compute.terminate_instance(a.instance_id, a.region, a.backend_data) + compute.api_client.delete_registry.assert_called_once_with("registry-a") + + def test_recovers_only_own_registry_after_lost_create_response(self, compute): + compute.api_client.create_registry.side_effect = DaytonaAPIError("response lost") + compute.api_client.get_registries.return_value = [ + {"name": "someone-else", "id": "other-id"}, + {"name": "dstack-test-unique-registry", "id": "recovered-id"}, + ] + + data = _run_job(compute, registry_auth=RegistryAuth(username="user", password="password")) + + assert json.loads(data.backend_data)["registry_id"] == "recovered-id" + compute.api_client.create_registry.assert_called_once() + compute.api_client.delete_registry.assert_not_called() + + def test_registry_failure_does_not_create_a_sandbox_or_delete_existing_credentials( + self, compute + ): + compute.api_client.create_registry.side_effect = DaytonaAPIError( + "forbidden", status_code=403 + ) + + with pytest.raises(DaytonaAPIError, match="forbidden"): + _run_job(compute, registry_auth=RegistryAuth(username="user", password="password")) + + compute.api_client.create_sandbox.assert_not_called() + compute.api_client.delete_registry.assert_not_called() + + def test_lost_sandbox_response_keeps_credentials_until_sandbox_is_gone(self, compute): + compute.api_client.create_sandbox.side_effect = DaytonaAPIError("response lost") + data = _run_job(compute, registry_auth=RegistryAuth(username="user", password="password")) + + with pytest.raises(NotYetTerminated): + compute.terminate_instance(data.instance_id, data.region, data.backend_data) + compute.api_client.delete_registry.assert_not_called() + compute.api_client.get_sandbox.return_value = None + compute.terminate_instance(data.instance_id, data.region, data.backend_data) + compute.api_client.delete_registry.assert_called_once_with("registry-id") + + def test_failed_registry_rollback_reports_manual_cleanup(self, compute, caplog): + error = DaytonaAPIError("invalid image", status_code=400) + compute.api_client.create_sandbox.side_effect = error + compute.api_client.delete_registry.side_effect = DaytonaAPIError( + "temporary outage", status_code=503 + ) + + with pytest.raises(DaytonaAPIError, match="invalid image") as exc: + _run_job( + compute, + registry_auth=RegistryAuth(username="user", password="password"), + ) + + assert exc.value is error + compute.api_client.delete_registry.assert_called_once_with("registry-id") + assert "registry-id" in caplog.text + assert "Delete it manually in Daytona" in caplog.text + + def test_registry_cleanup_survives_provisioning_and_server_restart(self, compute): + data = _run_job(compute, registry_auth=RegistryAuth(username="user", password="password")) + _update(compute, data) + data = JobProvisioningData.model_validate_json(data.model_dump_json()) + assert json.loads(data.backend_data) == { + "startup_command": None, + "registry_id": "registry-id", + } + restarted = DaytonaCompute(compute.config) + restarted.api_client = compute.api_client + compute.api_client.get_sandbox.return_value = None + compute.api_client.delete_registry.side_effect = [ + DaytonaAPIError("temporary outage"), + None, + ] + + with pytest.raises(DaytonaAPIError, match="temporary outage"): + restarted.terminate_instance(data.instance_id, data.region, data.backend_data) + restarted.terminate_instance(data.instance_id, data.region, data.backend_data) + + assert compute.api_client.delete_registry.call_count == 2 + + +class TestVolumes: + def test_mounts_selected_alternatives_and_multiple_volumes(self, compute): + volumes = [ + SimpleNamespace(name="data", volume_id="data-id"), + SimpleNamespace(name="cache", volume_id="cache-id"), + ] + _run_job( + compute, + volumes=volumes, + mount_points=[ + VolumeMountPoint(name=["another-provider", "data"], path="/data"), + VolumeMountPoint(name="cache", path="/cache"), + ], + ) + + assert compute.api_client.create_sandbox.call_args.args[0]["volumes"] == [ + {"volumeId": "data-id", "mountPath": "/data"}, + {"volumeId": "cache-id", "mountPath": "/cache"}, + ] + + def test_creates_volume_without_a_size_or_region_and_waits_until_ready(self, compute): + compute.api_client.get_volume.side_effect = [{"state": "creating"}, {"state": "ready"}] + with ( + patch(f"{COMPUTE_MODULE}.generate_unique_volume_name", return_value="unique-volume"), + patch(f"{COMPUTE_MODULE}.time.sleep") as sleep, + ): + data = compute.create_volume(SimpleNamespace()) + + compute.api_client.create_volume.assert_called_once_with("unique-volume") + sleep.assert_called_once_with(2) + assert data.volume_id == "volume-id" + assert data.size_gb is None + assert data.price == 0 + assert data.attachable is False + assert data.detachable is False + + def test_recovers_created_volume_by_name_after_lost_response(self, compute): + compute.api_client.create_volume.side_effect = DaytonaAPIError("response lost") + compute.api_client.get_volume_by_name.return_value = {"id": "recovered-id"} + with patch(f"{COMPUTE_MODULE}.generate_unique_volume_name", return_value="unique-volume"): + data = compute.create_volume(SimpleNamespace()) + + assert data.volume_id == "recovered-id" + compute.api_client.create_volume.assert_called_once() + compute.api_client.get_volume_by_name.assert_called_once_with("unique-volume") + + def test_creation_timeout_is_bounded_and_cleans_up(self, compute): + compute.api_client.get_volume.return_value = {"state": "creating"} + with ( + patch(f"{COMPUTE_MODULE}.generate_unique_volume_name", return_value="unique-volume"), + patch(f"{COMPUTE_MODULE}.time.monotonic", side_effect=[0, 61]), + ): + with pytest.raises(ComputeError, match="Timed out"): + compute.create_volume(SimpleNamespace()) + compute.api_client.delete_volume.assert_called_once_with("volume-id") + + def test_failed_volume_rollback_reports_manual_cleanup(self, compute, mocker, caplog): + mocker.patch(f"{COMPUTE_MODULE}.generate_unique_volume_name", return_value="unique-volume") + compute.api_client.get_volume.return_value = { + "state": "error", + "errorReason": "storage failed", + } + compute.api_client.delete_volume.side_effect = DaytonaAPIError( + "temporary outage", status_code=503 + ) + + with pytest.raises(ComputeError, match="storage failed"): + compute.create_volume(SimpleNamespace()) + + compute.api_client.delete_volume.assert_called_once_with("volume-id") + assert "volume-id" in caplog.text + assert "Delete it manually in Daytona" in caplog.text + + def test_registers_an_existing_volume_without_modifying_it(self, compute): + volume = SimpleNamespace(configuration=DaytonaVolumeConfiguration(volume_id="external-id")) + data = compute.register_volume(volume) + + assert data.volume_id == "external-id" + compute.api_client.get_volume.assert_called_once_with("external-id") + compute.api_client.create_volume.assert_not_called() + compute.api_client.delete_volume.assert_not_called() + + @pytest.mark.parametrize("state", [None, {"state": "deleted"}, {"state": "error"}]) + def test_rejects_unusable_external_volume_without_deleting_it(self, compute, state): + compute.api_client.get_volume.return_value = state + volume = SimpleNamespace(configuration=DaytonaVolumeConfiguration(volume_id="external-id")) + with pytest.raises(ComputeError): + compute.register_volume(volume) + compute.api_client.delete_volume.assert_not_called() + + def test_deletion_accepts_pending_request_without_repeating_delete(self, compute): + compute.api_client.get_volume.side_effect = [ + {"state": "ready"}, + {"state": "pending_delete"}, + {"state": "deleting"}, + {"state": "deleted"}, + None, + ] + volume = SimpleNamespace(volume_id="volume-id") + for _ in range(5): + compute.delete_volume(volume) + + compute.api_client.delete_volume.assert_called_once_with("volume-id") + + +class TestUpdateProvisioningData: + def test_queues_runner_and_returns_connection_data_once(self, compute): + data = _run_job(compute) + startup_command = json.loads(data.backend_data)["startup_command"] + compute.api_client.get_session.return_value = None + + _update(compute, data) + + compute.api_client.create_session.assert_called_once_with(TOOLBOX_URL, "dstack-runner") + compute.api_client.execute_session_command.assert_called_once_with( + TOOLBOX_URL, + "dstack-runner", + startup_command, + run_async=True, + ) + assert data.hostname == "localhost" + assert data.ssh_port == 10022 + assert data.ssh_proxy.hostname == "ssh.app.daytona.io" + assert data.ssh_proxy.port == 2222 + assert data.ssh_proxy.username == "ssh-token_123" + assert data.backend_data is None + # Reconstruct the persisted state, as a new server process would do. + data = JobProvisioningData.model_validate_json(data.model_dump_json()) + restarted = DaytonaCompute(compute.config) + restarted.api_client = compute.api_client + + _update(restarted, data) + _update(restarted, data) + + compute.api_client.get_session.assert_called_once() + compute.api_client.execute_session_command.assert_called_once() + compute.api_client.create_session.assert_called_once() + compute.api_client.create_ssh_access.assert_called_once_with( + "sandbox-id", expires_in_minutes=3650 * 24 * 60 + ) + + def test_executes_an_existing_empty_session_after_interrupted_bootstrap(self, compute): + data = _run_job(compute) + compute.api_client.get_session.return_value = { + "sessionId": "dstack-runner", + "commands": [], + } + + _update(compute, data) + + compute.api_client.create_session.assert_not_called() + compute.api_client.execute_session_command.assert_called_once() + compute.api_client.create_ssh_access.assert_called_once() + assert data.hostname == "localhost" + + def test_recovers_existing_startup_after_server_restart_before_connection_data_saved( + self, compute + ): + data = _run_job(compute) + compute.api_client.get_session.side_effect = [ + None, + {"commands": [{"id": "running", "command": "runner"}]}, + ] + access = compute.api_client.create_ssh_access.return_value + compute.api_client.create_ssh_access.side_effect = [ + DaytonaAPIError("response lost"), + access, + ] + + with pytest.raises(DaytonaAPIError, match="response lost"): + _update(compute, data) + assert data.hostname is None + data = JobProvisioningData.model_validate_json(data.model_dump_json()) + restarted = DaytonaCompute(compute.config) + restarted.api_client = compute.api_client + + _update(restarted, data) + + assert data.hostname == "localhost" + compute.api_client.create_session.assert_called_once() + compute.api_client.execute_session_command.assert_called_once() + + @pytest.mark.parametrize("sandbox", [None, {"id": "sandbox-id", "state": "building_snapshot"}]) + def test_waits_for_pending_or_not_yet_visible_sandbox(self, compute, sandbox): + data = _run_job(compute) + compute.api_client.get_sandbox.return_value = sandbox + + _update(compute, data) + + assert data.hostname is None + compute.api_client.get_toolbox_url.assert_not_called() + compute.api_client.create_ssh_access.assert_not_called() + + @pytest.mark.parametrize("state", ["error", "build_failed", "stopped", "destroying", "paused"]) + def test_reports_terminal_sandbox_failure(self, compute, state): + data = _run_job(compute) + compute.api_client.get_sandbox.return_value = { + "id": "sandbox-id", + "state": state, + "errorReason": "image pull failed", + } + + with pytest.raises(ProvisioningError, match="image pull failed"): + _update(compute, data) + + compute.api_client.create_ssh_access.assert_not_called() + + @pytest.mark.parametrize("exit_code", [0, 1]) + def test_reports_exited_runner_setup_instead_of_executing_it_again(self, compute, exit_code): + data = _run_job(compute) + compute.api_client.get_session.return_value = { + "commands": [{"id": "cmd", "command": "runner", "exitCode": exit_code}] + } + + with pytest.raises(ProvisioningError, match=f"exited with code {exit_code}"): + _update(compute, data) + + compute.api_client.execute_session_command.assert_not_called() + compute.api_client.create_ssh_access.assert_not_called() + + def test_recovers_when_async_exec_response_is_lost(self, compute): + data = _run_job(compute) + compute.api_client.get_session.side_effect = [ + {"commands": []}, + {"commands": [{"id": "running", "command": "runner"}]}, + ] + compute.api_client.execute_session_command.side_effect = DaytonaAPIError("response lost") + + with pytest.raises(DaytonaAPIError, match="response lost"): + _update(compute, data) + _update(compute, data) + + assert data.hostname == "localhost" + compute.api_client.execute_session_command.assert_called_once() + compute.api_client.create_ssh_access.assert_called_once() + + @pytest.mark.parametrize( + "command", + [ + "ssh ssh-token_123@ssh.app.daytona.io", # Default SSH port is supported. + "ssh -p 2222 ssh-token_123@ssh.app.daytona.io", + ], + ) + def test_accepts_native_ssh_command_argument_order(self, compute, command): + data = _run_job(compute) + compute.api_client.create_ssh_access.return_value["sshCommand"] = command + + _update(compute, data) + + assert data.ssh_proxy.port == (2222 if "-p" in command else 22) + + def test_uses_api_token_and_provider_supplied_ssh_gateway(self, compute): + data = _run_job(compute) + compute.api_client.create_ssh_access.return_value = { + "token": "current-token", + "sshCommand": "ssh previous-token@custom-gateway.example -p 2222", + } + + _update(compute, data) + + assert data.ssh_proxy.username == "current-token" + assert data.ssh_proxy.hostname == "custom-gateway.example" + assert data.ssh_proxy.port == 2222 + + @pytest.mark.parametrize( + "command", + [ + "ssh ssh.app.daytona.io", + "ssh ssh-token_123@", + "ssh ssh-token_123@ssh.app.daytona.io -p 0", + "ssh ssh-token_123@ssh.app.daytona.io -p 70000", + "ssh ssh-token_123@ssh.app.daytona.io -p invalid", + "echo ssh-token_123@ssh.app.daytona.io", + ], + ) + def test_rejects_invalid_native_ssh_details(self, compute, command): + data = _run_job(compute) + compute.api_client.create_ssh_access.return_value["sshCommand"] = command + + with pytest.raises(ProvisioningError, match="invalid SSH connection details"): + _update(compute, data) + + assert data.hostname is None + + def test_rejects_missing_ssh_token(self, compute): + data = _run_job(compute) + del compute.api_client.create_ssh_access.return_value["token"] + + with pytest.raises(ProvisioningError, match="invalid SSH connection details"): + _update(compute, data) + + assert data.hostname is None + + +class TestTerminateInstance: + def test_deletion_is_idempotent_while_provider_is_destroying(self, compute): + compute.api_client.get_sandbox.side_effect = [ + {"state": "started"}, + {"state": "destroying"}, + {"state": "destroyed"}, + None, + ] + + for _ in range(2): + with pytest.raises(NotYetTerminated): + compute.terminate_instance("sandbox-id", "us") + compute.terminate_instance("sandbox-id", "us") + compute.terminate_instance("sandbox-id", "us") + + compute.api_client.delete_sandbox.assert_called_once_with("sandbox-id") + + +class TestGetOffersByRequirements: + def test_offer_queries_use_only_api_key_accessible_endpoints(self, requests_mock): + # Organization details require dashboard authentication, even for a valid API key. + # No organization-details mock is registered so using that endpoint fails this test. + paths = [ + "/api-keys/current", + "/organizations/org-id/usage", + "/organizations/org-id/available-sandbox-classes", + ] + requests_mock.get(API_URL + paths[0], json={"organizationId": "org-id"}) + requests_mock.get(API_URL + paths[1], json={"regionUsage": [_cpu_usage(), _gpu_usage()]}) + requests_mock.get( + API_URL + paths[2], + json=[ + {"regionId": "us", "sandboxClass": "container", "gpuAvailable": False}, + {"regionId": "earth", "sandboxClass": "container", "gpuAvailable": True}, + ], + ) + compute = DaytonaCompute(DaytonaConfig(creds=DaytonaCreds(api_key="test-api-key"))) + + offers = _offers(compute, _item(), _item(gpu=True)) + + assert [offer.availability for offer in offers] == [InstanceAvailability.AVAILABLE] * 2 + assert [request.url for request in requests_mock.request_history] == [ + API_URL + path for path in paths + ] + assert all( + request.headers["Authorization"] == "Bearer test-api-key" + for request in requests_mock.request_history + ) + + def test_organization_id_is_cached_but_quotas_are_refreshed(self, compute): + assert _offers(compute, _item())[0].availability == InstanceAvailability.AVAILABLE + compute.api_client.get_organization_usage.return_value["regionUsage"] = [ + _cpu_usage(currentCpuUsage=100) + ] + + assert _offers(compute, _item())[0].availability == InstanceAvailability.NO_QUOTA + compute.api_client.get_current_api_key.assert_called_once_with() + assert compute.api_client.get_organization_usage.call_count == 2 + + def test_skips_quota_calls_without_catalog_offers(self, compute): + assert _offers(compute) == [] + + compute.api_client.get_current_api_key.assert_not_called() + compute.api_client.get_organization_usage.assert_not_called() + + def test_filters_configured_regions(self, compute): + compute.config.regions = ["earth"] + + offers = _offers(compute, _item(), _item(gpu=True)) + + assert [offer.region for offer in offers] == ["earth"] + + @pytest.mark.parametrize("resource", ["Cpu", "Memory", "Disk"]) + def test_cpu_aggregate_quotas_include_current_usage(self, compute, resource): + compute.api_client.get_organization_usage.return_value["regionUsage"] = [ + _cpu_usage(**{f"current{resource}Usage": 999}) + ] + + assert _offers(compute, _item())[0].availability == InstanceAvailability.NO_QUOTA + + @pytest.mark.parametrize("resource", ["Cpu", "Memory", "Disk"]) + def test_cpu_region_per_sandbox_limit_is_enforced(self, compute, resource): + compute.api_client.get_organization_usage.return_value["regionUsage"] = [ + _cpu_usage(**{f"max{resource}PerSandbox": 1}) + ] + + assert _offers(compute, _item())[0].availability == InstanceAvailability.NO_QUOTA + + def test_ignores_other_sandbox_classes_when_checking_usage(self, compute): + compute.api_client.get_organization_usage.return_value["regionUsage"].append( + _cpu_usage(sandboxClass="windows", totalCpuQuota=0) + ) + compute.api_client.get_available_sandbox_classes.return_value.append( + {"regionId": "earth", "sandboxClass": "windows", "gpuAvailable": False} + ) + + assert _offers(compute, _item())[0].availability == InstanceAvailability.AVAILABLE + assert _offers(compute, _item(gpu=True))[0].availability == InstanceAvailability.AVAILABLE + + @pytest.mark.parametrize("missing", ["usage", "sandbox_class"]) + def test_region_requires_both_quota_and_container_class(self, compute, missing): + if missing == "usage": + compute.api_client.get_organization_usage.return_value["regionUsage"] = [] + else: + compute.api_client.get_available_sandbox_classes.return_value = [] + + assert _offers(compute, _item())[0].availability == InstanceAvailability.NO_QUOTA + + @pytest.mark.parametrize("spot", [False, True]) + @pytest.mark.parametrize("total,used", [(4, 3), (0, 0)]) + def test_gpu_quota_uses_remaining_count_only_for_on_demand(self, compute, spot, total, used): + compute.api_client.get_organization_usage.return_value["regionUsage"] = [ + _gpu_usage(totalGpuQuota=total, currentGpuUsage=used) + ] + + expected = InstanceAvailability.AVAILABLE if spot else InstanceAvailability.NO_QUOTA + assert _offers(compute, _item(gpu=True, spot=spot))[0].availability == expected + + @pytest.mark.parametrize( + "allowed,expected", + [ + (None, InstanceAvailability.AVAILABLE), + ([], InstanceAvailability.NO_QUOTA), + (["RTX-PRO-6000"], InstanceAvailability.AVAILABLE), + (["H100"], InstanceAvailability.NO_QUOTA), + ], + ) + @pytest.mark.parametrize("source", ["usage", "sandbox_class"]) + @pytest.mark.parametrize("spot", [False, True]) + def test_gpu_allowlist_uses_api_gpu_name(self, compute, allowed, expected, source, spot): + if source == "usage": + compute.api_client.get_organization_usage.return_value["regionUsage"] = [ + _gpu_usage(allowedGpuTypes=allowed) + ] + else: + compute.api_client.get_available_sandbox_classes.return_value[1]["allowedGpuTypes"] = ( + allowed + ) + + assert _offers(compute, _item(gpu=True, spot=spot))[0].availability == expected + + @pytest.mark.parametrize("shape", [{"cpu": 33}, {"memory": 385}, {"disk": 1025}]) + @pytest.mark.parametrize("spot", [False, True]) + def test_gpu_resource_limit_is_multiplied_by_gpu_count(self, compute, shape, spot): + boundary = {name: value - 1 for name, value in shape.items()} + + assert ( + _offers(compute, _item(gpu=True, spot=spot, **boundary))[0].availability + == InstanceAvailability.AVAILABLE + ) + assert ( + _offers(compute, _item(gpu=True, spot=spot, **shape))[0].availability + == InstanceAvailability.NO_QUOTA + ) + + @pytest.mark.parametrize("resource", ["Cpu", "Memory", "Disk"]) + def test_gpu_null_limit_falls_back_to_absolute_region_per_sandbox_limit( + self, compute, resource + ): + usage = _gpu_usage(**{f"max{resource}PerGpu": None}) + sizes = {"Cpu": 4, "Memory": 16, "Disk": 101} + usage[f"max{resource}PerSandbox"] = sizes[resource] - 1 + compute.api_client.get_organization_usage.return_value["regionUsage"] = [usage] + + assert _offers(compute, _item(gpu=True))[0].availability == InstanceAvailability.NO_QUOTA + + def test_gpu_unknown_resource_limits_do_not_impose_guessed_global_defaults(self, compute): + compute.api_client.get_organization_usage.return_value["regionUsage"] = [ + _gpu_usage(maxCpuPerGpu=None, maxMemoryPerGpu=None, maxDiskPerGpu=None) + ] + + assert ( + _offers(compute, _item(gpu=True, cpu=32, memory=384, disk=1024))[0].availability + == InstanceAvailability.AVAILABLE + ) + + @pytest.mark.parametrize("spot", [False, True]) + def test_gpu_class_must_allow_gpu(self, compute, spot): + compute.api_client.get_available_sandbox_classes.return_value[1]["gpuAvailable"] = False + + assert ( + _offers(compute, _item(gpu=True, spot=spot))[0].availability + == InstanceAvailability.NO_QUOTA + ) diff --git a/src/tests/_internal/core/backends/daytona/test_configurator.py b/src/tests/_internal/core/backends/daytona/test_configurator.py new file mode 100644 index 0000000000..8e272dd376 --- /dev/null +++ b/src/tests/_internal/core/backends/daytona/test_configurator.py @@ -0,0 +1,96 @@ +import json +from unittest.mock import patch + +import pytest + +from dstack._internal.core.backends.daytona.api_client import API_URL +from dstack._internal.core.backends.daytona.configurator import DaytonaConfigurator +from dstack._internal.core.backends.daytona.models import ( + DaytonaBackendConfigWithCreds, + DaytonaCreds, +) +from dstack._internal.core.errors import ( + BackendInvalidCredentialsError, + ServerClientError, +) + + +@pytest.fixture +def client(): + with patch( + "dstack._internal.core.backends.daytona.configurator.DaytonaAPIClient", autospec=True + ) as client_class: + client = client_class.return_value + client.get_current_api_key.return_value = { + "organizationId": "test-org", + "permissions": ["write:sandboxes", "delete:sandboxes", "read:limits"], + } + client.get_shared_regions.return_value = [{"id": "us"}, {"id": "eu"}, {"id": "ap"}] + yield client + + +class TestDaytonaConfigurator: + def test_validate_config_accepts_only_required_permissions(self, client): + config = DaytonaBackendConfigWithCreds(creds=DaytonaCreds(api_key="test-key")) + + DaytonaConfigurator().validate_config(config, default_creds_enabled=False) + + client.get_current_api_key.assert_called_once_with() + client.get_shared_regions.assert_not_called() + + @pytest.mark.parametrize( + "permissions,missing", + [ + (["write:sandboxes", "delete:sandboxes"], "read:limits"), + (["write:sandboxes", "read:limits"], "delete:sandboxes"), + (["delete:sandboxes", "read:limits"], "write:sandboxes"), + ], + ) + def test_validate_config_reports_missing_permissions( + self, requests_mock, permissions, missing + ): + requests_mock.get(f"{API_URL}/api-keys/current", json={"permissions": permissions}) + config = DaytonaBackendConfigWithCreds(creds=DaytonaCreds(api_key="test-key")) + + with pytest.raises(BackendInvalidCredentialsError) as exc: + DaytonaConfigurator().validate_config(config, default_creds_enabled=False) + + assert exc.value.msg == f"Daytona API key is missing permissions: {missing}" + assert exc.value.fields == [["creds", "api_key"]] + assert requests_mock.call_count == 1 + + def test_validate_config_accepts_gpu_and_discovered_cpu_regions(self, client): + config = DaytonaBackendConfigWithCreds( + creds=DaytonaCreds(api_key="test-key"), regions=["earth", "eu", "ap"] + ) + + DaytonaConfigurator().validate_config(config, default_creds_enabled=False) + + client.get_current_api_key.assert_called_once_with() + client.get_shared_regions.assert_called_once_with() + + def test_validate_config_rejects_unknown_regions(self, client): + config = DaytonaBackendConfigWithCreds( + creds=DaytonaCreds(api_key="test-key"), regions=["unknown"] + ) + + with pytest.raises(ServerClientError) as exc_info: + DaytonaConfigurator().validate_config(config, default_creds_enabled=False) + + assert exc_info.value.fields == [["regions"]] + assert "unknown" in exc_info.value.msg + + def test_stored_config_keeps_credentials_separate(self): + config = DaytonaBackendConfigWithCreds( + creds=DaytonaCreds(api_key="test-secret"), regions=["earth", "us"] + ) + configurator = DaytonaConfigurator() + + record = configurator.create_backend(project_name="main", config=config) + + assert json.loads(record.config) == {"type": "daytona", "regions": ["earth", "us"]} + assert json.loads(record.auth) == {"type": "api_key", "api_key": "test-secret"} + assert configurator.get_backend_config_with_creds(record) == config + public_config = configurator.get_backend_config_without_creds(record).model_dump() + assert public_config == json.loads(record.config) + assert "test-secret" not in json.dumps(public_config) diff --git a/src/tests/_internal/core/backends/daytona/test_models.py b/src/tests/_internal/core/backends/daytona/test_models.py new file mode 100644 index 0000000000..3b9a823084 --- /dev/null +++ b/src/tests/_internal/core/backends/daytona/test_models.py @@ -0,0 +1,26 @@ +from pydantic import Field, TypeAdapter +from typing_extensions import Annotated + +from dstack._internal.core.backends.daytona.models import DaytonaBackendConfigWithCreds +from dstack._internal.core.backends.models import ( + AnyBackendFileConfigWithCreds, + BackendConfigWithCreds, +) + + +class TestDaytonaBackendConfig: + def test_api_and_server_file_configs_accept_daytona(self): + data = { + "type": "daytona", + "regions": ["earth", "eu"], + "creds": {"type": "api_key", "api_key": "test-key"}, + } + api_config = BackendConfigWithCreds.model_validate(data).root + file_config = TypeAdapter( + Annotated[AnyBackendFileConfigWithCreds, Field(discriminator="type")] + ).validate_python(data) + + assert isinstance(api_config, DaytonaBackendConfigWithCreds) + assert api_config == file_config + assert api_config.regions == ["earth", "eu"] + assert api_config.creds.api_key == "test-key" diff --git a/src/tests/_internal/core/models/test_volumes.py b/src/tests/_internal/core/models/test_volumes.py index 9f95b1bffe..46b692e2aa 100644 --- a/src/tests/_internal/core/models/test_volumes.py +++ b/src/tests/_internal/core/models/test_volumes.py @@ -1,13 +1,66 @@ +import json + import pytest from pydantic import ValidationError from dstack._internal.core.models.volumes import ( + DaytonaVolumeConfiguration, InstanceMountPoint, + VolumeConfiguration, + VolumeConfigurationWithSize, VolumeMountPoint, + VolumeProvisioningData, + VolumeSpec, parse_mount_point, ) +@pytest.mark.parametrize("backend", ["aws", "gcp", "runpod", "kubernetes"]) +class TestSizedVolumeConfiguration: + def test_managed_size_keeps_numeric_wire_format(self, backend): + spec = VolumeSpec.model_validate( + {"configuration": {"backend": backend, "region": "us", "size": "120GB"}} + ) + + assert isinstance(spec.configuration, VolumeConfigurationWithSize) + assert spec.configuration.size_gb == 120 + assert not spec.configuration.is_external + assert json.loads(spec.model_dump_json())["configuration"]["size"] == 120.0 + restored = VolumeSpec.model_validate_json(spec.model_dump_json()) + assert restored == spec + + +class TestDaytonaVolumeConfiguration: + @pytest.mark.parametrize("volume_id", [None, "existing-volume"]) + def test_roundtrip_without_region_or_size(self, volume_id): + spec = VolumeSpec.model_validate( + {"configuration": {"backend": "daytona", "name": "cache", "volume_id": volume_id}} + ) + + assert isinstance(spec.configuration, DaytonaVolumeConfiguration) + assert spec.configuration.is_external == (volume_id is not None) + assert "region" not in spec.configuration.model_dump() + assert "size" not in spec.configuration.model_dump() + assert VolumeSpec.model_validate_json(spec.model_dump_json()) == spec + + @pytest.mark.parametrize(("field", "value"), [("size", 100), ("region", "us")]) + def test_rejects_fixed_size_and_region(self, field, value): + with pytest.raises(ValidationError, match="Extra inputs are not permitted"): + VolumeConfiguration.model_validate({"backend": "daytona", field: value}) + + +class TestVolumeProvisioningData: + def test_existing_capacity_remains_an_integer(self): + data = VolumeProvisioningData.model_validate( + {"backend": "aws", "volume_id": "volume-id", "size_gb": 120} + ) + + payload = json.loads(data.model_dump_json()) + assert type(payload["size_gb"]) is int + assert payload["size_gb"] == 120 + assert VolumeProvisioningData.model_validate(payload).size_gb == 120 + + class TestVolumeMountPoint: def test_parse(self): assert VolumeMountPoint.parse("my-vol:/path/./to///dir/") == VolumeMountPoint( diff --git a/src/tests/_internal/server/routers/test_backends.py b/src/tests/_internal/server/routers/test_backends.py index 9aed7be000..824674c364 100644 --- a/src/tests/_internal/server/routers/test_backends.py +++ b/src/tests/_internal/server/routers/test_backends.py @@ -88,6 +88,7 @@ async def test_returns_backend_types(self, client: AsyncClient): "cloudrift", "crusoe", "datacrunch", + "daytona", "digitalocean", "gcp", "hotaisle", diff --git a/src/tests/_internal/server/services/jobs/test_jobs.py b/src/tests/_internal/server/services/jobs/test_jobs.py index 089b5b54bb..96b127e823 100644 --- a/src/tests/_internal/server/services/jobs/test_jobs.py +++ b/src/tests/_internal/server/services/jobs/test_jobs.py @@ -4,17 +4,26 @@ import pytest import dstack._internal.server.settings as server_settings +from dstack._internal.core.models.backends.base import BackendType from dstack._internal.core.models.common import RegistryAuth from dstack._internal.core.models.configurations import TaskConfiguration from dstack._internal.core.models.profiles import Profile from dstack._internal.core.models.repos.local import LocalRunRepoData from dstack._internal.core.models.resources import ResourcesSpec from dstack._internal.core.models.runs import JobSpec, RunSpec +from dstack._internal.core.models.volumes import DaytonaVolumeConfiguration from dstack._internal.server.services.docker import ImageConfig from dstack._internal.server.services.jobs import ( + _get_job_mount_point_attached_volume, get_job_specs_from_run_spec, job_spec_updatable_in_place, ) +from dstack._internal.server.testing.common import ( + get_job_provisioning_data, + get_volume, + get_volume_configuration, + get_volume_provisioning_data, +) @pytest.mark.parametrize( @@ -175,3 +184,38 @@ def test_does_not_mutate_the_new_spec(self): job_spec_updatable_in_place(old_job_spec, new_job_spec) assert new_job_spec.requirements.resources.cpu.arch is None + + +class TestGetJobMountPointAttachedVolume: + @pytest.mark.parametrize(("region", "gpu_count"), [("us", 0), ("earth", 1)]) + def test_global_volume_matches_cpu_and_gpu_regions(self, region, gpu_count): + volume = get_volume(configuration=DaytonaVolumeConfiguration()) + other_backend_volume = get_volume( + configuration=get_volume_configuration(backend=BackendType.AWS, region=region) + ) + provisioning_data = get_job_provisioning_data( + backend=BackendType.DAYTONA, region=region, gpu_count=gpu_count + ) + + assert ( + _get_job_mount_point_attached_volume([other_backend_volume, volume], provisioning_data) + == volume + ) + + def test_regional_volume_preserves_region_and_zone_matching(self): + volumes = [ + get_volume( + configuration=get_volume_configuration(backend=BackendType.AWS, region=region), + provisioning_data=get_volume_provisioning_data(availability_zone=zone), + ) + for region, zone in [ + ("eu-west-1", "eu-west-1a"), + ("us-east-1", "us-east-1b"), + ("us-east-1", "us-east-1a"), + ] + ] + provisioning_data = get_job_provisioning_data( + backend=BackendType.AWS, region="US-EAST-1", availability_zone="US-EAST-1A" + ) + + assert _get_job_mount_point_attached_volume(volumes, provisioning_data) == volumes[2] diff --git a/src/tests/_internal/server/services/test_instances.py b/src/tests/_internal/server/services/test_instances.py index defd1498d4..6b332a3e73 100644 --- a/src/tests/_internal/server/services/test_instances.py +++ b/src/tests/_internal/server/services/test_instances.py @@ -23,6 +23,7 @@ Profile, ) from dstack._internal.core.models.runs import JobStatus +from dstack._internal.core.models.volumes import DaytonaVolumeConfiguration from dstack._internal.server.models import InstanceModel from dstack._internal.server.schemas.runner import TaskListItem, TaskListResponse, TaskStatus from dstack._internal.server.services.runner.client import ShimClient @@ -90,6 +91,67 @@ async def test_includes_termination_reason_in_event_message_when_switching_direc class TestFilterInstances: + @pytest.mark.parametrize("backend", [BackendType.AWS, BackendType.GCP]) + @pytest.mark.parametrize( + ("instance_zone", "volume_zone", "matches"), + [ + (None, None, True), + (None, "us-1b", True), + ("us-1a", None, False), + ("us-1a", "us-1a", True), + ("US-1A", "us-1a", True), + ("us-1a", "us-1b", False), + ], + ) + def test_preserves_volume_availability_zone_matching( + self, backend, instance_zone, volume_zone, matches + ): + instance = InstanceModel( + unreachable=False, + health=HealthStatus.HEALTHY, + total_blocks=1, + job_provisioning_data=get_job_provisioning_data( + backend=backend, region="US-1", availability_zone=instance_zone + ).model_dump_json(), + ) + volume = get_volume( + configuration=get_volume_configuration(backend=backend, region="us-1"), + provisioning_data=get_volume_provisioning_data( + backend=backend, availability_zone=volume_zone + ), + ) + + result = instances_services.filter_instances( + instances=[instance], profile=Profile(), volumes=[[volume]] + ) + + assert result == ([instance] if matches else []) + + def test_global_volume_matches_cpu_and_gpu_regions(self): + instances = [ + InstanceModel( + unreachable=False, + health=HealthStatus.HEALTHY, + total_blocks=1, + job_provisioning_data=get_job_provisioning_data( + backend=backend, region=region, gpu_count=gpu_count + ).model_dump_json(), + ) + for backend, region, gpu_count in [ + (BackendType.DAYTONA, "us", 0), + (BackendType.DAYTONA, "earth", 1), + (BackendType.AWS, "us", 0), + ] + ] + + result = instances_services.filter_instances( + instances=instances, + profile=Profile(), + volumes=[[get_volume(configuration=DaytonaVolumeConfiguration())]], + ) + + assert result == instances[:2] + @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) async def test_returns_all_instances(self, test_db, session: AsyncSession): @@ -257,7 +319,13 @@ async def test_returns_volume_instances_without_region(self, test_db, session: A project=project, backend=BackendType.KUBERNETES, ) - instances = [aws_instance, kubernetes_instance] + other_cluster_instance = await create_instance( + session=session, + project=project, + backend=BackendType.KUBERNETES, + region="other-cluster", + ) + instances = [aws_instance, kubernetes_instance, other_cluster_instance] volume = get_volume( configuration=get_kubernetes_volume_configuration(), provisioning_data=get_volume_provisioning_data( diff --git a/src/tests/_internal/server/services/test_offers.py b/src/tests/_internal/server/services/test_offers.py index 017d6aa813..7ff0b4fe4a 100644 --- a/src/tests/_internal/server/services/test_offers.py +++ b/src/tests/_internal/server/services/test_offers.py @@ -6,6 +6,7 @@ from dstack._internal.core.models.profiles import Profile from dstack._internal.core.models.resources import ResourcesSpec from dstack._internal.core.models.runs import Requirements +from dstack._internal.core.models.volumes import DaytonaVolumeConfiguration from dstack._internal.server.services.offers import get_offers_by_requirements from dstack._internal.server.testing.common import ( get_instance_offer_with_availability, @@ -100,6 +101,31 @@ async def test_returns_volume_offers(self): m.assert_awaited_once() assert res == [(runpod_backend_mock, runpod_offer2)] + @pytest.mark.asyncio + async def test_global_volume_matches_cpu_and_gpu_regions(self): + daytona_backend = Mock(TYPE=BackendType.DAYTONA) + cpu_offer = get_instance_offer_with_availability(backend=BackendType.DAYTONA, region="us") + gpu_offer = get_instance_offer_with_availability( + backend=BackendType.DAYTONA, region="earth", gpu_count=1 + ) + daytona_backend.compute.return_value.get_offers.return_value = [cpu_offer, gpu_offer] + aws_backend = Mock(TYPE=BackendType.AWS) + aws_offer = get_instance_offer_with_availability(backend=BackendType.AWS, region="us") + aws_backend.compute.return_value.get_offers.return_value = [aws_offer] + + with patch( + "dstack._internal.server.services.backends.get_project_backends", + return_value=[daytona_backend, aws_backend], + ): + offers = await get_offers_by_requirements( + project=Mock(), + profile=Profile(), + requirements=Requirements(resources=ResourcesSpec()), + volumes=[[get_volume(configuration=DaytonaVolumeConfiguration())]], + ) + + assert offers == [(daytona_backend, cpu_offer), (daytona_backend, gpu_offer)] + @pytest.mark.asyncio async def test_returns_volume_offers_without_region(self): profile = Profile(name="test") @@ -117,7 +143,10 @@ async def test_returns_volume_offers_without_region(self): availability_zones=None, ) kubernetes_backend_mock.compute.return_value.get_offers.return_value = [ - kubernetes_offer + kubernetes_offer, + get_instance_offer_with_availability( + backend=BackendType.KUBERNETES, region="other-cluster" + ), ] m.return_value = [aws_backend_mock, kubernetes_backend_mock] res = await get_offers_by_requirements( diff --git a/src/tests/_internal/server/services/test_volumes.py b/src/tests/_internal/server/services/test_volumes.py index 82477812ab..b05455ed31 100644 --- a/src/tests/_internal/server/services/test_volumes.py +++ b/src/tests/_internal/server/services/test_volumes.py @@ -5,7 +5,12 @@ from dstack._internal.core.errors import ServerClientError from dstack._internal.core.models.backends.base import BackendType -from dstack._internal.core.models.volumes import AWSVolumeConfiguration, VolumeStatus +from dstack._internal.core.models.volumes import ( + AWSVolumeConfiguration, + DaytonaVolumeConfiguration, + VolumeStatus, + parse_volume_configuration, +) from dstack._internal.server.services.volumes import ( _get_volume_cost, _validate_volume_configuration, @@ -17,6 +22,27 @@ class TestValidateVolumeConfiguration: + def test_daytona_does_not_require_size_or_region(self): + _validate_volume_configuration(DaytonaVolumeConfiguration()) + + @pytest.mark.parametrize( + ("backend", "identifier_field"), + [ + ("aws", "volume_id"), + ("gcp", "volume_id"), + ("runpod", "volume_id"), + ("kubernetes", "claim_name"), + ], + ) + def test_regional_volume_requires_size_only_when_managed(self, backend, identifier_field): + configuration = {"backend": backend, "region": "us"} + with pytest.raises(ServerClientError, match="existing identifier or size"): + _validate_volume_configuration(parse_volume_configuration(configuration)) + + _validate_volume_configuration( + parse_volume_configuration({**configuration, identifier_field: "existing-volume"}) + ) + def test_external_volume_with_auto_cleanup_duration_raises_error(self): """External volumes (with volume_id) should not allow auto_cleanup_duration""" config = AWSVolumeConfiguration(