dstack 0.0.9__py3-none-any.whl → 0.20.7__py3-none-any.whl

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

@@ -0,0 +1,117 @@
+from abc import ABC, abstractmethod
+from typing import Any, ClassVar, Generic, List, Optional, TypeVar
+from uuid import UUID
+
+from dstack._internal.core.backends.base.backend import Backend
+from dstack._internal.core.backends.models import (
+    AnyBackendConfigWithCreds,
+    AnyBackendConfigWithoutCreds,
+)
+from dstack._internal.core.errors import BackendInvalidCredentialsError
+from dstack._internal.core.models.backends.base import BackendType
+from dstack._internal.core.models.common import CoreModel
+
+# Most clouds allow ~ 40-60 tags/labels per resource.
+# We'll introduce our own base limit that can be customized per backend if required.
+TAGS_MAX_NUM = 25
+
+BackendConfigWithoutCredsT = TypeVar(
+    "BackendConfigWithoutCredsT", bound=AnyBackendConfigWithoutCreds
+)
+BackendConfigWithCredsT = TypeVar("BackendConfigWithCredsT", bound=AnyBackendConfigWithCreds)
+
+
+class BackendRecord(CoreModel):
+    """
+    This model includes backend parameters to store in the DB.
+    """
+
+    # `config` stores text-encoded non-sensitive backend config parameters (e.g. json)
+    config: str
+    # `auth` stores text-encoded sensitive backend config parameters (e.g. json).
+    # Configurator should not encrypt/decrypt it. This is done by the caller.
+    auth: str
+
+
+class StoredBackendRecord(BackendRecord):
+    """
+    This model includes backend parameters stored in the DB.
+    """
+
+    # IDs of DB models.
+    # Can be used by externally-registered Configurator to work with the DB directly.
+    project_id: UUID
+    backend_id: UUID
+
+
+class Configurator(ABC, Generic[BackendConfigWithoutCredsT, BackendConfigWithCredsT]):
+    """
+    `Configurator` is responsible for configuring backends
+    and initializing `Backend` instances from backend configs.
+    Every backend must implement `Configurator` and register it
+    in `dstack._internal.core.backends.configurators`.
+    """
+
+    TYPE: ClassVar[BackendType]
+    # `BACKEND_CLASS` is used to introspect backend features without initializing it.
+    BACKEND_CLASS: ClassVar[type[Backend]]
+
+    @abstractmethod
+    def validate_config(self, config: BackendConfigWithCredsT, default_creds_enabled: bool):
+        """
+        Validates backend config including backend creds and other parameters.
+        Raises `ServerClientError` or its subclass if config is invalid.
+        If the backend supports default creds and not `default_creds_enabled`, should raise an error.
+        """
+        pass
+
+    @abstractmethod
+    def create_backend(self, project_name: str, config: BackendConfigWithCredsT) -> BackendRecord:
+        """
+        Sets up backend given backend config and returns
+        text-encoded config and creds to be stored in the DB.
+        It may perform backend initialization, create
+        cloud resources such as networks and managed identities, and
+        save additional configuration parameters.
+        It does not need to duplicate validation done by `validate_config()`
+        since the caller guarantees to call `validate_config()` first.
+        It may perform additional validation not possible in `validate_config()`
+        and raise `ServerClientError` or its subclass if config is invalid.
+        """
+        pass
+
+    @abstractmethod
+    def get_backend_config_with_creds(
+        self, record: StoredBackendRecord
+    ) -> BackendConfigWithCredsT:
+        """
+        Constructs `BackendConfig` with credentials included.
+        Used internally and when project admins need to see backend's creds.
+        """
+        pass
+
+    @abstractmethod
+    def get_backend_config_without_creds(
+        self, record: StoredBackendRecord
+    ) -> BackendConfigWithoutCredsT:
+        """
+        Constructs `BackendConfig` without sensitive information.
+        Used for API responses where creds should not be exposed.
+        """
+        pass
+
+    @abstractmethod
+    def get_backend(self, record: StoredBackendRecord) -> Backend:
+        """
+        Returns `Backend` instance from config and creds stored in `record`.
+        """
+        pass
+
+
+def raise_invalid_credentials_error(
+    fields: Optional[List[List[str]]] = None, details: Optional[Any] = None
+):
+    msg = BackendInvalidCredentialsError.msg
+    if details:
+        msg += f": {details}"
+    raise BackendInvalidCredentialsError(fields=fields, msg=msg)

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

@@ -0,0 +1,24 @@
+from pathlib import Path
+from typing import List
+
+from dstack._internal.core.models.common import CoreModel
+from dstack._internal.core.models.runs import Job
+from dstack._internal.core.models.volumes import Volume
+
+
+class JobConfiguration(CoreModel):
+    job: Job
+    volumes: List[Volume]
+
+
+def fill_data(values: dict, filename_field: str = "filename", data_field: str = "data") -> dict:
+    if values.get(data_field) is not None:
+        return values
+    if (filename := values.get(filename_field)) is None:
+        raise ValueError(f"Either `{filename_field}` or `{data_field}` must be specified")
+    try:
+        with open(Path(filename).expanduser()) as f:
+            values[data_field] = f.read()
+    except OSError:
+        raise ValueError(f"No such file {filename}")
+    return values

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

@@ -0,0 +1,232 @@
+from collections.abc import Iterable, Iterator
+from dataclasses import asdict
+from typing import Callable, List, Optional, TypeVar
+
+import gpuhunt
+from pydantic import parse_obj_as
+
+from dstack._internal.core.models.backends.base import BackendType
+from dstack._internal.core.models.instances import (
+    Disk,
+    Gpu,
+    InstanceOffer,
+    InstanceOfferWithAvailability,
+    InstanceType,
+    Resources,
+)
+from dstack._internal.core.models.resources import DEFAULT_DISK, CPUSpec, Memory, Range
+from dstack._internal.core.models.runs import Requirements
+from dstack._internal.utils.common import get_or_error
+
+# Offers not supported by all dstack versions are hidden behind one or more flags.
+# This list enables the flags that are currently supported.
+SUPPORTED_GPUHUNT_FLAGS = [
+    "oci-spot",
+    "lambda-arm",
+    "gcp-a4",
+    "gcp-g4",
+    "gcp-dws-calendar-mode",
+    "runpod-cluster",
+]
+
+
+def get_catalog_offers(
+    backend: BackendType,
+    locations: Optional[List[str]] = None,
+    requirements: Optional[Requirements] = None,
+    configurable_disk_size: Range[Memory] = Range[Memory](min=Memory.parse("1GB"), max=None),
+    extra_filter: Optional[Callable[[InstanceOffer], bool]] = None,
+    catalog: Optional[gpuhunt.Catalog] = None,
+) -> List[InstanceOffer]:
+    provider = backend.value
+    if backend == BackendType.DATACRUNCH:
+        provider = BackendType.VERDA.value  # Backward compatibility
+    if backend == BackendType.LAMBDA:
+        provider = "lambdalabs"
+    if backend == BackendType.AMDDEVCLOUD:
+        provider = "digitalocean"
+    q = requirements_to_query_filter(requirements)
+    q.provider = [provider]
+    offers = []
+
+    catalog = catalog if catalog is not None else gpuhunt.default_catalog()
+    for item in catalog.query(**asdict(q)):
+        if locations is not None and item.location not in locations:
+            continue
+        offer = catalog_item_to_offer(backend, item, requirements, configurable_disk_size)
+        if offer is None:
+            continue
+        if extra_filter is not None and not extra_filter(offer):
+            continue
+        offers.append(offer)
+    return offers
+
+
+def catalog_item_to_offer(
+    backend: BackendType,
+    item: gpuhunt.CatalogItem,
+    requirements: Optional[Requirements],
+    configurable_disk_size: Range[Memory],
+) -> Optional[InstanceOffer]:
+    gpus = []
+    if item.gpu_count > 0:
+        gpu = Gpu(
+            vendor=item.gpu_vendor, name=item.gpu_name, memory_mib=round(item.gpu_memory * 1024)
+        )
+        gpus = [gpu] * item.gpu_count
+    disk_size_mib = choose_disk_size_mib(
+        catalog_item_disk_size_gib=item.disk_size,
+        requirements_disk_size=requirements.resources.disk.size
+        if requirements and requirements.resources.disk
+        else None,
+        configurable_disk_size=configurable_disk_size,
+    )
+    if disk_size_mib is None:
+        return None
+    resources = Resources(
+        cpu_arch=item.cpu_arch,
+        cpus=item.cpu,
+        memory_mib=round(item.memory * 1024),
+        gpus=gpus,
+        spot=item.spot,
+        disk=Disk(size_mib=disk_size_mib),
+    )
+    return InstanceOffer(
+        backend=backend,
+        instance=InstanceType(
+            name=item.instance_name,
+            resources=resources,
+        ),
+        region=item.location,
+        price=item.price,
+        backend_data=item.provider_data,
+    )
+
+
+def offer_to_catalog_item(offer: InstanceOffer) -> gpuhunt.CatalogItem:
+    cpu_arch = offer.instance.resources.cpu_arch
+    if cpu_arch is None:
+        cpu_arch = gpuhunt.CPUArchitecture.X86
+    gpu_count = len(offer.instance.resources.gpus)
+    gpu_vendor = None
+    gpu_name = None
+    gpu_memory = None
+    if gpu_count > 0:
+        gpu = offer.instance.resources.gpus[0]
+        gpu_vendor = gpu.vendor
+        gpu_name = gpu.name
+        gpu_memory = gpu.memory_mib / 1024
+    return gpuhunt.CatalogItem(
+        provider=offer.backend.value,
+        instance_name=offer.instance.name,
+        location=offer.region,
+        price=offer.price,
+        cpu_arch=cpu_arch,
+        cpu=offer.instance.resources.cpus,
+        memory=offer.instance.resources.memory_mib / 1024,
+        gpu_count=gpu_count,
+        gpu_vendor=gpu_vendor,
+        gpu_name=gpu_name,
+        gpu_memory=gpu_memory,
+        spot=offer.instance.resources.spot,
+        disk_size=offer.instance.resources.disk.size_mib / 1024,
+    )
+
+
+def requirements_to_query_filter(req: Optional[Requirements]) -> gpuhunt.QueryFilter:
+    q = gpuhunt.QueryFilter(allowed_flags=SUPPORTED_GPUHUNT_FLAGS)
+    if req is None:
+        return q
+
+    q.max_price = req.max_price
+    q.spot = req.spot
+
+    res = req.resources
+    if res.cpu:
+        # TODO: Remove in 0.20. Use res.cpu directly
+        cpu = parse_obj_as(CPUSpec, res.cpu)
+        q.cpu_arch = cpu.arch
+        q.min_cpu = cpu.count.min
+        q.max_cpu = cpu.count.max
+    if res.memory:
+        q.min_memory = res.memory.min
+        q.max_memory = res.memory.max
+    if res.disk:
+        q.min_disk_size = res.disk.size.min
+        q.max_disk_size = res.disk.size.max
+
+    if res.gpu:
+        q.gpu_vendor = res.gpu.vendor
+        q.gpu_name = res.gpu.name
+        if res.gpu.memory:
+            q.min_gpu_memory = res.gpu.memory.min
+            q.max_gpu_memory = res.gpu.memory.max
+        if res.gpu.count:
+            q.min_gpu_count = res.gpu.count.min
+            q.max_gpu_count = res.gpu.count.max
+        if res.gpu.total_memory:
+            q.min_total_gpu_memory = res.gpu.total_memory.min
+            q.max_total_gpu_memory = res.gpu.total_memory.max
+        if res.gpu.compute_capability:
+            q.min_compute_capability = res.gpu.compute_capability
+
+    return q
+
+
+InstanceOfferT = TypeVar("InstanceOfferT", InstanceOffer, InstanceOfferWithAvailability)
+
+
+def filter_offers_by_requirements(
+    offers: Iterable[InstanceOfferT],
+    requirements: Optional[Requirements],
+) -> Iterator[InstanceOfferT]:
+    query_filter = requirements_to_query_filter(requirements)
+    for offer in offers:
+        catalog_item = offer_to_catalog_item(offer)
+        if gpuhunt.matches(catalog_item, q=query_filter):
+            yield offer
+
+
+def choose_disk_size_mib(
+    catalog_item_disk_size_gib: Optional[float],
+    requirements_disk_size: Optional[Range[Memory]],
+    configurable_disk_size: Range[Memory],
+) -> Optional[int]:
+    if catalog_item_disk_size_gib:
+        disk_size_gib = catalog_item_disk_size_gib
+    else:
+        disk_size_range = requirements_disk_size or DEFAULT_DISK.size
+        disk_size_range = disk_size_range.intersect(configurable_disk_size)
+        if disk_size_range is None:
+            return None
+        disk_size_gib = disk_size_range.min
+
+    return round(disk_size_gib * 1024)
+
+
+OfferModifier = Callable[[InstanceOfferWithAvailability], Optional[InstanceOfferWithAvailability]]
+
+
+def get_offers_disk_modifier(
+    configurable_disk_size: Range[Memory], requirements: Requirements
+) -> OfferModifier:
+    """
+    Returns a func that modifies offers disk by setting min value that satisfies both
+    `configurable_disk_size` and `requirements`.
+    """
+
+    def modifier(offer: InstanceOfferWithAvailability) -> Optional[InstanceOfferWithAvailability]:
+        requirements_disk_range = DEFAULT_DISK.size
+        if requirements.resources.disk is not None:
+            requirements_disk_range = requirements.resources.disk.size
+        disk_size_range = requirements_disk_range.intersect(configurable_disk_size)
+        if disk_size_range is None:
+            return None
+        offer_copy = offer.copy(deep=True)
+        offer_copy.instance.resources.disk = Disk(
+            size_mib=get_or_error(disk_size_range.min) * 1024
+        )
+        offer_copy.instance.resources.update_description()
+        return offer_copy
+
+    return modifier

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

@@ -0,0 +1,220 @@
+import os
+import re
+from typing import Any, Dict, List, Mapping, Optional, Union
+
+import requests
+from packaging import version
+from requests import Response
+
+from dstack._internal.core.errors import BackendError, BackendInvalidCredentialsError
+from dstack._internal.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+CLOUDRIFT_SERVER_ADDRESS = "https://api.cloudrift.ai"
+CLOUDRIFT_API_VERSION = "2025-05-29"
+
+
+class RiftClient:
+    def __init__(self, api_key: Optional[str] = None):
+        self.public_api_root = os.path.join(CLOUDRIFT_SERVER_ADDRESS, "api/v1")
+        self.api_key = api_key
+
+    def validate_api_key(self) -> bool:
+        """
+        Validates the API key by making a request to the server.
+        Returns True if the API key is valid, False otherwise.
+        """
+        try:
+            response = self._make_request("auth/me")
+            if isinstance(response, dict):
+                return "email" in response
+            return False
+        except BackendInvalidCredentialsError:
+            return False
+        except Exception as e:
+            logger.error(f"Error validating API key: {e}")
+            return False
+
+    def get_instance_types(self) -> List[Dict]:
+        request_data = {"selector": {"ByServiceAndLocation": {"services": ["vm"]}}}
+        response_data = self._make_request("instance-types/list", request_data)
+        if isinstance(response_data, dict):
+            return response_data.get("instance_types", [])
+        return []
+
+    def list_recipes(self) -> List[Dict]:
+        request_data = {}
+        response_data = self._make_request("recipes/list", request_data)
+        if isinstance(response_data, dict):
+            return response_data.get("groups", [])
+        return []
+
+    def get_vm_recipies(self) -> List[Dict]:
+        """
+        Retrieves a list of VM recipes from the CloudRift API.
+        Returns a list of dictionaries containing recipe information.
+        """
+        recipe_group = self.list_recipes()
+        vm_recipes = []
+        for group in recipe_group:
+            tags = group.get("tags", [])
+            has_vm = "vm" in map(str.lower, tags)
+            if group.get("name", "").lower() != "linux" or not has_vm:
+                continue
+
+            recipes = group.get("recipes", [])
+            for recipe in recipes:
+                details = recipe.get("details", {})
+                if details.get("VirtualMachine", False):
+                    vm_recipes.append(recipe)
+
+        return vm_recipes
+
+    def get_vm_image_url(self) -> Optional[str]:
+        recipes = self.get_vm_recipies()
+        ubuntu_images = []
+        for recipe in recipes:
+            has_nvidia_driver = "nvidia-driver" in recipe.get("tags", [])
+            if not has_nvidia_driver:
+                continue
+
+            recipe_name = recipe.get("name", "")
+            if "Ubuntu" not in recipe_name:
+                continue
+
+            url = recipe["details"].get("VirtualMachine", {}).get("image_url", None)
+            version_match = re.search(r".* (\d+\.\d+)", recipe_name)
+            if url and version_match and version_match.group(1):
+                ubuntu_version = version.parse(version_match.group(1))
+                ubuntu_images.append((ubuntu_version, url))
+
+        ubuntu_images.sort(key=lambda x: x[0])  # Sort by version
+        if ubuntu_images:
+            return ubuntu_images[-1][1]
+
+        return None
+
+    def deploy_instance(
+        self, instance_type: str, region: str, ssh_keys: List[str], cmd: str
+    ) -> List[str]:
+        image_url = self.get_vm_image_url()
+        if not image_url:
+            raise BackendError("No suitable VM image found.")
+
+        request_data = {
+            "config": {
+                "VirtualMachine": {
+                    "cloudinit_commands": cmd,
+                    "image_url": image_url,
+                    "ssh_key": {"PublicKeys": ssh_keys},
+                }
+            },
+            "selector": {
+                "ByInstanceTypeAndLocation": {
+                    "datacenters": [region],
+                    "instance_type": instance_type,
+                }
+            },
+            "with_public_ip": True,
+        }
+        logger.debug("Deploying instance with request data: %s", request_data)
+
+        response_data = self._make_request("instances/rent", request_data)
+        if isinstance(response_data, dict):
+            return response_data.get("instance_ids", [])
+        return []
+
+    def list_instances(self, instance_ids: Optional[List[str]] = None) -> List[Dict]:
+        request_data = {
+            "selector": {
+                "ByStatus": ["Initializing", "Active", "Deactivating"],
+            }
+        }
+        logger.debug("Listing instances with request data: %s", request_data)
+        response_data = self._make_request("instances/list", request_data)
+        if isinstance(response_data, dict):
+            return response_data.get("instances", [])
+
+        return []
+
+    def get_instance_by_id(self, instance_id: str) -> Optional[Dict]:
+        request_data = {"selector": {"ById": [instance_id]}}
+        logger.debug("Getting instance with request data: %s", request_data)
+        response_data = self._make_request("instances/list", request_data)
+        if isinstance(response_data, dict):
+            instances = response_data.get("instances", [])
+            if isinstance(instances, list) and len(instances) > 0:
+                return instances[0]
+
+        return None
+
+    def terminate_instance(self, instance_id: str) -> bool:
+        request_data = {"selector": {"ById": [instance_id]}}
+        logger.debug("Terminating instance with request data: %s", request_data)
+        response_data = self._make_request("instances/terminate", request_data)
+        if isinstance(response_data, dict):
+            logger.debug("Terminating instance with response: %s", response_data)
+            info = response_data.get("terminated", [])
+            is_terminated = len(info) > 0
+            if not is_terminated:
+                # check if the instance is already terminated
+                instance_info = self.get_instance_by_id(instance_id)
+                is_terminated = instance_info is None or instance_info.get("status") == "Inactive"
+                logger.debug(
+                    "Instance %s is already terminated: %s response: %s",
+                    instance_id,
+                    is_terminated,
+                    instance_info,
+                )
+            return is_terminated
+
+        return False
+
+    def _make_request(
+        self,
+        endpoint: str,
+        data: Optional[Mapping[str, Any]] = None,
+        method: str = "POST",
+        **kwargs,
+    ) -> Union[Mapping[str, Any], str, Response]:
+        headers = {}
+        if self.api_key is not None:
+            headers["X-API-Key"] = self.api_key
+
+        version = CLOUDRIFT_API_VERSION
+        full_url = f"{self.public_api_root}/{endpoint}"
+
+        try:
+            response = requests.request(
+                method,
+                full_url,
+                headers=headers,
+                json={"version": version, "data": data},
+                timeout=15,
+                **kwargs,
+            )
+
+            if not response.ok:
+                response.raise_for_status()
+            try:
+                response_json = response.json()
+                if isinstance(response_json, str):
+                    return response_json
+                if version is not None and version < response_json["version"]:
+                    logger.warning(
+                        "The API version %s is lower than the server version %s. ",
+                        version,
+                        response_json["version"],
+                    )
+                return response_json["data"]
+            except requests.exceptions.JSONDecodeError:
+                return response
+        except requests.HTTPError as e:
+            if e.response is not None and e.response.status_code in (
+                requests.codes.forbidden,
+                requests.codes.unauthorized,
+            ):
+                raise BackendInvalidCredentialsError(e.response.text)
+            raise

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

@@ -0,0 +1,16 @@
+from dstack._internal.core.backends.base.backend import Backend
+from dstack._internal.core.backends.cloudrift.compute import CloudRiftCompute
+from dstack._internal.core.backends.cloudrift.models import CloudRiftConfig
+from dstack._internal.core.models.backends.base import BackendType
+
+
+class CloudRiftBackend(Backend):
+    TYPE = BackendType.CLOUDRIFT
+    COMPUTE_CLASS = CloudRiftCompute
+
+    def __init__(self, config: CloudRiftConfig):
+        self.config = config
+        self._compute = CloudRiftCompute(self.config)
+
+    def compute(self) -> CloudRiftCompute:
+        return self._compute