Allow configuring subnetworks in GCP settings

mkdocs/docs/concepts/backends.md

@@ -615,6 +615,30 @@ gcloud projects list --format="json(projectId)"
         allowing all traffic within the VPC. This is needed for multi-node tasks to work.
         The default VPC already permits traffic within the VPC.
 
+        If the VPC has multiple subnets in a region, `dstack` uses any usable subnet.
+        Specify `subnetworks` to target a specific subnet in a region:
+
+        <div editor-title="~/.dstack/server/config.yml">
+
+        ```yaml
+        projects:
+        - name: main
+          backends:
+            - type: gcp
+              project_id: gcp-project-id
+              creds:
+                type: default
+
+              regions: [europe-west4]
+              vpc_name: my-custom-vpc
+              subnetworks:
+                europe-west4: my-custom-subnet
+        ```
+
+        </div>
+
+        Regions not mapped in `subnetworks` keep using any usable subnet of the VPC.
+
     === "Shared VPC"
 
         <div editor-title="~/.dstack/server/config.yml">
@@ -660,6 +684,8 @@ gcloud projects list --format="json(projectId)"
     ```
 
     Using private subnets assumes that both the `dstack` server and users can access the configured VPC's private subnets.
+    If the VPC has multiple subnets in a region and only some of them are reachable from the `dstack` server and users,
+    specify `subnetworks` to make `dstack` provision instances in a reachable subnet.
     Additionally, [Cloud NAT](https://cloud.google.com/nat/docs/overview) must be configured to provide access to external resources for provisioned instances.
 
 ### Lambda

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

@@ -273,8 +273,8 @@ def create_instance(
                 network=self.config.vpc_resource_name,
             )
         disk_size = round(instance_offer.instance.resources.disk.size_mib / 1024)
-        # Choose any usable subnet in a VPC.
-        # Configuring a specific subnet per region is not supported yet.
+        # Use the subnet configured in `subnetworks` for the region if any,
+        # otherwise choose any usable subnet in the VPC.
         subnetwork = self._get_vpc_subnet(instance_offer.region)
         extra_subnets = self._get_extra_subnets(
             region=instance_offer.region,
@@ -581,8 +581,8 @@ def create_gateway_replica(
         instance_name = generate_unique_gateway_instance_name(
             configuration, max_length=gcp_resources.MAX_RESOURCE_NAME_LEN
         )
-        # Choose any usable subnet in a VPC.
-        # Configuring a specific subnet per region is not supported yet.
+        # Use the subnet configured in `subnetworks` for the region if any,
+        # otherwise choose any usable subnet in the VPC.
         subnetwork = self._get_vpc_subnet(configuration.region)
 
         labels = {
@@ -968,10 +968,12 @@ def _list_usable_subnets(self) -> list[compute_v1.UsableSubnetwork]:
     def _get_vpc_subnet(self, region: str) -> Optional[str]:
         if self.config.vpc_name is None:
             return None
+        subnetworks = self.config.subnetworks
         return gcp_resources.get_vpc_subnet_or_error(
             vpc_name=self.config.vpc_name,
             region=region,
             usable_subnets=self._list_usable_subnets(),
+            subnetwork_name=subnetworks.get(region) if subnetworks else None,
         )
 
     @cachedmethod(

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

@@ -191,6 +191,17 @@ def _check_config_vpc(
         subnetworks_client: compute_v1.SubnetworksClient,
         routers_client: compute_v1.RoutersClient,
     ):
+        if config.subnetworks:
+            if config.vpc_name is None:
+                raise ServerClientError(
+                    "`vpc_name` must be specified when `subnetworks` is configured."
+                )
+            unknown_regions = set(config.subnetworks) - set(config.regions or DEFAULT_REGIONS)
+            if unknown_regions:
+                raise ServerClientError(
+                    f"`subnetworks` is configured for regions not in `regions`:"
+                    f" {sorted(unknown_regions)}"
+                )
         allocate_public_ip = config.public_ips if config.public_ips is not None else True
         nat_check = config.nat_check if config.nat_check is not None else True
         try:
@@ -200,6 +211,7 @@ def _check_config_vpc(
                 project_id=config.project_id,
                 regions=config.regions or DEFAULT_REGIONS,
                 vpc_name=config.vpc_name,
+                subnetworks=config.subnetworks,
                 shared_vpc_project_id=config.vpc_project_id,
                 allocate_public_ip=allocate_public_ip,
                 nat_check=nat_check,

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

@@ -38,6 +38,17 @@ class GCPBackendConfig(CoreModel):
         Optional[str],
         Field(description="The name of a custom VPC. If not specified, the default VPC is used"),
     ] = None
+    subnetworks: Annotated[
+        Optional[Dict[str, str]],
+        Field(
+            description=(
+                "The mapping from regions to names of subnetworks in the VPC specified by `vpc_name`."
+                " `dstack` provisions instances in the specified subnetwork in mapped regions"
+                " and in any usable subnetwork of the VPC in other regions."
+                " Requires `vpc_name` to be set"
+            )
+        ),
+    ] = None
     extra_vpcs: Annotated[
         Optional[List[str]],
         Field(

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

@@ -71,6 +71,7 @@ def check_vpc(
     regions: List[str],
     allocate_public_ip: bool,
     vpc_name: Optional[str] = None,
+    subnetworks: Optional[Dict[str, str]] = None,
     shared_vpc_project_id: Optional[str] = None,
     nat_check: bool = True,
 ):
@@ -88,6 +89,7 @@ def check_vpc(
                 vpc_name=vpc_name,
                 region=region,
                 usable_subnets=usable_subnets,
+                subnetwork_name=subnetworks.get(region) if subnetworks else None,
             )
     except google.api_core.exceptions.NotFound:
         raise ComputeError(f"Failed to find VPC project {vpc_project_id}")
@@ -305,12 +307,23 @@ def get_vpc_subnet_or_error(
     vpc_name: str,
     region: str,
     usable_subnets: list[compute_v1.UsableSubnetwork],
+    subnetwork_name: Optional[str] = None,
 ) -> str:
     """
-    Returns resource name of any usable subnet in a given VPC
-    (e.g. "projects/example-project/regions/europe-west4/subnetworks/example-subnet")
+    Returns resource name of a usable subnet in a given VPC
+    (e.g. "projects/example-project/regions/europe-west4/subnetworks/example-subnet").
+    If `subnetwork_name` is not specified, any usable subnet is returned.
     """
     vpc_subnets = get_vpc_subnets(vpc_name, region, usable_subnets)
+    if subnetwork_name is not None:
+        for subnet in vpc_subnets:
+            if subnet.split("/")[-1] == subnetwork_name:
+                return subnet
+        raise ComputeError(
+            f"Subnetwork {subnetwork_name} not found among usable subnetworks"
+            f" of VPC {vpc_name} in region {region}."
+            f" Available subnetworks: {[s.split('/')[-1] for s in vpc_subnets]}"
+        )
     if vpc_subnets:
         return vpc_subnets[0]
     raise ComputeError(

src/tests/_internal/core/backends/gcp/test_configurator.py

@@ -10,6 +10,7 @@
 from dstack._internal.core.errors import (
     BackendAuthError,
     BackendInvalidCredentialsError,
+    ServerClientError,
 )
 
 
@@ -27,6 +28,51 @@ def test_validate_config_valid(self):
             authenticate_mock.return_value = Mock(), Mock()
             GCPConfigurator().validate_config(config, default_creds_enabled=True)
 
+    def test_validate_config_valid_subnetworks(self):
+        config = GCPBackendConfigWithCreds(
+            creds=GCPServiceAccountCreds(data="valid", filename="-"),
+            project_id="valid-project",
+            regions=["us-west1", "europe-west4"],
+            vpc_name="my-vpc",
+            subnetworks={"us-west1": "my-subnet"},
+        )
+        with (
+            patch("dstack._internal.core.backends.gcp.auth.authenticate") as authenticate_mock,
+            patch("dstack._internal.core.backends.gcp.resources.check_vpc") as check_vpc_mock,
+        ):
+            authenticate_mock.return_value = Mock(), Mock()
+            GCPConfigurator().validate_config(config, default_creds_enabled=True)
+        assert check_vpc_mock.call_args.kwargs["subnetworks"] == {"us-west1": "my-subnet"}
+
+    def test_validate_config_subnetworks_without_vpc_name(self):
+        config = GCPBackendConfigWithCreds(
+            creds=GCPServiceAccountCreds(data="valid", filename="-"),
+            project_id="valid-project",
+            regions=["us-west1"],
+            subnetworks={"us-west1": "my-subnet"},
+        )
+        with (
+            patch("dstack._internal.core.backends.gcp.auth.authenticate") as authenticate_mock,
+            pytest.raises(ServerClientError, match="`vpc_name` must be specified"),
+        ):
+            authenticate_mock.return_value = Mock(), Mock()
+            GCPConfigurator().validate_config(config, default_creds_enabled=True)
+
+    def test_validate_config_subnetworks_region_not_in_regions(self):
+        config = GCPBackendConfigWithCreds(
+            creds=GCPServiceAccountCreds(data="valid", filename="-"),
+            project_id="valid-project",
+            regions=["us-west1"],
+            vpc_name="my-vpc",
+            subnetworks={"europe-west4": "my-subnet"},
+        )
+        with (
+            patch("dstack._internal.core.backends.gcp.auth.authenticate") as authenticate_mock,
+            pytest.raises(ServerClientError, match="regions not in `regions`"),
+        ):
+            authenticate_mock.return_value = Mock(), Mock()
+            GCPConfigurator().validate_config(config, default_creds_enabled=True)
+
     def test_validate_config_invalid_creds(self):
         config = GCPBackendConfigWithCreds(
             creds=GCPServiceAccountCreds(data="invalid", filename="-"),

src/tests/_internal/core/backends/gcp/test_resources.py

@@ -1,7 +1,87 @@
+import google.cloud.compute_v1 as compute_v1
 import pytest
 
 from dstack._internal.core.backends.gcp import resources as gcp_resources
-from dstack._internal.core.errors import BackendError
+from dstack._internal.core.errors import BackendError, ComputeError
+
+
+def _usable_subnet(
+    project: str, vpc: str, region: str, subnet: str
+) -> compute_v1.UsableSubnetwork:
+    return compute_v1.UsableSubnetwork(
+        network=(
+            f"https://www.googleapis.com/compute/v1/projects/{project}/global/networks/{vpc}"
+        ),
+        subnetwork=(
+            f"https://www.googleapis.com/compute/v1/projects/{project}"
+            f"/regions/{region}/subnetworks/{subnet}"
+        ),
+    )
+
+
+class TestGetVpcSubnetOrError:
+    def test_returns_first_subnet_when_name_not_specified(self):
+        usable_subnets = [
+            _usable_subnet("proj", "my-vpc", "us-west1", "subnet-a"),
+            _usable_subnet("proj", "my-vpc", "us-west1", "subnet-b"),
+        ]
+        subnet = gcp_resources.get_vpc_subnet_or_error(
+            vpc_name="my-vpc",
+            region="us-west1",
+            usable_subnets=usable_subnets,
+        )
+        assert subnet == "projects/proj/regions/us-west1/subnetworks/subnet-a"
+
+    def test_returns_subnet_matching_specified_name(self):
+        usable_subnets = [
+            _usable_subnet("proj", "my-vpc", "us-west1", "subnet-a"),
+            _usable_subnet("proj", "my-vpc", "us-west1", "subnet-b"),
+        ]
+        subnet = gcp_resources.get_vpc_subnet_or_error(
+            vpc_name="my-vpc",
+            region="us-west1",
+            usable_subnets=usable_subnets,
+            subnetwork_name="subnet-b",
+        )
+        assert subnet == "projects/proj/regions/us-west1/subnetworks/subnet-b"
+
+    def test_raises_when_specified_subnet_not_found(self):
+        usable_subnets = [
+            _usable_subnet("proj", "my-vpc", "us-west1", "subnet-a"),
+        ]
+        with pytest.raises(ComputeError, match=r"Available subnetworks: \['subnet-a'\]"):
+            gcp_resources.get_vpc_subnet_or_error(
+                vpc_name="my-vpc",
+                region="us-west1",
+                usable_subnets=usable_subnets,
+                subnetwork_name="missing",
+            )
+
+    def test_raises_when_specified_subnet_in_another_region(self):
+        usable_subnets = [
+            _usable_subnet("proj", "my-vpc", "us-west1", "subnet-a"),
+            _usable_subnet("proj", "my-vpc", "europe-west4", "subnet-b"),
+        ]
+        with pytest.raises(ComputeError, match="VPC my-vpc in region us-west1"):
+            gcp_resources.get_vpc_subnet_or_error(
+                vpc_name="my-vpc",
+                region="us-west1",
+                usable_subnets=usable_subnets,
+                subnetwork_name="subnet-b",
+            )
+
+    def test_matches_subnet_by_short_name_in_shared_vpc(self):
+        usable_subnets = [
+            _usable_subnet("host-proj", "shared-vpc", "us-west1", "subnet-a"),
+            _usable_subnet("host-proj", "shared-vpc", "us-west1", "subnet-b"),
+        ]
+        subnet = gcp_resources.get_vpc_subnet_or_error(
+            vpc_name="shared-vpc",
+            region="us-west1",
+            usable_subnets=usable_subnets,
+            subnetwork_name="subnet-b",
+        )
+        assert subnet == "projects/host-proj/regions/us-west1/subnetworks/subnet-b"
 
 
 class TestValidateLabels: