dstack 0.18.44__tar.gz → 0.19.0__tar.gz

{dstack-0.18.44 → dstack-0.19.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dstack
-Version: 0.18.44
+Version: 0.19.0
 Summary: dstack is an open-source orchestration engine for running AI workloads on any cloud or on-premises.
 Home-page: https://dstack.ai
 Author: Andrey Cheptsov

{dstack-0.18.44 → dstack-0.19.0}/src/dstack/_internal/cli/commands/gateway.py

@@ -13,6 +13,10 @@ from dstack._internal.cli.utils.common import (
 )
 from dstack._internal.cli.utils.gateway import get_gateways_table, print_gateways_table
 from dstack._internal.core.models.backends.base import BackendType
+from dstack._internal.core.models.gateways import GatewayConfiguration
+from dstack._internal.utils.logging import get_logger
+
+logger = get_logger(__name__)
 
 
 class GatewayCommand(APIBaseCommand):
@@ -41,7 +45,9 @@ class GatewayCommand(APIBaseCommand):
             )
 
         create_parser = subparsers.add_parser(
-            "create", help="Add a gateway", formatter_class=self._parser.formatter_class
+            "create",
+            help="Add a gateway. Deprecated in favor of `dstack apply` with gateway configuration.",
+            formatter_class=self._parser.formatter_class,
         )
         create_parser.set_defaults(subfunc=self._create)
         create_parser.add_argument(
@@ -100,10 +106,16 @@ class GatewayCommand(APIBaseCommand):
             pass
 
     def _create(self, args: argparse.Namespace):
+        logger.warning(
+            "`dstack gateway create` is deperecated in favor of `dstack apply` with gateway configurations."
+        )
         with console.status("Creating gateway..."):
-            gateway = self.api.client.gateways.create(
-                self.api.project, args.name, BackendType(args.backend), args.region
+            configuration = GatewayConfiguration(
+                name=args.name,
+                backend=BackendType(args.backend),
+                region=args.region,
             )
+            gateway = self.api.client.gateways.create(self.api.project, configuration)
             if args.set_default:
                 self.api.client.gateways.set_default(self.api.project, gateway.name)
             if args.domain:

{dstack-0.18.44 → dstack-0.19.0}/src/dstack/_internal/cli/commands/logs.py

@@ -1,6 +1,5 @@
 import argparse
 import sys
-from pathlib import Path
 
 from dstack._internal.cli.commands import APIBaseCommand
 from dstack._internal.cli.services.completion import RunNameCompleter
@@ -19,19 +18,6 @@ class LogsCommand(APIBaseCommand):
         self._parser.add_argument(
             "-d", "--diagnose", action="store_true", help="Show run diagnostic logs"
         )
-        self._parser.add_argument(
-            "-a",
-            "--attach",
-            action="store_true",
-            help="Set up an SSH tunnel and print logs as they follow",
-        )
-        self._parser.add_argument(
-            "--ssh-identity",
-            metavar="SSH_PRIVATE_KEY",
-            help="The private SSH key path for SSH tunneling",
-            type=Path,
-            dest="ssh_identity_file",
-        )
         self._parser.add_argument(
             "--replica",
             help="The replica number. Defaults to 0.",
@@ -51,14 +37,6 @@ class LogsCommand(APIBaseCommand):
         run = self.api.runs.get(args.run_name)
         if run is None:
             raise CLIError(f"Run {args.run_name} not found")
-        if not args.diagnose and args.attach:
-            if run.status.is_finished():
-                raise CLIError(f"Run {args.run_name} is finished")
-            else:
-                logger.warning(
-                    "`dstack logs --attach` is deprecated in favor of `dstack attach --logs`"
-                )
-                run.attach(args.ssh_identity_file)
         logs = run.logs(
             diagnose=args.diagnose,
             replica_num=args.replica,

{dstack-0.18.44 → dstack-0.19.0}/src/dstack/_internal/cli/commands/stats.py

@@ -2,7 +2,6 @@ import argparse
 import time
 from typing import Any, Dict, List, Optional, Union
 
-import requests
 from rich.live import Live
 from rich.table import Table
 
@@ -64,22 +63,14 @@ class StatsCommand(APIBaseCommand):
 
 def _get_run_jobs_metrics(api: Client, run: Run) -> List[JobMetrics]:
     metrics = []
-    try:
-        for job in run._run.jobs:
-            job_metrics = api.client.metrics.get_job_metrics(
-                project_name=api.project,
-                run_name=run.name,
-                replica_num=job.job_spec.replica_num,
-                job_num=job.job_spec.job_num,
-            )
-            metrics.append(job_metrics)
-    except requests.exceptions.HTTPError as e:
-        if e.response.status_code == 404:
-            raise CLIError(
-                "Metrics API is not supported for server versions before 0.18.18. "
-                "Update the server to use `dstack stats`."
-            )
-        raise
+    for job in run._run.jobs:
+        job_metrics = api.client.metrics.get_job_metrics(
+            project_name=api.project,
+            run_name=run.name,
+            replica_num=job.job_spec.replica_num,
+            job_num=job.job_spec.job_num,
+        )
+        metrics.append(job_metrics)
     return metrics
 
 

{dstack-0.18.44 → dstack-0.19.0}/src/dstack/_internal/cli/main.py

@@ -13,9 +13,7 @@ from dstack._internal.cli.commands.fleet import FleetCommand
 from dstack._internal.cli.commands.gateway import GatewayCommand
 from dstack._internal.cli.commands.init import InitCommand
 from dstack._internal.cli.commands.logs import LogsCommand
-from dstack._internal.cli.commands.pool import PoolCommand
 from dstack._internal.cli.commands.ps import PsCommand
-from dstack._internal.cli.commands.run import RunCommand
 from dstack._internal.cli.commands.server import ServerCommand
 from dstack._internal.cli.commands.stats import StatsCommand
 from dstack._internal.cli.commands.stop import StopCommand
@@ -40,7 +38,7 @@ def main():
     parser = argparse.ArgumentParser(
         description=(
             "Not sure where to start? Call [code]dstack init[/].\n"
-            "Define a [code].dstack.yml[/] configuration file and run it via [code]dstack run[/]\n"
+            "Define a [code].dstack.yml[/] configuration file and run it via [code]dstack apply[/]\n"
         ),
         formatter_class=RichHelpFormatter,
         epilog=(
@@ -65,11 +63,9 @@ def main():
     DeleteCommand.register(subparsers)
     FleetCommand.register(subparsers)
     GatewayCommand.register(subparsers)
-    PoolCommand.register(subparsers)
     InitCommand.register(subparsers)
     LogsCommand.register(subparsers)
     PsCommand.register(subparsers)
-    RunCommand.register(subparsers)
     ServerCommand.register(subparsers)
     StatsCommand.register(subparsers)
     StopCommand.register(subparsers)

{dstack-0.18.44 → dstack-0.19.0}/src/dstack/_internal/cli/services/configurators/fleet.py

@@ -3,7 +3,6 @@ import time
 from pathlib import Path
 from typing import List, Optional
 
-import requests
 from rich.table import Table
 
 from dstack._internal.cli.services.configurators.base import (
@@ -32,7 +31,6 @@ from dstack._internal.core.models.repos.base import Repo
 from dstack._internal.utils.common import local_time
 from dstack._internal.utils.logging import get_logger
 from dstack._internal.utils.ssh import convert_ssh_key_to_pem, generate_public_key, pkey_from_str
-from dstack.api._public import Client
 from dstack.api.utils import load_profile
 
 logger = get_logger(__name__)
@@ -60,7 +58,10 @@ class FleetConfigurator(ApplyEnvVarsConfiguratorMixin, BaseApplyConfigurator):
         _preprocess_spec(spec)
 
         with console.status("Getting apply plan..."):
-            plan = _get_plan(api=self.api, spec=spec)
+            plan = self.api.client.fleets.get_plan(
+                project_name=self.api.project,
+                spec=spec,
+            )
         _print_plan_header(plan)
 
         action_message = ""
@@ -234,42 +235,6 @@ def _resolve_ssh_key(ssh_key_path: Optional[str]) -> Optional[SSHKey]:
         exit()
 
 
-def _get_plan(api: Client, spec: FleetSpec) -> FleetPlan:
-    try:
-        return api.client.fleets.get_plan(
-            project_name=api.project,
-            spec=spec,
-        )
-    except requests.exceptions.HTTPError as e:
-        # Handle older server versions that do not have /get_plan for fleets
-        # TODO: Can be removed in 0.19
-        if e.response.status_code == 405:
-            logger.warning(
-                "Fleet apply plan is not fully supported before 0.18.17. "
-                "Update the server to view full-featured apply plan."
-            )
-            user = api.client.users.get_my_user()
-            spec.configuration_path = None
-            current_resource = None
-            if spec.configuration.name is not None:
-                try:
-                    current_resource = api.client.fleets.get(
-                        project_name=api.project, name=spec.configuration.name
-                    )
-                except ResourceNotExistsError:
-                    pass
-            return FleetPlan(
-                project_name=api.project,
-                user=user.username,
-                spec=spec,
-                current_resource=current_resource,
-                offers=[],
-                total_offers=0,
-                max_offer_price=0,
-            )
-        raise e
-
-
 def _print_plan_header(plan: FleetPlan):
     def th(s: str) -> str:
         return f"[bold]{s}[/bold]"

{dstack-0.18.44 → dstack-0.19.0}/src/dstack/_internal/cli/services/configurators/run.py

@@ -85,28 +85,11 @@ class BaseRunConfigurator(ApplyEnvVarsConfiguratorMixin, BaseApplyConfigurator):
             )
         profile = load_profile(Path.cwd(), configurator_args.profile)
         with console.status("Getting apply plan..."):
-            run_plan = self.api.runs.get_plan(
+            run_plan = self.api.runs.get_run_plan(
                 configuration=conf,
                 repo=repo,
                 configuration_path=configuration_path,
-                backends=profile.backends,
-                regions=profile.regions,
-                instance_types=profile.instance_types,
-                reservation=profile.reservation,
-                spot_policy=profile.spot_policy,
-                retry_policy=profile.retry_policy,
-                utilization_policy=profile.utilization_policy,
-                max_duration=profile.max_duration,
-                stop_duration=profile.stop_duration,
-                max_price=profile.max_price,
-                working_dir=conf.working_dir,
-                run_name=conf.name,
-                pool_name=profile.pool_name,
-                instance_name=profile.instance_name,
-                creation_policy=profile.creation_policy,
-                termination_policy=profile.termination_policy,
-                termination_policy_idle=profile.termination_idle_time,
-                idle_duration=profile.idle_duration,
+                profile=profile,
             )
 
         print_run_plan(run_plan, offers_limit=configurator_args.max_offers)
@@ -165,8 +148,8 @@ class BaseRunConfigurator(ApplyEnvVarsConfiguratorMixin, BaseApplyConfigurator):
 
         try:
             with console.status("Applying plan..."):
-                run = self.api.runs.exec_plan(
-                    run_plan, repo, reserve_ports=not command_args.detach
+                run = self.api.runs.apply_plan(
+                    run_plan=run_plan, repo=repo, reserve_ports=not command_args.detach
                 )
         except ServerClientError as e:
             raise CLIError(e.msg)
@@ -459,6 +442,14 @@ class DevEnvironmentConfigurator(RunWithPortsConfigurator):
                     "Fix by opening [code]Command Palette[/code], executing [code]Shell Command: "
                     "Install 'code' command in PATH[/code], and restarting terminal.[/]\n"
                 )
+        if conf.ide == "cursor" and conf.version is None:
+            conf.version = _detect_cursor_version()
+            if conf.version is None:
+                console.print(
+                    "[secondary]Unable to detect the Cursor version and pre-install extensions. "
+                    "Fix by opening [code]Command Palette[/code], executing [code]Shell Command: "
+                    "Install 'cursor' command in PATH[/code], and restarting terminal.[/]\n"
+                )
 
 
 class ServiceConfigurator(BaseRunConfigurator):
@@ -497,6 +488,16 @@ def _detect_vscode_version(exe: str = "code") -> Optional[str]:
     return None
 
 
+def _detect_cursor_version(exe: str = "cursor") -> Optional[str]:
+    try:
+        run = subprocess.run([exe, "--version"], capture_output=True)
+    except FileNotFoundError:
+        return None
+    if run.returncode == 0:
+        return run.stdout.decode().split("\n")[1].strip()
+    return None
+
+
 def _print_service_urls(run: Run) -> None:
     if run._run.run_spec.configuration.type != RunConfigurationType.SERVICE.value:
         return

{dstack-0.18.44 → dstack-0.19.0}/src/dstack/_internal/cli/services/profile.py

@@ -4,21 +4,18 @@ from typing import Union
 
 from dstack._internal.core.models.configurations import AnyRunConfiguration
 from dstack._internal.core.models.profiles import (
-    DEFAULT_INSTANCE_RETRY_DURATION,
-    DEFAULT_POOL_TERMINATION_IDLE_TIME,
     CreationPolicy,
     Profile,
-    ProfileRetryPolicy,
+    ProfileRetry,
     SpotPolicy,
-    TerminationPolicy,
     parse_duration,
     parse_max_duration,
 )
 
 
-def register_profile_args(parser: argparse.ArgumentParser, pool_add: bool = False):
+def register_profile_args(parser: argparse.ArgumentParser):
     """
-    Registers `parser` with `dstack run` and `dstack pool add`
+    Registers `parser` with `dstack apply` run configuration
     CLI arguments that override `profiles.yml` settings.
     """
     profile_group = parser.add_argument_group("Profile")
@@ -36,14 +33,13 @@ def register_profile_args(parser: argparse.ArgumentParser, pool_add: bool = Fals
         help="The maximum price per hour, in dollars",
         dest="max_price",
     )
-    if not pool_add:
-        profile_group.add_argument(
-            "--max-duration",
-            type=max_duration,
-            dest="max_duration",
-            help="The maximum duration of the run",
-            metavar="DURATION",
-        )
+    profile_group.add_argument(
+        "--max-duration",
+        type=max_duration,
+        dest="max_duration",
+        help="The maximum duration of the run",
+        metavar="DURATION",
+    )
     profile_group.add_argument(
         "-b",
         "--backend",
@@ -67,42 +63,28 @@ def register_profile_args(parser: argparse.ArgumentParser, pool_add: bool = Fals
         dest="instance_types",
         help="The cloud-specific instance types that will be tried for provisioning",
     )
-    if pool_add:
-        pools_group_exc = parser
-    else:
-        pools_group = parser.add_argument_group("Pools")
-        pools_group_exc = pools_group.add_mutually_exclusive_group()
-    pools_group_exc.add_argument(
-        "--pool",
-        dest="pool_name",
-        help="The name of the pool. If not set, the default pool will be used",
-    )
-    pools_group_exc.add_argument(
+
+    fleets_group = parser.add_argument_group("Fleets")
+    fleets_group_exc = fleets_group.add_mutually_exclusive_group()
+    fleets_group_exc.add_argument(
         "-R",
         "--reuse",
         dest="creation_policy_reuse",
         action="store_true",
-        help="Reuse instance from pool",
+        help="Reuse an existing instance from fleet (do not provision a new one)",
     )
-    pools_group_exc.add_argument(
+    fleets_group_exc.add_argument(
         "--dont-destroy",
         dest="dont_destroy",
         action="store_true",
-        help="Do not destroy instance after the run is finished",
+        help="Do not destroy instance after the run is finished (if the run provisions a new instance)",
     )
-    pools_group_exc.add_argument(
+    fleets_group_exc.add_argument(
         "--idle-duration",
         dest="idle_duration",
         type=str,
-        help="Time to wait before destroying the idle instance",
+        help="Time to wait before destroying the idle instance (if the run provisions a new instance)",
     )
-    if not pool_add:
-        pools_group_exc.add_argument(
-            "--instance",
-            dest="instance_name",
-            metavar="NAME",
-            help="Reuse instance from pool with name [code]NAME[/]",
-        )
 
     spot_group = parser.add_argument_group("Spot policy")
     spot_group_exc = spot_group.add_mutually_exclusive_group()
@@ -137,10 +119,8 @@ def register_profile_args(parser: argparse.ArgumentParser, pool_add: bool = Fals
 
     retry_group = parser.add_argument_group("Retry policy")
     retry_group_exc = retry_group.add_mutually_exclusive_group()
-    retry_group_exc.add_argument("--retry", action="store_const", dest="retry_policy", const=True)
-    retry_group_exc.add_argument(
-        "--no-retry", action="store_const", dest="retry_policy", const=False
-    )
+    retry_group_exc.add_argument("--retry", action="store_const", dest="retry", const=True)
+    retry_group_exc.add_argument("--no-retry", action="store_const", dest="retry", const=False)
     retry_group_exc.add_argument(
         "--retry-duration", type=retry_duration, dest="retry_duration", metavar="DURATION"
     )
@@ -149,7 +129,6 @@ def register_profile_args(parser: argparse.ArgumentParser, pool_add: bool = Fals
 def apply_profile_args(
     args: argparse.Namespace,
     profile_settings: Union[Profile, AnyRunConfiguration],
-    pool_add: bool = False,
 ):
     """
     Overrides `profile_settings` settings with arguments registered by `register_profile_args()`.
@@ -165,53 +144,25 @@ def apply_profile_args(
         profile_settings.instance_types = args.instance_types
     if args.max_price is not None:
         profile_settings.max_price = args.max_price
-    if not pool_add:
-        if args.max_duration is not None:
-            profile_settings.max_duration = args.max_duration
-
-    if args.pool_name:
-        profile_settings.pool_name = args.pool_name
+    if args.max_duration is not None:
+        profile_settings.max_duration = args.max_duration
 
     if args.idle_duration is not None:
-        profile_settings.termination_idle_time = args.idle_duration
-    if pool_add and args.idle_duration is None:
-        profile_settings.termination_idle_time = DEFAULT_POOL_TERMINATION_IDLE_TIME
-
-    if args.dont_destroy:
-        profile_settings.termination_policy = TerminationPolicy.DONT_DESTROY
-    if not pool_add:
-        if args.instance_name:
-            profile_settings.instance_name = args.instance_name
-        if args.creation_policy_reuse:
-            profile_settings.creation_policy = CreationPolicy.REUSE
+        profile_settings.idle_duration = args.idle_duration
+    elif args.dont_destroy:
+        profile_settings.idle_duration = False
+    if args.creation_policy_reuse:
+        profile_settings.creation_policy = CreationPolicy.REUSE
 
     if args.spot_policy is not None:
         profile_settings.spot_policy = args.spot_policy
-    if pool_add and args.spot_policy is None:  # ONDEMAND by default for `dstack pool add`
-        profile_settings.spot_policy = SpotPolicy.ONDEMAND
 
-    if not pool_add:
-        if args.retry_policy is not None:
-            if not profile_settings.retry_policy:
-                profile_settings.retry_policy = ProfileRetryPolicy()
-            profile_settings.retry_policy.retry = args.retry_policy
-        elif args.retry_duration is not None:
-            if not profile_settings.retry_policy:
-                profile_settings.retry_policy = ProfileRetryPolicy()
-            profile_settings.retry_policy.retry = True
-            profile_settings.retry_policy.duration = args.retry_duration
-    else:
-        if args.retry_policy is not None:
-            if not profile_settings.retry_policy:
-                profile_settings.retry_policy = ProfileRetryPolicy()
-            profile_settings.retry_policy.retry = args.retry_policy
-            if profile_settings.retry_policy.retry:
-                profile_settings.retry_policy.duration = DEFAULT_INSTANCE_RETRY_DURATION
-        elif args.retry_duration is not None:
-            if not profile_settings.retry_policy:
-                profile_settings.retry_policy = ProfileRetryPolicy()
-            profile_settings.retry_policy.retry = True
-            profile_settings.retry_policy.duration = args.retry_duration  # --retry-duration
+    if args.retry is not None:
+        profile_settings.retry = args.retry
+    elif args.retry_duration is not None:
+        profile_settings.retry = ProfileRetry(
+            duration=args.retry_duration,
+        )
 
 
 def max_duration(v: str) -> int:

{dstack-0.18.44 → dstack-0.19.0}/src/dstack/_internal/cli/utils/gateway.py

@@ -34,7 +34,7 @@ def get_gateways_table(
     for gateway in gateways:
         row = {
             "NAME": gateway.name,
-            "BACKEND": f"{gateway.backend.value} ({gateway.region})",
+            "BACKEND": f"{gateway.configuration.backend.value} ({gateway.configuration.region})",
             "HOSTNAME": gateway.hostname,
             "DOMAIN": gateway.wildcard_domain,
             "DEFAULT": "✓" if gateway.default else "",

dstack-0.19.0/src/dstack/_internal/core/backends/__init__.py

@@ -0,0 +1,59 @@
+from dstack._internal.core.backends.base.compute import (
+    ComputeWithCreateInstanceSupport,
+    ComputeWithGatewaySupport,
+    ComputeWithMultinodeSupport,
+    ComputeWithPlacementGroupSupport,
+    ComputeWithPrivateGatewaySupport,
+    ComputeWithReservationSupport,
+    ComputeWithVolumeSupport,
+)
+from dstack._internal.core.backends.base.configurator import Configurator
+from dstack._internal.core.backends.configurators import list_available_configurator_classes
+from dstack._internal.core.models.backends.base import BackendType
+
+
+def _get_backends_with_compute_feature(
+    configurator_classes: list[type[Configurator]],
+    compute_feature_class: type,
+) -> list[BackendType]:
+    backend_types = []
+    for configurator_class in configurator_classes:
+        compute_class = configurator_class.BACKEND_CLASS.COMPUTE_CLASS
+        if issubclass(compute_class, compute_feature_class):
+            backend_types.append(configurator_class.TYPE)
+    return backend_types
+
+
+_configurator_classes = list_available_configurator_classes()
+
+
+# The following backend lists do not include unavailable backends (i.e. backends missing deps).
+# TODO: Add LocalBackend to lists if it's enabled
+BACKENDS_WITH_CREATE_INSTANCE_SUPPORT = _get_backends_with_compute_feature(
+    configurator_classes=_configurator_classes,
+    compute_feature_class=ComputeWithCreateInstanceSupport,
+)
+BACKENDS_WITH_MULTINODE_SUPPORT = [BackendType.REMOTE] + _get_backends_with_compute_feature(
+    configurator_classes=_configurator_classes,
+    compute_feature_class=ComputeWithMultinodeSupport,
+)
+BACKENDS_WITH_PLACEMENT_GROUPS_SUPPORT = _get_backends_with_compute_feature(
+    configurator_classes=_configurator_classes,
+    compute_feature_class=ComputeWithPlacementGroupSupport,
+)
+BACKENDS_WITH_RESERVATION_SUPPORT = _get_backends_with_compute_feature(
+    configurator_classes=_configurator_classes,
+    compute_feature_class=ComputeWithReservationSupport,
+)
+BACKENDS_WITH_GATEWAY_SUPPORT = _get_backends_with_compute_feature(
+    configurator_classes=_configurator_classes,
+    compute_feature_class=ComputeWithGatewaySupport,
+)
+BACKENDS_WITH_PRIVATE_GATEWAY_SUPPORT = _get_backends_with_compute_feature(
+    configurator_classes=_configurator_classes,
+    compute_feature_class=ComputeWithPrivateGatewaySupport,
+)
+BACKENDS_WITH_VOLUMES_SUPPORT = _get_backends_with_compute_feature(
+    configurator_classes=_configurator_classes,
+    compute_feature_class=ComputeWithVolumeSupport,
+)

{dstack-0.18.44 → dstack-0.19.0}/src/dstack/_internal/core/backends/aws/auth.py

@@ -2,8 +2,8 @@ import boto3.session
 import botocore.exceptions
 from boto3.session import Session
 
+from dstack._internal.core.backends.aws.models import AnyAWSCreds, AWSAccessKeyCreds
 from dstack._internal.core.errors import BackendAuthError
-from dstack._internal.core.models.backends.aws import AnyAWSCreds, AWSAccessKeyCreds
 from dstack._internal.core.models.common import is_core_model_instance
 
 
@@ -29,12 +29,3 @@ def validate_credentials(session: Session):
         sts.get_caller_identity()
     except (botocore.exceptions.ClientError, botocore.exceptions.NoCredentialsError):
         raise BackendAuthError()
-
-
-def default_creds_available() -> bool:
-    session = boto3.session.Session()
-    try:
-        validate_credentials(session)
-    except BackendAuthError:
-        return False
-    return True

dstack-0.18.44/src/dstack/_internal/core/backends/aws/__init__.py → dstack-0.19.0/src/dstack/_internal/core/backends/aws/backend.py

@@ -1,14 +1,15 @@
 import botocore.exceptions
 
 from dstack._internal.core.backends.aws.compute import AWSCompute
-from dstack._internal.core.backends.aws.config import AWSConfig
-from dstack._internal.core.backends.base import Backend
+from dstack._internal.core.backends.aws.models import AWSConfig
+from dstack._internal.core.backends.base.backend import Backend
 from dstack._internal.core.errors import BackendInvalidCredentialsError
 from dstack._internal.core.models.backends.base import BackendType
 
 
 class AWSBackend(Backend):
-    TYPE: BackendType = BackendType.AWS
+    TYPE = BackendType.AWS
+    COMPUTE_CLASS = AWSCompute
 
     def __init__(self, config: AWSConfig):
         self.config = config