Make experimental GCP TPU support opt-in

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:
+
+    <div editor-title="~/.dstack/server/config.yml">
+
+    ```yaml
+    projects:
+    - name: main
+      backends:
+        - type: gcp
+          project_id: gcp-project-id
+          creds:
+            type: default
+          tpu: true
+    ```
+
+    </div>
+
+    Enable the Cloud TPU API in your GCP project:
+
+    ```shell
+    gcloud services enable tpu.googleapis.com --project=gcp-project-id
+    ```
+
 ??? info "VPC"
 
     === "VPC"

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"

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)

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:

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

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

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")