From 200faa55cae92a2faf1efc0d1e73232d4568301f Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Fri, 16 Oct 2026 12:09:17 +0200 Subject: [PATCH] Make experimental GCP TPU support opt-in --- mkdocs/docs/concepts/backends.md | 24 ++++++++ mkdocs/docs/examples/accelerators/tpu.md | 3 + .../_internal/core/backends/gcp/compute.py | 10 +++- .../_internal/core/backends/gcp/models.py | 17 ++++++ .../_internal/core/compatibility/backends.py | 3 + .../core/backends/gcp/test_compute.py | 55 +++++++++++++++++++ .../core/compatibility/test_backends.py | 19 ++++++- 7 files changed, 128 insertions(+), 3 deletions(-) create mode 100644 src/tests/_internal/core/backends/gcp/test_compute.py diff --git a/mkdocs/docs/concepts/backends.md b/mkdocs/docs/concepts/backends.md index 58cff86371..e891904bbb 100644 --- a/mkdocs/docs/concepts/backends.md +++ b/mkdocs/docs/concepts/backends.md @@ -627,6 +627,30 @@ gcloud projects list --format="json(projectId)" gcloud services enable compute.googleapis.com ``` +??? info "TPUs" + TPU support is experimental and disabled by default. To enable it, set `tpu: true` in the backend settings: + +
+ + ```yaml + projects: + - name: main + backends: + - type: gcp + project_id: gcp-project-id + creds: + type: default + tpu: true + ``` + +
+ + Enable the Cloud TPU API in your GCP project: + + ```shell + gcloud services enable tpu.googleapis.com --project=gcp-project-id + ``` + ??? info "VPC" === "VPC" diff --git a/mkdocs/docs/examples/accelerators/tpu.md b/mkdocs/docs/examples/accelerators/tpu.md index 8c4d1584bb..a9ccfad763 100644 --- a/mkdocs/docs/examples/accelerators/tpu.md +++ b/mkdocs/docs/examples/accelerators/tpu.md @@ -9,6 +9,9 @@ If you've configured the `gcp` backend in `dstack`, you can run dev environments Choose a TPU instance by specifying the TPU version and the number of cores (e.g. `v5litepod-8`) in the `gpu` property under `resources`, or request TPUs by specifying `tpu` as `vendor` ([see examples](../../guides/protips.md#gpu)). +TPU support is experimental and disabled by default. Before using these examples, set `tpu: true` +in your [GCP backend settings](../../concepts/backends.md#gcp). + Below are a few examples on using TPUs for deployment and fine-tuning. !!! info "Multi-host TPUs" diff --git a/src/dstack/_internal/core/backends/gcp/compute.py b/src/dstack/_internal/core/backends/gcp/compute.py index 243bc13caa..9cd795e688 100644 --- a/src/dstack/_internal/core/backends/gcp/compute.py +++ b/src/dstack/_internal/core/backends/gcp/compute.py @@ -176,7 +176,9 @@ def get_all_offers_with_availability( ) -> List[InstanceOfferWithAvailability]: regions = get_or_error(self.config.regions) zones_by_key: Dict[Tuple, List[str]] = {} - catalog_item_filter = _make_catalog_item_filter(regions, zones_by_key) + catalog_item_filter = _make_catalog_item_filter( + regions, zones_by_key, allow_tpu=self.config.allow_tpu + ) offers = get_catalog_offers( backend=BackendType.GCP, catalog_item_filter=catalog_item_filter, @@ -1488,9 +1490,11 @@ def _offer_dedup_key( def _make_catalog_item_filter( regions: List[str], zones_by_key: Dict[Tuple, List[str]], + *, + allow_tpu: bool, ) -> Callable[[gpuhunt.CatalogItem], bool]: """ - Returns a filter that checks region, instance support, and deduplicates + Returns a filter that checks region, TPU opt-in, instance support, and deduplicates per-zone items into per-region offers. Zones are collected in `zones_by_key` so the caller can attach them to offers later. """ @@ -1500,6 +1504,8 @@ def _filter(item: gpuhunt.CatalogItem) -> bool: region = item.location[:-2] if region not in regions: return False + if not allow_tpu and item.gpu_vendor == gpuhunt.AcceleratorVendor.GOOGLE: + return False if not _is_supported_gcp_instance(item.instance_name, item.gpu_name): return False key = _offer_dedup_key(item.instance_name, item.spot, item.gpu_name, region) diff --git a/src/dstack/_internal/core/backends/gcp/models.py b/src/dstack/_internal/core/backends/gcp/models.py index 9fca5bd8f7..15eabb7222 100644 --- a/src/dstack/_internal/core/backends/gcp/models.py +++ b/src/dstack/_internal/core/backends/gcp/models.py @@ -6,6 +6,8 @@ from dstack._internal.core.backends.base.models import fill_data from dstack._internal.core.models.common import CoreModel +GCP_TPU_DEFAULT = False + class GCPServiceAccountCreds(CoreModel): type: Annotated[Literal["service_account"], Field(description="The type of credentials")] = ( @@ -34,6 +36,15 @@ class GCPBackendConfig(CoreModel): regions: Annotated[ Optional[List[str]], Field(description="The list of GCP regions. Omit to use all regions") ] = None + tpu: Annotated[ + Optional[bool], + Field( + description=( + "Whether TPU offers can be suggested in addition to CPU and GPU VMs (experimental)." + f" Defaults to `{str(GCP_TPU_DEFAULT).lower()}`" + ) + ), + ] = None vpc_name: Annotated[ Optional[str], Field(description="The name of a custom VPC. If not specified, the default VPC is used"), @@ -155,6 +166,12 @@ class GCPStoredConfig(GCPBackendConfig): class GCPConfig(GCPStoredConfig): creds: Annotated[AnyGCPCreds, Field(discriminator="type")] + @property + def allow_tpu(self) -> bool: + if self.tpu is not None: + return self.tpu + return GCP_TPU_DEFAULT + @property def allocate_public_ips(self) -> bool: if self.public_ips is not None: diff --git a/src/dstack/_internal/core/compatibility/backends.py b/src/dstack/_internal/core/compatibility/backends.py index a4bb7f21a5..b742e3c155 100644 --- a/src/dstack/_internal/core/compatibility/backends.py +++ b/src/dstack/_internal/core/compatibility/backends.py @@ -1,3 +1,4 @@ +from dstack._internal.core.backends.gcp.models import GCPBackendConfigWithCreds from dstack._internal.core.backends.hotaisle.models import HotAisleBackendConfigWithCreds from dstack._internal.core.backends.models import AnyBackendConfigWithCreds from dstack._internal.core.models.common import IncludeExcludeDictType @@ -10,6 +11,8 @@ def get_backend_config_excludes(config: AnyBackendConfigWithCreds) -> IncludeExc clients backward-compatibility with older servers. """ excludes: IncludeExcludeDictType = {} + if isinstance(config, GCPBackendConfigWithCreds) and config.tpu is None: + excludes["tpu"] = True if isinstance(config, HotAisleBackendConfigWithCreds) and config.bare_metal is None: excludes["bare_metal"] = True return excludes diff --git a/src/tests/_internal/core/backends/gcp/test_compute.py b/src/tests/_internal/core/backends/gcp/test_compute.py new file mode 100644 index 0000000000..6675d5df53 --- /dev/null +++ b/src/tests/_internal/core/backends/gcp/test_compute.py @@ -0,0 +1,55 @@ +import pytest +from google.auth.credentials import AnonymousCredentials +from gpuhunt import AcceleratorVendor, CatalogItem + +from dstack._internal.core.backends.gcp.compute import GCPCompute +from dstack._internal.core.backends.gcp.models import GCPConfig, GCPDefaultCreds + + +class TestGetAllOffersWithAvailability: + @pytest.mark.parametrize( + "tpu_config", [{}, {"tpu": False}, {"tpu": True}], ids=["default", "disabled", "enabled"] + ) + def test_returns_tpu_offers_only_if_enabled(self, mocker, tpu_config): + mocker.patch( + "dstack._internal.core.backends.gcp.auth.authenticate", + return_value=(AnonymousCredentials(), None), + ) + compute = GCPCompute( + GCPConfig( + project_id="test-project", + regions=["us-central1"], + creds=GCPDefaultCreds(), + **tpu_config, + ) + ) + mocker.patch.object(compute.regions_client, "list", return_value=[]) + items = [ + CatalogItem( + provider="gcp", + instance_name=name, + location="us-central1-a", + price=1.0, + cpu=4, + memory=16, + gpu_count=int(vendor is not None), + gpu_vendor=vendor, + gpu_name=gpu_name, + gpu_memory=16 if vendor is not None else None, + spot=False, + disk_size=None, + ) + for name, vendor, gpu_name in [ + ("n2-standard-4", None, None), + ("g2-standard-4", AcceleratorVendor.NVIDIA, "L4"), + ("v5litepod-8", AcceleratorVendor.GOOGLE, "v5litepod-8"), + ] + ] + mocker.patch("gpuhunt.default_catalog").return_value.query.return_value = items + + offers = compute.get_all_offers_with_availability(unallocated_resources=False) + + expected_names = ["n2-standard-4", "g2-standard-4"] + if tpu_config.get("tpu"): + expected_names.append("v5litepod-8") + assert [offer.instance.name for offer in offers] == expected_names diff --git a/src/tests/_internal/core/compatibility/test_backends.py b/src/tests/_internal/core/compatibility/test_backends.py index 20e55a002c..dfc3273ee2 100644 --- a/src/tests/_internal/core/compatibility/test_backends.py +++ b/src/tests/_internal/core/compatibility/test_backends.py @@ -1,17 +1,34 @@ import json +import pytest + +from dstack._internal.core.backends.gcp.models import GCPBackendConfigWithCreds, GCPDefaultCreds from dstack._internal.core.backends.hotaisle.models import ( HotAisleAPIKeyCreds, HotAisleBackendConfigWithCreds, ) +from dstack._internal.core.backends.models import AnyBackendConfigWithCreds from dstack._internal.core.compatibility.backends import get_backend_config_excludes -def _request_body(config: HotAisleBackendConfigWithCreds) -> dict: +def _request_body(config: AnyBackendConfigWithCreds) -> dict: return json.loads(config.model_dump_json(exclude=get_backend_config_excludes(config))) class TestGetBackendConfigExcludes: + def test_excludes_unset_gcp_tpu(self): + config = GCPBackendConfigWithCreds(project_id="test-project", creds=GCPDefaultCreds()) + + assert "tpu" not in _request_body(config) + + @pytest.mark.parametrize("tpu", [False, True]) + def test_keeps_set_gcp_tpu(self, tpu): + config = GCPBackendConfigWithCreds( + project_id="test-project", creds=GCPDefaultCreds(), tpu=tpu + ) + + assert _request_body(config)["tpu"] is tpu + def test_excludes_unset_hotaisle_bare_metal(self): config = HotAisleBackendConfigWithCreds( team_handle="test-team", creds=HotAisleAPIKeyCreds(api_key="test-key")