todo archit

feat: TPC

PiperOrigin-RevId: 877503411

Author: vertex-sdk-bot

google/cloud/aiplatform/constants/base.py

@@ -19,53 +19,52 @@
 
 
 DEFAULT_REGION = "us-central1"
-SUPPORTED_REGIONS = frozenset(
-    {
-        "africa-south1",
-        "asia-east1",
-        "asia-east2",
-        "asia-northeast1",
-        "asia-northeast2",
-        "asia-northeast3",
-        "asia-south1",
-        "asia-south2",
-        "asia-southeast1",
-        "asia-southeast2",
-        "australia-southeast1",
-        "australia-southeast2",
-        "europe-central2",
-        "europe-north1",
-        "europe-north2",
-        "europe-southwest1",
-        "europe-west1",
-        "europe-west2",
-        "europe-west3",
-        "europe-west4",
-        "europe-west6",
-        "europe-west8",
-        "europe-west9",
-        "europe-west12",
-        "global",
-        "me-central1",
-        "me-central2",
-        "me-west1",
-        "northamerica-northeast1",
-        "northamerica-northeast2",
-        "southamerica-east1",
-        "southamerica-west1",
-        "us-central1",
-        "us-east1",
-        "us-east4",
-        "us-east5",
-        "us-east7",
-        "us-south1",
-        "us-west1",
-        "us-west2",
-        "us-west3",
-        "us-west4",
-        "us-west8",
-    }
-)
+SUPPORTED_REGIONS = frozenset({
+    "africa-south1",
+    "asia-east1",
+    "asia-east2",
+    "asia-northeast1",
+    "asia-northeast2",
+    "asia-northeast3",
+    "asia-south1",
+    "asia-south2",
+    "asia-southeast1",
+    "asia-southeast2",
+    "australia-southeast1",
+    "australia-southeast2",
+    "europe-central2",
+    "europe-north1",
+    "europe-north2",
+    "europe-southwest1",
+    "europe-west1",
+    "europe-west2",
+    "europe-west3",
+    "europe-west4",
+    "europe-west6",
+    "europe-west8",
+    "europe-west9",
+    "europe-west12",
+    "global",
+    "me-central1",
+    "me-central2",
+    "me-west1",
+    "northamerica-northeast1",
+    "northamerica-northeast2",
+    "southamerica-east1",
+    "southamerica-west1",
+    "us-central1",
+    "us-east1",
+    "us-east4",
+    "us-east5",
+    "us-east7",
+    "us-south1",
+    "us-west1",
+    "us-west2",
+    "us-west3",
+    "us-west4",
+    "us-west8",
+    "u-us-prp1",
+})
 
 API_BASE_PATH = "aiplatform.googleapis.com"
 PREDICTION_API_BASE_PATH = API_BASE_PATH

google/cloud/aiplatform/initializer.py

@@ -133,6 +133,7 @@ def __init__(self):
         self._request_metadata = None
         self._resource_type = None
         self._async_rest_credentials = None
+        self._universe_domain = None
 
     def init(
         self,
@@ -153,6 +154,7 @@ def init(
         api_key: Optional[str] = None,
         api_transport: Optional[str] = None,
         request_metadata: Optional[Sequence[Tuple[str, str]]] = None,
+        universe_domain: Optional[str] = None,
     ):
         """Updates common initialization parameters with provided options.
 
@@ -220,6 +222,8 @@ def init(
                 beta state (preview).
             request_metadata:
                 Optional. Additional gRPC metadata to send with every client request.
+            universe_domain (str):
+                Optional. The universe domain.
         Raises:
             ValueError:
                 If experiment_description is provided but experiment is not.
@@ -291,6 +295,8 @@ def init(
             self._request_metadata = request_metadata
         if api_key is not None:
             self._api_key = api_key
+        if universe_domain is not None:
+            self._universe_domain = universe_domain
         self._resource_type = None
 
         # Finally, perform secondary state updates
@@ -348,6 +354,11 @@ def api_key(self) -> Optional[str]:
         """API Key, if provided."""
         return self._api_key
 
+    @property
+    def universe_domain(self) -> Optional[str]:
+        """Default universe domain, if provided."""
+        return self._universe_domain
+
     @property
     def project(self) -> str:
         """Default project."""
@@ -382,7 +393,11 @@ def location(self) -> str:
 
         location = os.getenv("GOOGLE_CLOUD_REGION") or os.getenv("CLOUD_ML_REGION")
         if location:
-            utils.validate_region(location)
+            utils.validate_region(
+                location,
+                api_endpoint=self.api_endpoint,
+                universe_domain=self.universe_domain,
+            )
             return location
 
         return constants.DEFAULT_REGION
@@ -449,6 +464,7 @@ def get_client_options(
         api_base_path_override: Optional[str] = None,
         api_key: Optional[str] = None,
         api_path_override: Optional[str] = None,
+        universe_domain: Optional[str] = None,
     ) -> client_options.ClientOptions:
         """Creates GAPIC client_options using location and type.
 
@@ -461,6 +477,7 @@ def get_client_options(
             api_base_path_override (str): Optional. Override default API base path.
             api_key (str): Optional. API key to use for the client.
             api_path_override (str): Optional. Override default api path.
+            universe_domain (str): Optional. Override default universe domain.
         Returns:
             clients_options (google.api_core.client_options.ClientOptions):
                 A ClientOptions object set with regionalized API endpoint, i.e.
@@ -491,7 +508,11 @@ def get_client_options(
             region = location_override or self.location
             region = region.lower()
 
-            utils.validate_region(region)
+            utils.validate_region(
+                region,
+                api_endpoint=self.api_endpoint,
+                universe_domain=universe_domain or self.universe_domain,
+            )
 
             service_base_path = api_base_path_override or (
                 constants.PREDICTION_API_BASE_PATH
@@ -508,9 +529,14 @@ def get_client_options(
         # Project/location take precedence over api_key
         if api_key and not self._project:
             return client_options.ClientOptions(
-                api_endpoint=api_endpoint, api_key=api_key
+                api_endpoint=api_endpoint,
+                api_key=api_key,
+                universe_domain=universe_domain or self.universe_domain,
             )
-        return client_options.ClientOptions(api_endpoint=api_endpoint)
+        return client_options.ClientOptions(
+            api_endpoint=api_endpoint,
+            universe_domain=universe_domain or self.universe_domain,
+        )
 
     def common_location_path(
         self, project: Optional[str] = None, location: Optional[str] = None
@@ -524,7 +550,11 @@ def common_location_path(
             resource_parent: Formatted parent resource string.
         """
         if location:
-            utils.validate_region(location)
+            utils.validate_region(
+                location,
+                api_endpoint=self.api_endpoint,
+                universe_domain=self.universe_domain,
+            )
 
         return "/".join(
             [
@@ -546,6 +576,7 @@ def create_client(
         api_path_override: Optional[str] = None,
         appended_user_agent: Optional[List[str]] = None,
         appended_gapic_version: Optional[str] = None,
+        universe_domain: Optional[str] = None,
     ) -> _TVertexAiServiceClientWithOverride:
         """Instantiates a given VertexAiServiceClient with optional
         overrides.
@@ -565,6 +596,8 @@ def create_client(
                 separated by spaces.
             appended_gapic_version (str):
                 Optional. GAPIC version suffix appended in the client info.
+            universe_domain (str):
+                Optional. universe domain override.
         Returns:
             client: Instantiated Vertex AI Service client with optional overrides
         """
@@ -607,6 +640,7 @@ def create_client(
                 api_key=api_key,
                 api_base_path_override=api_base_path_override,
                 api_path_override=api_path_override,
+                universe_domain=universe_domain,
             ),
             "client_info": client_info,
         }

google/cloud/aiplatform/utils/__init__.py

@@ -299,11 +299,18 @@ def validate_labels(labels: Dict[str, str]):
             )
 
 
-def validate_region(region: str) -> bool:
+def validate_region(
+    region: str,
+    api_endpoint: Optional[str] = None,
+    universe_domain: Optional[str] = None,
+) -> bool:
     """Validates region against supported regions.
 
     Args:
         region: region to validate
+        api_endpoint: Optional API endpoint.
+        universe_domain: Optional universe domain.
+
     Returns:
         bool: True if no errors raised
     Raises:
@@ -316,9 +323,16 @@ def validate_region(region: str) -> bool:
 
     region = region.lower()
     if region not in constants.SUPPORTED_REGIONS:
-        raise ValueError(
-            f"Unsupported region for Vertex AI, select from {constants.SUPPORTED_REGIONS}"
-        )
+        if not (
+            api_endpoint
+            or universe_domain
+            or initializer.global_config.api_endpoint
+            or initializer.global_config.universe_domain
+        ):
+            raise ValueError(
+                "Unsupported region for Vertex AI, select from"
+                f" {constants.SUPPORTED_REGIONS}"
+            )
 
     return True
 

tests/unit/aiplatform/test_initializer.py

@@ -47,6 +47,10 @@
 _TEST_NETWORK = "projects/12345/global/networks/myVPC"
 _TEST_SERVICE_ACCOUNT = "test-service-account@test-project.iam.gserviceaccount.com"
 
+_TEST_LOCATION_TPC = "u-us-prp1"
+_TEST_ENDPOINT_TPC = "u-us-prp1-aiplatform.apis-tpczero.goog"
+_TEST_UNIVERSE_TPC = "apis-tpczero.goog"
+
 # tensorboard
 _TEST_TENSORBOARD_ID = "1028944691210842416"
 _TEST_TENSORBOARD_NAME = f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/tensorboards/{_TEST_TENSORBOARD_ID}"
@@ -547,6 +551,71 @@ def test_create_client_with_request_metadata_prediction_service(self):
             for metadata_key in ["global_param", "request_param"]:
                 assert metadata_key in headers
 
+    def test_init_with_universe_domain(self):
+        """Test that aiplatform.init supports universe_domain."""
+        initializer.global_config.init(
+            project=_TEST_PROJECT, universe_domain=_TEST_UNIVERSE_TPC
+        )
+        assert initializer.global_config.universe_domain == _TEST_UNIVERSE_TPC
+
+    def test_get_client_options_with_universe_domain(self):
+        """Test that ClientOptions correctly inherits universe_domain."""
+        initializer.global_config.init(
+            project=_TEST_PROJECT,
+            location="us-central1",
+            universe_domain=_TEST_UNIVERSE_TPC,
+        )
+        client_options = initializer.global_config.get_client_options()
+        assert client_options.universe_domain == _TEST_UNIVERSE_TPC
+
+    def test_init_tpc_success(self):
+        """Test that aiplatform.init succeeds with TPC parameters."""
+        # Should NOT raise ValueError
+        initializer.global_config.init(
+            project=_TEST_PROJECT,
+            location=_TEST_LOCATION_TPC,
+            api_endpoint=_TEST_ENDPOINT_TPC,
+            universe_domain=_TEST_UNIVERSE_TPC,
+        )
+        assert initializer.global_config.location == _TEST_LOCATION_TPC
+
+    def test_get_impersonated_credentials_tpc(self):
+        """Test impersonated credentials for TPC."""
+        info = {
+            "type": "impersonated_service_account",
+            "source_credentials": {
+                "type": "external_account_authorized_user",
+                "token_url": "https://sts.apis-tpczero.goog/v1/oauth/token",
+            },
+            "service_account_impersonation_url": (
+                "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/test-sa@test.iam.gserviceaccount.com:generateAccessToken"
+            ),
+            "universe_domain": _TEST_UNIVERSE_TPC,
+        }
+
+        with mock.patch(
+            "google.auth._default._get_external_account_authorized_user_credentials"
+        ) as mock_get_ext:
+            mock_source_creds = mock.Mock()
+            mock_get_ext.return_value = (mock_source_creds, None)
+
+            with mock.patch(
+                "google.auth.impersonated_credentials.Credentials"
+            ) as mock_imp_creds:
+                # pylint: disable=protected-access
+                google.auth._default._get_impersonated_service_account_credentials(
+                    "fake-file", info, ["scope"]
+                )
+
+                # Verify that iam_endpoint_override was set correctly for TPC
+                _, kwargs = mock_imp_creds.call_args
+                target_principal = "test-sa@test.iam.gserviceaccount.com"
+                expected_iam_endpoint = (
+                    f"https://iamcredentials.{_TEST_UNIVERSE_TPC}/v1/projects/"
+                    f"-/serviceAccounts/{target_principal}:generateAccessToken"
+                )
+                assert kwargs["iam_endpoint_override"] == expected_iam_endpoint
+
 
 class TestThreadPool:
     def teardown_method(self):