From 96a9fcf5bb0aca0e42f39428a041d0c70f7f9bc3 Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Mon, 27 Jul 2026 15:44:08 -0700 Subject: [PATCH] feat: Support us mREP endpoint for Batch Prediction, Training, and Online Prediction. PiperOrigin-RevId: 954870469 --- google/cloud/aiplatform/constants/base.py | 5 ++++ google/cloud/aiplatform/initializer.py | 14 ++++++++---- google/cloud/aiplatform/utils/__init__.py | 28 +++++++++++++++++++++-- tests/unit/aiplatform/test_initializer.py | 24 +++++++++++++++++++ tests/unit/aiplatform/test_utils.py | 23 +++++++++++++++++++ 5 files changed, 87 insertions(+), 7 deletions(-) diff --git a/google/cloud/aiplatform/constants/base.py b/google/cloud/aiplatform/constants/base.py index 3abf032874..7bf78013b6 100644 --- a/google/cloud/aiplatform/constants/base.py +++ b/google/cloud/aiplatform/constants/base.py @@ -67,6 +67,11 @@ } ) +# Multi-regional (mREP) jurisdictions. These are served on dedicated REP hosts +# (aiplatform..rep.googleapis.com) rather than the locational +# -aiplatform.googleapis.com form. +MREP_JURISDICTIONS = frozenset({"us"}) + API_BASE_PATH = "aiplatform.googleapis.com" PREDICTION_API_BASE_PATH = API_BASE_PATH diff --git a/google/cloud/aiplatform/initializer.py b/google/cloud/aiplatform/initializer.py index cffb76d035..13d7ce0217 100644 --- a/google/cloud/aiplatform/initializer.py +++ b/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: diff --git a/google/cloud/aiplatform/utils/__init__.py b/google/cloud/aiplatform/utils/__init__.py index 4abe2b07b1..0ee79d0c37 100644 --- a/google/cloud/aiplatform/utils/__init__.py +++ b/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. diff --git a/tests/unit/aiplatform/test_initializer.py b/tests/unit/aiplatform/test_initializer.py index 8b7347efd5..fe12dbdb99 100644 --- a/tests/unit/aiplatform/test_initializer.py +++ b/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, diff --git a/tests/unit/aiplatform/test_utils.py b/tests/unit/aiplatform/test_utils.py index 71b5689f82..0688f78c5a 100644 --- a/tests/unit/aiplatform/test_utils.py +++ b/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: