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

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

@@ -0,0 +1,266 @@
+from collections.abc import Iterable
+from typing import Dict, List, Optional
+
+from verda import VerdaClient
+from verda.exceptions import APIException
+from verda.instances import Instance
+
+from dstack._internal.core.backends.base.backend import Compute
+from dstack._internal.core.backends.base.compute import (
+    ComputeWithAllOffersCached,
+    ComputeWithCreateInstanceSupport,
+    ComputeWithPrivilegedSupport,
+    generate_unique_instance_name,
+    get_shim_commands,
+)
+from dstack._internal.core.backends.base.offers import (
+    OfferModifier,
+    get_catalog_offers,
+    get_offers_disk_modifier,
+)
+from dstack._internal.core.backends.verda.models import VerdaConfig
+from dstack._internal.core.errors import NoCapacityError
+from dstack._internal.core.models.backends.base import BackendType
+from dstack._internal.core.models.instances import (
+    InstanceAvailability,
+    InstanceConfiguration,
+    InstanceOffer,
+    InstanceOfferWithAvailability,
+)
+from dstack._internal.core.models.placement import PlacementGroup
+from dstack._internal.core.models.resources import Memory, Range
+from dstack._internal.core.models.runs import JobProvisioningData, Requirements
+from dstack._internal.utils.logging import get_logger
+from dstack._internal.utils.ssh import get_public_key_fingerprint
+
+logger = get_logger("verda.compute")
+
+MAX_INSTANCE_NAME_LEN = 60
+
+IMAGE_SIZE = Memory.parse("50GB")
+
+CONFIGURABLE_DISK_SIZE = Range[Memory](min=IMAGE_SIZE, max=None)
+
+
+class VerdaCompute(
+    ComputeWithAllOffersCached,
+    ComputeWithCreateInstanceSupport,
+    ComputeWithPrivilegedSupport,
+    Compute,
+):
+    def __init__(self, config: VerdaConfig, backend_type: BackendType):
+        super().__init__()
+        self.config = config
+        self.client = VerdaClient(
+            client_id=self.config.creds.client_id,
+            client_secret=self.config.creds.client_secret,
+        )
+        self.backend_type = backend_type
+
+    def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability]:
+        offers = get_catalog_offers(
+            backend=self.backend_type,
+            locations=self.config.regions,
+        )
+        offers_with_availability = self._get_offers_with_availability(offers)
+        return offers_with_availability
+
+    def get_offers_modifiers(self, requirements: Requirements) -> Iterable[OfferModifier]:
+        return [get_offers_disk_modifier(CONFIGURABLE_DISK_SIZE, requirements)]
+
+    def _get_offers_with_availability(
+        self, offers: List[InstanceOffer]
+    ) -> List[InstanceOfferWithAvailability]:
+        raw_availabilities: List[Dict] = self.client.instances.get_availabilities()
+
+        region_availabilities = {}
+        for location in raw_availabilities:
+            location_code = location["location_code"]
+            availabilities = location["availabilities"]
+            for name in availabilities:
+                key = (name, location_code)
+                region_availabilities[key] = InstanceAvailability.AVAILABLE
+
+        availability_offers = []
+        for offer in offers:
+            key = (offer.instance.name, offer.region)
+            availability = region_availabilities.get(key, InstanceAvailability.NOT_AVAILABLE)
+            availability_offers.append(
+                InstanceOfferWithAvailability(**offer.dict(), availability=availability)
+            )
+
+        return availability_offers
+
+    def create_instance(
+        self,
+        instance_offer: InstanceOfferWithAvailability,
+        instance_config: InstanceConfiguration,
+        placement_group: Optional[PlacementGroup],
+    ) -> JobProvisioningData:
+        instance_name = generate_unique_instance_name(
+            instance_config, max_length=MAX_INSTANCE_NAME_LEN
+        )
+        public_keys = instance_config.get_public_keys()
+        ssh_ids = []
+        for ssh_public_key in public_keys:
+            ssh_ids.append(
+                # verda allows you to use the same name
+                _get_or_create_ssh_key(
+                    client=self.client,
+                    name=f"dstack-{instance_config.instance_name}.key",
+                    public_key=ssh_public_key,
+                )
+            )
+
+        commands = get_shim_commands()
+        startup_script = " ".join([" && ".join(commands)])
+        script_name = f"dstack-{instance_config.instance_name}.sh"
+        startup_script_ids = _get_or_create_startup_scrpit(
+            client=self.client,
+            name=script_name,
+            script=startup_script,
+        )
+
+        disk_size = round(instance_offer.instance.resources.disk.size_mib / 1024)
+        image_id = _get_vm_image_id(instance_offer)
+
+        logger.debug(
+            "Deploying Verda instance",
+            {
+                "instance_type": instance_offer.instance.name,
+                "ssh_key_ids": ssh_ids,
+                "startup_script_id": startup_script_ids,
+                "hostname": instance_name,
+                "description": instance_name,
+                "image": image_id,
+                "disk_size": disk_size,
+                "location": instance_offer.region,
+            },
+        )
+        instance = _deploy_instance(
+            client=self.client,
+            instance_type=instance_offer.instance.name,
+            ssh_key_ids=ssh_ids,
+            startup_script_id=startup_script_ids,
+            hostname=instance_name,
+            description=instance_name,
+            image=image_id,
+            disk_size=disk_size,
+            is_spot=instance_offer.instance.resources.spot,
+            location=instance_offer.region,
+        )
+        return JobProvisioningData(
+            backend=instance_offer.backend,
+            instance_type=instance_offer.instance,
+            instance_id=instance.id,
+            hostname=None,
+            internal_ip=None,
+            region=instance.location,
+            price=instance_offer.price,
+            username="root",
+            ssh_port=22,
+            dockerized=True,
+            ssh_proxy=None,
+            backend_data=None,
+        )
+
+    def terminate_instance(
+        self, instance_id: str, region: str, backend_data: Optional[str] = None
+    ):
+        try:
+            self.client.instances.action(id_list=[instance_id], action="delete")
+        except APIException as e:
+            if e.message in [
+                "Invalid instance id",
+                "Can't discontinue a discontinued instance",
+            ]:
+                logger.debug("Skipping instance %s termination. Instance not found.", instance_id)
+                return
+            raise
+
+    def update_provisioning_data(
+        self,
+        provisioning_data: JobProvisioningData,
+        project_ssh_public_key: str,
+        project_ssh_private_key: str,
+    ):
+        instance = _get_instance_by_id(self.client, provisioning_data.instance_id)
+        if instance is not None and instance.status == "running":
+            provisioning_data.hostname = instance.ip
+
+
+def _get_vm_image_id(instance_offer: InstanceOfferWithAvailability) -> str:
+    # https://api.verda.com/v1/images
+    if len(instance_offer.instance.resources.gpus) > 0 and instance_offer.instance.resources.gpus[
+        0
+    ].name in ["V100", "A6000"]:
+        # Ubuntu 22.04 + CUDA 12.0 + Docker
+        return "2088da25-bb0d-41cc-a191-dccae45d96fd"
+    # Ubuntu 24.04 + CUDA 12.8 Open + Docker
+    return "77777777-4f48-4249-82b3-f199fb9b701b"
+
+
+def _get_or_create_ssh_key(client: VerdaClient, name: str, public_key: str) -> str:
+    fingerprint = get_public_key_fingerprint(public_key)
+    keys = client.ssh_keys.get()
+    found_keys = [key for key in keys if fingerprint == get_public_key_fingerprint(key.public_key)]
+    if found_keys:
+        key = found_keys[0]
+        return key.id
+    key = client.ssh_keys.create(name, public_key)
+    return key.id
+
+
+def _get_or_create_startup_scrpit(client: VerdaClient, name: str, script: str) -> str:
+    scripts = client.startup_scripts.get()
+    found_scripts = [startup_script for startup_script in scripts if script == startup_script]
+    if found_scripts:
+        startup_script = found_scripts[0]
+        return startup_script.id
+
+    startup_script = client.startup_scripts.create(name, script)
+    return startup_script.id
+
+
+def _get_instance_by_id(
+    client: VerdaClient,
+    instance_id: str,
+) -> Optional[Instance]:
+    try:
+        return client.instances.get_by_id(instance_id)
+    except APIException as e:
+        if e.message == "Invalid instance id":
+            return None
+        raise
+
+
+def _deploy_instance(
+    client: VerdaClient,
+    instance_type: str,
+    image: str,
+    ssh_key_ids: List[str],
+    hostname: str,
+    description: str,
+    startup_script_id: str,
+    disk_size: int,
+    is_spot: bool,
+    location: str,
+) -> Instance:
+    try:
+        instance = client.instances.create(
+            instance_type=instance_type,
+            image=image,
+            ssh_key_ids=ssh_key_ids,
+            hostname=hostname,
+            description=description,
+            startup_script_id=startup_script_id,
+            pricing="FIXED_PRICE",
+            is_spot=is_spot,
+            location=location,
+            os_volume={"name": "OS volume", "size": disk_size},
+        )
+    except APIException as e:
+        # FIXME: Catch only no capacity errors
+        raise NoCapacityError(f"Verda API error: {e.message}")
+
+    return instance

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

@@ -0,0 +1,73 @@
+import json
+
+from verda import VerdaClient
+from verda.exceptions import APIException
+
+from dstack._internal.core.backends.base.configurator import (
+    BackendRecord,
+    Configurator,
+    raise_invalid_credentials_error,
+)
+from dstack._internal.core.backends.verda.backend import VerdaBackend
+from dstack._internal.core.backends.verda.models import (
+    VerdaBackendConfig,
+    VerdaBackendConfigWithCreds,
+    VerdaConfig,
+    VerdaCreds,
+    VerdaStoredConfig,
+)
+from dstack._internal.core.models.backends.base import (
+    BackendType,
+)
+
+
+class VerdaConfigurator(
+    Configurator[
+        VerdaBackendConfig,
+        VerdaBackendConfigWithCreds,
+    ]
+):
+    TYPE = BackendType.VERDA
+    BACKEND_CLASS = VerdaBackend
+
+    def validate_config(self, config: VerdaBackendConfigWithCreds, default_creds_enabled: bool):
+        self._validate_creds(config.creds)
+
+    def create_backend(
+        self, project_name: str, config: VerdaBackendConfigWithCreds
+    ) -> BackendRecord:
+        return BackendRecord(
+            config=VerdaStoredConfig(
+                **VerdaBackendConfig.__response__.parse_obj(config).dict()
+            ).json(),
+            auth=VerdaCreds.parse_obj(config.creds).json(),
+        )
+
+    def get_backend_config_with_creds(self, record: BackendRecord) -> VerdaBackendConfigWithCreds:
+        config = self._get_config(record)
+        return VerdaBackendConfigWithCreds.__response__.parse_obj(config)
+
+    def get_backend_config_without_creds(self, record: BackendRecord) -> VerdaBackendConfig:
+        config = self._get_config(record)
+        return VerdaBackendConfig.__response__.parse_obj(config)
+
+    def get_backend(self, record: BackendRecord) -> VerdaBackend:
+        config = self._get_config(record)
+        return VerdaBackend(config=config)
+
+    def _get_config(self, record: BackendRecord) -> VerdaConfig:
+        return VerdaConfig.__response__(
+            **json.loads(record.config),
+            creds=VerdaCreds.parse_raw(record.auth),
+        )
+
+    def _validate_creds(self, creds: VerdaCreds):
+        try:
+            VerdaClient(
+                client_id=creds.client_id,
+                client_secret=creds.client_secret,
+            )
+        except APIException as e:
+            if e.code == "unauthorized_request":
+                raise_invalid_credentials_error(fields=[["creds", "api_key"]])
+            raise

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

@@ -0,0 +1,38 @@
+from typing import Annotated, List, Literal, Optional, Union
+
+from pydantic import Field
+
+from dstack._internal.core.models.common import CoreModel
+
+
+class VerdaAPIKeyCreds(CoreModel):
+    type: Annotated[Literal["api_key"], Field(description="The type of credentials")] = "api_key"
+    client_id: Annotated[str, Field(description="The client ID")]
+    client_secret: Annotated[str, Field(description="The client secret")]
+
+
+AnyVerdaCreds = VerdaAPIKeyCreds
+VerdaCreds = AnyVerdaCreds
+
+
+class VerdaBackendConfig(CoreModel):
+    type: Annotated[Literal["verda", "datacrunch"], Field(description="The type of backend")]
+    regions: Annotated[
+        Optional[List[str]],
+        Field(description="The list of Verda regions. Omit to use all regions"),
+    ] = None
+
+
+class VerdaBackendConfigWithCreds(VerdaBackendConfig):
+    creds: Annotated[AnyVerdaCreds, Field(description="The credentials")]
+
+
+AnyVerdaBackendConfig = Union[VerdaBackendConfig, VerdaBackendConfigWithCreds]
+
+
+class VerdaStoredConfig(VerdaBackendConfig):
+    pass
+
+
+class VerdaConfig(VerdaStoredConfig):
+    creds: AnyVerdaCreds

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

@@ -0,0 +1,116 @@
+import base64
+from typing import Any, Optional
+
+import requests
+from requests import Response
+
+from dstack._internal.core.errors import BackendError, BackendInvalidCredentialsError
+
+API_URL = "https://api.vultr.com/v2"
+
+
+class VultrApiClient:
+    def __init__(self, api_key: str):
+        self.api_key = api_key
+
+    def validate_api_key(self) -> bool:
+        try:
+            self._make_request("GET", "/ssh-keys")
+        except BackendInvalidCredentialsError:
+            return False
+        return True
+
+    def get_instance(self, instance_id: str, plan_type: str) -> dict:
+        if plan_type == "bare-metal":
+            response = self._make_request("GET", f"/bare-metals/{instance_id}")
+            return response.json()["bare_metal"]
+        else:
+            response = self._make_request("GET", f"/instances/{instance_id}")
+            return response.json()["instance"]
+
+    def get_vpc_for_region(self, region: str) -> Optional[dict]:
+        response = self._make_request("GET", "/vpcs?per_page=500")
+        vpcs = response.json().get("vpcs", [])
+        if vpcs:
+            for vpc in vpcs:
+                if vpc["description"] == f"dstack-vpc-{region}":
+                    return vpc
+        return None
+
+    def create_vpc(self, region: str) -> dict:
+        data = {"region": region, "description": f"dstack-vpc-{region}"}
+        response = self._make_request("POST", "/vpcs", data=data)
+        return response.json()["vpc"]
+
+    def launch_instance(self, region: str, plan: str, label: str, user_data: str, vpc_id: str):
+        # For Bare-metals
+        if "vbm" in plan:
+            # "Docker on Ubuntu 22.04" is required for bare-metals.
+            data = {
+                "region": region,
+                "plan": plan,
+                "label": label,
+                "image_id": "docker",
+                "user_data": base64.b64encode(user_data.encode()).decode(),
+                "attach_vpc": [vpc_id],
+            }
+            resp = self._make_request("POST", "/bare-metals", data)
+            return resp.json()["bare_metal"]["id"]
+        # For VMs
+        elif "vcg" in plan:
+            # Ubuntu 22.04 will be installed. For gpu VMs, docker is preinstalled.
+            data = {
+                "region": region,
+                "plan": plan,
+                "label": label,
+                "os_id": 1743,
+                "user_data": base64.b64encode(user_data.encode()).decode(),
+                "attach_vpc": [vpc_id],
+            }
+            resp = self._make_request("POST", "/instances", data)
+            return resp.json()["instance"]["id"]
+        else:
+            data = {
+                "region": region,
+                "plan": plan,
+                "label": label,
+                "image_id": "docker",
+                "user_data": base64.b64encode(user_data.encode()).decode(),
+                "attach_vpc": [vpc_id],
+            }
+            resp = self._make_request("POST", "/instances", data)
+            return resp.json()["instance"]["id"]
+
+    def terminate_instance(self, instance_id: str, plan_type: str):
+        if plan_type == "bare-metal":
+            # Terminate bare-metal instance
+            endpoint = f"/bare-metals/{instance_id}"
+        else:
+            # Terminate virtual machine instance
+            endpoint = f"/instances/{instance_id}"
+        self._make_request("DELETE", endpoint)
+
+    def _make_request(self, method: str, path: str, data: Any = None) -> Response:
+        try:
+            response = requests.request(
+                method=method,
+                url=API_URL + path,
+                json=data,
+                headers={"Authorization": f"Bearer {self.api_key}"},
+                timeout=30,
+            )
+            response.raise_for_status()
+            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)
+            if e.response is not None and e.response.status_code in (
+                requests.codes.bad_request,
+                requests.codes.internal_server_error,
+                requests.codes.not_found,
+            ):
+                raise BackendError(e.response.text)
+            raise

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

@@ -0,0 +1,16 @@
+from dstack._internal.core.backends.base.backend import Backend
+from dstack._internal.core.backends.vultr.compute import VultrCompute
+from dstack._internal.core.backends.vultr.models import VultrConfig
+from dstack._internal.core.models.backends.base import BackendType
+
+
+class VultrBackend(Backend):
+    TYPE = BackendType.VULTR
+    COMPUTE_CLASS = VultrCompute
+
+    def __init__(self, config: VultrConfig):
+        self.config = config
+        self._compute = VultrCompute(self.config)
+
+    def compute(self) -> VultrCompute:
+        return self._compute

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

@@ -0,0 +1,167 @@
+import json
+import re
+from typing import List, Optional
+
+import requests
+
+from dstack._internal.core.backends.base.backend import Compute
+from dstack._internal.core.backends.base.compute import (
+    ComputeWithAllOffersCached,
+    ComputeWithCreateInstanceSupport,
+    ComputeWithMultinodeSupport,
+    ComputeWithPrivilegedSupport,
+    generate_unique_instance_name,
+    get_user_data,
+)
+from dstack._internal.core.backends.base.offers import get_catalog_offers
+from dstack._internal.core.backends.vultr.api_client import VultrApiClient
+from dstack._internal.core.backends.vultr.models import VultrConfig
+from dstack._internal.core.errors import BackendError, ProvisioningError
+from dstack._internal.core.models.backends.base import BackendType
+from dstack._internal.core.models.instances import (
+    InstanceAvailability,
+    InstanceConfiguration,
+    InstanceOffer,
+    InstanceOfferWithAvailability,
+)
+from dstack._internal.core.models.placement import PlacementGroup
+from dstack._internal.core.models.runs import JobProvisioningData
+from dstack._internal.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+MAX_INSTANCE_NAME_LEN = 64
+
+
+class VultrCompute(
+    ComputeWithAllOffersCached,
+    ComputeWithCreateInstanceSupport,
+    ComputeWithPrivilegedSupport,
+    ComputeWithMultinodeSupport,
+    Compute,
+):
+    def __init__(self, config: VultrConfig):
+        super().__init__()
+        self.config = config
+        self.api_client = VultrApiClient(config.creds.api_key)
+
+    def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability]:
+        offers = get_catalog_offers(
+            backend=BackendType.VULTR,
+            requirements=None,
+            locations=self.config.regions or None,
+            extra_filter=_supported_instances,
+        )
+        offers = [
+            InstanceOfferWithAvailability(
+                **offer.dict(), availability=InstanceAvailability.AVAILABLE
+            )
+            for offer in offers
+        ]
+        return offers
+
+    def create_instance(
+        self,
+        instance_offer: InstanceOfferWithAvailability,
+        instance_config: InstanceConfiguration,
+        placement_group: Optional[PlacementGroup],
+    ) -> JobProvisioningData:
+        instance_name = generate_unique_instance_name(
+            instance_config, max_length=MAX_INSTANCE_NAME_LEN
+        )
+        # create vpc
+        vpc = self.api_client.get_vpc_for_region(instance_offer.region)
+        if not vpc:
+            vpc = self.api_client.create_vpc(instance_offer.region)
+
+        subnet = vpc["v4_subnet"]
+        subnet_mask = vpc["v4_subnet_mask"]
+
+        instance_id = self.api_client.launch_instance(
+            region=instance_offer.region,
+            label=instance_name,
+            plan=instance_offer.instance.name,
+            user_data=get_user_data(
+                authorized_keys=instance_config.get_public_keys(),
+                firewall_allow_from_subnets=[f"{subnet}/{subnet_mask}"],
+            ),
+            vpc_id=vpc["id"],
+        )
+
+        launched_instance = 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,
+            ssh_port=22,
+            username="root",
+            ssh_proxy=None,
+            dockerized=True,
+            backend_data=json.dumps(
+                {
+                    "plan_type": "bare-metal"
+                    if "vbm" in instance_offer.instance.name
+                    else "vm_instance"
+                }
+            ),
+        )
+        return launched_instance
+
+    def terminate_instance(
+        self, instance_id: str, region: str, backend_data: Optional[str] = None
+    ) -> None:
+        plan_type = json.loads(backend_data)["plan_type"]
+        try:
+            self.api_client.terminate_instance(instance_id=instance_id, plan_type=plan_type)
+        except requests.HTTPError as e:
+            raise BackendError(e.response.text)
+
+    def update_provisioning_data(
+        self,
+        provisioning_data: JobProvisioningData,
+        project_ssh_public_key: str,
+        project_ssh_private_key: str,
+    ):
+        plan_type = json.loads(provisioning_data.backend_data)["plan_type"]
+        instance_data = self.api_client.get_instance(provisioning_data.instance_id, plan_type)
+        # Access specific fields
+        instance_status = instance_data["status"]
+        instance_main_ip = instance_data["main_ip"]
+        instance_internal_ip = instance_data["internal_ip"]
+        if instance_status == "active":
+            provisioning_data.hostname = instance_main_ip
+            provisioning_data.internal_ip = instance_internal_ip
+        if instance_status == "failed":
+            raise ProvisioningError("VM entered FAILED state")
+
+
+def _supported_instances(offer: InstanceOffer) -> bool:
+    # The vbm-4c-32gb plan does not support VPC, so it is excluded.
+    if offer.instance.name == "vbm-4c-32gb":
+        return False
+    if offer.instance.resources.spot:
+        return False
+    for family in [
+        # Bare Metal - GPU
+        r"vbm-\d+c-\d+gb-\d+-(a100|h100|l40|mi300x)-gpu",
+        # Bare Metal - AMD CPU
+        r"vbm-\d+c-\d+gb-amd",
+        # Bare Metal - Intel CPU
+        r"vbm-\d+c-\d+gb(-v\d+)?",
+        # Cloud GPU
+        r"vcg-(a16|a40|l40s|a100)-\d+c-\d+g-\d+vram",
+        # Cloud Compute - Regular Performance
+        r"vc2-\d+c-\d+gb(-sc1)?",
+        # Cloud Compute - High Frequency
+        r"vhf-\d+c-\d+gb(-sc1)?",
+        # Cloud Compute - High Performance
+        r"vhp-\d+c-\d+gb-(intel|amd)(-sc1)?",
+        # Optimized Cloud Compute
+        r"voc-[cgms]-\d+c-\d+gb-\d+s-amd(-sc1)?",
+    ]:
+        if re.fullmatch(family, offer.instance.name):
+            return True
+    return False