Support A4 instances with the B200 GPU on GCP (#3100)

This implementation allows provisioning both
individual A4 instances and clusters, but clusters
do not yet support high-speed networking, since it
requires a
[different network setup](https://cloud.google.com/ai-hypercomputer/docs/create/create-vm#setup-network).

Author: jvstme

src/dstack/_internal/core/backends/base/offers.py

@@ -22,6 +22,7 @@
 SUPPORTED_GPUHUNT_FLAGS = [
     "oci-spot",
     "lambda-arm",
+    "gcp-a4",
 ]
 
 

src/dstack/_internal/core/backends/gcp/compute.py

@@ -867,8 +867,8 @@ def _has_gpu_quota(quotas: Dict[str, float], resources: Resources) -> bool:
     gpu = resources.gpus[0]
     if _is_tpu(gpu.name):
         return True
-    if gpu.name == "H100":
-        # H100 and H100_MEGA quotas are not returned by `regions_client.list`
+    if gpu.name in ["B200", "H100"]:
+        # B200, H100 and H100_MEGA quotas are not returned by `regions_client.list`
         return True
     quota_name = f"NVIDIA_{gpu.name}_GPUS"
     if gpu.name == "A100" and gpu.memory_mib == 80 * 1024:

src/dstack/_internal/core/backends/gcp/resources.py

@@ -19,6 +19,7 @@
 DSTACK_GATEWAY_TAG = "dstack-gateway-instance"
 
 supported_accelerators = [
+    {"accelerator_name": "nvidia-b200", "gpu_name": "B200", "memory_mb": 1024 * 180},
     {"accelerator_name": "nvidia-a100-80gb", "gpu_name": "A100", "memory_mb": 1024 * 80},
     {"accelerator_name": "nvidia-tesla-a100", "gpu_name": "A100", "memory_mb": 1024 * 40},
     {"accelerator_name": "nvidia-l4", "gpu_name": "L4", "memory_mb": 1024 * 24},
@@ -476,5 +477,6 @@ def instance_type_supports_persistent_disk(instance_type_name: str) -> bool:
             "n4-",
             "h3-",
             "v6e",
+            "a4-",
         ]
     )

src/dstack/_internal/server/background/tasks/common.py

@@ -19,4 +19,6 @@ def get_provisioning_timeout(backend_type: BackendType, instance_type_name: str)
         return timedelta(minutes=20)
     if backend_type == BackendType.VULTR and instance_type_name.startswith("vbm"):
         return timedelta(minutes=55)
+    if backend_type == BackendType.GCP and instance_type_name == "a4-highgpu-8g":
+        return timedelta(minutes=16)
     return timedelta(minutes=10)