dstack 0.18.44__py3-none-any.whl → 0.19.0__py3-none-any.whl

dstack/_internal/core/backends/cudo/config.py

@@ -1,9 +0,0 @@
-from dstack._internal.core.backends.base.config import BackendConfig
-from dstack._internal.core.models.backends.cudo import (
-    AnyCudoCreds,
-    CudoStoredConfig,
-)
-
-
-class CudoConfig(CudoStoredConfig, BackendConfig):
-    creds: AnyCudoCreds

dstack/_internal/core/backends/datacrunch/config.py

@@ -1,9 +0,0 @@
-from dstack._internal.core.backends.base.config import BackendConfig
-from dstack._internal.core.models.backends.datacrunch import (
-    AnyDataCrunchCreds,
-    DataCrunchStoredConfig,
-)
-
-
-class DataCrunchConfig(DataCrunchStoredConfig, BackendConfig):
-    creds: AnyDataCrunchCreds

dstack/_internal/core/backends/gcp/config.py

@@ -1,22 +0,0 @@
-from dstack._internal.core.backends.base.config import BackendConfig
-from dstack._internal.core.models.backends.gcp import AnyGCPCreds, GCPStoredConfig
-
-
-class GCPConfig(GCPStoredConfig, BackendConfig):
-    creds: AnyGCPCreds
-
-    @property
-    def allocate_public_ips(self) -> bool:
-        if self.public_ips is not None:
-            return self.public_ips
-        return True
-
-    @property
-    def vpc_resource_name(self) -> str:
-        vpc_name = self.vpc_name
-        if vpc_name is None:
-            vpc_name = "default"
-        project_id = self.project_id
-        if self.vpc_project_id is not None:
-            project_id = self.vpc_project_id
-        return f"projects/{project_id}/global/networks/{vpc_name}"

dstack/_internal/core/backends/kubernetes/config.py

@@ -1,6 +0,0 @@
-from dstack._internal.core.backends.base.config import BackendConfig
-from dstack._internal.core.models.backends.kubernetes import KubernetesStoredConfig
-
-
-class KubernetesConfig(KubernetesStoredConfig, BackendConfig):
-    pass

dstack/_internal/core/backends/lambdalabs/config.py

@@ -1,9 +0,0 @@
-from dstack._internal.core.backends.base.config import BackendConfig
-from dstack._internal.core.models.backends.lambdalabs import (
-    AnyLambdaCreds,
-    LambdaStoredConfig,
-)
-
-
-class LambdaConfig(LambdaStoredConfig, BackendConfig):
-    creds: AnyLambdaCreds

dstack/_internal/core/backends/nebius/__init__.py

@@ -1,15 +0,0 @@
-from dstack._internal.core.backends.base import Backend
-from dstack._internal.core.backends.nebius.compute import NebiusCompute
-from dstack._internal.core.backends.nebius.config import NebiusConfig
-from dstack._internal.core.models.backends.base import BackendType
-
-
-class NebiusBackend(Backend):
-    TYPE: BackendType = BackendType.NEBIUS
-
-    def __init__(self, config: NebiusConfig):
-        self.config = config
-        self._compute = NebiusCompute(self.config)
-
-    def compute(self) -> NebiusCompute:
-        return self._compute

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

@@ -1,319 +0,0 @@
-import time
-from typing import Dict, List, Optional
-
-import jwt
-import requests
-
-from dstack._internal.core.backends.nebius.types import (
-    ClientError,
-    ConflictError,
-    ForbiddenError,
-    NebiusError,
-    NotFoundError,
-    ResourcesSpec,
-    ServiceAccount,
-)
-from dstack._internal.utils.logging import get_logger
-
-logger = get_logger("nebius")
-API_URL = "api.ai.nebius.cloud"
-REQUEST_TIMEOUT = 15
-
-
-class NebiusAPIClient:
-    # Reference: https://nebius.ai/docs/api-design-guide/
-    def __init__(self, service_account: ServiceAccount):
-        self.service_account = service_account
-        self.s = requests.Session()
-        self.expires_at = 0
-
-    def get_token(self):
-        now = int(time.time())
-        if now + 60 < self.expires_at:
-            return
-        logger.debug("Refreshing IAM token")
-        expires_at = now + 3600
-        payload = {
-            "aud": self.url("iam", "/tokens"),
-            "iss": self.service_account["service_account_id"],
-            "iat": now,
-            "exp": expires_at,
-        }
-        jwt_token = jwt.encode(
-            payload,
-            self.service_account["private_key"],
-            algorithm="PS256",
-            headers={"kid": self.service_account["id"]},
-        )
-
-        resp = requests.post(payload["aud"], json={"jwt": jwt_token}, timeout=REQUEST_TIMEOUT)
-        resp.raise_for_status()
-        iam_token = resp.json()["iamToken"]
-        self.s.headers["Authorization"] = f"Bearer {iam_token}"
-        self.expires_at = expires_at
-
-    def compute_zones_list(self) -> List[dict]:
-        logger.debug("Fetching compute zones")
-        self.get_token()
-        resp = self.s.get(self.url("compute", "/zones"), timeout=REQUEST_TIMEOUT)
-        self.raise_for_status(resp)
-        return resp.json()["zones"]
-
-    def resource_manager_folders_create(self, cloud_id: str, name: str, **kwargs) -> dict:
-        logger.debug("Creating folder %s", name)
-        self.get_token()
-        resp = self.s.post(
-            self.url("resource-manager", "/folders"),
-            json=omit_none(
-                cloudId=cloud_id,
-                name=name,
-                **kwargs,
-            ),
-            timeout=REQUEST_TIMEOUT,
-        )
-        self.raise_for_status(resp)
-        return resp.json()
-
-    def vpc_networks_create(self, folder_id: str, name: str, **kwargs) -> dict:
-        logger.debug("Creating network %s in %s", name, folder_id)
-        self.get_token()
-        resp = self.s.post(
-            self.url("vpc", "/networks"),
-            json=omit_none(
-                folderId=folder_id,
-                name=name,
-                **kwargs,
-            ),
-            timeout=REQUEST_TIMEOUT,
-        )
-        self.raise_for_status(resp)
-        return resp.json()
-
-    def vpc_networks_list(self, folder_id: str, filter: Optional[str] = None) -> List[dict]:
-        logger.debug("Fetching networks in %s", folder_id)
-        return self.list(
-            "vpc",
-            "networks",
-            params=dict(
-                folderId=folder_id,
-                filter=filter,
-            ),
-        )
-
-    def vpc_subnets_create(
-        self,
-        folder_id: str,
-        name: str,
-        network_id: str,
-        zone: str,
-        cird_blocks: List[str],
-        **kwargs,
-    ) -> dict:
-        logger.debug("Creating subnet %s in %s", name, network_id)
-        self.get_token()
-        resp = self.s.post(
-            self.url("vpc", "/subnets"),
-            json=omit_none(
-                folderId=folder_id,
-                name=name,
-                networkId=network_id,
-                zoneId=zone,
-                v4CidrBlocks=cird_blocks,
-                **kwargs,
-            ),
-            timeout=REQUEST_TIMEOUT,
-        )
-        self.raise_for_status(resp)
-        return resp.json()
-
-    def vpc_subnets_list(self, folder_id: str, filter: Optional[str] = None) -> List[dict]:
-        logger.debug("Fetching subnets in %s", folder_id)
-        return self.list(
-            "vpc",
-            "subnets",
-            params=dict(
-                folderId=folder_id,
-                filter=filter,
-            ),
-        )
-
-    def vpc_security_groups_create(
-        self, folder_id: str, name: str, network_id: str, rule_specs: List[dict], **kwargs
-    ) -> dict:
-        logger.debug("Creating security group %s in %s", name, folder_id)
-        self.get_token()
-        resp = self.s.post(
-            self.url("vpc", "/securityGroups"),
-            json=omit_none(
-                folderId=folder_id,
-                name=name,
-                networkId=network_id,
-                ruleSpecs=rule_specs,
-                **kwargs,
-            ),
-            timeout=REQUEST_TIMEOUT,
-        )
-        self.raise_for_status(resp)
-        return resp.json()
-
-    def vpc_security_groups_list(self, folder_id: str, filter: Optional[str] = None) -> List[dict]:
-        logger.debug("Fetching security groups in %s", folder_id)
-        return self.list(
-            "vpc",
-            "securityGroups",
-            params=dict(
-                folderId=folder_id,
-                filter=filter,
-            ),
-        )
-
-    def vpc_security_groups_delete(self, security_group_id: str):
-        logger.debug("Deleting security group %s", security_group_id)
-        self.get_token()
-        resp = self.s.delete(
-            self.url("vpc", f"/securityGroups/{security_group_id}"), timeout=REQUEST_TIMEOUT
-        )
-        self.raise_for_status(resp)
-
-    def compute_instances_create(
-        self,
-        folder_id: str,
-        name: str,
-        zone_id: str,
-        platform_id: str,
-        resources_spec: ResourcesSpec,
-        metadata: Optional[Dict[str, str]],
-        disk_size_gb: int,
-        image_id: str,
-        subnet_id: str,
-        security_group_ids: List[str],
-        **kwargs,
-    ) -> dict:
-        # Reference: https://nebius.ai/docs/api-design-guide/compute/v1/api-ref/Instance/create
-        logger.debug("Creating instance %s (%s) in %s", name, platform_id, folder_id)
-        self.get_token()
-        resp = self.s.post(
-            self.url("compute", "/instances"),
-            json=omit_none(
-                folderId=folder_id,
-                name=name,
-                zoneId=zone_id,
-                platformId=platform_id,
-                resourcesSpec=resources_spec,
-                metadata=metadata,
-                boot_disk_spec=dict(
-                    autoDelete=True,
-                    diskSpec=dict(
-                        typeId="network-ssd",
-                        size=disk_size_gb * 1024 * 1024 * 1024,
-                        imageId=image_id,
-                    ),
-                ),
-                networkInterfaceSpecs=[
-                    dict(
-                        subnetId=subnet_id,
-                        primaryV4AddressSpec=dict(
-                            oneToOneNatSpec=dict(
-                                ipVersion="IPV4",
-                            ),
-                        ),
-                        securityGroupIds=security_group_ids,
-                    )
-                ],
-                **kwargs,
-            ),
-            timeout=REQUEST_TIMEOUT,
-        )
-        self.raise_for_status(resp)
-        return resp.json()
-
-    def compute_instances_list(
-        self, folder_id: str, filter: Optional[str] = None, order_by: Optional[str] = None
-    ) -> List[dict]:
-        logger.debug("Fetching instances in %s", folder_id)
-        return self.list(
-            "compute",
-            "instances",
-            params=dict(
-                folderId=folder_id,
-                filter=filter,
-                orderBy=order_by,
-            ),
-        )
-
-    def compute_instances_delete(self, instance_id: str):
-        logger.debug("Deleting instance %s", instance_id)
-        self.get_token()
-        resp = self.s.delete(
-            self.url("compute", f"/instances/{instance_id}"), timeout=REQUEST_TIMEOUT
-        )
-        self.raise_for_status(resp)
-
-    def compute_instances_get(self, instance_id: str, full: bool = False) -> dict:
-        logger.debug("Fetching instance %s", instance_id)
-        self.get_token()
-        resp = self.s.get(
-            self.url("compute", f"/instances/{instance_id}"),
-            params=dict(
-                view="FULL" if full else "BASIC",
-            ),
-            timeout=REQUEST_TIMEOUT,
-        )
-        self.raise_for_status(resp)
-        return resp.json()
-
-    def compute_images_list(
-        self, folder_id: str, filter: Optional[str] = None, order_by: Optional[str] = None
-    ):
-        logger.debug("Fetching images in %s", folder_id)
-        return self.list(
-            "compute",
-            "images",
-            params=dict(
-                folderId=folder_id,
-                filter=filter,
-                orderBy=order_by,
-            ),
-        )
-
-    def list(self, service: str, resource: str, params: dict, page_size: int = 1000) -> List[dict]:
-        page_token = None
-        output = []
-        while True:
-            self.get_token()
-            resp = self.s.get(
-                self.url(service, f"/{resource}"),
-                params=omit_none(
-                    pageSize=page_size,
-                    pageToken=page_token,
-                    **params,
-                ),
-                timeout=REQUEST_TIMEOUT,
-            )
-            self.raise_for_status(resp)
-            data = resp.json()
-            output += data.get(resource, [])
-            page_token = data.get("nextPageToken")
-            if not page_token:
-                break
-        return output
-
-    def url(self, service: str, path: str, version="v1") -> str:
-        return f"https://{service}.{API_URL.rstrip('/')}/{service}/{version}/{path.lstrip('/')}"
-
-    def raise_for_status(self, resp: requests.Response):
-        if resp.status_code == 400:
-            raise NebiusError(resp.text)
-        if resp.status_code == 401:
-            raise ClientError(resp.text)
-        if resp.status_code == 403:
-            raise ForbiddenError(resp.text)
-        if resp.status_code == 404:
-            raise NotFoundError(resp.text)
-        if resp.status_code == 409:
-            raise ConflictError(resp.text)
-        resp.raise_for_status()
-
-
-def omit_none(**kwargs) -> dict:
-    return {k: v for k, v in kwargs.items() if v is not None}

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

@@ -1,220 +0,0 @@
-import json
-import re
-import time
-from typing import List, Optional
-
-import dstack.version as version
-from dstack._internal import settings
-from dstack._internal.core.backends.base import Compute
-from dstack._internal.core.backends.base.compute import get_job_instance_name, get_user_data
-from dstack._internal.core.backends.base.offers import get_catalog_offers
-from dstack._internal.core.backends.nebius.api_client import NebiusAPIClient
-from dstack._internal.core.backends.nebius.config import NebiusConfig
-from dstack._internal.core.backends.nebius.types import (
-    ForbiddenError,
-    NotFoundError,
-    ResourcesSpec,
-)
-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,
-    InstanceOfferWithAvailability,
-    SSHKey,
-)
-from dstack._internal.core.models.resources import Memory, Range
-from dstack._internal.core.models.runs import Job, JobProvisioningData, Requirements, Run
-from dstack._internal.core.models.volumes import Volume
-
-MEGABYTE = 1024**2
-INSTANCE_PULL_INTERVAL = 10
-# TODO: find out the actual lower bound considering dstack image size, 50GB is made up
-CONFIGURABLE_DISK_SIZE = Range[Memory](min=Memory.parse("50GB"), max=Memory.parse("4TB"))
-
-
-class NebiusCompute(Compute):
-    def __init__(self, config: NebiusConfig):
-        super().__init__()
-        self.config = config
-        self.api_client = NebiusAPIClient(json.loads(self.config.creds.data))
-
-    def get_offers(
-        self, requirements: Optional[Requirements] = None
-    ) -> List[InstanceOfferWithAvailability]:
-        offers = get_catalog_offers(
-            backend=BackendType.NEBIUS,
-            locations=self.config.regions,
-            requirements=requirements,
-            configurable_disk_size=CONFIGURABLE_DISK_SIZE,
-        )
-        # TODO(egor-s) quotas
-        return [
-            InstanceOfferWithAvailability(
-                **offer.dict(), availability=InstanceAvailability.UNKNOWN
-            )
-            for offer in offers
-        ]
-
-    def create_instance(
-        self,
-        instance_offer: InstanceOfferWithAvailability,
-        instance_config: InstanceConfiguration,
-    ) -> JobProvisioningData:
-        cuda = len(instance_offer.instance.resources.gpus) > 0
-        security_group_id = self._get_security_group_id(project_name=instance_config.project_name)
-        subnet_id = self._get_subnet_id(zone=instance_offer.region)
-        image_id = self._get_image_id(cuda=cuda)
-
-        try:
-            disk_size = round(instance_offer.instance.resources.disk.size_mib / 1024)
-            resp = self.api_client.compute_instances_create(
-                folder_id=self.config.folder_id,
-                name=instance_config.instance_name,
-                zone_id=instance_offer.region,
-                platform_id=instance_offer.instance.name,
-                resources_spec=ResourcesSpec(
-                    memory=int(instance_offer.instance.resources.memory_mib * MEGABYTE),
-                    cores=instance_offer.instance.resources.cpus,
-                    coreFraction=100,
-                    gpus=len(instance_offer.instance.resources.gpus),
-                ),
-                metadata={
-                    "user-data": get_user_data(authorized_keys=instance_config.get_public_keys())
-                },
-                disk_size_gb=disk_size,
-                image_id=image_id,
-                subnet_id=subnet_id,
-                security_group_ids=[security_group_id],
-                labels=self._get_labels(project=instance_config.project_name),
-            )
-        except ForbiddenError as e:
-            if instance_offer.instance.name in e.args[0]:
-                raise NoCapacityError(json.loads(e.args[0])["message"])
-            raise
-        instance_id = resp["metadata"]["instanceId"]
-        try:
-            while True:
-                instance = self.api_client.compute_instances_get(instance_id)
-                if "primaryV4Address" in instance["networkInterfaces"][0]:
-                    break
-                time.sleep(INSTANCE_PULL_INTERVAL)
-        except Exception:
-            self.terminate_instance(instance_id, instance_offer.region)
-            raise
-        return JobProvisioningData(
-            backend=instance_offer.backend,
-            instance_type=instance_offer.instance,
-            instance_id=instance_id,
-            hostname=instance["networkInterfaces"][0]["primaryV4Address"]["oneToOneNat"][
-                "address"
-            ],
-            internal_ip=None,
-            region=instance_offer.region,
-            price=instance_offer.price,
-            username="ubuntu",
-            ssh_port=22,
-            dockerized=True,
-            ssh_proxy=None,
-            backend_data=None,
-        )
-
-    def run_job(
-        self,
-        run: Run,
-        job: Job,
-        instance_offer: InstanceOfferWithAvailability,
-        project_ssh_public_key: str,
-        project_ssh_private_key: str,
-        volumes: List[Volume],
-    ) -> JobProvisioningData:
-        instance_config = InstanceConfiguration(
-            project_name=run.project_name,
-            instance_name=get_job_instance_name(run, job),  # TODO: generate name
-            ssh_keys=[
-                SSHKey(public=project_ssh_public_key.strip()),
-            ],
-            user=run.user,
-        )
-        return self.create_instance(instance_offer, instance_config)
-
-    def terminate_instance(
-        self, instance_id: str, region: str, backend_data: Optional[str] = None
-    ):
-        try:
-            self.api_client.compute_instances_delete(instance_id)
-        except NotFoundError:
-            pass
-
-    def _get_security_group_id(self, project_name: str) -> str:
-        name = project_name
-        security_groups = self.api_client.vpc_security_groups_list(
-            folder_id=self.config.folder_id,
-            filter=f'name="{name}"',
-        )
-        if security_groups:
-            return security_groups[0]["id"]
-        resp = self.api_client.vpc_security_groups_create(
-            folder_id=self.config.folder_id,
-            name=name,
-            network_id=self.config.network_id,
-            rule_specs=[
-                {
-                    "description": "SSH access",
-                    "direction": "INGRESS",
-                    "ports": {"fromPort": 22, "toPort": 22},
-                    "protocolName": "ANY",
-                    "cidrBlocks": {"v4CidrBlocks": ["0.0.0.0/0"]},
-                },
-                {
-                    "description": "Project intranet",
-                    "direction": "INGRESS",
-                    "protocolName": "ANY",
-                    "predefinedTarget": "self_security_group",
-                },
-                {
-                    "description": "Internet access",
-                    "direction": "EGRESS",
-                    "protocolName": "ANY",
-                    "cidrBlocks": {"v4CidrBlocks": ["0.0.0.0/0"]},
-                },
-            ],
-            description="For job instance, by dstack",
-            labels=self._get_labels(project=project_name),
-        )
-        return resp["response"]["id"]
-
-    def _get_subnet_id(self, zone: str, name: Optional[str] = None) -> str:
-        name = name or f"default-{zone}"
-        subnets = self.api_client.vpc_subnets_list(folder_id=self.config.folder_id)
-        for subnet in subnets:
-            if subnet["name"] == name:
-                return subnet["id"]
-        n = len(subnets)
-        resp = self.api_client.vpc_subnets_create(
-            folder_id=self.config.folder_id,
-            name=name,
-            network_id=self.config.network_id,
-            zone=zone,
-            cird_blocks=[f"10.{n}.0.0/16"],
-            labels=self._get_labels(),
-        )
-        return resp["response"]["id"]
-
-    def _get_image_id(self, cuda: bool) -> str:
-        image_name = re.sub(r"[^a-z0-9-]", "-", f"dstack-{version.base_image}")
-        if cuda:
-            image_name += "-cuda"
-        images = self.api_client.compute_images_list(
-            folder_id="bjel82ie37qos4pc6guk", filter=f'name="{image_name}"'
-        )
-        return images[0]["id"]
-
-    def _get_labels(self, **kwargs) -> dict:
-        labels = {
-            "owner": "dstack",
-            **kwargs,
-        }
-        if settings.DSTACK_VERSION is not None:
-            labels["dstack-version"] = settings.DSTACK_VERSION.replace(".", "-")
-        return labels

dstack/_internal/core/backends/nebius/config.py

@@ -1,6 +0,0 @@
-from dstack._internal.core.backends.base.config import BackendConfig
-from dstack._internal.core.models.backends.nebius import AnyNebiusCreds, NebiusStoredConfig
-
-
-class NebiusConfig(NebiusStoredConfig, BackendConfig):
-    creds: AnyNebiusCreds

dstack/_internal/core/backends/nebius/types.py

@@ -1,37 +0,0 @@
-from typing import TypedDict
-
-
-class ServiceAccount(TypedDict):
-    id: str
-    service_account_id: str
-    created_at: str
-    key_algorithm: str
-    public_key: str
-    private_key: str
-
-
-class ResourcesSpec(TypedDict):
-    memory: int
-    cores: int
-    coreFraction: int
-    gpus: int
-
-
-class NebiusError(Exception):
-    pass
-
-
-class ClientError(NebiusError):
-    pass
-
-
-class ForbiddenError(NebiusError):
-    pass
-
-
-class NotFoundError(NebiusError):
-    pass
-
-
-class ConflictError(NebiusError):
-    pass

dstack/_internal/core/backends/oci/config.py

@@ -1,6 +0,0 @@
-from dstack._internal.core.backends.base.config import BackendConfig
-from dstack._internal.core.models.backends.oci import AnyOCICreds, OCIStoredConfig
-
-
-class OCIConfig(OCIStoredConfig, BackendConfig):
-    creds: AnyOCICreds

dstack/_internal/core/backends/runpod/config.py

@@ -1,17 +0,0 @@
-from dstack._internal.core.backends.base.config import BackendConfig
-from dstack._internal.core.models.backends.runpod import (
-    AnyRunpodCreds,
-    RunpodStoredConfig,
-)
-
-RUNPOD_COMMUNITY_CLOUD_DEFAULT = True
-
-
-class RunpodConfig(RunpodStoredConfig, BackendConfig):
-    creds: AnyRunpodCreds
-
-    @property
-    def allow_community_cloud(self) -> bool:
-        if self.community_cloud is not None:
-            return self.community_cloud
-        return RUNPOD_COMMUNITY_CLOUD_DEFAULT