dstack 0.18.43__py3-none-any.whl → 0.19.0rc1__py3-none-any.whl

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

@@ -25,6 +25,7 @@ from dstack._internal.core.models.gateways import (
 from dstack._internal.core.models.instances import (
     InstanceConfiguration,
     InstanceOfferWithAvailability,
+    SSHKey,
 )
 from dstack._internal.core.models.placement import PlacementGroup, PlacementGroupProvisioningData
 from dstack._internal.core.models.runs import Job, JobProvisioningData, Requirements, Run
@@ -46,6 +47,11 @@ DSTACK_RUNNER_BINARY_PATH = f"/usr/local/bin/{DSTACK_RUNNER_BINARY_NAME}"
 
 
 class Compute(ABC):
+    """
+    A base class for all compute implementations with minimal features.
+    If a compute supports additional features, it must also subclass `ComputeWith*` classes.
+    """
+
     def __init__(self):
         self._offers_cache_lock = threading.Lock()
         self._offers_cache = TTLCache(maxsize=5, ttl=30)
@@ -54,6 +60,11 @@ class Compute(ABC):
     def get_offers(
         self, requirements: Optional[Requirements] = None
     ) -> List[InstanceOfferWithAvailability]:
+        """
+        Returns offers with availability matching `requirements`.
+        If the provider is added to gpuhunt, typically gets offers using `base.offers.get_catalog_offers()`
+        and extends them with availability info.
+        """
         pass
 
     @abstractmethod
@@ -86,6 +97,47 @@ class Compute(ABC):
         """
         pass
 
+    def update_provisioning_data(
+        self,
+        provisioning_data: JobProvisioningData,
+        project_ssh_public_key: str,
+        project_ssh_private_key: str,
+    ):
+        """
+        This method is called if `JobProvisioningData` returned from `run_job()`/`create_instance()`
+        is not complete, e.g. missing `hostname` or `ssh_port`.
+        It can be used if getting complete provisioning data takes a long of time.
+        It should not wait but return immediately.
+        If it raises `ProvisioningError`, there will be no further attempts to update the provisioning data,
+        and the run will be terminated.
+        """
+        pass
+
+    def _get_offers_cached_key(self, requirements: Optional[Requirements] = None) -> int:
+        # Requirements is not hashable, so we use a hack to get arguments hash
+        if requirements is None:
+            return hash(None)
+        return hash(requirements.json())
+
+    @cachedmethod(
+        cache=lambda self: self._offers_cache,
+        key=_get_offers_cached_key,
+        lock=lambda self: self._offers_cache_lock,
+    )
+    def get_offers_cached(
+        self, requirements: Optional[Requirements] = None
+    ) -> List[InstanceOfferWithAvailability]:
+        return self.get_offers(requirements)
+
+
+class ComputeWithCreateInstanceSupport(ABC):
+    """
+    Must be subclassed and implemented to support fleets (instance creation without running a job).
+    Typically, a compute that runs VMs would implement it,
+    and a compute that runs containers would not.
+    """
+
+    @abstractmethod
     def create_instance(
         self,
         instance_offer: InstanceOfferWithAvailability,
@@ -96,24 +148,77 @@ class Compute(ABC):
         If required to wait to get the IP address or SSH port, return partially filled `JobProvisioningData`
         and implement `update_provisioning_data()`.
         """
-        raise NotImplementedError()
+        pass
 
-    def update_provisioning_data(
+    def run_job(
         self,
-        provisioning_data: JobProvisioningData,
+        run: Run,
+        job: Job,
+        instance_offer: InstanceOfferWithAvailability,
         project_ssh_public_key: str,
         project_ssh_private_key: str,
-    ):
+        volumes: List[Volume],
+    ) -> JobProvisioningData:
         """
-        This method is called if `JobProvisioningData` returned from `run_job()`/`create_instance()`
-        is not complete, e.g. missing `hostname` or `ssh_port`.
-        It can be used if getting complete provisioning data takes a long of time.
-        It should not wait but return immediately.
-        If it raises `ProvisioningError`, there will be no further attempts to update the provisioning data,
-        and the run will be terminated.
+        The default `run_job()` implementation for all backends that support `create_instance()`.
+        Override only if custom `run_job()` behavior is required.
         """
-        pass
+        instance_config = InstanceConfiguration(
+            project_name=run.project_name,
+            instance_name=get_job_instance_name(run, job),
+            user=run.user,
+            ssh_keys=[SSHKey(public=project_ssh_public_key.strip())],
+            volumes=volumes,
+            reservation=run.run_spec.configuration.reservation,
+        )
+        instance_offer = instance_offer.copy()
+        self._restrict_instance_offer_az_to_volumes_az(instance_offer, volumes)
+        return self.create_instance(instance_offer, instance_config)
+
+    def _restrict_instance_offer_az_to_volumes_az(
+        self,
+        instance_offer: InstanceOfferWithAvailability,
+        volumes: List[Volume],
+    ):
+        if len(volumes) == 0:
+            return
+        volume = volumes[0]
+        if (
+            volume.provisioning_data is not None
+            and volume.provisioning_data.availability_zone is not None
+        ):
+            if instance_offer.availability_zones is None:
+                instance_offer.availability_zones = [volume.provisioning_data.availability_zone]
+            instance_offer.availability_zones = [
+                z
+                for z in instance_offer.availability_zones
+                if z == volume.provisioning_data.availability_zone
+            ]
+
+
+class ComputeWithMultinodeSupport:
+    """
+    Must be subclassed to support multinode tasks and cluster fleets.
+    Instances provisioned in the same project/region must be interconnected.
+    """
+
+    pass
+
+
+class ComputeWithReservationSupport:
+    """
+    Must be subclassed to support provisioning from reservations.
+    """
+
+    pass
+
 
+class ComputeWithPlacementGroupSupport(ABC):
+    """
+    Must be subclassed and implemented to support placement groups.
+    """
+
+    @abstractmethod
     def create_placement_group(
         self,
         placement_group: PlacementGroup,
@@ -121,8 +226,9 @@ class Compute(ABC):
         """
         Creates a placement group.
         """
-        raise NotImplementedError()
+        pass
 
+    @abstractmethod
     def delete_placement_group(
         self,
         placement_group: PlacementGroup,
@@ -131,8 +237,15 @@ class Compute(ABC):
         Deletes a placement group.
         If the group does not exist, it should not raise errors but return silently.
         """
-        raise NotImplementedError()
+        pass
+
 
+class ComputeWithGatewaySupport(ABC):
+    """
+    Must be subclassed and imlemented to support gateways.
+    """
+
+    @abstractmethod
     def create_gateway(
         self,
         configuration: GatewayComputeConfiguration,
@@ -140,8 +253,9 @@ class Compute(ABC):
         """
         Creates a gateway instance.
         """
-        raise NotImplementedError()
+        pass
 
+    @abstractmethod
     def terminate_gateway(
         self,
         instance_id: str,
@@ -152,21 +266,39 @@ class Compute(ABC):
         Terminates a gateway instance. Generally, it passes the call to `terminate_instance()`,
         but may perform additional work such as deleting a load balancer when a gateway has one.
         """
-        raise NotImplementedError()
+        pass
 
+
+class ComputeWithPrivateGatewaySupport:
+    """
+    Must be subclassed to support private gateways.
+    `create_gateway()` must be able to create private gateways.
+    """
+
+    pass
+
+
+class ComputeWithVolumeSupport(ABC):
+    """
+    Must be subclassed and implemented to support volumes.
+    """
+
+    @abstractmethod
     def register_volume(self, volume: Volume) -> VolumeProvisioningData:
         """
         Returns VolumeProvisioningData for an existing volume.
         Used to add external volumes to dstack.
         """
-        raise NotImplementedError()
+        pass
 
+    @abstractmethod
     def create_volume(self, volume: Volume) -> VolumeProvisioningData:
         """
         Creates a new volume.
         """
         raise NotImplementedError()
 
+    @abstractmethod
     def delete_volume(self, volume: Volume):
         """
         Deletes a volume.
@@ -176,13 +308,17 @@ class Compute(ABC):
     def attach_volume(self, volume: Volume, instance_id: str) -> VolumeAttachmentData:
         """
         Attaches a volume to the instance.
-        If the volume is not found, it should raise `ComputeError()` instead of a thrid-party exception.
+        If the volume is not found, it should raise `ComputeError()`.
+        Implement only if compute may return `VolumeProvisioningData.attachable`.
+        Otherwise, volumes should be attached by `run_job()`.
         """
         raise NotImplementedError()
 
     def detach_volume(self, volume: Volume, instance_id: str, force: bool = False):
         """
         Detaches a volume from the instance.
+        Implement only if compute may return `VolumeProvisioningData.detachable`.
+        Otherwise, volumes should be detached on instance termination.
         """
         raise NotImplementedError()
 
@@ -195,22 +331,6 @@ class Compute(ABC):
         """
         return True
 
-    def _get_offers_cached_key(self, requirements: Optional[Requirements] = None) -> int:
-        # Requirements is not hashable, so we use a hack to get arguments hash
-        if requirements is None:
-            return hash(None)
-        return hash(requirements.json())
-
-    @cachedmethod(
-        cache=lambda self: self._offers_cache,
-        key=_get_offers_cached_key,
-        lock=lambda self: self._offers_cache_lock,
-    )
-    def get_offers_cached(
-        self, requirements: Optional[Requirements] = None
-    ) -> List[InstanceOfferWithAvailability]:
-        return self.get_offers(requirements)
-
 
 def get_job_instance_name(run: Run, job: Job) -> str:
     return job.job_spec.job_name

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

@@ -0,0 +1,105 @@
+from abc import ABC, abstractmethod
+from typing import Any, ClassVar, List, Optional
+from uuid import UUID
+
+from dstack._internal.core.backends.base.backend import Backend
+from dstack._internal.core.backends.models import (
+    AnyBackendConfig,
+    AnyBackendConfigWithCreds,
+)
+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
+
+
+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):
+    """
+    `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: AnyBackendConfigWithCreds, 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: AnyBackendConfigWithCreds
+    ) -> 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(
+        self, record: StoredBackendRecord, include_creds: bool
+    ) -> AnyBackendConfig:
+        """
+        Constructs `BackendConfig` to be returned in API responses.
+        Project admins may need to see backend's creds. In this case `include_creds` will be `True`.
+        Otherwise, no sensitive information should be included.
+        """
+        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,14 @@
+from pathlib import Path
+
+
+def fill_data(values: dict):
+    if values.get("data") is not None:
+        return values
+    if "filename" not in values:
+        raise ValueError()
+    try:
+        with open(Path(values["filename"]).expanduser()) as f:
+            values["data"] = f.read()
+    except OSError:
+        raise ValueError(f"No such file {values['filename']}")
+    return values

dstack/_internal/core/backends/configurators.py

@@ -0,0 +1,138 @@
+from typing import List, Optional, Type, Union
+
+from dstack._internal.core.backends.base.configurator import Configurator
+from dstack._internal.core.models.backends.base import BackendType
+
+_CONFIGURATOR_CLASSES: List[Type[Configurator]] = []
+
+
+try:
+    from dstack._internal.core.backends.aws.configurator import AWSConfigurator
+
+    _CONFIGURATOR_CLASSES.append(AWSConfigurator)
+except ImportError:
+    pass
+
+try:
+    from dstack._internal.core.backends.azure.configurator import AzureConfigurator
+
+    _CONFIGURATOR_CLASSES.append(AzureConfigurator)
+except ImportError:
+    pass
+
+try:
+    from dstack._internal.core.backends.cudo.configurator import (
+        CudoConfigurator,
+    )
+
+    _CONFIGURATOR_CLASSES.append(CudoConfigurator)
+except ImportError:
+    pass
+
+try:
+    from dstack._internal.core.backends.datacrunch.configurator import (
+        DataCrunchConfigurator,
+    )
+
+    _CONFIGURATOR_CLASSES.append(DataCrunchConfigurator)
+except ImportError:
+    pass
+
+try:
+    from dstack._internal.core.backends.gcp.configurator import GCPConfigurator
+
+    _CONFIGURATOR_CLASSES.append(GCPConfigurator)
+except ImportError:
+    pass
+
+try:
+    from dstack._internal.core.backends.kubernetes.configurator import (
+        KubernetesConfigurator,
+    )
+
+    _CONFIGURATOR_CLASSES.append(KubernetesConfigurator)
+except ImportError:
+    pass
+
+try:
+    from dstack._internal.core.backends.lambdalabs.configurator import (
+        LambdaConfigurator,
+    )
+
+    _CONFIGURATOR_CLASSES.append(LambdaConfigurator)
+except ImportError:
+    pass
+
+try:
+    from dstack._internal.core.backends.oci.configurator import OCIConfigurator
+
+    _CONFIGURATOR_CLASSES.append(OCIConfigurator)
+except ImportError:
+    pass
+
+try:
+    from dstack._internal.core.backends.runpod.configurator import RunpodConfigurator
+
+    _CONFIGURATOR_CLASSES.append(RunpodConfigurator)
+except ImportError:
+    pass
+
+try:
+    from dstack._internal.core.backends.tensordock.configurator import (
+        TensorDockConfigurator,
+    )
+
+    _CONFIGURATOR_CLASSES.append(TensorDockConfigurator)
+except ImportError:
+    pass
+
+try:
+    from dstack._internal.core.backends.vastai.configurator import VastAIConfigurator
+
+    _CONFIGURATOR_CLASSES.append(VastAIConfigurator)
+except ImportError:
+    pass
+
+try:
+    from dstack._internal.core.backends.vultr.configurator import VultrConfigurator
+
+    _CONFIGURATOR_CLASSES.append(VultrConfigurator)
+except ImportError:
+    pass
+
+
+_BACKEND_TYPE_TO_CONFIGURATOR_CLASS_MAP = {c.TYPE: c for c in _CONFIGURATOR_CLASSES}
+_BACKEND_TYPES = [c.TYPE for c in _CONFIGURATOR_CLASSES]
+
+
+def get_configurator(backend_type: Union[BackendType, str]) -> Optional[Configurator]:
+    """
+    Returns an available `Configurator` for a given `backend_type`.
+    """
+    backend_type = BackendType(backend_type)
+    configurator_class = _BACKEND_TYPE_TO_CONFIGURATOR_CLASS_MAP.get(backend_type)
+    if configurator_class is None:
+        return None
+    return configurator_class()
+
+
+def list_available_backend_types() -> List[BackendType]:
+    """
+    Lists all backend types available on the server.
+    """
+    return _BACKEND_TYPES
+
+
+def list_available_configurator_classes() -> List[type[Configurator]]:
+    """
+    Lists all backend configurator classes available on the server.
+    """
+    return _CONFIGURATOR_CLASSES
+
+
+def register_configurator(configurator: Type[Configurator]):
+    """
+    A hook to for registering new configurators without importing them.
+    Can be used to extend dstack functionality.
+    """
+    _BACKEND_TYPE_TO_CONFIGURATOR_CLASS_MAP[configurator.TYPE] = configurator

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

@@ -1,15 +0,0 @@
-from dstack._internal.core.backends.base import Backend
-from dstack._internal.core.backends.cudo.compute import CudoCompute
-from dstack._internal.core.backends.cudo.config import CudoConfig
-from dstack._internal.core.models.backends.base import BackendType
-
-
-class CudoBackend(Backend):
-    TYPE: BackendType = BackendType.CUDO
-
-    def __init__(self, config: CudoConfig):
-        self.config = config
-        self._compute = CudoCompute(self.config)
-
-    def compute(self) -> CudoCompute:
-        return self._compute

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

@@ -0,0 +1,16 @@
+from dstack._internal.core.backends.base.backend import Backend
+from dstack._internal.core.backends.cudo.compute import CudoCompute
+from dstack._internal.core.backends.cudo.models import CudoConfig
+from dstack._internal.core.models.backends.base import BackendType
+
+
+class CudoBackend(Backend):
+    TYPE = BackendType.CUDO
+    COMPUTE_CLASS = CudoCompute
+
+    def __init__(self, config: CudoConfig):
+        self.config = config
+        self._compute = CudoCompute(self.config)
+
+    def compute(self) -> CudoCompute:
+        return self._compute

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

@@ -2,25 +2,23 @@ from typing import List, Optional
 
 import requests
 
-from dstack._internal.core.backends.base import Compute
+from dstack._internal.core.backends.base.backend import Compute
 from dstack._internal.core.backends.base.compute import (
+    ComputeWithCreateInstanceSupport,
     generate_unique_instance_name,
-    get_job_instance_name,
     get_shim_commands,
 )
 from dstack._internal.core.backends.base.offers import get_catalog_offers
 from dstack._internal.core.backends.cudo.api_client import CudoApiClient
-from dstack._internal.core.backends.cudo.config import CudoConfig
+from dstack._internal.core.backends.cudo.models import CudoConfig
 from dstack._internal.core.errors import BackendError, NoCapacityError, ProvisioningError
 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.runs import Job, JobProvisioningData, Requirements, Run
-from dstack._internal.core.models.volumes import Volume
+from dstack._internal.core.models.runs import JobProvisioningData, Requirements
 from dstack._internal.utils.logging import get_logger
 
 logger = get_logger(__name__)
@@ -29,7 +27,10 @@ logger = get_logger(__name__)
 MAX_RESOURCE_NAME_LEN = 30
 
 
-class CudoCompute(Compute):
+class CudoCompute(
+    ComputeWithCreateInstanceSupport,
+    Compute,
+):
     def __init__(self, config: CudoConfig):
         super().__init__()
         self.config = config
@@ -51,25 +52,6 @@ class CudoCompute(Compute):
         ]
         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],
-    ) -> JobProvisioningData:
-        instance_config = InstanceConfiguration(
-            project_name=run.project_name,
-            instance_name=get_job_instance_name(run, job),
-            ssh_keys=[
-                SSHKey(public=project_ssh_public_key.strip()),
-            ],
-            user=run.user,
-        )
-        return self.create_instance(instance_offer, instance_config)
-
     def create_instance(
         self,
         instance_offer: InstanceOfferWithAvailability,