Add extra_vpcs to GCP config

Author: r4victor

.github/workflows/gcp-a3mega-image.yml

@@ -2,7 +2,6 @@ name: Build GCP A3 Mega VM image
 
 on:
   - workflow_dispatch
-  - push
 
 env:
   PACKER_VERSION: "1.9.2"

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

@@ -1,7 +1,7 @@
 import concurrent.futures
 import json
 from collections import defaultdict
-from typing import Callable, Dict, List, Literal, Optional
+from typing import Callable, Dict, List, Literal, Optional, Tuple
 
 import google.api_core.exceptions
 import google.cloud.compute_v1 as compute_v1
@@ -192,6 +192,12 @@ def create_instance(
             config=self.config,
             region=instance_offer.region,
         )
+        extra_subnets = _get_extra_subnets(
+            subnetworks_client=self.subnetworks_client,
+            config=self.config,
+            region=instance_offer.region,
+            instance_type_name=instance_offer.instance.name,
+        )
         placement_policy = None
         if instance_config.placement_group_name is not None:
             placement_policy = gcp_resources.get_placement_policy_resource_name(
@@ -300,6 +306,7 @@ def create_instance(
                 service_account=self.config.vm_service_account,
                 network=self.config.vpc_resource_name,
                 subnetwork=subnetwork,
+                extra_subnetworks=extra_subnets,
                 allocate_public_ip=allocate_public_ip,
                 placement_policy=placement_policy,
             )
@@ -741,21 +748,6 @@ def detach_volume(self, volume: Volume, instance_id: str, force: bool = False):
         )
 
 
-def _get_vpc_subnet(
-    subnetworks_client: compute_v1.SubnetworksClient,
-    config: GCPConfig,
-    region: str,
-) -> Optional[str]:
-    if config.vpc_name is None:
-        return None
-    return gcp_resources.get_vpc_subnet_or_error(
-        subnetworks_client=subnetworks_client,
-        vpc_project_id=config.vpc_project_id or config.project_id,
-        vpc_name=config.vpc_name,
-        region=region,
-    )
-
-
 def _supported_instances_and_zones(
     regions: List[str],
 ) -> Optional[Callable[[InstanceOffer], bool]]:
@@ -814,6 +806,47 @@ def _unique_instance_name(instance: InstanceType) -> str:
     return f"{name}-{gpu.name}-{gpu.memory_mib}"
 
 
+def _get_vpc_subnet(
+    subnetworks_client: compute_v1.SubnetworksClient,
+    config: GCPConfig,
+    region: str,
+) -> Optional[str]:
+    if config.vpc_name is None:
+        return None
+    return gcp_resources.get_vpc_subnet_or_error(
+        subnetworks_client=subnetworks_client,
+        vpc_project_id=config.vpc_project_id or config.project_id,
+        vpc_name=config.vpc_name,
+        region=region,
+    )
+
+
+def _get_extra_subnets(
+    subnetworks_client: compute_v1.SubnetworksClient,
+    config: GCPConfig,
+    region: str,
+    instance_type_name: str,
+) -> List[Tuple[str, str]]:
+    if config.extra_vpcs is None:
+        return []
+    if instance_type_name != "a3-megagpu-8g":
+        return []
+    extra_subnets = []
+    for vpc_name in config.extra_vpcs:
+        subnet = gcp_resources.get_vpc_subnet_or_error(
+            subnetworks_client=subnetworks_client,
+            vpc_project_id=config.vpc_project_id or config.project_id,
+            vpc_name=vpc_name,
+            region=region,
+        )
+        vpc_resource_name = gcp_resources.vpc_name_to_vpc_resource_name(
+            project_id=config.vpc_project_id or config.project_id,
+            vpc_name=vpc_name,
+        )
+        extra_subnets.append((vpc_resource_name, subnet))
+    return extra_subnets[:8]
+
+
 def _get_image_id(instance_type_name: str, cuda: bool) -> str:
     if instance_type_name == "a3-megagpu-8g":
         image_name = "dstack-a3mega-2"

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

@@ -199,3 +199,5 @@ def _check_config_vpc(
             )
         except BackendError as e:
             raise ServerClientError(e.args[0])
+        # Not checking config.extra_vpc so that users are not required to configure subnets for all regions
+        # but only for regions they intend to use. Validation will be done on provisioning.

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

@@ -33,7 +33,19 @@ class GCPBackendConfig(CoreModel):
     regions: Annotated[
         Optional[List[str]], Field(description="The list of GCP regions. Omit to use all regions")
     ] = None
-    vpc_name: Annotated[Optional[str], Field(description="The name of a custom VPC")] = None
+    vpc_name: Annotated[
+        Optional[str],
+        Field(description="The name of a custom VPC. If not specified, the default VPC is used"),
+    ] = None
+    extra_vpcs: Annotated[
+        Optional[List[str]],
+        Field(
+            description=(
+                "The names of additional VPCs used for GPUDirect. Specify eight VPCs to maximize bandwidth."
+                " Each VPC must have a subnet and a firewall rule allowing internal traffic across all subnets"
+            )
+        ),
+    ] = None
     vpc_project_id: Annotated[
         Optional[str],
         Field(description="The shared VPC hosted project ID. Required for shared VPC only"),

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

@@ -1,6 +1,6 @@
 import concurrent.futures
 import re
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Tuple
 
 import google.api_core.exceptions
 import google.cloud.compute_v1 as compute_v1
@@ -116,6 +116,7 @@ def create_instance_struct(
     service_account: Optional[str] = None,
     network: str = "global/networks/default",
     subnetwork: Optional[str] = None,
+    extra_subnetworks: Optional[List[Tuple[str, str]]] = None,
     allocate_public_ip: bool = True,
     placement_policy: Optional[str] = None,
 ) -> compute_v1.Instance:
@@ -126,6 +127,7 @@ def create_instance_struct(
         network=network,
         subnetwork=subnetwork,
         allocate_public_ip=allocate_public_ip,
+        extra_subnetworks=extra_subnetworks,
     )
 
     disk = compute_v1.AttachedDisk()
@@ -184,6 +186,7 @@ def _get_network_interfaces(
     network: str,
     subnetwork: Optional[str],
     allocate_public_ip: bool,
+    extra_subnetworks: Optional[List[Tuple[str, str]]],
 ) -> List[compute_v1.NetworkInterface]:
     network_interface = compute_v1.NetworkInterface()
     network_interface.network = network
@@ -199,11 +202,11 @@ def _get_network_interfaces(
         network_interface.access_configs = []
 
     network_interfaces = [network_interface]
-    for i in range(1, 9):
+    for network, subnetwork in extra_subnetworks or []:
         network_interfaces.append(
             compute_v1.NetworkInterface(
-                network=f"projects/dstack/global/networks/dstack-test-data-net-{i}",
-                subnetwork=f"projects/dstack/regions/europe-west4/subnetworks/dstack-test-data-sub-{i}",
+                network=network,
+                subnetwork=subnetwork,
             )
         )
     return network_interfaces
@@ -420,6 +423,10 @@ def full_resource_name_to_name(full_resource_name: str) -> str:
     return full_resource_name.split("/")[-1]
 
 
+def vpc_name_to_vpc_resource_name(project_id: str, vpc_name: str) -> str:
+    return f"projects/{project_id}/global/networks/{vpc_name}"
+
+
 def get_placement_policy_resource_name(
     project_id: str,
     region: str,