diff --git a/frontend/src/types/backend.d.ts b/frontend/src/types/backend.d.ts index e18b89b006..216d60829d 100644 --- a/frontend/src/types/backend.d.ts +++ b/frontend/src/types/backend.d.ts @@ -9,6 +9,7 @@ declare type TBackendType = | 'gcp' | 'kubernetes' | 'lambda' + | 'lium' | 'local' | 'nebius' | 'remote' diff --git a/mkdocs/docs/concepts/backends.md b/mkdocs/docs/concepts/backends.md index 58cff86371..ba30e07ed0 100644 --- a/mkdocs/docs/concepts/backends.md +++ b/mkdocs/docs/concepts/backends.md @@ -1506,6 +1506,44 @@ projects: * Creating volumes: `read:volumes`, `write:volumes`, and `delete:volumes` * Registering existing volumes: `read:volumes` +### Lium + +The Lium backend is experimental. Create an API key under **Access → API Keys** in the +[Lium dashboard](https://lium.io/ssh-keys), then configure the backend: + +
+ +```yaml +projects: + - name: main + backends: + - type: lium + creds: + type: api_key + api_key: YOUR_LIUM_API_KEY +``` + +
+ +To restrict provisioning to specific regions, set `regions`, for example `regions: [us-ca, de-he]`. +Omit it to use all available regions. + +Private Docker images are supported through the run configuration's +[`registry_auth`](../reference/dstack.yml/task.md#registry_auth) property. + +!!! info "Storage and supported features" + Local disk is unencrypted and is deleted with the pod. The backend does not support + GPU blocks, instance volumes, or multi-node runs. + + Lium splits local disk between `/root` (two thirds) and the container filesystem + (one third). The total allocation is not available in a single directory. + + Persistent [network volumes](volumes.md#network-volumes) use S3-backed storage and support + jobs running as `root` only. You can attach one volume per pod at the path specified in + the run configuration. Storage grows with usage, and volumes can be attached in any Lium + region. Omit `size` and `region` from the volume configuration. + Do not rely on local-disk POSIX semantics or concurrent-writer consistency. + ### 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 6aa26ead5a..68aee12c9d 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`, `kubernetes`, and `daytona` backends. +> Network volumes are currently supported for the `aws`, `gcp`, `runpod`, `kubernetes`, `daytona`, and `lium` backends. ### Apply a configuration @@ -52,6 +52,22 @@ If you use this configuration, `dstack` will create a new volume based on the sp name: my-volume ``` +??? info "Lium" + Lium volumes use S3-backed storage and grow with usage. Omit `size` and `region`: + + ```yaml + type: volume + backend: lium + name: my-volume + ``` + + A run can mount one volume. Lium forces files on the volume to be owned by `root`, + so `dstack` rejects volume mounts for non-root jobs. Jobs without network volumes + can run as non-root users normally. + + Lium bills storage and read traffic separately from compute. Use these volumes for + datasets and checkpoints; do not assume local-disk performance or POSIX semantics. + ??? 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`. @@ -287,7 +303,7 @@ If you've registered an existing volume, it will be de-registered with `dstack` 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 `daytona`, `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`, `lium`, `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 58883b4744..0952f57cb1 100644 --- a/mkdocs/docs/reference/dstack.yml/volume.md +++ b/mkdocs/docs/reference/dstack.yml/volume.md @@ -34,6 +34,17 @@ The `volume` configuration type allows creating, registering, and updating [volu backend: required: true +=== "Lium" + + Lium volumes grow with usage and do not take `size` or `region`. + A run can mount one volume and must run as `root`. + + #SCHEMA# dstack._internal.core.models.volumes.LiumVolumeConfiguration + 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 538f703a89..00cf2f1957 100644 --- a/mkdocs/docs/reference/server/config.yml.md +++ b/mkdocs/docs/reference/server/config.yml.md @@ -159,6 +159,23 @@ to configure [backends](../../concepts/backends.md) and other [server-level sett type: required: true +##### `projects[n].backends[type=lium]` { #lium data-toc-label="lium" } + +#SCHEMA# dstack._internal.core.backends.lium.models.LiumBackendConfigWithCreds + overrides: + show_root_heading: false + type: + required: true + item_id_prefix: lium- + +###### `projects[n].backends[type=lium].creds` { #lium-creds data-toc-label="creds" } + +#SCHEMA# dstack._internal.core.backends.lium.models.LiumAPIKeyCreds + overrides: + show_root_heading: false + type: + required: true + ##### `projects[n].backends[type=nebius]` { #nebius data-toc-label="nebius" } #SCHEMA# dstack._internal.core.backends.nebius.models.NebiusBackendFileConfigWithCreds diff --git a/pyproject.toml b/pyproject.toml index 1ddb8ef87d..7e18c00c33 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,7 +33,8 @@ dependencies = [ "python-multipart>=0.0.16", "filelock", "psutil", - "gpuhunt==0.1.33", + # TODO: Merge/release gpuhunt#259 and pin the released version. + "gpuhunt @ git+https://github.com/Datura-ai/gpuhunt.git@lium-provider", "argcomplete>=3.5.0", "ignore-python>=0.2.0", "apscheduler<4", @@ -57,6 +58,10 @@ dstack = "dstack._internal.cli.main:main" [tool.hatch.version] path = "src/dstack/version.py" +[tool.hatch.metadata] +# TODO: Remove after pinning the released gpuhunt package. +allow-direct-references = true + [tool.hatch.build.targets.sdist] artifacts = [ "src/dstack/_internal/server/statics/**", diff --git a/src/dstack/_internal/core/backends/configurators.py b/src/dstack/_internal/core/backends/configurators.py index 7bad991340..826073914c 100644 --- a/src/dstack/_internal/core/backends/configurators.py +++ b/src/dstack/_internal/core/backends/configurators.py @@ -112,6 +112,13 @@ except ImportError: pass +try: + from dstack._internal.core.backends.lium.configurator import LiumConfigurator + + _CONFIGURATOR_CLASSES.append(LiumConfigurator) +except ImportError: + pass + try: from dstack._internal.core.backends.nebius.configurator import ( NebiusConfigurator, diff --git a/src/dstack/_internal/core/backends/lium/__init__.py b/src/dstack/_internal/core/backends/lium/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/dstack/_internal/core/backends/lium/api_client.py b/src/dstack/_internal/core/backends/lium/api_client.py new file mode 100644 index 0000000000..c7ab749a97 --- /dev/null +++ b/src/dstack/_internal/core/backends/lium/api_client.py @@ -0,0 +1,176 @@ +import re +from typing import Any, Optional +from urllib.parse import quote + +import requests + +from dstack._internal.core.errors import ComputeError + +API_URL = "https://lium.io/api" +REQUEST_TIMEOUT_SECONDS = 30 + + +class LiumAPIError(ComputeError): + def __init__( + self, + message: str, + status_code: Optional[int] = None, + code: Optional[str] = None, + idempotent_replayed: bool = False, + ): + super().__init__(message) + self.status_code = status_code + self.code = code + self.idempotent_replayed = idempotent_replayed + + +class LiumAPIClient: + def __init__(self, api_key: str): + self.api_key = api_key + + def get_user(self) -> dict[str, Any]: + return _object(self._make_request("GET", "/users/me")) + + def get_version(self) -> dict[str, Any]: + return _object(self._make_request("GET", "/version")) + + def get_pods(self) -> list[dict[str, Any]]: + return _objects(self._make_request("GET", "/pods")) + + def get_account_events(self, since: Optional[str] = None) -> list[dict[str, Any]]: + # The account event log retains pod names and IDs after deletion. + params: dict[str, Any] = {"limit": 1000} + if since is not None: + params["since"] = since + return _objects(self._make_request("GET", "/users/me/events", params=params)) + + def get_pod(self, pod_id: str) -> Optional[dict[str, Any]]: + try: + return _object(self._make_request("GET", f"/pods/{quote(pod_id, safe='')}")) + except LiumAPIError as e: + if e.status_code == 404: + return None + raise + + def create_template(self, payload: dict[str, Any]) -> dict[str, Any]: + return _object(self._make_request("POST", "/templates", json=payload)) + + def get_templates(self) -> list[dict[str, Any]]: + return _objects(self._make_request("GET", "/templates")) + + def delete_template(self, template_id: str) -> None: + self._delete(f"/templates/{quote(template_id, safe='')}") + + def create_docker_credential(self, payload: dict[str, Any]) -> dict[str, Any]: + return _object(self._make_request("POST", "/docker-credentials/", json=payload)) + + def delete_docker_credential(self, credential_id: str) -> None: + self._delete(f"/docker-credentials/{quote(credential_id, safe='')}") + + def rent_executor( + self, executor_id: str, payload: dict[str, Any], idempotency_key: str + ) -> dict[str, Any]: + # Retries must keep both the key and the request body unchanged. The caller + # reconciles uncertain outcomes through the pod list before another rental. + return _object( + self._make_request( + "POST", + f"/executors/{quote(executor_id, safe='')}/rent", + json=payload, + headers={"Idempotency-Key": idempotency_key}, + ) + ) + + def delete_pod(self, pod_id: str) -> None: + self._delete(f"/pods/{quote(pod_id, safe='')}") + + def create_volume(self, name: str) -> dict[str, Any]: + return _object(self._make_request("POST", "/volumes", json={"name": name})) + + def get_volumes(self) -> list[dict[str, Any]]: + return _objects(self._make_request("GET", "/volumes")) + + def get_volume(self, volume_id: str) -> Optional[dict[str, Any]]: + try: + return _object(self._make_request("GET", f"/volumes/{quote(volume_id, safe='')}")) + except LiumAPIError as e: + if e.status_code == 404: + return None + raise + + def delete_volume(self, volume_id: str) -> None: + self._delete(f"/volumes/{quote(volume_id, safe='')}") + + def _delete(self, path: str) -> None: + try: + self._make_request("DELETE", path) + except LiumAPIError as e: + if e.status_code != 404: + raise + + def _make_request( + self, + method: str, + path: str, + *, + json: Optional[dict[str, Any]] = None, + headers: Optional[dict[str, str]] = None, + params: Optional[dict[str, Any]] = None, + ) -> Any: + try: + response = requests.request( + method=method, + url=API_URL + path, + headers={"X-API-Key": self.api_key, **(headers or {})}, + json=json, + params=params, + timeout=REQUEST_TIMEOUT_SECONDS, + allow_redirects=False, + ) + except requests.Timeout: + raise LiumAPIError("Lium API request timed out") from None + except requests.RequestException: + raise LiumAPIError("Lium API request failed") from None + if not 200 <= response.status_code < 300: + code = None + try: + data = response.json() + except ValueError: + data = None + if isinstance(data, dict) and isinstance(data.get("error"), dict): + value = data["error"].get("code") + if ( + isinstance(value, str) + and re.fullmatch(r"[a-z][a-z0-9_]{0,63}", value) + and self.api_key not in value + ): + code = value + # Show only a short machine code; provider messages may reflect credentials. + message = f"Lium API returned HTTP {response.status_code}" + if code is not None: + message += f" ({code})" + raise LiumAPIError( + message, + status_code=response.status_code, + code=code, + idempotent_replayed=response.headers.get("Idempotent-Replayed", "").lower() + == "true", + ) + if method == "DELETE" or not response.content: + return None + try: + return response.json() + except ValueError: + raise LiumAPIError("Lium API returned a non-JSON response") from None + + +def _object(data: Any) -> dict[str, Any]: + if not isinstance(data, dict): + raise LiumAPIError("Lium 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 LiumAPIError("Lium API returned an invalid list response") + return data diff --git a/src/dstack/_internal/core/backends/lium/backend.py b/src/dstack/_internal/core/backends/lium/backend.py new file mode 100644 index 0000000000..07869b02f3 --- /dev/null +++ b/src/dstack/_internal/core/backends/lium/backend.py @@ -0,0 +1,16 @@ +from dstack._internal.core.backends.base.backend import Backend +from dstack._internal.core.backends.lium.compute import LiumCompute +from dstack._internal.core.backends.lium.models import LiumConfig +from dstack._internal.core.models.backends.base import BackendType + + +class LiumBackend(Backend): + TYPE = BackendType.LIUM + COMPUTE_CLASS = LiumCompute + + def __init__(self, config: LiumConfig): + self.config = config + self._compute = LiumCompute(self.config) + + def compute(self) -> LiumCompute: + return self._compute diff --git a/src/dstack/_internal/core/backends/lium/compute.py b/src/dstack/_internal/core/backends/lium/compute.py new file mode 100644 index 0000000000..4cec05c8ee --- /dev/null +++ b/src/dstack/_internal/core/backends/lium/compute.py @@ -0,0 +1,591 @@ +import re +import shlex +import uuid +from datetime import datetime, timedelta +from pathlib import PurePosixPath +from typing import Any, List, Optional + +import gpuhunt +import paramiko +from gpuhunt.providers.lium import LiumProvider + +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.lium.api_client import LiumAPIClient, LiumAPIError +from dstack._internal.core.backends.lium.models import LiumConfig +from dstack._internal.core.consts import DSTACK_RUNNER_SSH_PORT +from dstack._internal.core.errors import ( + ComputeError, + NoCapacityError, + NotYetTerminated, + ProvisioningError, +) +from dstack._internal.core.models.backends.base import BackendType +from dstack._internal.core.models.common import CoreModel, validate_json_extra_ignore +from dstack._internal.core.models.instances import ( + InstanceAvailability, + 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 ( + LiumVolumeConfiguration, + Volume, + VolumeMountPoint, + VolumeProvisioningData, +) +from dstack._internal.utils.common import get_current_datetime, get_or_error +from dstack._internal.utils.logging import get_logger +from dstack._internal.utils.ssh import convert_ssh_key_to_pem, pkey_from_str + +logger = get_logger(__name__) + +MAX_INSTANCE_NAME_LEN = 63 +SSH_TIMEOUT_SECONDS = 15 +BOOTSTRAP_DIR = "/root/.dstack-lium" +TERMINAL_POD_STATES = { + "STOPPED", + "FAILED", + "CREATION_FAILED", + "BROKEN", + "REBOOT_FAILED", + "DELETING", + "DELETED", +} +NO_CAPACITY_CODES = { + "node_not_found", + "node_rent_in_progress", + "node_not_verified", + "node_unavailable", + "node_paused_by_provider", + "provider_banned", + "node_below_bandwidth_floor", + "gpu_count_invalid", + "gpu_split_not_allowed", + "no_executor_matches_spec", +} + + +class LiumCompute(ComputeWithFilteredOffersCached, ComputeWithVolumeSupport, Compute): + def __init__(self, config: LiumConfig): + super().__init__() + self.config = config + self.api_client = LiumAPIClient(config.creds.api_key) + self._catalog = gpuhunt.Catalog(balance_resources=False, auto_reload=False) + self._catalog.add_provider(LiumProvider()) + + def get_offers_by_requirements( + self, requirements: Requirements, full_offers: bool, unallocated_resources: bool + ) -> List[InstanceOfferWithAvailability]: + offers = get_catalog_offers( + backend=BackendType.LIUM, + locations=self.config.regions or None, + requirements=requirements, + catalog=self._catalog, + ) + return [ + offer.with_availability( + availability=InstanceAvailability.AVAILABLE, + 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 + if not image_name or re.search(r"\s", image_name): + raise ComputeError("Invalid Lium container image name") + volume_commands = _get_volume_commands(run, job, volumes) + 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), + pre_runner_commands=volume_commands, + ) + data = LiumInstanceBackendData( + pod_name=instance_name, + executor_id=instance_offer.instance.name, + startup_command=" && ".join(commands), + ) + if job.job_spec.registry_auth is not None: + auth = job.job_spec.registry_auth + try: + credential = self.api_client.create_docker_credential( + {"docker_username": auth.username, "docker_password": auth.password} + ) + data.credential_id = _resource_id(credential) + except LiumAPIError as e: + if not _is_definitive_error(e): + # Credentials have neither a name nor idempotent creation. Never + # recover/delete someone else's record by matching only its username. + raise ComputeError( + "Lium registry credential creation could not be confirmed. " + "Check Docker credentials in Lium before retrying." + ) from e + raise + try: + data.template_id = self._create_template( + { + "name": instance_name, + # Follow the official SDK: Lium splits the complete image reference. + "docker_image": image_name, + "docker_image_tag": "", + "docker_image_digest": "", + "entrypoint": "/bin/sh", + "startup_commands": "-c 'exec sleep infinity'", + "internal_ports": [22, DSTACK_RUNNER_SSH_PORT], + "volumes": ["/root"], + "is_private": True, + "supports_docker": False, + "supports_volume_encryption": False, + "docker_credential_id": data.credential_id, + } + ) + except Exception: + self._cleanup_failed_resources(data) + raise + data.rent_payload = { + "pod_name": instance_name, + "template_id": data.template_id, + "user_public_key": project_ssh_public_key, + "gpu_count": len(instance_offer.instance.resources.gpus), + "initial_port_count": 2, + "enable_jupyter": False, + # Encrypted /root can become inaccessible after a provider host reboot. + "enable_volume_encryption": False, + } + data.idempotency_key = str(uuid.uuid4()) + data.rent_started_at = get_current_datetime() + if volumes: + data.rent_payload["volume_id"] = get_or_error(volumes[0].volume_id) + pod_id = self._rent(data) + return JobProvisioningData( + backend=BackendType.LIUM, + instance_type=instance_offer.instance, + instance_id=pod_id or instance_name, + region=instance_offer.region, + price=instance_offer.price, + username="root", + dockerized=False, + backend_data=data.model_dump_json(), + ) + + def update_provisioning_data( + self, + provisioning_data: JobProvisioningData, + project_ssh_public_key: str, + project_ssh_private_key: str, + ) -> None: + data = validate_json_extra_ignore( + LiumInstanceBackendData, get_or_error(provisioning_data.backend_data) + ) + if data.rent_failed: + raise ProvisioningError("Lium rental failed while reconciling its creation response") + if data.rent_payload is not None: + pod = self._find_pending_pod(data) + if pod is None: + if not self._can_replay_rent(data): + return + pod_id = self._rent(data, reconcile=True) + provisioning_data.backend_data = data.model_dump_json() + if pod_id is None: + return + provisioning_data.instance_id = pod_id + pod = self.api_client.get_pod(pod_id) + else: + provisioning_data.instance_id = _resource_id(pod) + data.rent_payload = None + data.idempotency_key = None + provisioning_data.backend_data = data.model_dump_json() + else: + pod = self.api_client.get_pod(provisioning_data.instance_id) + if pod is None: + raise ProvisioningError("Lium pod was not found") + status = str(pod.get("status", "")).upper() + if status in TERMINAL_POD_STATES: + raise ProvisioningError(f"Lium pod entered {status} state") + if status != "RUNNING": + return + ssh = _get_provider_ssh(pod) + try: + _bootstrap_runner(ssh, project_ssh_private_key, data.startup_command) + except (paramiko.SSHException, OSError): + # Lium marks RUNNING before its separate SSH bootstrap finishes. + logger.debug("Waiting for Lium pod %s SSH", provisioning_data.instance_id) + return + ports = pod.get("ports_mapping") or {} + runner_port = ports.get(str(DSTACK_RUNNER_SSH_PORT)) + if runner_port is not None: + provisioning_data.hostname = ssh.hostname + provisioning_data.ssh_port = _port(runner_port) + provisioning_data.ssh_proxy = None + else: + provisioning_data.hostname = "localhost" + provisioning_data.ssh_port = DSTACK_RUNNER_SSH_PORT + provisioning_data.ssh_proxy = ssh + + def terminate_instance( + self, instance_id: str, region: str, backend_data: Optional[str] = None + ) -> None: + data = validate_json_extra_ignore(LiumInstanceBackendData, get_or_error(backend_data)) + if data.rent_failed: + self._cleanup_resources(data) + return + if data.rent_payload is not None: + pod = self._find_pending_pod(data) + if pod is None: + # Replaying a rent during cancellation could create a new paid pod. + # An empty list cannot prove an in-flight rent will never commit. + raise NotYetTerminated( + f"Lium rental outcome is uncertain for {data.pod_name}; " + "reconcile it in Lium before removing the instance" + ) + instance_id = _resource_id(pod) + else: + pod = self.api_client.get_pod(instance_id) + if pod is None or pod.get("status") == "DELETED": + self._cleanup_resources(data) + return + if pod.get("status") != "DELETING": + try: + self.api_client.delete_pod(instance_id) + except LiumAPIError as e: + if e.code != "pod_state_conflict": + raise + current = self.api_client.get_pod(instance_id) + if current is not None and current.get("status") != "DELETING": + raise + raise NotYetTerminated("Lium pod deletion is in progress") + + def register_volume(self, volume: Volume) -> VolumeProvisioningData: + assert isinstance(volume.configuration, LiumVolumeConfiguration) + volume_id = get_or_error(volume.configuration.volume_id) + if self.api_client.get_volume(volume_id) is None: + raise ComputeError(f"Lium volume {volume_id} not found") + 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 LiumAPIError as e: + if _is_definitive_error(e): + raise + data = next((v for v in self.api_client.get_volumes() if v.get("name") == name), None) + if data is None: + raise + return _volume_provisioning_data(_resource_id(data)) + + def delete_volume(self, volume: Volume) -> None: + if volume.volume_id is not None: + self.api_client.delete_volume(volume.volume_id) + + def _create_template(self, payload: dict[str, Any]) -> str: + try: + template = self.api_client.create_template(payload) + return _resource_id(template) + except LiumAPIError as e: + if _is_definitive_error(e): + raise + for template in self.api_client.get_templates(): + if template.get("name") == payload["name"]: + return _resource_id(template) + logger.warning( + "Lium template %s creation could not be confirmed; check templates in Lium", + payload["name"], + ) + raise + + def _rent(self, data: "LiumInstanceBackendData", *, reconcile: bool = False) -> Optional[str]: + try: + response = self.api_client.rent_executor( + data.executor_id, + get_or_error(data.rent_payload), + get_or_error(data.idempotency_key), + ) + pod_id = _resource_id(response, key="pod_id") + except LiumAPIError as e: + if _is_definitive_error(e) and (not reconcile or e.idempotent_replayed): + data.rent_failed = True + data.rent_payload = None + data.idempotency_key = None + self._cleanup_failed_resources(data) + if reconcile: + # Persist the failed outcome before the next poll raises. + return None + if e.code in NO_CAPACITY_CODES: + raise NoCapacityError(f"Lium node is unavailable ({e.code})") from e + raise + logger.warning( + "Lium pod %s creation response is uncertain; reconciling by name", + data.pod_name, + ) + return None + data.rent_payload = None + data.idempotency_key = None + return pod_id + + def _find_pending_pod(self, data: "LiumInstanceBackendData") -> Optional[dict[str, Any]]: + for pod in self.api_client.get_pods(): + if pod.get("pod_name") == data.pod_name and pod.get("executor_id") == data.executor_id: + return pod + since = ( + (data.rent_started_at - timedelta(minutes=1)).isoformat() + if data.rent_started_at is not None + else None + ) + # Account events retain pod IDs after they disappear from the live list. + # This also makes cancellation recovery durable across calls, without + # process-local state or replaying a rental during deletion. + for event in self.api_client.get_account_events(since=since): + if event.get("pod_name") != data.pod_name: + continue + pod_id = event.get("pod_id") + if not isinstance(pod_id, str) or not pod_id: + continue + pod = self.api_client.get_pod(pod_id) + deleted = event.get("to_status") == "DELETED" or ( + event.get("sub_event_type") == "pod-delete.success" + ) + if pod is None and deleted: + return {"id": pod_id, "status": "DELETED"} + if pod is not None and pod.get("executor_id") == data.executor_id: + return pod + return None + + def _can_replay_rent(self, data: "LiumInstanceBackendData") -> bool: + if data.rent_started_at is None: + return False + if get_current_datetime() - data.rent_started_at >= timedelta(hours=24): + return False + # Older Lium releases may ignore the key and rent another pod on retry. + features = self.api_client.get_version().get("features") + return isinstance(features, list) and "rent_idempotency" in features + + def _cleanup_failed_resources(self, data: "LiumInstanceBackendData") -> None: + try: + self._cleanup_resources(data) + except Exception: + logger.exception( + "Failed to delete Lium template %s or credential %s after provisioning failed. " + "Delete them manually in Lium.", + data.template_id, + data.credential_id, + ) + + def _cleanup_resources(self, data: "LiumInstanceBackendData") -> None: + if data.template_id is not None: + self.api_client.delete_template(data.template_id) + if data.credential_id is not None: + self.api_client.delete_docker_credential(data.credential_id) + + +class LiumInstanceBackendData(CoreModel): + pod_name: str + executor_id: str + startup_command: str + template_id: Optional[str] = None + credential_id: Optional[str] = None + rent_payload: Optional[dict[str, Any]] = None + idempotency_key: Optional[str] = None + rent_started_at: Optional[datetime] = None + rent_failed: bool = False + + +def _get_volume_commands(run: Run, job: Job, volumes: List[Volume]) -> list[str]: + if not volumes: + return [] + if len(volumes) > 1: + raise ComputeError("Lium supports mounting only one network volume per pod") + user = job.job_spec.user + if user is not None and user.uid != 0 and user.username != "root": + raise ComputeError( + "Lium network volumes support root jobs only. " + "Lium forces file ownership to root, so non-root jobs cannot update files." + ) + mount_points = job.job_spec.volumes + if mount_points is None: + mount_points = run.run_spec.configuration.volumes + for mount in mount_points: + if not isinstance(mount, VolumeMountPoint): + continue + names = [mount.name] if isinstance(mount.name, str) else mount.name + if volumes[0].name in names: + path = PurePosixPath(mount.path) + break + else: + raise ComputeError("Lium volume has no mount path in the job configuration") + # Mounting over bootstrap/system directories would break SSH or the runner. + if ( + path == PurePosixPath("/") + or path == PurePosixPath(BOOTSTRAP_DIR) + or path in PurePosixPath(BOOTSTRAP_DIR).parents + or any(path.is_relative_to(p) for p in ("/proc", "/sys", "/dev")) + or path != PurePosixPath("/mnt") + and path.is_relative_to("/mnt") + ): + raise ComputeError(f"Lium volume cannot be mounted at {path}") + commands = ["grep -q ' /mnt .* - fuse.s3fs ' /proc/self/mountinfo"] + if str(path) != "/mnt": + target = shlex.quote(str(path)) + commands += [f"mkdir -p -- {target}", f"mount --bind /mnt {target}"] + return commands + + +def _volume_provisioning_data(volume_id: str) -> VolumeProvisioningData: + return VolumeProvisioningData( + backend=BackendType.LIUM, + volume_id=volume_id, + size_gb=None, + # Storage and read traffic are billed by usage, not as a fixed monthly price. + price=None, + attachable=False, + detachable=False, + ) + + +def _resource_id(resource: dict[str, Any], key: str = "id") -> str: + value = resource.get(key) + if not isinstance(value, str) or not value: + raise LiumAPIError(f"Lium response has no {key}") + return value + + +def _is_definitive_error(error: LiumAPIError) -> bool: + return ( + error.status_code is not None + and error.status_code < 500 + and error.status_code != 408 + and error.code not in {"idempotency_in_progress", "idempotency_outcome_unknown"} + ) + + +def _get_provider_ssh(pod: dict[str, Any]) -> SSHConnectionParams: + # Parse connection data only; never execute the provider's SSH command. + try: + parts = shlex.split(pod["ssh_connect_cmd"]) + target = next(part for part in parts[1:] if "@" in part) + user, hostname = target.rsplit("@", 1) + port = _port(parts[parts.index("-p") + 1]) if "-p" in parts else 22 + if parts[0] != "ssh" or user != "root" or not hostname: + raise ValueError + except (KeyError, IndexError, StopIteration, TypeError, ValueError) as e: + raise ProvisioningError("Lium returned invalid root SSH connection details") from e + return SSHConnectionParams(hostname=hostname.strip("[]"), port=port, username="root") + + +def _port(value: Any) -> int: + try: + port = int(value) + if isinstance(value, bool) or not 1 <= port <= 65535: + raise ValueError + return port + except (TypeError, ValueError) as e: + raise ProvisioningError("Lium returned an invalid SSH port") from e + + +def _bootstrap_runner(ssh: SSHConnectionParams, private_key: str, startup_command: str) -> None: + with paramiko.SSHClient() as client: + client.set_missing_host_key_policy(paramiko.AutoAddPolicy()) + client.connect( + hostname=ssh.hostname, + port=ssh.port, + username="root", + pkey=pkey_from_str(convert_ssh_key_to_pem(private_key)), + timeout=SSH_TIMEOUT_SECONDS, + banner_timeout=SSH_TIMEOUT_SECONDS, + auth_timeout=SSH_TIMEOUT_SECONDS, + allow_agent=False, + look_for_keys=False, + ) + _, stdout, _ = client.exec_command( + _get_bootstrap_command(startup_command), timeout=SSH_TIMEOUT_SECONDS + ) + output = stdout.read().decode().strip() + if stdout.channel.recv_exit_status() != 0: + raise ProvisioningError("Failed to prepare the Lium runner bootstrap") + if output != "running": + raise ProvisioningError("Lium runner bootstrap exited; check /root/.dstack-lium/log") + + +def _get_bootstrap_command(startup_command: str) -> str: + runner_path = f"{BOOTSTRAP_DIR}/runner.sh" + monitor_path = f"{BOOTSTRAP_DIR}/monitor.sh" + exit_path = f"{BOOTSTRAP_DIR}/exit-code" + # SSH replaces the image's environment and working directory. Restore both + # from PID 1 without requiring Python in the image. Keep ENV NUL-delimited; + # the argument count check rejects oversized environments instead of using + # only some variables. + # The bootstrap files stay private, but runner-created user homes and SSH + # files must remain readable after the runner drops to the image USER. + image_command = _get_image_env_and_cwd_command(runner_path) + monitor = f"""umask 022 +{image_command} +status=$? +printf "%s\\n" "$status" > {exit_path} +""" + runner_script = startup_command + "\n" + command = f"""umask 077 +command -v xargs env tr wc nohup > /dev/null || exit 1 +mkdir -p {BOOTSTRAP_DIR} || exit 1 +if test -f {exit_path}; then + printf exited +elif test -d {BOOTSTRAP_DIR}/started; then + printf running +elif mkdir {BOOTSTRAP_DIR}/started 2>/dev/null; then + printf %s {shlex.quote(runner_script)} > {runner_path} || exit 1 + printf %s {shlex.quote(monitor)} > {monitor_path} || exit 1 + nohup /bin/sh {monitor_path} > {BOOTSTRAP_DIR}/log 2>&1 < /dev/null & + printf running +else + printf running +fi""" + return shlex.join(["/bin/sh", "-c", command]) + + +def _get_image_env_and_cwd_command(runner_path: str) -> str: + # Resolve the real directory so the non-root job does not need to traverse + # root-only /proc/1/cwd or inherit root SSH's /root directory. + script = f"""test "$#" -eq "$0" || exit 255 +cd -P /proc/1/cwd || exit 255 +exec env -i -- "$@" /bin/sh {shlex.quote(runner_path)}""" + command = shlex.join( + [ + "xargs", + "-0", + "-n", + "2147483647", + "/bin/sh", + "-c", + script, + ] + ) + # xargs may split large input into multiple commands. Exit 255 stops it on the + # first partial batch, before starting a runner with only some image variables. + return ( + "env_count=$(tr -cd '\\000' < /proc/1/environ | wc -c) &&\n" + 'test "$env_count" -gt 0 &&\n' + f'{command} "$env_count" < /proc/1/environ' + ) diff --git a/src/dstack/_internal/core/backends/lium/configurator.py b/src/dstack/_internal/core/backends/lium/configurator.py new file mode 100644 index 0000000000..85b27ea447 --- /dev/null +++ b/src/dstack/_internal/core/backends/lium/configurator.py @@ -0,0 +1,56 @@ +import json + +from dstack._internal.core.backends.base.configurator import BackendRecord, Configurator +from dstack._internal.core.backends.lium.api_client import LiumAPIClient, LiumAPIError +from dstack._internal.core.backends.lium.backend import LiumBackend +from dstack._internal.core.backends.lium.models import ( + LiumBackendConfig, + LiumBackendConfigWithCreds, + LiumConfig, + LiumCreds, + LiumStoredConfig, +) +from dstack._internal.core.errors import BackendInvalidCredentialsError +from dstack._internal.core.models.backends.base import BackendType +from dstack._internal.core.models.common import validate_extra_ignore, validate_json_extra_ignore + + +class LiumConfigurator(Configurator[LiumBackendConfig, LiumBackendConfigWithCreds]): + TYPE = BackendType.LIUM + BACKEND_CLASS = LiumBackend + + def validate_config(self, config: LiumBackendConfigWithCreds, default_creds_enabled: bool): + try: + LiumAPIClient(config.creds.api_key).get_user() + except LiumAPIError as e: + if e.status_code == 401: + raise BackendInvalidCredentialsError(fields=[["creds", "api_key"]]) from e + raise + + def create_backend( + self, project_name: str, config: LiumBackendConfigWithCreds + ) -> BackendRecord: + return BackendRecord( + config=LiumStoredConfig( + **validate_extra_ignore(LiumBackendConfig, config).model_dump() + ).model_dump_json(), + auth=LiumCreds.model_validate(config.creds).model_dump_json(), + ) + + def get_backend_config_with_creds(self, record: BackendRecord) -> LiumBackendConfigWithCreds: + return validate_extra_ignore(LiumBackendConfigWithCreds, self._get_config(record)) + + def get_backend_config_without_creds(self, record: BackendRecord) -> LiumBackendConfig: + return validate_extra_ignore(LiumBackendConfig, self._get_config(record)) + + def get_backend(self, record: BackendRecord) -> LiumBackend: + return LiumBackend(config=self._get_config(record)) + + def _get_config(self, record: BackendRecord) -> LiumConfig: + return validate_extra_ignore( + LiumConfig, + { + **json.loads(record.config), + "creds": validate_json_extra_ignore(LiumCreds, record.auth), + }, + ) diff --git a/src/dstack/_internal/core/backends/lium/models.py b/src/dstack/_internal/core/backends/lium/models.py new file mode 100644 index 0000000000..3125449d70 --- /dev/null +++ b/src/dstack/_internal/core/backends/lium/models.py @@ -0,0 +1,41 @@ +from typing import Annotated, Literal, Optional, Union + +from pydantic import Field + +from dstack._internal.core.models.common import CoreModel + + +class LiumAPIKeyCreds(CoreModel): + type: Annotated[Literal["api_key"], Field(description="The type of credentials")] = "api_key" + api_key: Annotated[str, Field(description="The Lium API key", min_length=1)] + + +AnyLiumCreds = LiumAPIKeyCreds +LiumCreds = AnyLiumCreds + + +class LiumBackendConfig(CoreModel): + type: Annotated[Literal["lium"], Field(description="The type of backend")] = "lium" + regions: Annotated[ + Optional[list[str]], + Field( + description=( + "The list of Lium regions, such as `us-ca` or `de-he`. Omit to use all regions" + ) + ), + ] = None + + +class LiumBackendConfigWithCreds(LiumBackendConfig): + creds: Annotated[AnyLiumCreds, Field(description="The credentials")] + + +AnyLiumBackendConfig = Union[LiumBackendConfig, LiumBackendConfigWithCreds] + + +class LiumStoredConfig(LiumBackendConfig): + pass + + +class LiumConfig(LiumStoredConfig): + creds: AnyLiumCreds diff --git a/src/dstack/_internal/core/backends/models.py b/src/dstack/_internal/core/backends/models.py index 94d9dcd094..c18a9d247d 100644 --- a/src/dstack/_internal/core/backends/models.py +++ b/src/dstack/_internal/core/backends/models.py @@ -59,6 +59,10 @@ LambdaBackendConfig, LambdaBackendConfigWithCreds, ) +from dstack._internal.core.backends.lium.models import ( + LiumBackendConfig, + LiumBackendConfigWithCreds, +) from dstack._internal.core.backends.nebius.models import ( NebiusBackendConfig, NebiusBackendConfigWithCreds, @@ -113,6 +117,7 @@ JarvisLabsBackendConfig, KubernetesBackendConfig, LambdaBackendConfig, + LiumBackendConfig, NebiusBackendConfig, OCIBackendConfig, RunpodBackendConfig, @@ -143,6 +148,7 @@ JarvisLabsBackendConfigWithCreds, KubernetesBackendConfigWithCreds, LambdaBackendConfigWithCreds, + LiumBackendConfigWithCreds, OCIBackendConfigWithCreds, NebiusBackendConfigWithCreds, RunpodBackendConfigWithCreds, @@ -188,6 +194,7 @@ class BackendConfigWithCreds(RootModel[AnyBackendConfigWithCredsTagged]): JarvisLabsBackendFileConfigWithCreds, KubernetesBackendFileConfigWithCreds, LambdaBackendConfigWithCreds, + LiumBackendConfigWithCreds, OCIBackendConfigWithCreds, NebiusBackendFileConfigWithCreds, RunpodBackendConfigWithCreds, diff --git a/src/dstack/_internal/core/models/backends/base.py b/src/dstack/_internal/core/models/backends/base.py index 6ee2931b18..9ace18f860 100644 --- a/src/dstack/_internal/core/models/backends/base.py +++ b/src/dstack/_internal/core/models/backends/base.py @@ -19,6 +19,7 @@ class BackendType(str, enum.Enum): JARVISLABS (BackendType): JarvisLabs KUBERNETES (BackendType): Kubernetes LAMBDA (BackendType): Lambda Cloud + LIUM (BackendType): Lium NEBIUS (BackendType): Nebius AI Cloud OCI (BackendType): Oracle Cloud Infrastructure RUNPOD (BackendType): Runpod Cloud @@ -46,6 +47,7 @@ class BackendType(str, enum.Enum): JARVISLABS = "jarvislabs" KUBERNETES = "kubernetes" LAMBDA = "lambda" + LIUM = "lium" REMOTE = "remote" NEBIUS = "nebius" OCI = "oci" diff --git a/src/dstack/_internal/core/models/volumes.py b/src/dstack/_internal/core/models/volumes.py index 92b3a650ee..089faa31c3 100644 --- a/src/dstack/_internal/core/models/volumes.py +++ b/src/dstack/_internal/core/models/volumes.py @@ -188,12 +188,19 @@ class DaytonaVolumeConfiguration(VolumeConfigurationWithVolumeID): ) +class LiumVolumeConfiguration(VolumeConfigurationWithVolumeID): + backend: Annotated[Literal[BackendType.LIUM], Field(description="The volume backend")] = ( + BackendType.LIUM + ) + + AnyVolumeConfiguration = Union[ AWSVolumeConfiguration, GCPVolumeConfiguration, RunpodVolumeConfiguration, KubernetesVolumeConfiguration, DaytonaVolumeConfiguration, + LiumVolumeConfiguration, ] diff --git a/src/tests/_internal/core/backends/lium/__init__.py b/src/tests/_internal/core/backends/lium/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/tests/_internal/core/backends/lium/test_api_client.py b/src/tests/_internal/core/backends/lium/test_api_client.py new file mode 100644 index 0000000000..52605a3f6f --- /dev/null +++ b/src/tests/_internal/core/backends/lium/test_api_client.py @@ -0,0 +1,238 @@ +import pytest +import requests + +from dstack._internal.core.backends.lium.api_client import API_URL, LiumAPIClient, LiumAPIError + + +class TestLiumAPIClient: + def test_identity_uses_api_key_and_timeout(self, requests_mock): + requests_mock.get(f"{API_URL}/users/me", json={"id": "user-id"}) + + assert LiumAPIClient("test-key").get_user() == {"id": "user-id"} + + assert requests_mock.last_request.headers["X-API-Key"] == "test-key" + assert requests_mock.last_request.timeout == 30 + + def test_version_reports_provider_features(self, requests_mock): + version = {"features": ["rent_idempotency"]} + requests_mock.get(f"{API_URL}/version", json=version) + + assert LiumAPIClient("test-key").get_version() == version + + def test_rental_retry_preserves_key_and_body(self, requests_mock): + requests_mock.post(f"{API_URL}/executors/node-id/rent", json={"pod_id": "pod-id"}) + client = LiumAPIClient("test-key") + payload = {"pod_name": "dstack-job", "template_id": "template-id", "gpu_count": 2} + + assert client.rent_executor("node-id", payload, "rent-key") == {"pod_id": "pod-id"} + client.rent_executor("node-id", payload, "rent-key") + + first, second = requests_mock.request_history + assert first.json() == payload + assert first.body == second.body + assert first.headers["Idempotency-Key"] == second.headers["Idempotency-Key"] == "rent-key" + + @pytest.mark.parametrize( + ["method", "path", "payload"], + [ + ("create_template", "/templates", {"docker_image": "registry.example/image:tag"}), + ( + "create_docker_credential", + "/docker-credentials/", + {"docker_username": "user", "docker_password": "private-password"}, + ), + ], + ) + def test_create_posts_payload_once(self, requests_mock, method, path, payload): + requests_mock.post(API_URL + path, json={"id": "resource-id"}) + + assert getattr(LiumAPIClient("test-key"), method)(payload) == {"id": "resource-id"} + + assert requests_mock.last_request.json() == payload + assert requests_mock.call_count == 1 + + @pytest.mark.parametrize("error", [requests.ReadTimeout, requests.ConnectionError]) + def test_rental_transport_error_does_not_retry_or_expose_credentials( + self, requests_mock, error + ): + requests_mock.post(f"{API_URL}/executors/node-id/rent", exc=error("secret")) + + with pytest.raises(LiumAPIError) as exc: + LiumAPIClient("secret").rent_executor("node-id", {"pod_name": "job"}, "rent-key") + + assert exc.value.status_code is None + assert "secret" not in str(exc.value) + assert exc.value.__suppress_context__ + assert requests_mock.call_count == 1 + + @pytest.mark.parametrize("status", [401, 403, 409, 429, 503]) + def test_http_error_preserves_status_and_code_without_body(self, requests_mock, status): + requests_mock.post( + f"{API_URL}/executors/node-id/rent", + status_code=status, + json={"error": {"code": "idempotency_outcome_unknown", "message": "secret"}}, + ) + + with pytest.raises(LiumAPIError) as exc: + LiumAPIClient("secret").rent_executor("node-id", {"pod_name": "job"}, "rent-key") + + assert str(exc.value) == (f"Lium API returned HTTP {status} (idempotency_outcome_unknown)") + assert exc.value.status_code == status + assert exc.value.code == "idempotency_outcome_unknown" + assert requests_mock.call_count == 1 + + @pytest.mark.parametrize( + "code", ["node_unavailable\nsecret", "x" * 65, {"detail": "secret"}, "lium_api_secret"] + ) + def test_error_omits_malformed_codes_and_reflected_api_key(self, requests_mock, code): + requests_mock.get( + f"{API_URL}/users/me", + status_code=403, + json={"error": {"code": code, "message": "private credentials"}}, + ) + + with pytest.raises(LiumAPIError) as exc: + LiumAPIClient("lium_api_secret").get_user() + + assert str(exc.value) == "Lium API returned HTTP 403" + assert exc.value.code is None + + def test_get_pod_returns_none_only_for_missing_pod(self, requests_mock): + requests_mock.get(f"{API_URL}/pods/pod-id", json={"id": "pod-id", "status": "PENDING"}) + requests_mock.get(f"{API_URL}/pods/missing", status_code=404) + requests_mock.get(f"{API_URL}/pods/forbidden", status_code=403) + client = LiumAPIClient("test-key") + + assert client.get_pod("pod-id") == {"id": "pod-id", "status": "PENDING"} + assert client.get_pod("missing") is None + with pytest.raises(LiumAPIError) as exc: + client.get_pod("forbidden") + assert exc.value.status_code == 403 + + @pytest.mark.parametrize("since", [None, "2026-10-16T08:00:00Z"]) + def test_account_events_include_deleted_pods_and_request_maximum_history( + self, requests_mock, since + ): + event = { + "pod_id": "pod-id", + "pod_name": "dstack-job", + "event_type": "pod-lifecycle", + "to_status": "DELETED", + } + requests_mock.get(f"{API_URL}/users/me/events", json=[event]) + + assert LiumAPIClient("test-key").get_account_events(since) == [event] + + params = {"limit": ["1000"]} + if since is not None: + params["since"] = [since.lower()] + assert requests_mock.last_request.qs == params + + @pytest.mark.parametrize( + ["headers", "expected"], + [ + ({}, False), + ({"Idempotent-Replayed": "false"}, False), + ({"Idempotent-Replayed": "true"}, True), + ], + ) + def test_refusal_exposes_idempotent_replay_header(self, requests_mock, headers, expected): + requests_mock.post( + f"{API_URL}/executors/node-id/rent", + status_code=400, + headers=headers, + json={"error": {"code": "node_unavailable"}}, + ) + + with pytest.raises(LiumAPIError) as exc: + LiumAPIClient("test-key").rent_executor("node-id", {"pod_name": "job"}, "rent-key") + + assert exc.value.idempotent_replayed is expected + assert requests_mock.call_count == 1 + + def test_volume_create_and_lookup(self, requests_mock): + volume = {"id": "volume-id", "name": "checkpoints"} + requests_mock.post(f"{API_URL}/volumes", json=volume) + requests_mock.get(f"{API_URL}/volumes/volume-id", json=volume) + requests_mock.get(f"{API_URL}/volumes/missing", status_code=404) + requests_mock.get( + f"{API_URL}/volumes/forbidden", + status_code=403, + json={"success": False, "message": "secret"}, + ) + client = LiumAPIClient("test-key") + + assert client.create_volume("checkpoints") == volume + assert requests_mock.last_request.json() == {"name": "checkpoints"} + assert client.get_volume("volume-id") == volume + assert client.get_volume("missing") is None + with pytest.raises(LiumAPIError) as exc: + client.get_volume("forbidden") + assert exc.value.status_code == 403 + assert exc.value.code is None + assert "secret" not in str(exc.value) + + @pytest.mark.parametrize( + ["method", "path"], + [ + ("get_pods", "/pods"), + ("get_templates", "/templates"), + ], + ) + def test_lists_require_object_items(self, requests_mock, method, path): + requests_mock.get( + API_URL + path, [{"json": [{"id": "resource-id"}]}, {"json": ["invalid"]}] + ) + call = getattr(LiumAPIClient("test-key"), method) + + assert call() == [{"id": "resource-id"}] + with pytest.raises(LiumAPIError, match="invalid list response"): + call() + + @pytest.mark.parametrize( + ["method", "path"], + [ + ("delete_pod", "/pods/resource-id"), + ("delete_template", "/templates/resource-id"), + ("delete_docker_credential", "/docker-credentials/resource-id"), + ("delete_volume", "/volumes/resource-id"), + ], + ) + @pytest.mark.parametrize("status", [200, 204, 404]) + def test_delete_is_idempotent(self, requests_mock, method, path, status): + requests_mock.delete(API_URL + path, status_code=status) + + getattr(LiumAPIClient("test-key"), method)("resource-id") + + assert requests_mock.call_count == 1 + + def test_delete_error_is_not_silenced(self, requests_mock): + requests_mock.delete(f"{API_URL}/pods/pod-id", status_code=503, text="secret") + + with pytest.raises(LiumAPIError) as exc: + LiumAPIClient("test-key").delete_pod("pod-id") + + assert exc.value.status_code == 503 + assert str(exc.value) == "Lium API returned HTTP 503" + + def test_redirect_does_not_forward_api_key(self, requests_mock): + requests_mock.get( + f"{API_URL}/users/me", + status_code=302, + headers={"Location": "https://untrusted.example/identity"}, + ) + + with pytest.raises(LiumAPIError) as exc: + LiumAPIClient("test-key").get_user() + + assert exc.value.status_code == 302 + assert requests_mock.call_count == 1 + + @pytest.mark.parametrize("body", ["not JSON", '"not an object"', "null"]) + def test_invalid_create_response_does_not_retry(self, requests_mock, body): + requests_mock.post(f"{API_URL}/templates", text=body) + + with pytest.raises(LiumAPIError): + LiumAPIClient("test-key").create_template({"docker_image": "ubuntu:24.04"}) + + assert requests_mock.call_count == 1 diff --git a/src/tests/_internal/core/backends/lium/test_compute.py b/src/tests/_internal/core/backends/lium/test_compute.py new file mode 100644 index 0000000000..a2616f7c5f --- /dev/null +++ b/src/tests/_internal/core/backends/lium/test_compute.py @@ -0,0 +1,665 @@ +import json +import shlex +import stat +import subprocess +from datetime import datetime, timedelta, timezone +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import gpuhunt +import paramiko +import pytest + +from dstack._internal.core.backends.lium.api_client import LiumAPIError +from dstack._internal.core.backends.lium.compute import ( + LiumCompute, + LiumInstanceBackendData, + _bootstrap_runner, + _get_bootstrap_command, + _get_image_env_and_cwd_command, + _get_provider_ssh, +) +from dstack._internal.core.backends.lium.models import LiumConfig, LiumCreds +from dstack._internal.core.errors import ( + ComputeError, + NoCapacityError, + 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.unix import UnixUser +from dstack._internal.utils.crypto import generate_rsa_key_pair_bytes + +COMPUTE_MODULE = "dstack._internal.core.backends.lium.compute" +PUBLIC_KEY = ( + "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAINOmx0T+hBRaJ6jCi21ZYe2NW3EZS8e0Mdwl+yZJt+kD project" +) +STARTUP_COMMANDS = ["install runner", "runner --key 'a key with spaces'"] + + +@pytest.fixture +def compute(): + with patch(f"{COMPUTE_MODULE}.LiumAPIClient", autospec=True) as client_class: + compute = LiumCompute(LiumConfig(creds=LiumCreds(api_key="test-api-key"))) + client = client_class.return_value + client.create_template.return_value = {"id": "template-id"} + client.create_docker_credential.return_value = {"id": "credential-id"} + client.rent_executor.return_value = {"success": True, "pod_id": "pod-id"} + client.get_pod.return_value = _pod() + client.get_pods.return_value = [] + client.get_account_events.return_value = [] + client.get_version.return_value = {"features": ["rent_idempotency"]} + client.get_templates.return_value = [] + yield compute + + +def _pod(**overrides): + return { + "id": "pod-id", + "pod_name": "dstack-test-unique", + "executor_id": "executor-id", + "status": "RUNNING", + "ssh_connect_cmd": "ssh root@192.0.2.10 -p 40022", + "ports_mapping": {"22": 40022, "10022": 40100}, + **overrides, + } + + +def _item(): + return gpuhunt.CatalogItem( + provider="lium", + instance_name="executor-id", + location="us-ia", + price=0.42, + cpu=6, + memory=42, + disk_size=100, + gpu_count=1, + gpu_name="A6000", + gpu_vendor=gpuhunt.AcceleratorVendor.NVIDIA, + gpu_memory=48, + spot=False, + provider_data={"rent_url": "https://lium.io/browse-pods/executor-id"}, + ) + + +def _offer(compute): + with patch.object(compute._catalog, "query", return_value=[_item()]): + return compute.get_offers_by_requirements( + Requirements(resources=ResourcesSpec()), + full_offers=False, + unallocated_resources=False, + )[0] + + +def _run_job(compute, **overrides): + job_spec = SimpleNamespace( + image_name="pytorch/pytorch:2.11.0-cuda12.8-cudnn9-runtime", + registry_auth=None, + volumes=[], + user=None, + ) + for name in ("image_name", "registry_auth", "user", "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(compute), + "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="dstack-test-unique", + ), + 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") + + +def _terminate(compute, data): + compute.terminate_instance(data.instance_id, data.region, data.backend_data) + + +class TestRunJob: + def test_uses_offer_and_configures_root_ssh_bootstrap(self, compute): + offer = _offer(compute) + user = UnixUser.parse("10001:10001") + + data = _run_job(compute, instance_offer=offer, user=user) + + assert offer.availability == InstanceAvailability.AVAILABLE + assert offer.instance_runtime == InstanceRuntime.RUNNER + assert data.backend == BackendType.LIUM + assert data.instance_id == "pod-id" + assert data.price == 0.42 + assert data.instance_type.resources.memory_mib == 42 * 1024 + assert data.username == "root" + assert data.dockerized is False + assert data.hostname is None + assert data.ssh_port is None + template = compute.api_client.create_template.call_args.args[0] + assert template["entrypoint"] == "/bin/sh" + assert shlex.split(template["startup_commands"]) == ["-c", "exec sleep infinity"] + assert template["internal_ports"] == [22, 10022] + assert template["supports_volume_encryption"] is False + assert "user" not in template + executor, payload, key = compute.api_client.rent_executor.call_args.args + assert executor == "executor-id" + assert payload["gpu_count"] == 1 + assert payload["template_id"] == "template-id" + assert payload["user_public_key"] == PUBLIC_KEY + assert payload["enable_volume_encryption"] is False + assert "termination_hours" not in payload + assert key + backend = LiumInstanceBackendData.model_validate_json(data.backend_data) + assert backend.rent_payload is None + assert backend.startup_command == " && ".join(STARTUP_COMMANDS) + + @pytest.mark.parametrize( + "image", + ["ubuntu", "registry.example:5000/team/image:tag", "registry.example/image@sha256:abcd"], + ) + def test_passes_full_image_reference_using_sdk_convention(self, compute, image): + _run_job(compute, image_name=image) + + template = compute.api_client.create_template.call_args.args[0] + assert template["docker_image"] == image + assert template["docker_image_tag"] == "" + assert template["docker_image_digest"] == "" + + def test_private_registry_credentials_are_selected_and_not_stored(self, compute): + data = _run_job( + compute, registry_auth=RegistryAuth(username="user", password="secret-password") + ) + + compute.api_client.create_docker_credential.assert_called_once_with( + {"docker_username": "user", "docker_password": "secret-password"} + ) + template = compute.api_client.create_template.call_args.args[0] + assert template["docker_credential_id"] == "credential-id" + assert "secret-password" not in data.backend_data + assert json.loads(data.backend_data)["credential_id"] == "credential-id" + + def test_uncertain_credential_creation_never_recovers_by_username(self, compute): + compute.api_client.create_docker_credential.side_effect = LiumAPIError("timeout") + + with pytest.raises(ComputeError, match="Check Docker credentials"): + _run_job(compute, registry_auth=RegistryAuth(username="user", password="secret")) + + compute.api_client.delete_docker_credential.assert_not_called() + compute.api_client.create_template.assert_not_called() + compute.api_client.rent_executor.assert_not_called() + + def test_recovers_lost_template_response_by_unique_name(self, compute): + compute.api_client.create_template.side_effect = LiumAPIError("timeout") + compute.api_client.get_templates.return_value = [ + {"id": "someone-elses-template", "name": "other"}, + {"id": "recovered-template", "name": "dstack-test-unique"}, + ] + + data = _run_job(compute) + + assert json.loads(data.backend_data)["template_id"] == "recovered-template" + assert compute.api_client.rent_executor.call_args.args[1]["template_id"] == ( + "recovered-template" + ) + + def test_template_failure_cleans_up_owned_credential(self, compute): + compute.api_client.create_template.side_effect = LiumAPIError("bad image", 400) + + with pytest.raises(LiumAPIError): + _run_job(compute, registry_auth=RegistryAuth(username="user", password="secret")) + + compute.api_client.delete_docker_credential.assert_called_once_with("credential-id") + compute.api_client.rent_executor.assert_not_called() + + def test_capacity_failure_cleans_up_template_then_credential(self, compute): + compute.api_client.rent_executor.side_effect = LiumAPIError( + "unavailable", 409, "node_paused_by_provider" + ) + cleanup = Mock() + cleanup.attach_mock(compute.api_client.delete_template, "template") + cleanup.attach_mock(compute.api_client.delete_docker_credential, "credential") + + with pytest.raises(NoCapacityError): + _run_job(compute, registry_auth=RegistryAuth(username="user", password="secret")) + + assert [call[0] for call in cleanup.mock_calls] == ["template", "credential"] + + def test_uncertain_rent_retains_cleanup_ids_and_original_request(self, compute): + compute.api_client.rent_executor.side_effect = LiumAPIError("timeout") + + data = _run_job(compute) + + assert data.instance_id == "dstack-test-unique" + backend = LiumInstanceBackendData.model_validate_json(data.backend_data) + assert backend.template_id == "template-id" + assert backend.rent_payload == compute.api_client.rent_executor.call_args.args[1] + assert backend.idempotency_key == compute.api_client.rent_executor.call_args.args[2] + compute.api_client.delete_template.assert_not_called() + + +class TestUpdateProvisioningData: + @pytest.mark.parametrize("state", ["PENDING", "CREATING", "PULLING"]) + def test_waits_for_running_pod(self, compute, state): + data = _run_job(compute) + compute.api_client.get_pod.return_value = _pod(status=state) + + with patch(f"{COMPUTE_MODULE}._bootstrap_runner") as bootstrap: + _update(compute, data) + + bootstrap.assert_not_called() + assert data.hostname is None + + @pytest.mark.parametrize("state", ["FAILED", "STOPPED", "CREATION_FAILED", "DELETING"]) + def test_terminal_state_fails_provisioning(self, compute, state): + data = _run_job(compute) + compute.api_client.get_pod.return_value = _pod(status=state) + + with pytest.raises(ProvisioningError, match=state): + _update(compute, data) + + def test_sets_direct_runner_ssh_after_bootstrap(self, compute): + data = _run_job(compute) + + with patch(f"{COMPUTE_MODULE}._bootstrap_runner") as bootstrap: + _update(compute, data) + + ssh, key, command = bootstrap.call_args.args + assert (ssh.hostname, ssh.port, ssh.username) == ("192.0.2.10", 40022, "root") + assert key == "private-key" + assert command == " && ".join(STARTUP_COMMANDS) + assert (data.hostname, data.ssh_port) == ("192.0.2.10", 40100) + assert data.ssh_proxy is None + + def test_uses_provider_proxy_when_runner_port_is_not_mapped(self, compute): + data = _run_job(compute) + compute.api_client.get_pod.return_value = _pod(ports_mapping={"22": 40022}) + + with patch(f"{COMPUTE_MODULE}._bootstrap_runner"): + _update(compute, data) + + assert (data.hostname, data.ssh_port) == ("localhost", 10022) + assert (data.ssh_proxy.hostname, data.ssh_proxy.port) == ("192.0.2.10", 40022) + + def test_running_does_not_imply_provider_ssh_is_ready(self, compute): + data = _run_job(compute) + + with patch(f"{COMPUTE_MODULE}._bootstrap_runner", side_effect=paramiko.SSHException()): + _update(compute, data) + + assert data.hostname is None + + def test_recovers_uncertain_rent_without_replaying_if_pod_exists(self, compute): + compute.api_client.rent_executor.side_effect = LiumAPIError("timeout") + data = _run_job(compute) + compute.api_client.get_pods.return_value = [ + _pod(id="unrelated", executor_id="wrong-executor"), + _pod(), + ] + + with patch(f"{COMPUTE_MODULE}._bootstrap_runner"): + _update(compute, data) + + assert data.instance_id == "pod-id" + assert json.loads(data.backend_data)["rent_payload"] is None + assert compute.api_client.rent_executor.call_count == 1 + + def test_replays_exact_request_and_key_only_during_provisioning(self, compute): + compute.api_client.rent_executor.side_effect = [ + LiumAPIError("timeout"), + {"success": True, "pod_id": "pod-id"}, + ] + data = _run_job(compute) + + with patch(f"{COMPUTE_MODULE}._bootstrap_runner"): + _update(compute, data) + + first, second = compute.api_client.rent_executor.call_args_list + assert first == second + assert data.instance_id == "pod-id" + + def test_definitive_replay_failure_can_be_terminated_without_rerenting(self, compute): + compute.api_client.rent_executor.side_effect = [ + LiumAPIError("timeout"), + LiumAPIError("unavailable", 400, "node_unavailable", idempotent_replayed=True), + ] + data = _run_job(compute) + + _update(compute, data) + # The server only saves the result when update_provisioning_data returns. + # A later failing call must not be needed to persist the definitive outcome. + persisted = data.model_dump_json() + next_poll = JobProvisioningData.model_validate_json(persisted) + with pytest.raises(ProvisioningError, match="rental failed"): + _update(compute, next_poll) + _terminate(compute, JobProvisioningData.model_validate_json(persisted)) + + assert json.loads(data.backend_data)["rent_failed"] is True + assert compute.api_client.rent_executor.call_count == 2 + compute.api_client.delete_pod.assert_not_called() + + @pytest.mark.parametrize("status", [401, 403, 422, 429]) + def test_unrecorded_replay_rejection_keeps_original_rental_uncertain(self, compute, status): + compute.api_client.rent_executor.side_effect = [ + LiumAPIError("timeout"), + LiumAPIError("rejected before idempotency", status), + ] + data = _run_job(compute) + original_backend_data = data.backend_data + + _update(compute, data) + + assert data.backend_data == original_backend_data + assert compute.api_client.rent_executor.call_count == 2 + compute.api_client.delete_template.assert_not_called() + compute.api_client.delete_docker_credential.assert_not_called() + + def test_expired_idempotency_key_is_never_replayed(self, compute): + compute.api_client.rent_executor.side_effect = LiumAPIError("timeout") + data = _run_job(compute) + backend = LiumInstanceBackendData.model_validate_json(data.backend_data) + backend.rent_started_at = datetime.now(timezone.utc) - timedelta(days=1) + data.backend_data = backend.model_dump_json() + + _update(compute, data) + + assert compute.api_client.rent_executor.call_count == 1 + compute.api_client.get_version.assert_not_called() + compute.api_client.delete_template.assert_not_called() + + def test_backend_without_idempotency_feature_is_not_retried(self, compute): + compute.api_client.rent_executor.side_effect = LiumAPIError("timeout") + compute.api_client.get_version.return_value = {"features": []} + data = _run_job(compute) + + _update(compute, data) + + assert compute.api_client.rent_executor.call_count == 1 + compute.api_client.delete_template.assert_not_called() + + +class TestTerminateInstance: + def test_keeps_cleanup_resources_until_pod_disappears(self, compute): + data = _run_job(compute, registry_auth=RegistryAuth(username="user", password="secret")) + + with pytest.raises(NotYetTerminated): + _terminate(compute, data) + + compute.api_client.delete_pod.assert_called_once_with("pod-id") + compute.api_client.delete_template.assert_not_called() + compute.api_client.delete_docker_credential.assert_not_called() + compute.api_client.get_pod.return_value = None + _terminate(compute, data) + compute.api_client.delete_template.assert_called_once_with("template-id") + compute.api_client.delete_docker_credential.assert_called_once_with("credential-id") + + def test_deleting_state_does_not_repeat_delete(self, compute): + data = _run_job(compute) + compute.api_client.get_pod.return_value = _pod(status="DELETING") + + with pytest.raises(NotYetTerminated): + _terminate(compute, data) + + compute.api_client.delete_pod.assert_not_called() + + def test_reconciles_delete_state_conflict(self, compute): + data = _run_job(compute) + compute.api_client.get_pod.side_effect = [_pod(), _pod(status="DELETING")] + compute.api_client.delete_pod.side_effect = LiumAPIError( + "conflict", 400, "pod_state_conflict" + ) + + with pytest.raises(NotYetTerminated): + _terminate(compute, data) + + compute.api_client.delete_template.assert_not_called() + + def test_uncertain_outcome_never_rerents_during_termination(self, compute): + compute.api_client.rent_executor.side_effect = LiumAPIError("timeout") + data = _run_job(compute) + + with pytest.raises(NotYetTerminated, match="uncertain"): + _terminate(compute, data) + + assert compute.api_client.rent_executor.call_count == 1 + compute.api_client.delete_template.assert_not_called() + compute.api_client.delete_pod.assert_not_called() + + def test_uncertain_rental_found_by_name_is_deleted(self, compute): + compute.api_client.rent_executor.side_effect = LiumAPIError("timeout") + data = _run_job(compute) + compute.api_client.get_pods.return_value = [_pod()] + + with pytest.raises(NotYetTerminated): + _terminate(compute, data) + + compute.api_client.delete_pod.assert_called_once_with("pod-id") + + @pytest.mark.parametrize( + "deleted_event", [{"to_status": "DELETED"}, {"sub_event_type": "pod-delete.success"}] + ) + def test_deleted_rental_is_recovered_from_durable_account_events(self, compute, deleted_event): + compute.api_client.rent_executor.side_effect = LiumAPIError("timeout") + data = _run_job(compute) + compute.api_client.get_pods.return_value = [_pod()] + with pytest.raises(NotYetTerminated): + _terminate(compute, data) + # A fresh backend receives exactly the same persisted data. The deleted + # pod is no longer listed, but its account event still identifies it. + compute.api_client.get_pods.return_value = [] + compute.api_client.get_pod.return_value = None + compute.api_client.get_account_events.return_value = [ + {"pod_id": "pod-id", "pod_name": "dstack-test-unique", **deleted_event} + ] + next_compute = LiumCompute(compute.config) + + _terminate(next_compute, JobProvisioningData.model_validate_json(data.model_dump_json())) + + compute.api_client.get_pod.assert_called_once_with("pod-id") + compute.api_client.delete_template.assert_called_once_with("template-id") + assert compute.api_client.rent_executor.call_count == 1 + + def test_unrelated_events_do_not_resolve_uncertain_outcome(self, compute): + compute.api_client.rent_executor.side_effect = LiumAPIError("timeout") + data = _run_job(compute) + compute.api_client.get_account_events.return_value = [ + {"pod_id": "other-id", "pod_name": "other-pod", "to_status": "DELETED"} + ] + + with pytest.raises(NotYetTerminated, match="uncertain"): + _terminate(compute, data) + + compute.api_client.get_pod.assert_not_called() + compute.api_client.delete_template.assert_not_called() + + def test_creation_event_without_visible_pod_is_not_proof_of_deletion(self, compute): + compute.api_client.rent_executor.side_effect = LiumAPIError("timeout") + data = _run_job(compute) + compute.api_client.get_pod.return_value = None + compute.api_client.get_account_events.return_value = [ + {"pod_id": "pod-id", "pod_name": "dstack-test-unique", "to_status": "PENDING"} + ] + + with pytest.raises(NotYetTerminated, match="uncertain"): + _terminate(compute, data) + + compute.api_client.delete_template.assert_not_called() + + +class TestBootstrapRunner: + def test_accepts_server_generated_pkcs8_project_key(self): + # Real key generation/conversion earns its subprocess cost here: mocking + # pkey_from_str hid incompatibility with the server's actual PKCS8 format. + private_key, public_key = generate_rsa_key_pair_bytes() + assert private_key.startswith(b"-----BEGIN PRIVATE KEY-----") + with patch(f"{COMPUTE_MODULE}.paramiko.SSHClient") as ssh_class: + client = ssh_class.return_value.__enter__.return_value + stdout = Mock() + stdout.read.return_value = b"running" + stdout.channel.recv_exit_status.return_value = 0 + client.exec_command.return_value = (Mock(), stdout, Mock()) + + _bootstrap_runner(_get_provider_ssh(_pod()), private_key.decode(), "runner") + + pkey = client.connect.call_args.kwargs["pkey"] + assert pkey.get_name() == "ssh-rsa" + assert pkey.get_base64() == public_key.decode().split()[1] + + def test_private_bootstrap_does_not_restrict_runner_created_files(self, tmp_path): + # Execute the generated shell so a restrictive inherited umask cannot + # silently break non-root access to runner-created job/SSH files. + directory = tmp_path / "bootstrap" + job_directory = tmp_path / "job" + startup = f"mkdir {shlex.quote(str(job_directory))} && touch {job_directory}/key" + with ( + patch(f"{COMPUTE_MODULE}.BOOTSTRAP_DIR", str(directory)), + patch( + f"{COMPUTE_MODULE}._get_image_env_and_cwd_command", + side_effect=lambda path: shlex.join(["/bin/sh", path]), + ), + ): + command = shlex.split(_get_bootstrap_command(startup)) + # The production command detaches; join the immediate shell child in this + # test instead of polling or sleeping for the exit marker. + command[-1] += "\nwait" + + subprocess.run(command, capture_output=True, check=True, timeout=5) + + assert stat.S_IMODE(directory.stat().st_mode) == 0o700 + assert stat.S_IMODE((directory / "runner.sh").stat().st_mode) == 0o600 + assert stat.S_IMODE(job_directory.stat().st_mode) == 0o755 + assert stat.S_IMODE((job_directory / "key").stat().st_mode) == 0o644 + assert (directory / "exit-code").read_text() == "0\n" + + @pytest.mark.parametrize("output,status", [(b"running", 0), (b"exited", 0), (b"", 1)]) + def test_checks_existing_bootstrap_status(self, output, status): + ssh = _get_provider_ssh(_pod()) + with ( + patch(f"{COMPUTE_MODULE}.paramiko.SSHClient") as ssh_class, + patch(f"{COMPUTE_MODULE}.pkey_from_str", return_value="parsed-key"), + ): + client = ssh_class.return_value.__enter__.return_value + stdout = Mock() + stdout.read.return_value = output + stdout.channel.recv_exit_status.return_value = status + client.exec_command.return_value = (Mock(), stdout, Mock()) + if output == b"running": + _bootstrap_runner(ssh, "private-key", "runner") + else: + with pytest.raises(ProvisioningError): + _bootstrap_runner(ssh, "private-key", "runner") + assert client.connect.call_args.kwargs["allow_agent"] is False + assert client.connect.call_args.kwargs["look_for_keys"] is False + assert client.connect.call_args.kwargs["pkey"] == "parsed-key" + + @pytest.mark.parametrize("marker,output", [("started", "running"), ("exit-code", "exited")]) + def test_repeated_bootstrap_does_not_rewrite_or_launch(self, tmp_path, marker, output): + # Exercise the real shell branching, which mocks cannot validate. + directory = tmp_path / "bootstrap" + directory.mkdir() + if marker == "started": + (directory / marker).mkdir() + else: + (directory / marker).write_text("7\n") + with patch(f"{COMPUTE_MODULE}.BOOTSTRAP_DIR", str(directory)): + command = _get_bootstrap_command("exit 99") + + result = subprocess.run(shlex.split(command), capture_output=True, check=True) + + assert result.stdout.decode() == output + assert not (directory / "runner.sh").exists() + assert not (directory / "monitor.sh").exists() + + def test_restores_image_environment_and_physical_working_directory(self, tmp_path): + # Black-box check of the NUL-delimited xargs/env boundary; interpreting this + # data as shell syntax would execute substitutions or split newline values. + environment = { + "VALUE": "spaces and 'quotes'\n$(touch unwanted) `touch unwanted2`", + "EMPTY": "", + } + env_path = tmp_path / "environ" + env_path.write_bytes( + b"\0".join(f"{key}={value}".encode() for key, value in environment.items()) + b"\0" + ) + image_directory = tmp_path / "image workdir" + image_directory.mkdir() + cwd_link = tmp_path / "cwd" + cwd_link.symlink_to(image_directory, target_is_directory=True) + runner = tmp_path / "runner.sh" + runner.write_text( + 'printf "%s\\0%s\\0%s\\0%s" "$VALUE" "$EMPTY" "${SSH_ONLY-unset}" "$PWD"\n' + ) + command = ( + _get_image_env_and_cwd_command(str(runner)) + .replace("/proc/1/environ", shlex.quote(str(env_path))) + .replace("/proc/1/cwd", shlex.quote(str(cwd_link))) + ) + + result = subprocess.run( + ["/bin/sh", "-c", command], + capture_output=True, + check=True, + cwd=tmp_path, + env={"SSH_ONLY": "must not leak"}, + ) + + assert result.stdout == ( + environment["VALUE"].encode() + + b"\0\0unset\0" + + str(image_directory.resolve()).encode() + ) + assert not (tmp_path / "unwanted").exists() + assert not (tmp_path / "unwanted2").exists() + assert not (image_directory / "unwanted").exists() + assert not (image_directory / "unwanted2").exists() + + def test_missing_image_working_directory_does_not_start_runner(self, tmp_path): + env_path = tmp_path / "environ" + env_path.write_bytes(b"VALUE=example\0") + runner = tmp_path / "runner.sh" + runner.write_text("printf started\n") + command = ( + _get_image_env_and_cwd_command(str(runner)) + .replace("/proc/1/environ", shlex.quote(str(env_path))) + .replace("/proc/1/cwd", shlex.quote(str(tmp_path / "missing"))) + ) + + result = subprocess.run(["/bin/sh", "-c", command], capture_output=True) + + assert result.returncode != 0 + assert result.stdout == b"" + + def test_oversized_environment_fails_without_starting_partial_runner(self, tmp_path): + env_path = tmp_path / "environ" + env_path.write_bytes(b"\0".join(f"VAR_{n}=value".encode() for n in range(100)) + b"\0") + runner = tmp_path / "runner.sh" + runner.write_text("printf started\n") + command = ( + _get_image_env_and_cwd_command(str(runner)) + .replace("/proc/1/environ", shlex.quote(str(env_path))) + .replace("/proc/1/cwd", shlex.quote(str(tmp_path))) + ) + # Force a small limit to test the same xargs failure on Linux and macOS. + command = command.replace("xargs -0", "xargs -s 512 -0") + + result = subprocess.run(["/bin/sh", "-c", command], capture_output=True) + + assert result.returncode != 0 + assert result.stdout == b"" diff --git a/src/tests/_internal/core/backends/lium/test_configurator.py b/src/tests/_internal/core/backends/lium/test_configurator.py new file mode 100644 index 0000000000..b62555101e --- /dev/null +++ b/src/tests/_internal/core/backends/lium/test_configurator.py @@ -0,0 +1,52 @@ +import json +from unittest.mock import patch + +import pytest + +from dstack._internal.core.backends.lium.api_client import API_URL, LiumAPIError +from dstack._internal.core.backends.lium.configurator import LiumConfigurator +from dstack._internal.core.backends.lium.models import LiumBackendConfigWithCreds, LiumCreds +from dstack._internal.core.errors import BackendInvalidCredentialsError + + +class TestLiumConfigurator: + def test_validate_config(self): + config = LiumBackendConfigWithCreds(creds=LiumCreds(api_key="test-key")) + with patch( + "dstack._internal.core.backends.lium.configurator.LiumAPIClient", autospec=True + ) as client: + LiumConfigurator().validate_config(config, default_creds_enabled=False) + + client.assert_called_once_with("test-key") + client.return_value.get_user.assert_called_once_with() + + def test_invalid_credentials(self, requests_mock): + requests_mock.get(f"{API_URL}/users/me", status_code=401) + config = LiumBackendConfigWithCreds(creds=LiumCreds(api_key="invalid")) + + with pytest.raises(BackendInvalidCredentialsError) as exc: + LiumConfigurator().validate_config(config, default_creds_enabled=False) + + assert exc.value.fields == [["creds", "api_key"]] + + def test_server_error_is_not_invalid_credentials(self, requests_mock): + requests_mock.get(f"{API_URL}/users/me", status_code=503) + config = LiumBackendConfigWithCreds(creds=LiumCreds(api_key="test-key")) + + with pytest.raises(LiumAPIError): + LiumConfigurator().validate_config(config, default_creds_enabled=False) + + def test_stored_config_keeps_credentials_separate(self): + config = LiumBackendConfigWithCreds( + creds=LiumCreds(api_key="test-secret"), regions=["us-ia", "de-he"] + ) + configurator = LiumConfigurator() + + record = configurator.create_backend(project_name="main", config=config) + + assert json.loads(record.config) == {"type": "lium", "regions": ["us-ia", "de-he"]} + 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/lium/test_models.py b/src/tests/_internal/core/backends/lium/test_models.py new file mode 100644 index 0000000000..96d27b4b2f --- /dev/null +++ b/src/tests/_internal/core/backends/lium/test_models.py @@ -0,0 +1,44 @@ +import pytest +from pydantic import Field, TypeAdapter, ValidationError +from typing_extensions import Annotated + +from dstack._internal.core.backends.lium.models import LiumBackendConfigWithCreds +from dstack._internal.core.backends.models import ( + AnyBackendFileConfigWithCreds, + BackendConfigWithCreds, +) +from dstack._internal.core.models.volumes import LiumVolumeConfiguration, VolumeConfiguration + + +class TestLiumBackendConfig: + def test_api_and_server_file_configs_accept_lium(self): + data = { + "type": "lium", + "regions": ["us-ia", "de-he"], + "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, LiumBackendConfigWithCreds) + assert api_config == file_config + assert api_config.regions == ["us-ia", "de-he"] + + +class TestLiumVolumeConfiguration: + @pytest.mark.parametrize("volume_id", [None, "volume-id"]) + def test_accepts_unbounded_volumes_without_region(self, volume_id): + config = VolumeConfiguration.model_validate( + {"backend": "lium", "volume_id": volume_id} + ).root + + assert isinstance(config, LiumVolumeConfiguration) + assert config.is_external == (volume_id is not None) + assert config.external_volume_id == volume_id + + @pytest.mark.parametrize("field,value", [("size", "100GB"), ("region", "us-ia")]) + def test_rejects_unsupported_size_and_region(self, field, value): + with pytest.raises(ValidationError): + VolumeConfiguration.model_validate({"backend": "lium", field: value}) diff --git a/src/tests/_internal/core/backends/lium/test_volumes.py b/src/tests/_internal/core/backends/lium/test_volumes.py new file mode 100644 index 0000000000..ac52eb8042 --- /dev/null +++ b/src/tests/_internal/core/backends/lium/test_volumes.py @@ -0,0 +1,114 @@ +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import pytest + +from dstack._internal.core.backends.lium.api_client import LiumAPIError +from dstack._internal.core.backends.lium.compute import LiumCompute, _get_volume_commands +from dstack._internal.core.backends.lium.models import LiumConfig, LiumCreds +from dstack._internal.core.errors import ComputeError +from dstack._internal.core.models.unix import UnixUser +from dstack._internal.core.models.volumes import LiumVolumeConfiguration, VolumeMountPoint + +COMPUTE_MODULE = "dstack._internal.core.backends.lium.compute" + + +@pytest.fixture +def compute(): + compute = LiumCompute(LiumConfig(creds=LiumCreds(api_key="test-key"))) + compute.api_client = Mock() + return compute + + +class TestVolumeLifecycle: + def test_create_unbounded_volume(self, compute): + compute.api_client.create_volume.return_value = {"id": "volume-id"} + with patch(f"{COMPUTE_MODULE}.generate_unique_volume_name", return_value="unique-name"): + data = compute.create_volume(Mock()) + + compute.api_client.create_volume.assert_called_once_with("unique-name") + assert data.volume_id == "volume-id" + assert data.size_gb is None + assert data.price is None + assert not data.attachable + assert not data.detachable + + def test_recovers_lost_create_response_by_unique_name(self, compute): + compute.api_client.create_volume.side_effect = LiumAPIError("timed out") + compute.api_client.get_volumes.return_value = [ + {"id": "other-id", "name": "other"}, + {"id": "volume-id", "name": "unique-name"}, + ] + with patch(f"{COMPUTE_MODULE}.generate_unique_volume_name", return_value="unique-name"): + data = compute.create_volume(Mock()) + + assert data.volume_id == "volume-id" + assert compute.api_client.create_volume.call_count == 1 + + def test_register_existing_volume(self, compute): + volume = SimpleNamespace(configuration=LiumVolumeConfiguration(volume_id="existing")) + compute.api_client.get_volume.return_value = {"id": "existing"} + + data = compute.register_volume(volume) + + assert data.volume_id == "existing" + compute.api_client.create_volume.assert_not_called() + + def test_register_missing_volume(self, compute): + compute.api_client.get_volume.return_value = None + with pytest.raises(ComputeError, match="not found"): + compute.register_volume( + SimpleNamespace(configuration=LiumVolumeConfiguration(volume_id="missing")) + ) + + def test_delete_volume(self, compute): + compute.delete_volume(SimpleNamespace(volume_id="volume-id")) + compute.api_client.delete_volume.assert_called_once_with("volume-id") + + +class TestVolumeMount: + @pytest.mark.parametrize("user", [None, UnixUser(uid=0), UnixUser(username="root")]) + def test_root_custom_path(self, user): + commands = _volume_commands(user=user, path="/checkpoints") + + assert commands[0] == "grep -q ' /mnt .* - fuse.s3fs ' /proc/self/mountinfo" + assert commands[1:] == ["mkdir -p -- /checkpoints", "mount --bind /mnt /checkpoints"] + + @pytest.mark.parametrize("user", [UnixUser(uid=10001), UnixUser(username="app")]) + def test_nonroot_rejected(self, user): + with pytest.raises(ComputeError, match="root jobs only"): + _volume_commands(user=user) + + def test_native_mount_path_needs_no_bind(self): + assert len(_volume_commands(path="/mnt")) == 1 + + def test_quotes_mount_path(self): + commands = _volume_commands(path="/checkpoint dir; echo unsafe") + assert commands[-1] == "mount --bind /mnt '/checkpoint dir; echo unsafe'" + + @pytest.mark.parametrize("path", ["/", "/root", "/root/.dstack-lium", "/proc/a", "/mnt/a"]) + def test_rejects_paths_that_hide_bootstrap_or_nest_source(self, path): + with pytest.raises(ComputeError, match="cannot be mounted"): + _volume_commands(path=path) + + def test_two_volumes_rejected(self): + with pytest.raises(ComputeError, match="only one"): + _get_volume_commands(Mock(), Mock(), [Mock(), Mock()]) + + def test_selects_volume_from_alternatives(self): + job = SimpleNamespace( + job_spec=SimpleNamespace( + user=None, + volumes=[VolumeMountPoint(name=["elsewhere", "cache"], path="/data")], + ) + ) + assert _get_volume_commands(Mock(), job, [SimpleNamespace(name="cache")])[-1] == ( + "mount --bind /mnt /data" + ) + + +def _volume_commands(user=None, path="/checkpoints"): + job = SimpleNamespace( + job_spec=SimpleNamespace(user=user, volumes=[VolumeMountPoint(name="cache", path=path)]) + ) + return _get_volume_commands(Mock(), job, [SimpleNamespace(name="cache")]) diff --git a/src/tests/_internal/server/routers/test_backends.py b/src/tests/_internal/server/routers/test_backends.py index 824674c364..665f1fff67 100644 --- a/src/tests/_internal/server/routers/test_backends.py +++ b/src/tests/_internal/server/routers/test_backends.py @@ -95,6 +95,7 @@ async def test_returns_backend_types(self, client: AsyncClient): "jarvislabs", "kubernetes", "lambda", + "lium", "nebius", "oci", "runpod",