diff --git a/mkdocs/docs/concepts/backends.md b/mkdocs/docs/concepts/backends.md index fa88e9ff61..1c00e4ddd6 100644 --- a/mkdocs/docs/concepts/backends.md +++ b/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: + +
+ + ```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 + ``` + +
+ + Regions not mapped in `subnetworks` keep using any usable subnet of the VPC. + === "Shared VPC"
@@ -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 diff --git a/src/dstack/_internal/core/backends/gcp/compute.py b/src/dstack/_internal/core/backends/gcp/compute.py index 5b30226004..45f46466c2 100644 --- a/src/dstack/_internal/core/backends/gcp/compute.py +++ b/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( diff --git a/src/dstack/_internal/core/backends/gcp/configurator.py b/src/dstack/_internal/core/backends/gcp/configurator.py index 1621756cc9..ffb5b2cb61 100644 --- a/src/dstack/_internal/core/backends/gcp/configurator.py +++ b/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, diff --git a/src/dstack/_internal/core/backends/gcp/models.py b/src/dstack/_internal/core/backends/gcp/models.py index 325cab6293..9fca5bd8f7 100644 --- a/src/dstack/_internal/core/backends/gcp/models.py +++ b/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( diff --git a/src/dstack/_internal/core/backends/gcp/resources.py b/src/dstack/_internal/core/backends/gcp/resources.py index ffbbda3158..30b21ddc9c 100644 --- a/src/dstack/_internal/core/backends/gcp/resources.py +++ b/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( diff --git a/src/tests/_internal/core/backends/gcp/test_configurator.py b/src/tests/_internal/core/backends/gcp/test_configurator.py index 635b11f2cf..abf0252559 100644 --- a/src/tests/_internal/core/backends/gcp/test_configurator.py +++ b/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="-"), diff --git a/src/tests/_internal/core/backends/gcp/test_resources.py b/src/tests/_internal/core/backends/gcp/test_resources.py index 47fe52438d..21a8263254 100644 --- a/src/tests/_internal/core/backends/gcp/test_resources.py +++ b/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: