dstack 0.19.27__py3-none-any.whl → 0.19.29__py3-none-any.whl

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

@@ -1,17 +1,18 @@
 from concurrent.futures import ThreadPoolExecutor
 from functools import cached_property
-from typing import List, Optional
+from typing import Callable, List, Optional
 
 import oci
 
 from dstack._internal.core.backends.base.compute import (
     Compute,
+    ComputeWithAllOffersCached,
     ComputeWithCreateInstanceSupport,
     ComputeWithMultinodeSupport,
     generate_unique_instance_name,
     get_user_data,
 )
-from dstack._internal.core.backends.base.offers import get_catalog_offers
+from dstack._internal.core.backends.base.offers import get_catalog_offers, get_offers_disk_modifier
 from dstack._internal.core.backends.oci import resources
 from dstack._internal.core.backends.oci.models import OCIConfig
 from dstack._internal.core.backends.oci.region import make_region_clients_map
@@ -47,6 +48,7 @@ CONFIGURABLE_DISK_SIZE = Range[Memory](min=Memory.parse("50GB"), max=Memory.pars
 
 
 class OCICompute(
+    ComputeWithAllOffersCached,
     ComputeWithCreateInstanceSupport,
     ComputeWithMultinodeSupport,
     Compute,
@@ -60,14 +62,10 @@ class OCICompute(
     def shapes_quota(self) -> resources.ShapesQuota:
         return resources.ShapesQuota.load(self.regions, self.config.compartment_id)
 
-    def get_offers(
-        self, requirements: Optional[Requirements] = None
-    ) -> List[InstanceOfferWithAvailability]:
+    def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability]:
         offers = get_catalog_offers(
             backend=BackendType.OCI,
             locations=self.config.regions,
-            requirements=requirements,
-            configurable_disk_size=CONFIGURABLE_DISK_SIZE,
             extra_filter=_supported_instances,
         )
 
@@ -96,6 +94,11 @@ class OCICompute(
 
         return offers_with_availability
 
+    def get_offers_modifier(
+        self, requirements: Requirements
+    ) -> Callable[[InstanceOfferWithAvailability], Optional[InstanceOfferWithAvailability]]:
+        return get_offers_disk_modifier(CONFIGURABLE_DISK_SIZE, requirements)
+
     def terminate_instance(
         self, instance_id: str, region: str, backend_data: Optional[str] = None
     ) -> None:

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

@@ -1,17 +1,18 @@
 import json
 import uuid
 from datetime import timedelta
-from typing import List, Optional
+from typing import Callable, List, Optional
 
 from dstack._internal.core.backends.base.backend import Compute
 from dstack._internal.core.backends.base.compute import (
+    ComputeWithAllOffersCached,
     ComputeWithVolumeSupport,
     generate_unique_instance_name,
     generate_unique_volume_name,
     get_docker_commands,
     get_job_instance_name,
 )
-from dstack._internal.core.backends.base.offers import get_catalog_offers
+from dstack._internal.core.backends.base.offers import get_catalog_offers, get_offers_disk_modifier
 from dstack._internal.core.backends.runpod.api_client import RunpodApiClient
 from dstack._internal.core.backends.runpod.models import RunpodConfig
 from dstack._internal.core.consts import DSTACK_RUNNER_SSH_PORT
@@ -27,6 +28,7 @@ from dstack._internal.core.models.instances import (
     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, VolumeProvisioningData
 from dstack._internal.utils.common import get_current_datetime
@@ -39,8 +41,12 @@ MAX_RESOURCE_NAME_LEN = 60
 
 CONTAINER_REGISTRY_AUTH_CLEANUP_INTERVAL = 60 * 60 * 24  # 24 hour
 
+# RunPod does not seem to have any limits on the disk size.
+CONFIGURABLE_DISK_SIZE = Range[Memory](min=Memory.parse("1GB"), max=None)
+
 
 class RunpodCompute(
+    ComputeWithAllOffersCached,
     ComputeWithVolumeSupport,
     Compute,
 ):
@@ -51,13 +57,11 @@ class RunpodCompute(
         self.config = config
         self.api_client = RunpodApiClient(config.creds.api_key)
 
-    def get_offers(
-        self, requirements: Optional[Requirements] = None
-    ) -> List[InstanceOfferWithAvailability]:
+    def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability]:
         offers = get_catalog_offers(
             backend=BackendType.RUNPOD,
             locations=self.config.regions or None,
-            requirements=requirements,
+            requirements=None,
             extra_filter=lambda o: _is_secure_cloud(o.region) or self.config.allow_community_cloud,
         )
         offers = [
@@ -68,6 +72,11 @@ class RunpodCompute(
         ]
         return offers
 
+    def get_offers_modifier(
+        self, requirements: Requirements
+    ) -> Callable[[InstanceOfferWithAvailability], Optional[InstanceOfferWithAvailability]]:
+        return get_offers_disk_modifier(CONFIGURABLE_DISK_SIZE, requirements)
+
     def run_job(
         self,
         run: Run,

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

@@ -2,6 +2,7 @@ from typing import List, Optional
 
 from dstack._internal.core.backends.base.backend import Compute
 from dstack._internal.core.backends.base.compute import (
+    ComputeWithAllOffersCached,
     ComputeWithCreateInstanceSupport,
     ComputeWithGatewaySupport,
     ComputeWithMultinodeSupport,
@@ -28,6 +29,7 @@ logger = get_logger(__name__)
 
 class {{ backend_name }}Compute(
     # TODO: Choose ComputeWith* classes to extend and implement
+    # ComputeWithAllOffersCached,
     # ComputeWithCreateInstanceSupport,
     # ComputeWithMultinodeSupport,
     # ComputeWithReservationSupport,
@@ -42,7 +44,7 @@ class {{ backend_name }}Compute(
         self.config = config
 
     def get_offers(
-        self, requirements: Optional[Requirements] = None
+        self, requirements: Requirements
     ) -> List[InstanceOfferWithAvailability]:
         # If the provider is added to gpuhunt, you'd typically get offers
         # using `get_catalog_offers()` and extend them with availability info.

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

@@ -39,9 +39,7 @@ class TensorDockCompute(
         self.config = config
         self.api_client = TensorDockAPIClient(config.creds.api_key, config.creds.api_token)
 
-    def get_offers(
-        self, requirements: Optional[Requirements] = None
-    ) -> List[InstanceOfferWithAvailability]:
+    def get_offers(self, requirements: Requirements) -> List[InstanceOfferWithAvailability]:
         offers = get_catalog_offers(
             backend=BackendType.TENSORDOCK,
             requirements=requirements,

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

@@ -4,6 +4,8 @@ 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"

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

@@ -5,6 +5,7 @@ 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,
 )
@@ -30,7 +31,10 @@ logger = get_logger(__name__)
 MAX_INSTANCE_NAME_LEN = 60
 
 
-class VastAICompute(Compute):
+class VastAICompute(
+    ComputeWithFilteredOffersCached,
+    Compute,
+):
     def __init__(self, config: VastAIConfig):
         super().__init__()
         self.config = config
@@ -49,8 +53,8 @@ class VastAICompute(Compute):
             )
         )
 
-    def get_offers(
-        self, requirements: Optional[Requirements] = None
+    def get_offers_by_requirements(
+        self, requirements: Requirements
     ) -> List[InstanceOfferWithAvailability]:
         offers = get_catalog_offers(
             backend=BackendType.VASTAI,

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

@@ -6,6 +6,7 @@ import requests
 
 from dstack._internal.core.backends.base.backend import Compute
 from dstack._internal.core.backends.base.compute import (
+    ComputeWithAllOffersCached,
     ComputeWithCreateInstanceSupport,
     ComputeWithMultinodeSupport,
     generate_unique_instance_name,
@@ -23,7 +24,7 @@ from dstack._internal.core.models.instances import (
     InstanceOfferWithAvailability,
 )
 from dstack._internal.core.models.placement import PlacementGroup
-from dstack._internal.core.models.runs import JobProvisioningData, Requirements
+from dstack._internal.core.models.runs import JobProvisioningData
 from dstack._internal.utils.logging import get_logger
 
 logger = get_logger(__name__)
@@ -32,6 +33,7 @@ MAX_INSTANCE_NAME_LEN = 64
 
 
 class VultrCompute(
+    ComputeWithAllOffersCached,
     ComputeWithCreateInstanceSupport,
     ComputeWithMultinodeSupport,
     Compute,
@@ -41,12 +43,10 @@ class VultrCompute(
         self.config = config
         self.api_client = VultrApiClient(config.creds.api_key)
 
-    def get_offers(
-        self, requirements: Optional[Requirements] = None
-    ) -> List[InstanceOfferWithAvailability]:
+    def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability]:
         offers = get_catalog_offers(
             backend=BackendType.VULTR,
-            requirements=requirements,
+            requirements=None,
             locations=self.config.regions or None,
             extra_filter=_supported_instances,
         )

dstack/_internal/core/compatibility/runs.py

@@ -31,6 +31,8 @@ def get_apply_plan_excludes(plan: ApplyRunPlanInput) -> Optional[IncludeExcludeD
         current_resource_excludes["status_message"] = True
         if current_resource.deployment_num == 0:
             current_resource_excludes["deployment_num"] = True
+        if current_resource.fleet is None:
+            current_resource_excludes["fleet"] = True
         apply_plan_excludes["current_resource"] = current_resource_excludes
         current_resource_excludes["run_spec"] = get_run_spec_excludes(current_resource.run_spec)
         job_submissions_excludes: IncludeExcludeDictType = {}

dstack/_internal/core/models/common.py

@@ -1,10 +1,10 @@
 import re
 from enum import Enum
-from typing import Any, Callable, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable, Mapping, Optional, Union
 
 import orjson
 from pydantic import Field
-from pydantic_duality import DualBaseModel
+from pydantic_duality import generate_dual_base_model
 from typing_extensions import Annotated
 
 from dstack._internal.utils.json_utils import pydantic_orjson_dumps
@@ -17,46 +17,73 @@ IncludeExcludeDictType = dict[
 IncludeExcludeType = Union[IncludeExcludeSetType, IncludeExcludeDictType]
 
 
+class CoreConfig:
+    json_loads = orjson.loads
+    json_dumps = pydantic_orjson_dumps
+
+
+# All dstack models inherit from pydantic-duality's DualBaseModel.
 # DualBaseModel creates two classes for the model:
 # one with extra = "forbid" (CoreModel/CoreModel.__request__),
 # and another with extra = "ignore" (CoreModel.__response__).
-# This allows to use the same model both for a strict parsing of the user input and
-# for a permissive parsing of the server responses.
-class CoreModel(DualBaseModel):
-    class Config:
-        json_loads = orjson.loads
-        json_dumps = pydantic_orjson_dumps
-
-    def json(
-        self,
-        *,
-        include: Optional[IncludeExcludeType] = None,
-        exclude: Optional[IncludeExcludeType] = None,
-        by_alias: bool = False,
-        skip_defaults: Optional[bool] = None,  # ignore as it's deprecated
-        exclude_unset: bool = False,
-        exclude_defaults: bool = False,
-        exclude_none: bool = False,
-        encoder: Optional[Callable[[Any], Any]] = None,
-        models_as_dict: bool = True,  # does not seems to be needed by dstack or dependencies
-        **dumps_kwargs: Any,
-    ) -> str:
-        """
-        Override `json()` method so that it calls `dict()`.
-        Allows changing how models are serialized by overriding `dict()` only.
-        By default, `json()` won't call `dict()`, so changes applied in `dict()` won't take place.
-        """
-        data = self.dict(
-            by_alias=by_alias,
-            include=include,
-            exclude=exclude,
-            exclude_unset=exclude_unset,
-            exclude_defaults=exclude_defaults,
-            exclude_none=exclude_none,
-        )
-        if self.__custom_root_type__:
-            data = data["__root__"]
-        return self.__config__.json_dumps(data, default=encoder, **dumps_kwargs)
+# This allows to use the same model both for strict parsing of the user input and
+# for permissive parsing of the server responses.
+#
+# We define a func to generate CoreModel dynamically that can be used
+# to define custom Config for both __request__ and __response__ models.
+# Note: Defining config in the model class directly overrides
+# pydantic-duality's base config, breaking __response__.
+def generate_dual_core_model(
+    custom_config: Union[type, Mapping],
+) -> "type[CoreModel]":
+    class CoreModel(generate_dual_base_model(custom_config)):
+        def json(
+            self,
+            *,
+            include: Optional[IncludeExcludeType] = None,
+            exclude: Optional[IncludeExcludeType] = None,
+            by_alias: bool = False,
+            skip_defaults: Optional[bool] = None,  # ignore as it's deprecated
+            exclude_unset: bool = False,
+            exclude_defaults: bool = False,
+            exclude_none: bool = False,
+            encoder: Optional[Callable[[Any], Any]] = None,
+            models_as_dict: bool = True,  # does not seems to be needed by dstack or dependencies
+            **dumps_kwargs: Any,
+        ) -> str:
+            """
+            Override `json()` method so that it calls `dict()`.
+            Allows changing how models are serialized by overriding `dict()` only.
+            By default, `json()` won't call `dict()`, so changes applied in `dict()` won't take place.
+            """
+            data = self.dict(
+                by_alias=by_alias,
+                include=include,
+                exclude=exclude,
+                exclude_unset=exclude_unset,
+                exclude_defaults=exclude_defaults,
+                exclude_none=exclude_none,
+            )
+            if self.__custom_root_type__:
+                data = data["__root__"]
+            return self.__config__.json_dumps(data, default=encoder, **dumps_kwargs)
+
+    return CoreModel
+
+
+if TYPE_CHECKING:
+
+    class CoreModel(generate_dual_base_model(CoreConfig)):
+        pass
+else:
+    CoreModel = generate_dual_core_model(CoreConfig)
+
+
+class FrozenConfig(CoreConfig):
+    frozen = True
+
+
+FrozenCoreModel = generate_dual_core_model(FrozenConfig)
 
 
 class Duration(int):
@@ -93,7 +120,7 @@ class Duration(int):
         raise ValueError(f"Cannot parse the duration {v}")
 
 
-class RegistryAuth(CoreModel):
+class RegistryAuth(FrozenCoreModel):
     """
     Credentials for pulling a private Docker image.
 
@@ -105,9 +132,6 @@ class RegistryAuth(CoreModel):
     username: Annotated[str, Field(description="The username")]
     password: Annotated[str, Field(description="The password or access token")]
 
-    class Config(CoreModel.Config):
-        frozen = True
-
 
 class ApplyAction(str, Enum):
     CREATE = "create"  # resource is to be created or overridden

dstack/_internal/core/models/configurations.py

@@ -10,12 +10,23 @@ from pydantic import Field, ValidationError, conint, constr, root_validator, val
 from typing_extensions import Self
 
 from dstack._internal.core.errors import ConfigurationError
-from dstack._internal.core.models.common import CoreModel, Duration, RegistryAuth
+from dstack._internal.core.models.common import (
+    CoreConfig,
+    CoreModel,
+    Duration,
+    RegistryAuth,
+    generate_dual_core_model,
+)
 from dstack._internal.core.models.envs import Env
 from dstack._internal.core.models.files import FilePathMapping
 from dstack._internal.core.models.fleets import FleetConfiguration
 from dstack._internal.core.models.gateways import GatewayConfiguration
-from dstack._internal.core.models.profiles import ProfileParams, parse_duration, parse_off_duration
+from dstack._internal.core.models.profiles import (
+    ProfileParams,
+    ProfileParamsConfig,
+    parse_duration,
+    parse_off_duration,
+)
 from dstack._internal.core.models.resources import Range, ResourcesSpec
 from dstack._internal.core.models.services import AnyModel, OpenAIChatModel
 from dstack._internal.core.models.unix import UnixUser
@@ -276,7 +287,20 @@ class HTTPHeaderSpec(CoreModel):
     ]
 
 
-class ProbeConfig(CoreModel):
+class ProbeConfigConfig(CoreConfig):
+    @staticmethod
+    def schema_extra(schema: Dict[str, Any]):
+        add_extra_schema_types(
+            schema["properties"]["timeout"],
+            extra_types=[{"type": "string"}],
+        )
+        add_extra_schema_types(
+            schema["properties"]["interval"],
+            extra_types=[{"type": "string"}],
+        )
+
+
+class ProbeConfig(generate_dual_core_model(ProbeConfigConfig)):
     type: Literal["http"]  # expect other probe types in the future, namely `exec`
     url: Annotated[
         Optional[str], Field(description=f"The URL to request. Defaults to `{DEFAULT_PROBE_URL}`")
@@ -331,18 +355,6 @@ class ProbeConfig(CoreModel):
         ),
     ] = None
 
-    class Config(CoreModel.Config):
-        @staticmethod
-        def schema_extra(schema: Dict[str, Any]):
-            add_extra_schema_types(
-                schema["properties"]["timeout"],
-                extra_types=[{"type": "string"}],
-            )
-            add_extra_schema_types(
-                schema["properties"]["interval"],
-                extra_types=[{"type": "string"}],
-            )
-
     @validator("timeout", pre=True)
     def parse_timeout(cls, v: Optional[Union[int, str]]) -> Optional[int]:
         if v is None:
@@ -381,6 +393,19 @@ class ProbeConfig(CoreModel):
         return values
 
 
+class BaseRunConfigurationConfig(CoreConfig):
+    @staticmethod
+    def schema_extra(schema: Dict[str, Any]):
+        add_extra_schema_types(
+            schema["properties"]["volumes"]["items"],
+            extra_types=[{"type": "string"}],
+        )
+        add_extra_schema_types(
+            schema["properties"]["files"]["items"],
+            extra_types=[{"type": "string"}],
+        )
+
+
 class BaseRunConfiguration(CoreModel):
     type: Literal["none"]
     name: Annotated[
@@ -484,18 +509,6 @@ class BaseRunConfiguration(CoreModel):
     # deprecated since 0.18.31; task, service -- no effect; dev-environment -- executed right before `init`
     setup: CommandsList = []
 
-    class Config(CoreModel.Config):
-        @staticmethod
-        def schema_extra(schema: Dict[str, Any]):
-            add_extra_schema_types(
-                schema["properties"]["volumes"]["items"],
-                extra_types=[{"type": "string"}],
-            )
-            add_extra_schema_types(
-                schema["properties"]["files"]["items"],
-                extra_types=[{"type": "string"}],
-            )
-
     @validator("python", pre=True, always=True)
     def convert_python(cls, v, values) -> Optional[PythonVersion]:
         if v is not None and values.get("image"):
@@ -621,20 +634,25 @@ class DevEnvironmentConfigurationParams(CoreModel):
         return None
 
 
+class DevEnvironmentConfigurationConfig(
+    ProfileParamsConfig,
+    BaseRunConfigurationConfig,
+):
+    @staticmethod
+    def schema_extra(schema: Dict[str, Any]):
+        ProfileParamsConfig.schema_extra(schema)
+        BaseRunConfigurationConfig.schema_extra(schema)
+
+
 class DevEnvironmentConfiguration(
     ProfileParams,
     BaseRunConfiguration,
     ConfigurationWithPortsParams,
     DevEnvironmentConfigurationParams,
+    generate_dual_core_model(DevEnvironmentConfigurationConfig),
 ):
     type: Literal["dev-environment"] = "dev-environment"
 
-    class Config(ProfileParams.Config, BaseRunConfiguration.Config):
-        @staticmethod
-        def schema_extra(schema: Dict[str, Any]):
-            ProfileParams.Config.schema_extra(schema)
-            BaseRunConfiguration.Config.schema_extra(schema)
-
     @validator("entrypoint")
     def validate_entrypoint(cls, v: Optional[str]) -> Optional[str]:
         if v is not None:
@@ -646,20 +664,38 @@ class TaskConfigurationParams(CoreModel):
     nodes: Annotated[int, Field(description="Number of nodes", ge=1)] = 1
 
 
+class TaskConfigurationConfig(
+    ProfileParamsConfig,
+    BaseRunConfigurationConfig,
+):
+    @staticmethod
+    def schema_extra(schema: Dict[str, Any]):
+        ProfileParamsConfig.schema_extra(schema)
+        BaseRunConfigurationConfig.schema_extra(schema)
+
+
 class TaskConfiguration(
     ProfileParams,
     BaseRunConfiguration,
     ConfigurationWithCommandsParams,
     ConfigurationWithPortsParams,
     TaskConfigurationParams,
+    generate_dual_core_model(TaskConfigurationConfig),
 ):
     type: Literal["task"] = "task"
 
-    class Config(ProfileParams.Config, BaseRunConfiguration.Config):
-        @staticmethod
-        def schema_extra(schema: Dict[str, Any]):
-            ProfileParams.Config.schema_extra(schema)
-            BaseRunConfiguration.Config.schema_extra(schema)
+
+class ServiceConfigurationParamsConfig(CoreConfig):
+    @staticmethod
+    def schema_extra(schema: Dict[str, Any]):
+        add_extra_schema_types(
+            schema["properties"]["replicas"],
+            extra_types=[{"type": "integer"}, {"type": "string"}],
+        )
+        add_extra_schema_types(
+            schema["properties"]["model"],
+            extra_types=[{"type": "string"}],
+        )
 
 
 class ServiceConfigurationParams(CoreModel):
@@ -719,18 +755,6 @@ class ServiceConfigurationParams(CoreModel):
         Field(description="List of probes used to determine job health"),
     ] = []
 
-    class Config(CoreModel.Config):
-        @staticmethod
-        def schema_extra(schema: Dict[str, Any]):
-            add_extra_schema_types(
-                schema["properties"]["replicas"],
-                extra_types=[{"type": "integer"}, {"type": "string"}],
-            )
-            add_extra_schema_types(
-                schema["properties"]["model"],
-                extra_types=[{"type": "string"}],
-            )
-
     @validator("port")
     def convert_port(cls, v) -> PortMapping:
         if isinstance(v, int):
@@ -797,25 +821,27 @@ class ServiceConfigurationParams(CoreModel):
         return v
 
 
+class ServiceConfigurationConfig(
+    ProfileParamsConfig,
+    BaseRunConfigurationConfig,
+    ServiceConfigurationParamsConfig,
+):
+    @staticmethod
+    def schema_extra(schema: Dict[str, Any]):
+        ProfileParamsConfig.schema_extra(schema)
+        BaseRunConfigurationConfig.schema_extra(schema)
+        ServiceConfigurationParamsConfig.schema_extra(schema)
+
+
 class ServiceConfiguration(
     ProfileParams,
     BaseRunConfiguration,
     ConfigurationWithCommandsParams,
     ServiceConfigurationParams,
+    generate_dual_core_model(ServiceConfigurationConfig),
 ):
     type: Literal["service"] = "service"
 
-    class Config(
-        ProfileParams.Config,
-        BaseRunConfiguration.Config,
-        ServiceConfigurationParams.Config,
-    ):
-        @staticmethod
-        def schema_extra(schema: Dict[str, Any]):
-            ProfileParams.Config.schema_extra(schema)
-            BaseRunConfiguration.Config.schema_extra(schema)
-            ServiceConfigurationParams.Config.schema_extra(schema)
-
 
 AnyRunConfiguration = Union[DevEnvironmentConfiguration, TaskConfiguration, ServiceConfiguration]
 
@@ -876,7 +902,7 @@ class DstackConfiguration(CoreModel):
         Field(discriminator="type"),
     ]
 
-    class Config(CoreModel.Config):
+    class Config(CoreConfig):
         json_loads = orjson.loads
         json_dumps = pydantic_orjson_dumps_with_indent