From c82a429f6bd054ba39cb3807bdbc0e68a3be2c83 Mon Sep 17 00:00:00 2001 From: Jvst Me Date: Mon, 27 Jul 2026 18:57:37 +0200 Subject: [PATCH] Add AWS `experimental_instance_types` setting Add an AWS backend setting that accepts a 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`). Today, this setting can be used to test the `p5en.48xlarge` instance type, which wasn't yet tested internally at `dstack`. ```yaml experimental_instance_types: [p5en.48xlarge] ``` --- src/dstack/_internal/core/backends/aws/compute.py | 15 ++++++++++++--- src/dstack/_internal/core/backends/aws/models.py | 10 ++++++++++ .../_internal/server/routers/test_backends.py | 1 + 3 files changed, 23 insertions(+), 3 deletions(-) diff --git a/src/dstack/_internal/core/backends/aws/compute.py b/src/dstack/_internal/core/backends/aws/compute.py index 61620165c0..b28cb0e104 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 d76a4c9ab3..95459bebe6 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 9b548ab297..910c62dbfe 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()), }