diff --git a/src/dstack/_internal/core/backends/aws/compute.py b/src/dstack/_internal/core/backends/aws/compute.py index 61620165c..b28cb0e10 100644 --- a/src/dstack/_internal/core/backends/aws/compute.py +++ b/src/dstack/_internal/core/backends/aws/compute.py @@ -1,7 +1,8 @@ import threading -from collections.abc import Iterable +from collections.abc import Container, Iterable from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass, field +from functools import partial from typing import Any, Callable, Dict, List, Optional, Tuple import boto3 @@ -141,6 +142,10 @@ def __init__( ) else: # default creds self.session = boto3.Session() + self._supported_instances = partial( + _supported_instances, + experimental_instance_types=set(self.config.experimental_instance_types or []), + ) # 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. @@ -163,7 +168,7 @@ def get_all_offers_with_availability( offers = get_catalog_offers( backend=BackendType.AWS, locations=self.config.regions, - extra_filter=_supported_instances, + extra_filter=self._supported_instances, ) regions = list(set(i.region for i in offers)) regions_to_quotas = self._get_regions_to_quotas(self.session, regions) @@ -1243,7 +1248,11 @@ def _get_regions_to_zones(session: boto3.Session, regions: List[str]) -> Dict[st return regions_to_zones -def _supported_instances(offer: InstanceOffer) -> bool: +def _supported_instances( + offer: InstanceOffer, experimental_instance_types: Container[str] +) -> bool: + if offer.instance.name in experimental_instance_types: + return True for family in [ "m7i.", "c7i.", diff --git a/src/dstack/_internal/core/backends/aws/models.py b/src/dstack/_internal/core/backends/aws/models.py index d76a4c9ab..95459bebe 100644 --- a/src/dstack/_internal/core/backends/aws/models.py +++ b/src/dstack/_internal/core/backends/aws/models.py @@ -106,6 +106,16 @@ class AWSBackendConfig(CoreModel): description="The mapping of instance categories (CPU, NVIDIA GPU) to AMI configurations" ), ] = None + experimental_instance_types: Annotated[ + Optional[List[str]], + Field( + description=( + "The list of instance type names to allow provisioning in addition to" + " the standard supported instance families. Only works for instance types" + " included in `dstack`'s pricing catalog (`gpuhunt`)" + ) + ), + ] = None class AWSBackendConfigWithCreds(AWSBackendConfig): diff --git a/src/tests/_internal/server/routers/test_backends.py b/src/tests/_internal/server/routers/test_backends.py index 9b548ab29..910c62dbf 100644 --- a/src/tests/_internal/server/routers/test_backends.py +++ b/src/tests/_internal/server/routers/test_backends.py @@ -828,6 +828,7 @@ async def test_returns_config_info(self, test_db, session: AsyncSession, client: "iam_instance_profile": None, "tags": None, "os_images": None, + "experimental_instance_types": None, "creds": json.loads(backend.auth.get_plaintext_or_error()), }