feat: Support us mREP endpoint for Batch Prediction, Training, and Online Prediction.

PiperOrigin-RevId: 953605490

Author: vertex-sdk-bot

google/cloud/aiplatform/constants/base.py

@@ -67,6 +67,11 @@
     }
 )
 
+# Multi-regional (mREP) jurisdictions. These are served on dedicated REP hosts
+# (aiplatform.<geo>.rep.googleapis.com) rather than the locational
+# <region>-aiplatform.googleapis.com form.
+MREP_JURISDICTIONS = frozenset({"us"})
+
 API_BASE_PATH = "aiplatform.googleapis.com"
 PREDICTION_API_BASE_PATH = API_BASE_PATH
 

google/cloud/aiplatform/initializer.py

@@ -493,11 +493,15 @@ def get_client_options(
                 else constants.API_BASE_PATH
             )
 
-            api_endpoint = (
-                f"{region}-{service_base_path}"
-                if not api_path_override
-                else api_path_override
-            )
+            if api_path_override:
+                api_endpoint = api_path_override
+            elif ".rep." in service_base_path:
+                # Already an mREP host (e.g. via api_base_path_override); use as-is.
+                api_endpoint = service_base_path
+            elif utils.is_mrep_location(region):
+                api_endpoint = utils.mrep_endpoint(service_base_path, region)
+            else:
+                api_endpoint = f"{region}-{service_base_path}"
 
         # Project/location take precedence over api_key
         if api_key and not self._project:

google/cloud/aiplatform/utils/__init__.py

@@ -314,14 +314,38 @@ def validate_region(region: str) -> bool:
         )
 
     region = region.lower()
-    if region not in constants.SUPPORTED_REGIONS:
+    if region not in constants.SUPPORTED_REGIONS and not is_mrep_location(region):
         raise ValueError(
-            f"Unsupported region for Vertex AI, select from {constants.SUPPORTED_REGIONS}"
+            "Unsupported region for Vertex AI, select from "
+            f"{constants.SUPPORTED_REGIONS} or from the multi-regional "
+            f"jurisdictions {constants.MREP_JURISDICTIONS}"
         )
 
     return True
 
 
+def is_mrep_location(location: str) -> bool:
+    """Returns whether location is a multi-regional (mREP) jurisdiction."""
+    return location.lower() in constants.MREP_JURISDICTIONS
+
+
+def mrep_endpoint(service_base_path: str, location: str) -> str:
+    """Returns the mREP host for a jurisdiction.
+
+    Args:
+        service_base_path: the base service host, e.g. "aiplatform.googleapis.com".
+        location: the multi-regional jurisdiction, e.g. "us".
+
+    Returns:
+        The mREP host, e.g. "aiplatform.us.rep.googleapis.com". A base host
+        without a domain (no ".") is returned unchanged.
+    """
+    if "." not in service_base_path:
+        return service_base_path
+    service, base_domain = service_base_path.split(".", 1)
+    return f"{service}.{location}.rep.{base_domain}"
+
+
 def validate_accelerator_type(accelerator_type: str) -> bool:
     """Validates user provided accelerator_type string for training and
     prediction.

tests/unit/aiplatform/test_initializer.py

@@ -310,6 +310,30 @@ def test_create_client_with_global_location(self):
         assert isinstance(client, utils.PredictionClientWithOverride)
         assert client._transport._host == f"https://{constants.API_BASE_PATH}"
 
+    def test_create_client_with_mrep_location(self):
+        initializer.global_config.init(project=_TEST_PROJECT, location="us")
+        client = initializer.global_config.create_client(
+            client_class=utils.ModelClientWithOverride
+        )
+        assert initializer.global_config.location == "us"
+        assert client._transport._host == "aiplatform.us.rep.googleapis.com:443"
+
+    def test_create_client_with_mrep_location_prediction_client(self):
+        initializer.global_config.init(project=_TEST_PROJECT, location="us")
+        client = initializer.global_config.create_client(
+            client_class=utils.PredictionClientWithOverride,
+            prediction_client=True,
+        )
+        assert client._transport._host == "aiplatform.us.rep.googleapis.com:443"
+
+    def test_create_client_with_mrep_location_and_rep_base_path_override(self):
+        initializer.global_config.init(project=_TEST_PROJECT, location="us")
+        client = initializer.global_config.create_client(
+            client_class=utils.ModelClientWithOverride,
+            api_base_path_override="aiplatform.us.rep.googleapis.com",
+        )
+        assert client._transport._host == "aiplatform.us.rep.googleapis.com:443"
+
     def test_create_client_with_global_location_and_api_endpoint(self):
         initializer.global_config.init(
             project=_TEST_PROJECT,

tests/unit/aiplatform/test_utils.py

@@ -142,6 +142,29 @@ def test_invalid_region_does_not_raise_with_valid_region():
     aiplatform.utils.validate_region(region="us-central1")
 
 
+def test_validate_region_does_not_raise_with_mrep_jurisdiction():
+    aiplatform.utils.validate_region(region="us")
+
+
+def test_is_mrep_location_true_for_jurisdiction():
+    assert aiplatform.utils.is_mrep_location("us")
+
+
+def test_is_mrep_location_false_for_locational_region():
+    assert not aiplatform.utils.is_mrep_location("us-central1")
+
+
+def test_mrep_endpoint_builds_rep_host():
+    assert (
+        aiplatform.utils.mrep_endpoint("aiplatform.googleapis.com", "us")
+        == "aiplatform.us.rep.googleapis.com"
+    )
+
+
+def test_mrep_endpoint_returns_base_without_domain_unchanged():
+    assert aiplatform.utils.mrep_endpoint("localhost", "us") == "localhost"
+
+
 @pytest.fixture
 def copy_tree_mock():
     with mock.patch("shutil.copytree") as copy_tree_mock: