todo:archit

feat: One-line summary of change for external release notes.

PiperOrigin-RevId: 877503411

Author: vertex-sdk-bot

google/cloud/aiplatform/constants/base.py

@@ -19,53 +19,53 @@
 
 
 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",
-    }
-)
+DEFAULT_UNIVERSE_DOMAIN = "googleapis.com"
+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,14 +508,29 @@ 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
                 if prediction_client
                 else constants.API_BASE_PATH
             )
 
+            current_universe_domain = (
+                universe_domain
+                or self.universe_domain
+                or constants.DEFAULT_UNIVERSE_DOMAIN
+            )
+
+            if not api_base_path_override and current_universe_domain != "googleapis.com":
+                service_base_path = service_base_path.replace(
+                    "googleapis.com", current_universe_domain
+                )
+
             api_endpoint = (
                 f"{region}-{service_base_path}"
                 if not api_path_override
@@ -508,9 +540,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 +561,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 +587,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 +607,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 +651,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/system/aiplatform/test_tpc.py

@@ -0,0 +1,63 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from google.cloud import aiplatform
+from tests.system.aiplatform import e2e_base
+
+_TEST_PROJECT = "test-project"
+_TEST_LOCATION_TPC = "u-us-prp1"
+_TEST_ENDPOINT_TPC = "u-us-prp1-aiplatform.apis-tpczero.goog"
+_TEST_UNIVERSE_TPC = "apis-tpczero.goog"
+
+
+class TestTpcInitializer(e2e_base.TestEndToEnd):
+    """Tests TPC support in the initializer without monkeypatching."""
+
+    _temp_prefix = "test_tpc_initializer_"
+
+    def test_init_tpc_sets_global_config(self):
+        # This verifies that our changes to validate_region allow u-us-prp1
+        # and that universe_domain is correctly stored.
+        aiplatform.init(
+            project=_TEST_PROJECT,
+            location=_TEST_LOCATION_TPC,
+            universe_domain=_TEST_UNIVERSE_TPC,
+        )
+
+        assert aiplatform.initializer.global_config.location == _TEST_LOCATION_TPC
+        assert aiplatform.initializer.global_config.universe_domain == _TEST_UNIVERSE_TPC
+
+    def test_tpc_client_creation_plumbing(self):
+        """Verifies that clients created after TPC init have correct TPC endpoints."""
+        aiplatform.init(
+            project=_TEST_PROJECT,
+            location=_TEST_LOCATION_TPC,
+            universe_domain=_TEST_UNIVERSE_TPC,
+        )
+
+        # Instantiate a client (e.g., Endpoint)
+        # We don't call any methods that trigger real RPCs.
+        ds = aiplatform.Endpoint.list()
+
+        # Check the underlying client's endpoint
+        # The SDK should have constructed the endpoint using the TPC location and universe domain.
+        client = aiplatform.initializer.global_config.create_client(
+            client_class=aiplatform.utils.EndpointClientWithOverride
+        )
+
+        expected_host = f"{_TEST_LOCATION_TPC}-aiplatform.{_TEST_UNIVERSE_TPC}:443"
+        assert client._transport._host == expected_host