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

dstack/_internal/core/backends/template/configurator.py.jinja

@@ -0,0 +1,69 @@
+import json
+
+from dstack._internal.core.backends.base.configurator import (
+    BackendRecord,
+    Configurator,
+    raise_invalid_credentials_error,
+)
+from dstack._internal.core.backends.{{ backend_name|lower }}.backend import {{ backend_name }}Backend
+from dstack._internal.core.backends.{{ backend_name|lower }}.models import (
+    Any{{ backend_name }}Creds,
+    {{ backend_name }}BackendConfig,
+    {{ backend_name }}BackendConfigWithCreds,
+    {{ backend_name }}Config,
+    {{ backend_name }}Creds,
+    {{ backend_name }}StoredConfig,
+)
+from dstack._internal.core.models.backends.base import (
+    BackendType,
+)
+
+
+class {{ backend_name }}Configurator(
+    Configurator[
+        {{ backend_name }}BackendConfig,
+        {{ backend_name }}BackendConfigWithCreds,
+    ]
+):
+    TYPE = BackendType.{{ backend_name|upper }}
+    BACKEND_CLASS = {{ backend_name }}Backend
+
+    def validate_config(
+        self, config: {{ backend_name }}BackendConfigWithCreds, default_creds_enabled: bool
+    ):
+        self._validate_creds(config.creds)
+        # TODO: If possible, validate config.regions and any other config parameters
+
+    def create_backend(
+        self, project_name: str, config: {{ backend_name }}BackendConfigWithCreds
+    ) -> BackendRecord:
+        return BackendRecord(
+            config={{ backend_name }}StoredConfig(
+                **{{ backend_name }}BackendConfig.__response__.parse_obj(config).dict()
+            ).json(),
+            auth={{ backend_name }}Creds.parse_obj(config.creds).json(),
+        )
+
+    def get_backend_config_with_creds(self, record: BackendRecord) -> {{ backend_name }}BackendConfigWithCreds:
+        config = self._get_config(record)
+        return {{ backend_name }}BackendConfigWithCreds.__response__.parse_obj(config)
+
+    def get_backend_config_without_creds(self, record: BackendRecord) -> {{ backend_name }}BackendConfig:
+        config = self._get_config(record)
+        return {{ backend_name }}BackendConfig.__response__.parse_obj(config)
+
+    def get_backend(self, record: BackendRecord) -> {{ backend_name }}Backend:
+        config = self._get_config(record)
+        return {{ backend_name }}Backend(config=config)
+
+    def _get_config(self, record: BackendRecord) -> {{ backend_name }}Config:
+        return {{ backend_name }}Config.__response__(
+            **json.loads(record.config),
+            creds={{ backend_name }}Creds.parse_raw(record.auth),
+        )
+
+    def _validate_creds(self, creds: Any{{ backend_name }}Creds):
+        # TODO: Implement API key or other creds validation
+        # if valid:
+        #     return
+        raise_invalid_credentials_error(fields=[["creds", "api_key"]])

dstack/_internal/core/backends/template/models.py.jinja

@@ -0,0 +1,62 @@
+from typing import Annotated, List, Literal, Optional, Union
+
+from pydantic import Field
+
+from dstack._internal.core.models.common import CoreModel
+
+
+# The template uses "api_key" creds as the most popular creds type.
+# TODO: Adjust it or add additional creds models if necessary.
+class {{ backend_name }}APIKeyCreds(CoreModel):
+    type: Annotated[Literal["api_key"], Field(description="The type of credentials")] = "api_key"
+    api_key: Annotated[str, Field(description="The API key")]
+
+
+Any{{ backend_name }}Creds = {{ backend_name }}APIKeyCreds
+{{ backend_name }}Creds = Any{{ backend_name }}Creds
+
+
+class {{ backend_name }}BackendConfig(CoreModel):
+    """
+    The backend config used in the API, server/config.yml, `{{ backend_name }}Configurator`.
+    It also serves as a base class for other backend config models.
+    Should not include creds.
+    """
+
+    type: Annotated[
+        Literal["{{ backend_name|lower }}"],
+        Field(description="The type of backend"),
+    ] = "{{ backend_name|lower }}"
+    regions: Annotated[
+        Optional[List[str]],
+        Field(description="The list of {{ backend_name }} regions. Omit to use all regions"),
+    ] = None
+    # TODO: Add additional backend parameters if necessary
+
+
+class {{ backend_name }}BackendConfigWithCreds({{ backend_name }}BackendConfig):
+    """
+    Same as `{{ backend_name }}BackendConfig` but also includes creds.
+    """
+
+    creds: Annotated[Any{{ backend_name }}Creds, Field(description="The credentials")]
+
+
+Any{{ backend_name }}BackendConfig = Union[{{ backend_name }}BackendConfig, {{ backend_name }}BackendConfigWithCreds]
+
+
+class {{ backend_name }}StoredConfig({{ backend_name }}BackendConfig):
+    """
+    The backend config used for config parameters in the DB.
+    Can extend `{{ backend_name }}BackendConfig` with additional parameters.
+    """
+
+    pass
+
+
+class {{ backend_name }}Config({{ backend_name }}StoredConfig):
+    """
+    The backend config used by `{{ backend_name }}Backend` and `{{ backend_name }}Compute`.
+    """
+
+    creds: Any{{ backend_name }}Creds

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

@@ -0,0 +1,40 @@
+from typing import Annotated, List, Literal, Optional, Union
+
+from pydantic import Field
+
+from dstack._internal.core.models.common import CoreModel
+
+# TODO: TensorDock is deprecated and will be removed in the future
+
+
+class TensorDockAPIKeyCreds(CoreModel):
+    type: Annotated[Literal["api_key"], Field(description="The type of credentials")] = "api_key"
+    api_key: Annotated[str, Field(description="The API key")]
+    api_token: Annotated[str, Field(description="The API token")]
+
+
+AnyTensorDockCreds = TensorDockAPIKeyCreds
+TensorDockCreds = AnyTensorDockCreds
+
+
+class TensorDockBackendConfig(CoreModel):
+    type: Annotated[Literal["tensordock"], Field(description="The type of backend")] = "tensordock"
+    regions: Annotated[
+        Optional[List[str]],
+        Field(description="The list of TensorDock regions. Omit to use all regions"),
+    ] = None
+
+
+class TensorDockBackendConfigWithCreds(TensorDockBackendConfig):
+    creds: Annotated[AnyTensorDockCreds, Field(description="The credentials")]
+
+
+AnyTensorDockBackendConfig = Union[TensorDockBackendConfig, TensorDockBackendConfigWithCreds]
+
+
+class TensorDockStoredConfig(TensorDockBackendConfig):
+    pass
+
+
+class TensorDockConfig(TensorDockStoredConfig):
+    creds: AnyTensorDockCreds

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

@@ -0,0 +1,143 @@
+import threading
+import time
+from typing import List, Optional, Union
+
+import requests
+from requests.adapters import HTTPAdapter, Retry
+
+import dstack._internal.server.services.docker as docker
+from dstack._internal.core.consts import DSTACK_RUNNER_SSH_PORT
+from dstack._internal.core.errors import NoCapacityError
+from dstack._internal.core.models.common import RegistryAuth
+
+
+class VastAIAPIClient:
+    def __init__(self, api_key: str):
+        self.api_url = "https://console.vast.ai/api/v0".rstrip("/")
+        self.api_key = api_key
+        self.s = requests.Session()  # TODO: set adequate timeout everywhere the session is used
+        retries = Retry(
+            total=5,
+            backoff_factor=1,
+            status_forcelist=[429, 504],
+        )
+        self.s.mount(prefix=self._url("/instances/"), adapter=HTTPAdapter(max_retries=retries))
+        self.lock = threading.Lock()
+        self.instances_cache_ts: float = 0
+        self.instances_cache: List[dict] = []
+
+    def get_bundle(self, bundle_id: Union[str, int]) -> Optional[dict]:
+        resp = self.s.post(self._url("/bundles/"), json={"id": {"eq": bundle_id}})
+        resp.raise_for_status()
+        data = resp.json()
+        offers = data["offers"]
+        return offers[0] if offers else None
+
+    def create_instance(
+        self,
+        instance_name: str,
+        bundle_id: Union[str, int],
+        image_name: str,
+        onstart: str,
+        disk_size: int,
+        registry_auth: Optional[RegistryAuth] = None,
+    ) -> dict:
+        """
+        Args:
+            instance_name: instance label
+            bundle_id: desired host
+            image_name: docker image name
+            onstart: commands to run on start
+            registry_auth: registry auth credentials for private images
+
+        Raises:
+            NoCapacityError: if instance cannot be created
+
+        Returns:
+            create instance response
+        """
+        image_login = None
+        if registry_auth:
+            registry = docker.parse_image_name(image_name).registry or "docker.io"
+            image_login = f"-u {registry_auth.username} -p {registry_auth.password} {registry}"
+        payload = {
+            "client_id": "me",
+            "image": image_name,
+            "disk": disk_size,
+            "label": instance_name,
+            "env": {
+                f"-p {DSTACK_RUNNER_SSH_PORT}:{DSTACK_RUNNER_SSH_PORT}": "1",
+            },
+            "user": "root",
+            "onstart": "/bin/sh",
+            "args": ["-c", onstart],
+            "runtype": "args",
+            "image_login": image_login,
+            "python_utf8": False,
+            "lang_utf8": False,
+            "use_jupyter_lab": False,
+            "jupyter_dir": None,
+            "create_from": None,
+            "force": False,
+        }
+        resp = self.s.put(self._url(f"/asks/{bundle_id}/"), json=payload)
+        if resp.status_code != 200 or not (data := resp.json())["success"]:
+            raise NoCapacityError(resp.text)
+        self._invalidate_cache()
+        return data
+
+    def destroy_instance(self, instance_id: Union[str, int]) -> bool:
+        """
+        Args:
+            instance_id: instance to destroy
+
+        Returns:
+            True if instance was destroyed successfully
+        """
+        resp = self.s.delete(self._url(f"/instances/{instance_id}/"))
+        if resp.status_code != 200 or not resp.json()["success"]:
+            return False
+        self._invalidate_cache()
+        return True
+
+    def get_instances(self, cache_ttl: float = 3.0) -> List[dict]:
+        with self.lock:
+            if time.time() - self.instances_cache_ts > cache_ttl:
+                resp = self.s.get(self._url("/instances/"))
+                resp.raise_for_status()
+                data = resp.json()
+                self.instances_cache_ts = time.time()
+                self.instances_cache = data["instances"]
+            return self.instances_cache
+
+    def get_instance(self, instance_id: Union[str, int]) -> Optional[dict]:
+        instances = self.get_instances()
+        for instance in instances:
+            if instance["id"] == int(instance_id):
+                return instance
+        return None
+
+    def request_logs(self, instance_id: Union[str, int]) -> dict:
+        resp = self.s.put(
+            self._url(f"/instances/request_logs/{instance_id}/"), json={"tail": "1000"}
+        )
+        resp.raise_for_status()
+        data = resp.json()
+        if not data["success"]:
+            raise requests.HTTPError(data)
+        return data
+
+    def auth_test(self) -> bool:
+        try:
+            self.get_instances()
+            return True
+        except requests.HTTPError:
+            return False
+
+    def _url(self, path):
+        return f"{self.api_url}/{path.lstrip('/')}?api_key={self.api_key}"
+
+    def _invalidate_cache(self):
+        with self.lock:
+            self.instances_cache_ts = 0
+            self.instances_cache = []

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

@@ -0,0 +1,16 @@
+from dstack._internal.core.backends.base.backend import Backend
+from dstack._internal.core.backends.vastai.compute import VastAICompute
+from dstack._internal.core.backends.vastai.models import VastAIConfig
+from dstack._internal.core.models.backends.base import BackendType
+
+
+class VastAIBackend(Backend):
+    TYPE = BackendType.VASTAI
+    COMPUTE_CLASS = VastAICompute
+
+    def __init__(self, config: VastAIConfig):
+        self.config = config
+        self._compute = VastAICompute(self.config)
+
+    def compute(self) -> VastAICompute:
+        return self._compute

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

@@ -0,0 +1,141 @@
+from typing import List, Optional
+
+import gpuhunt
+from gpuhunt.providers.vastai import VastAIProvider
+
+from dstack._internal.core.backends.base.backend import Compute
+from dstack._internal.core.backends.base.compute import (
+    ComputeWithFilteredOffersCached,
+    generate_unique_instance_name_for_job,
+    get_docker_commands,
+)
+from dstack._internal.core.backends.base.offers import get_catalog_offers
+from dstack._internal.core.backends.vastai.api_client import VastAIAPIClient
+from dstack._internal.core.backends.vastai.models import VastAIConfig
+from dstack._internal.core.consts import DSTACK_RUNNER_SSH_PORT
+from dstack._internal.core.errors import ProvisioningError
+from dstack._internal.core.models.backends.base import BackendType
+from dstack._internal.core.models.instances import (
+    InstanceAvailability,
+    InstanceOfferWithAvailability,
+    InstanceRuntime,
+)
+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 Volume
+from dstack._internal.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+# Undocumented but names of len 60 work
+MAX_INSTANCE_NAME_LEN = 60
+
+
+class VastAICompute(
+    ComputeWithFilteredOffersCached,
+    Compute,
+):
+    def __init__(self, config: VastAIConfig):
+        super().__init__()
+        self.config = config
+        self.api_client = VastAIAPIClient(config.creds.api_key)
+        self.catalog = gpuhunt.Catalog(balance_resources=False, auto_reload=False)
+        self.catalog.add_provider(
+            VastAIProvider(
+                extra_filters={
+                    "direct_port_count": {"gte": 1},
+                    "reliability2": {"gte": 0.9},
+                    "inet_down": {"gt": 128},
+                    "verified": {"eq": True},
+                    "cuda_max_good": {"gte": 12.8},
+                    "compute_cap": {"gte": 600},
+                }
+            )
+        )
+
+    def get_offers_by_requirements(
+        self, requirements: Requirements
+    ) -> List[InstanceOfferWithAvailability]:
+        offers = get_catalog_offers(
+            backend=BackendType.VASTAI,
+            locations=self.config.regions or None,
+            requirements=requirements,
+            # TODO(egor-s): spots currently not supported
+            extra_filter=lambda offer: not offer.instance.resources.spot,
+            catalog=self.catalog,
+        )
+        offers = [
+            InstanceOfferWithAvailability(
+                **offer.dict(),
+                availability=InstanceAvailability.AVAILABLE,
+                instance_runtime=InstanceRuntime.RUNNER,
+            )
+            for offer in offers
+        ]
+        return 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],
+    ) -> JobProvisioningData:
+        instance_name = generate_unique_instance_name_for_job(
+            run, job, max_length=MAX_INSTANCE_NAME_LEN
+        )
+        assert run.run_spec.ssh_key_pub is not None
+        commands = get_docker_commands(
+            [run.run_spec.ssh_key_pub.strip(), project_ssh_public_key.strip()]
+        )
+        resp = self.api_client.create_instance(
+            instance_name=instance_name,
+            bundle_id=instance_offer.instance.name,
+            image_name=job.job_spec.image_name,
+            onstart=" && ".join(commands),
+            disk_size=round(instance_offer.instance.resources.disk.size_mib / 1024),
+            registry_auth=job.job_spec.registry_auth,
+        )
+        instance_id = resp["new_contract"]
+        return JobProvisioningData(
+            backend=instance_offer.backend,
+            instance_type=instance_offer.instance,
+            instance_id=instance_id,
+            hostname=None,
+            internal_ip=None,
+            region=instance_offer.region,
+            price=instance_offer.price,
+            username="root",
+            ssh_port=None,
+            dockerized=False,
+            ssh_proxy=None,
+            backend_data=None,
+        )
+
+    def terminate_instance(
+        self, instance_id: str, region: str, backend_data: Optional[str] = None
+    ):
+        self.api_client.destroy_instance(instance_id)
+
+    def update_provisioning_data(
+        self,
+        provisioning_data: JobProvisioningData,
+        project_ssh_public_key: str,
+        project_ssh_private_key: str,
+    ):
+        resp = self.api_client.get_instance(provisioning_data.instance_id)
+        if resp is not None:
+            if resp["actual_status"] == "running":
+                provisioning_data.hostname = resp["public_ipaddr"].strip()
+                provisioning_data.ssh_port = int(
+                    resp["ports"][f"{DSTACK_RUNNER_SSH_PORT}/tcp"][0]["HostPort"]
+                )
+            if (
+                resp["actual_status"] == "created"
+                and ": OCI runtime create failed:" in resp["status_msg"]
+            ):
+                raise ProvisioningError(resp["status_msg"])

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

@@ -0,0 +1,69 @@
+import json
+
+from dstack._internal.core.backends.base.configurator import (
+    BackendRecord,
+    Configurator,
+    raise_invalid_credentials_error,
+)
+from dstack._internal.core.backends.vastai import api_client
+from dstack._internal.core.backends.vastai.backend import VastAIBackend
+from dstack._internal.core.backends.vastai.models import (
+    VastAIBackendConfig,
+    VastAIBackendConfigWithCreds,
+    VastAIConfig,
+    VastAICreds,
+    VastAIStoredConfig,
+)
+from dstack._internal.core.models.backends.base import (
+    BackendType,
+)
+
+REGIONS = []
+
+
+class VastAIConfigurator(
+    Configurator[
+        VastAIBackendConfig,
+        VastAIBackendConfigWithCreds,
+    ]
+):
+    TYPE = BackendType.VASTAI
+    BACKEND_CLASS = VastAIBackend
+
+    def validate_config(self, config: VastAIBackendConfigWithCreds, default_creds_enabled: bool):
+        self._validate_vastai_creds(config.creds.api_key)
+
+    def create_backend(
+        self, project_name: str, config: VastAIBackendConfigWithCreds
+    ) -> BackendRecord:
+        if config.regions is None:
+            config.regions = REGIONS
+        return BackendRecord(
+            config=VastAIStoredConfig(
+                **VastAIBackendConfig.__response__.parse_obj(config).dict()
+            ).json(),
+            auth=VastAICreds.parse_obj(config.creds).json(),
+        )
+
+    def get_backend_config_with_creds(self, record: BackendRecord) -> VastAIBackendConfigWithCreds:
+        config = self._get_config(record)
+        return VastAIBackendConfigWithCreds.__response__.parse_obj(config)
+
+    def get_backend_config_without_creds(self, record: BackendRecord) -> VastAIBackendConfig:
+        config = self._get_config(record)
+        return VastAIBackendConfig.__response__.parse_obj(config)
+
+    def get_backend(self, record: BackendRecord) -> VastAIBackend:
+        config = self._get_config(record)
+        return VastAIBackend(config=config)
+
+    def _get_config(self, record: BackendRecord) -> VastAIConfig:
+        return VastAIConfig.__response__(
+            **json.loads(record.config),
+            creds=VastAICreds.parse_raw(record.auth),
+        )
+
+    def _validate_vastai_creds(self, api_key: str):
+        client = api_client.VastAIAPIClient(api_key=api_key)
+        if not client.auth_test():
+            raise_invalid_credentials_error(fields=[["creds", "api_key"]])

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

@@ -0,0 +1,37 @@
+from typing import Annotated, List, Literal, Optional, Union
+
+from pydantic import Field
+
+from dstack._internal.core.models.common import CoreModel
+
+
+class VastAIAPIKeyCreds(CoreModel):
+    type: Annotated[Literal["api_key"], Field(description="The type of credentials")] = "api_key"
+    api_key: Annotated[str, Field(description="The API key")]
+
+
+AnyVastAICreds = VastAIAPIKeyCreds
+VastAICreds = AnyVastAICreds
+
+
+class VastAIBackendConfig(CoreModel):
+    type: Annotated[Literal["vastai"], Field(description="The type of backend")] = "vastai"
+    regions: Annotated[
+        Optional[List[str]],
+        Field(description="The list of VastAI regions. Omit to use all regions"),
+    ] = None
+
+
+class VastAIBackendConfigWithCreds(VastAIBackendConfig):
+    creds: Annotated[AnyVastAICreds, Field(description="The credentials")]
+
+
+AnyVastAIBackendConfig = Union[VastAIBackendConfig, VastAIBackendConfigWithCreds]
+
+
+class VastAIStoredConfig(VastAIBackendConfig):
+    pass
+
+
+class VastAIConfig(VastAIStoredConfig):
+    creds: AnyVastAICreds

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

@@ -0,0 +1,16 @@
+from dstack._internal.core.backends.base.backend import Backend
+from dstack._internal.core.backends.verda.compute import VerdaCompute
+from dstack._internal.core.backends.verda.models import VerdaConfig
+from dstack._internal.core.models.backends.base import BackendType
+
+
+class VerdaBackend(Backend):
+    TYPE = BackendType.VERDA
+    COMPUTE_CLASS = VerdaCompute
+
+    def __init__(self, config: VerdaConfig):
+        self.config = config
+        self._compute = VerdaCompute(self.config, self.TYPE)
+
+    def compute(self) -> VerdaCompute:
+        return self._compute