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

dstack/_internal/cli/commands/__init__.py

@@ -1,20 +1,22 @@
 import argparse
 import os
+import shlex
 from abc import ABC, abstractmethod
-from typing import List, Optional
+from typing import ClassVar, Optional
 
 from rich_argparse import RichHelpFormatter
 
 from dstack._internal.cli.services.completion import ProjectNameCompleter
-from dstack._internal.cli.utils.common import configure_logging
+from dstack._internal.core.errors import CLIError
 from dstack.api import Client
 
 
 class BaseCommand(ABC):
-    NAME: str = "name the command"
-    DESCRIPTION: str = "describe the command"
-    DEFAULT_HELP: bool = True
-    ALIASES: Optional[List[str]] = None
+    NAME: ClassVar[str] = "name the command"
+    DESCRIPTION: ClassVar[str] = "describe the command"
+    DEFAULT_HELP: ClassVar[bool] = True
+    ALIASES: ClassVar[Optional[list[str]]] = None
+    ACCEPT_EXTRA_ARGS: ClassVar[bool] = False
 
     def __init__(self, parser: argparse.ArgumentParser):
         self._parser = parser
@@ -50,7 +52,8 @@ class BaseCommand(ABC):
 
     @abstractmethod
     def _command(self, args: argparse.Namespace):
-        pass
+        if not self.ACCEPT_EXTRA_ARGS and args.extra_args:
+            raise CLIError(f"Unrecognized arguments: {shlex.join(args.extra_args)}")
 
 
 class APIBaseCommand(BaseCommand):
@@ -65,5 +68,5 @@ class APIBaseCommand(BaseCommand):
         ).completer = ProjectNameCompleter()  # type: ignore[attr-defined]
 
     def _command(self, args: argparse.Namespace):
-        configure_logging()
+        super()._command(args)
         self.api = Client.from_config(project_name=args.project)

dstack/_internal/cli/commands/apply.py

@@ -1,4 +1,5 @@
 import argparse
+import shlex
 
 from argcomplete import FilesCompleter  # type: ignore[attr-defined]
 
@@ -19,6 +20,7 @@ class ApplyCommand(APIBaseCommand):
     NAME = "apply"
     DESCRIPTION = "Apply a configuration"
     DEFAULT_HELP = False
+    ACCEPT_EXTRA_ARGS = True
 
     def _register(self):
         super()._register()
@@ -84,13 +86,14 @@ class ApplyCommand(APIBaseCommand):
             configurator_class = get_apply_configurator_class(configuration.type)
             configurator = configurator_class(api_client=self.api)
             configurator_parser = configurator.get_parser()
-            known, unknown = configurator_parser.parse_known_args(args.unknown)
+            configurator_args, unknown_args = configurator_parser.parse_known_args(args.extra_args)
+            if unknown_args:
+                raise CLIError(f"Unrecognized arguments: {shlex.join(unknown_args)}")
             configurator.apply_configuration(
                 conf=configuration,
                 configuration_path=configuration_path,
                 command_args=args,
-                configurator_args=known,
-                unknown_args=unknown,
+                configurator_args=configurator_args,
             )
         except KeyboardInterrupt:
             console.print("\nOperation interrupted by user. Exiting...")

dstack/_internal/cli/commands/completion.py

@@ -1,3 +1,5 @@
+import argparse
+
 import argcomplete
 
 from dstack._internal.cli.commands import BaseCommand
@@ -15,6 +17,6 @@ class CompletionCommand(BaseCommand):
             choices=["bash", "zsh"],
         )
 
-    def _command(self, args):
+    def _command(self, args: argparse.Namespace):
         super()._command(args)
         print(argcomplete.shellcode(["dstack"], shell=args.shell))  # type: ignore[attr-defined]

dstack/_internal/cli/commands/config.py

@@ -40,6 +40,7 @@ class ConfigCommand(BaseCommand):
         )
 
     def _command(self, args: argparse.Namespace):
+        super()._command(args)
         config_manager = ConfigManager()
         if args.remove:
             config_manager.delete_project(args.project)

dstack/_internal/cli/commands/init.py

@@ -9,7 +9,7 @@ from dstack._internal.cli.services.repos import (
     is_git_repo_url,
     register_init_repo_args,
 )
-from dstack._internal.cli.utils.common import configure_logging, confirm_ask, console, warn
+from dstack._internal.cli.utils.common import confirm_ask, console, warn
 from dstack._internal.core.errors import ConfigurationError
 from dstack._internal.core.models.repos.remote import RemoteRepo
 from dstack._internal.core.services.configs import ConfigManager
@@ -52,7 +52,7 @@ class InitCommand(BaseCommand):
         )
 
     def _command(self, args: argparse.Namespace):
-        configure_logging()
+        super()._command(args)
 
         repo_path: Optional[Path] = None
         repo_url: Optional[str] = None

dstack/_internal/cli/commands/offer.py

@@ -99,7 +99,7 @@ class OfferCommand(APIBaseCommand):
         conf = TaskConfiguration(commands=[":"])
 
         configurator = OfferConfigurator(api_client=self.api)
-        configurator.apply_args(conf, args, [])
+        configurator.apply_args(conf, args)
         profile = load_profile(Path.cwd(), profile_name=args.profile)
 
         run_spec = RunSpec(

dstack/_internal/cli/commands/project.py

@@ -67,6 +67,7 @@ class ProjectCommand(BaseCommand):
         set_default_parser.set_defaults(subfunc=self._set_default)
 
     def _command(self, args: argparse.Namespace):
+        super()._command(args)
         if not hasattr(args, "subfunc"):
             args.subfunc = self._list
         args.subfunc(args)

dstack/_internal/cli/commands/server.py

@@ -1,5 +1,5 @@
+import argparse
 import os
-from argparse import Namespace
 
 from dstack._internal import settings
 from dstack._internal.cli.commands import BaseCommand
@@ -53,7 +53,7 @@ class ServerCommand(BaseCommand):
         )
         self._parser.add_argument("--token", type=str, help="The admin user token")
 
-    def _command(self, args: Namespace):
+    def _command(self, args: argparse.Namespace):
         super()._command(args)
 
         if not UVICORN_INSTALLED:

dstack/_internal/cli/main.py

@@ -83,7 +83,7 @@ def main():
     argcomplete.autocomplete(parser, always_complete_options=False)
 
     args, unknown_args = parser.parse_known_args()
-    args.unknown = unknown_args
+    args.extra_args = unknown_args
 
     try:
         check_for_updates()

dstack/_internal/cli/services/configurators/base.py

@@ -1,7 +1,7 @@
 import argparse
 import os
 from abc import ABC, abstractmethod
-from typing import Generic, List, TypeVar, Union, cast
+from typing import ClassVar, Generic, List, TypeVar, Union, cast
 
 from dstack._internal.cli.services.args import env_var
 from dstack._internal.core.errors import ConfigurationError
@@ -18,7 +18,7 @@ ApplyConfigurationT = TypeVar("ApplyConfigurationT", bound=AnyApplyConfiguration
 
 
 class BaseApplyConfigurator(ABC, Generic[ApplyConfigurationT]):
-    TYPE: ApplyConfigurationType
+    TYPE: ClassVar[ApplyConfigurationType]
 
     def __init__(self, api_client: Client):
         self.api = api_client
@@ -30,7 +30,6 @@ class BaseApplyConfigurator(ABC, Generic[ApplyConfigurationT]):
         configuration_path: str,
         command_args: argparse.Namespace,
         configurator_args: argparse.Namespace,
-        unknown_args: List[str],
     ):
         """
         Implements `dstack apply` for a given configuration type.
@@ -40,7 +39,6 @@ class BaseApplyConfigurator(ABC, Generic[ApplyConfigurationT]):
             configuration_path: The path to the configuration file.
             command_args: The args parsed by `dstack apply`.
             configurator_args: The known args parsed by `cls.get_parser()`.
-            unknown_args: The unknown args after parsing by `cls.get_parser()`.
         """
         pass
 

dstack/_internal/cli/services/configurators/fleet.py

@@ -1,7 +1,7 @@
 import argparse
 import time
 from pathlib import Path
-from typing import List, Optional
+from typing import Optional
 
 from rich.table import Table
 
@@ -46,7 +46,7 @@ logger = get_logger(__name__)
 
 
 class FleetConfigurator(ApplyEnvVarsConfiguratorMixin, BaseApplyConfigurator[FleetConfiguration]):
-    TYPE: ApplyConfigurationType = ApplyConfigurationType.FLEET
+    TYPE = ApplyConfigurationType.FLEET
 
     def apply_configuration(
         self,
@@ -54,9 +54,8 @@ class FleetConfigurator(ApplyEnvVarsConfiguratorMixin, BaseApplyConfigurator[Fle
         configuration_path: str,
         command_args: argparse.Namespace,
         configurator_args: argparse.Namespace,
-        unknown_args: List[str],
     ):
-        self.apply_args(conf, configurator_args, unknown_args)
+        self.apply_args(conf, configurator_args)
         profile = load_profile(Path.cwd(), None)
         spec = FleetSpec(
             configuration=conf,
@@ -309,7 +308,7 @@ class FleetConfigurator(ApplyEnvVarsConfiguratorMixin, BaseApplyConfigurator[Fle
         )
         cls.register_env_args(configuration_group)
 
-    def apply_args(self, conf: FleetConfiguration, args: argparse.Namespace, unknown: List[str]):
+    def apply_args(self, conf: FleetConfiguration, args: argparse.Namespace):
         if args.name:
             conf.name = args.name
         self.apply_env_vars(conf.env, args)

dstack/_internal/cli/services/configurators/gateway.py

@@ -1,6 +1,5 @@
 import argparse
 import time
-from typing import List
 
 from rich.table import Table
 
@@ -27,7 +26,7 @@ from dstack.api._public import Client
 
 
 class GatewayConfigurator(BaseApplyConfigurator[GatewayConfiguration]):
-    TYPE: ApplyConfigurationType = ApplyConfigurationType.GATEWAY
+    TYPE = ApplyConfigurationType.GATEWAY
 
     def apply_configuration(
         self,
@@ -35,9 +34,8 @@ class GatewayConfigurator(BaseApplyConfigurator[GatewayConfiguration]):
         configuration_path: str,
         command_args: argparse.Namespace,
         configurator_args: argparse.Namespace,
-        unknown_args: List[str],
     ):
-        self.apply_args(conf, configurator_args, unknown_args)
+        self.apply_args(conf, configurator_args)
         spec = GatewaySpec(
             configuration=conf,
             configuration_path=configuration_path,
@@ -179,7 +177,7 @@ class GatewayConfigurator(BaseApplyConfigurator[GatewayConfiguration]):
             help="The gateway name",
         )
 
-    def apply_args(self, conf: GatewayConfiguration, args: argparse.Namespace, unknown: List[str]):
+    def apply_args(self, conf: GatewayConfiguration, args: argparse.Namespace):
         if args.name:
             conf.name = args.name
 

dstack/_internal/cli/services/configurators/run.py

@@ -1,4 +1,5 @@
 import argparse
+import shlex
 import subprocess
 import sys
 import time
@@ -35,6 +36,7 @@ from dstack._internal.core.models.configurations import (
     LEGACY_REPO_DIR,
     AnyRunConfiguration,
     ApplyConfigurationType,
+    ConfigurationWithCommandsParams,
     ConfigurationWithPortsParams,
     DevEnvironmentConfiguration,
     PortMapping,
@@ -80,20 +82,17 @@ class BaseRunConfigurator(
     ApplyEnvVarsConfiguratorMixin,
     BaseApplyConfigurator[RunConfigurationT],
 ):
-    TYPE: ApplyConfigurationType
-
     def apply_configuration(
         self,
         conf: RunConfigurationT,
         configuration_path: str,
         command_args: argparse.Namespace,
         configurator_args: argparse.Namespace,
-        unknown_args: List[str],
     ):
         if configurator_args.repo and configurator_args.no_repo:
             raise CLIError("Either --repo or --no-repo can be specified")
 
-        self.apply_args(conf, configurator_args, unknown_args)
+        self.apply_args(conf, configurator_args)
         self.validate_gpu_vendor_and_image(conf)
         self.validate_cpu_arch_and_image(conf)
 
@@ -395,7 +394,7 @@ class BaseRunConfigurator(
         )
         register_init_repo_args(repo_group)
 
-    def apply_args(self, conf: RunConfigurationT, args: argparse.Namespace, unknown: List[str]):
+    def apply_args(self, conf: RunConfigurationT, args: argparse.Namespace):
         apply_profile_args(args, conf)
         if args.run_name:
             conf.name = args.run_name
@@ -408,16 +407,6 @@ class BaseRunConfigurator(
 
         self.apply_env_vars(conf.env, args)
         self.interpolate_env(conf)
-        self.interpolate_run_args(conf.setup, unknown)
-
-    def interpolate_run_args(self, value: List[str], unknown):
-        run_args = " ".join(unknown)
-        interpolator = VariablesInterpolator({"run": {"args": run_args}}, skip=["secrets"])
-        try:
-            for i in range(len(value)):
-                value[i] = interpolator.interpolate_or_error(value[i])
-        except InterpolatorError as e:
-            raise ConfigurationError(e.args[0])
 
     def interpolate_env(self, conf: RunConfigurationT):
         env_dict = conf.env.as_dict()
@@ -701,18 +690,50 @@ class RunWithPortsConfiguratorMixin:
             conf.ports = list(_merge_ports(conf.ports, args.ports).values())
 
 
-class TaskConfigurator(RunWithPortsConfiguratorMixin, BaseRunConfigurator):
+class RunWithCommandsConfiguratorMixin:
+    @classmethod
+    def register_commands_args(cls, parser: argparse.ArgumentParser):
+        parser.add_argument(
+            "run_args",
+            help=(
+                "Run arguments. Available in the configuration [code]commands[/code] as"
+                " [code]${{ run.args }}[/code]."
+                " Use [code]--[/code] to separate run options from [code]dstack[/code] options"
+            ),
+            nargs="*",
+            metavar="RUN_ARGS",
+        )
+
+    def apply_commands_args(
+        self,
+        conf: ConfigurationWithCommandsParams,
+        args: argparse.Namespace,
+    ):
+        commands = conf.commands
+        run_args = shlex.join(args.run_args)
+        interpolator = VariablesInterpolator({"run": {"args": run_args}}, skip=["secrets"])
+        try:
+            for i, command in enumerate(commands):
+                commands[i] = interpolator.interpolate_or_error(command)
+        except InterpolatorError as e:
+            raise ConfigurationError(e.args[0])
+
+
+class TaskConfigurator(
+    RunWithPortsConfiguratorMixin, RunWithCommandsConfiguratorMixin, BaseRunConfigurator
+):
     TYPE = ApplyConfigurationType.TASK
 
     @classmethod
     def register_args(cls, parser: argparse.ArgumentParser):
         super().register_args(parser)
         cls.register_ports_args(parser)
+        cls.register_commands_args(parser)
 
-    def apply_args(self, conf: TaskConfiguration, args: argparse.Namespace, unknown: List[str]):
-        super().apply_args(conf, args, unknown)
+    def apply_args(self, conf: TaskConfiguration, args: argparse.Namespace):
+        super().apply_args(conf, args)
         self.apply_ports_args(conf, args)
-        self.interpolate_run_args(conf.commands, unknown)
+        self.apply_commands_args(conf, args)
 
 
 class DevEnvironmentConfigurator(RunWithPortsConfiguratorMixin, BaseRunConfigurator):
@@ -723,10 +744,8 @@ class DevEnvironmentConfigurator(RunWithPortsConfiguratorMixin, BaseRunConfigura
         super().register_args(parser)
         cls.register_ports_args(parser)
 
-    def apply_args(
-        self, conf: DevEnvironmentConfiguration, args: argparse.Namespace, unknown: List[str]
-    ):
-        super().apply_args(conf, args, unknown)
+    def apply_args(self, conf: DevEnvironmentConfiguration, args: argparse.Namespace):
+        super().apply_args(conf, args)
         self.apply_ports_args(conf, args)
         if conf.ide == "vscode" and conf.version is None:
             conf.version = _detect_vscode_version()
@@ -746,12 +765,17 @@ class DevEnvironmentConfigurator(RunWithPortsConfiguratorMixin, BaseRunConfigura
                 )
 
 
-class ServiceConfigurator(BaseRunConfigurator):
+class ServiceConfigurator(RunWithCommandsConfiguratorMixin, BaseRunConfigurator):
     TYPE = ApplyConfigurationType.SERVICE
 
-    def apply_args(self, conf: ServiceConfiguration, args: argparse.Namespace, unknown: List[str]):
-        super().apply_args(conf, args, unknown)
-        self.interpolate_run_args(conf.commands, unknown)
+    @classmethod
+    def register_args(cls, parser: argparse.ArgumentParser):
+        super().register_args(parser)
+        cls.register_commands_args(parser)
+
+    def apply_args(self, conf: TaskConfiguration, args: argparse.Namespace):
+        super().apply_args(conf, args)
+        self.apply_commands_args(conf, args)
 
 
 def _merge_ports(conf: List[PortMapping], args: List[PortMapping]) -> Dict[int, PortMapping]:

dstack/_internal/cli/services/configurators/volume.py

@@ -1,6 +1,5 @@
 import argparse
 import time
-from typing import List
 
 from rich.table import Table
 
@@ -26,7 +25,7 @@ from dstack.api._public import Client
 
 
 class VolumeConfigurator(BaseApplyConfigurator[VolumeConfiguration]):
-    TYPE: ApplyConfigurationType = ApplyConfigurationType.VOLUME
+    TYPE = ApplyConfigurationType.VOLUME
 
     def apply_configuration(
         self,
@@ -34,9 +33,8 @@ class VolumeConfigurator(BaseApplyConfigurator[VolumeConfiguration]):
         configuration_path: str,
         command_args: argparse.Namespace,
         configurator_args: argparse.Namespace,
-        unknown_args: List[str],
     ):
-        self.apply_args(conf, configurator_args, unknown_args)
+        self.apply_args(conf, configurator_args)
         spec = VolumeSpec(
             configuration=conf,
             configuration_path=configuration_path,
@@ -167,7 +165,7 @@ class VolumeConfigurator(BaseApplyConfigurator[VolumeConfiguration]):
             help="The volume name",
         )
 
-    def apply_args(self, conf: VolumeConfiguration, args: argparse.Namespace, unknown: List[str]):
+    def apply_args(self, conf: VolumeConfiguration, args: argparse.Namespace):
         if args.name:
             conf.name = args.name
 

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

@@ -1,6 +1,6 @@
 import threading
 from concurrent.futures import ThreadPoolExecutor, as_completed
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Callable, Dict, List, Optional, Tuple
 
 import boto3
 import botocore.client
@@ -18,6 +18,7 @@ from dstack._internal.core.backends.aws.models import (
 )
 from dstack._internal.core.backends.base.compute import (
     Compute,
+    ComputeWithAllOffersCached,
     ComputeWithCreateInstanceSupport,
     ComputeWithGatewaySupport,
     ComputeWithMultinodeSupport,
@@ -32,7 +33,7 @@ from dstack._internal.core.backends.base.compute import (
     get_user_data,
     merge_tags,
 )
-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.errors import (
     ComputeError,
     NoCapacityError,
@@ -87,6 +88,7 @@ def _ec2client_cache_methodkey(self, ec2_client, *args, **kwargs):
 
 
 class AWSCompute(
+    ComputeWithAllOffersCached,
     ComputeWithCreateInstanceSupport,
     ComputeWithMultinodeSupport,
     ComputeWithReservationSupport,
@@ -109,6 +111,8 @@ class AWSCompute(
         # Caches to avoid redundant API calls when provisioning many instances
         # get_offers is already cached but we still cache its sub-functions
         # with more aggressive/longer caches.
+        self._offers_post_filter_cache_lock = threading.Lock()
+        self._offers_post_filter_cache = TTLCache(maxsize=10, ttl=180)
         self._get_regions_to_quotas_cache_lock = threading.Lock()
         self._get_regions_to_quotas_execution_lock = threading.Lock()
         self._get_regions_to_quotas_cache = TTLCache(maxsize=10, ttl=300)
@@ -125,43 +129,11 @@ class AWSCompute(
         self._get_image_id_and_username_cache_lock = threading.Lock()
         self._get_image_id_and_username_cache = TTLCache(maxsize=100, ttl=600)
 
-    def get_offers(
-        self, requirements: Optional[Requirements] = None
-    ) -> List[InstanceOfferWithAvailability]:
-        filter = _supported_instances
-        if requirements and requirements.reservation:
-            region_to_reservation = {}
-            for region in self.config.regions:
-                reservation = aws_resources.get_reservation(
-                    ec2_client=self.session.client("ec2", region_name=region),
-                    reservation_id=requirements.reservation,
-                    instance_count=1,
-                )
-                if reservation is not None:
-                    region_to_reservation[region] = reservation
-
-            def _supported_instances_with_reservation(offer: InstanceOffer) -> bool:
-                # Filter: only instance types supported by dstack
-                if not _supported_instances(offer):
-                    return False
-                # Filter: Spot instances can't be used with reservations
-                if offer.instance.resources.spot:
-                    return False
-                region = offer.region
-                reservation = region_to_reservation.get(region)
-                # Filter: only instance types matching the capacity reservation
-                if not bool(reservation and offer.instance.name == reservation["InstanceType"]):
-                    return False
-                return True
-
-            filter = _supported_instances_with_reservation
-
+    def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability]:
         offers = get_catalog_offers(
             backend=BackendType.AWS,
             locations=self.config.regions,
-            requirements=requirements,
-            configurable_disk_size=CONFIGURABLE_DISK_SIZE,
-            extra_filter=filter,
+            extra_filter=_supported_instances,
         )
         regions = list(set(i.region for i in offers))
         with self._get_regions_to_quotas_execution_lock:
@@ -185,6 +157,49 @@ class AWSCompute(
             )
         return availability_offers
 
+    def get_offers_modifier(
+        self, requirements: Requirements
+    ) -> Callable[[InstanceOfferWithAvailability], Optional[InstanceOfferWithAvailability]]:
+        return get_offers_disk_modifier(CONFIGURABLE_DISK_SIZE, requirements)
+
+    def _get_offers_cached_key(self, requirements: Requirements) -> int:
+        # Requirements is not hashable, so we use a hack to get arguments hash
+        return hash(requirements.json())
+
+    @cachedmethod(
+        cache=lambda self: self._offers_post_filter_cache,
+        key=_get_offers_cached_key,
+        lock=lambda self: self._offers_post_filter_cache_lock,
+    )
+    def get_offers_post_filter(
+        self, requirements: Requirements
+    ) -> Optional[Callable[[InstanceOfferWithAvailability], bool]]:
+        if requirements.reservation:
+            region_to_reservation = {}
+            for region in get_or_error(self.config.regions):
+                reservation = aws_resources.get_reservation(
+                    ec2_client=self.session.client("ec2", region_name=region),
+                    reservation_id=requirements.reservation,
+                    instance_count=1,
+                )
+                if reservation is not None:
+                    region_to_reservation[region] = reservation
+
+            def reservation_filter(offer: InstanceOfferWithAvailability) -> bool:
+                # Filter: Spot instances can't be used with reservations
+                if offer.instance.resources.spot:
+                    return False
+                region = offer.region
+                reservation = region_to_reservation.get(region)
+                # Filter: only instance types matching the capacity reservation
+                if not bool(reservation and offer.instance.name == reservation["InstanceType"]):
+                    return False
+                return True
+
+            return reservation_filter
+
+        return None
+
     def terminate_instance(
         self, instance_id: str, region: str, backend_data: Optional[str] = None
     ) -> None:

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

@@ -2,7 +2,7 @@ import base64
 import enum
 import re
 from concurrent.futures import ThreadPoolExecutor, as_completed
-from typing import Dict, List, Optional, Tuple
+from typing import Callable, Dict, List, Optional, Tuple
 
 from azure.core.credentials import TokenCredential
 from azure.core.exceptions import ResourceExistsError, ResourceNotFoundError
@@ -39,6 +39,7 @@ from dstack._internal.core.backends.azure import utils as azure_utils
 from dstack._internal.core.backends.azure.models import AzureConfig
 from dstack._internal.core.backends.base.compute import (
     Compute,
+    ComputeWithAllOffersCached,
     ComputeWithCreateInstanceSupport,
     ComputeWithGatewaySupport,
     ComputeWithMultinodeSupport,
@@ -48,7 +49,7 @@ from dstack._internal.core.backends.base.compute import (
     get_user_data,
     merge_tags,
 )
-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.errors import ComputeError, NoCapacityError
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.gateways import (
@@ -73,6 +74,7 @@ CONFIGURABLE_DISK_SIZE = Range[Memory](min=Memory.parse("30GB"), max=Memory.pars
 
 
 class AzureCompute(
+    ComputeWithAllOffersCached,
     ComputeWithCreateInstanceSupport,
     ComputeWithMultinodeSupport,
     ComputeWithGatewaySupport,
@@ -89,14 +91,10 @@ class AzureCompute(
             credential=credential, subscription_id=config.subscription_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.AZURE,
             locations=self.config.regions,
-            requirements=requirements,
-            configurable_disk_size=CONFIGURABLE_DISK_SIZE,
             extra_filter=_supported_instances,
         )
         offers_with_availability = _get_offers_with_availability(
@@ -106,6 +104,11 @@ class AzureCompute(
         )
         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 create_instance(
         self,
         instance_offer: InstanceOfferWithAvailability,