feat: support user provided api endpoint.

PiperOrigin-RevId: 583223550

google/cloud/aiplatform/initializer.py

@@ -103,6 +103,7 @@ def __init__(self):
  self._encryption_spec_key_name = None
  self._network = None
  self._service_account = None
+ self._api_endpoint = None
 
  def init(
  self,
@@ -119,6 +120,7 @@ def init(
  encryption_spec_key_name: Optional[str] = None,
  network: Optional[str] = None,
  service_account: Optional[str] = None,
+ api_endpoint: Optional[str] = None,
  ):
  """Updates common initialization parameters with provided options.
 
@@ -174,11 +176,17 @@ def init(
  PipelineJob, HyperparameterTuningJob, CustomTrainingJob,
  CustomPythonPackageTrainingJob, CustomContainerTrainingJob,
  ModelEvaluationJob.
+ api_endpoint (str):
+ Optional. The desired API endpoint,
+ e.g., us-central1-aiplatform.googleapis.com
  Raises:
  ValueError:
  If experiment_description is provided but experiment is not.
  """
 
+ if api_endpoint is not None:
+ self._api_endpoint = api_endpoint
+
  if experiment_description and experiment is None:
  raise ValueError(
  "Experiment needs to be set in `init` in order to add experiment descriptions."
@@ -252,6 +260,11 @@ def get_encryption_spec(
  )
  return encryption_spec
 
+ @property
+ def api_endpoint(self) -> Optional[str]:
+ """Default API endpoint, if provided."""
+ return self._api_endpoint
+
  @property
  def project(self) -> str:
  """Default project."""
@@ -351,27 +364,32 @@ def get_client_options(
  { "api_endpoint": "us-central1-aiplatform.googleapis.com" } or
  { "api_endpoint": "asia-east1-aiplatform.googleapis.com" }
  """
- if not (self.location or location_override):
- raise ValueError(
- "No location found. Provide or initialize SDK with a location."
- )
 
- region = location_override or self.location
- region = region.lower()
+ api_endpoint = self.api_endpoint
 
- utils.validate_region(region)
+ if api_endpoint is None:
+ if not (self.location or location_override):
+ raise ValueError(
+ "No location found. Provide or initialize SDK with a location."
+ )
 
- service_base_path = api_base_path_override or (
- constants.PREDICTION_API_BASE_PATH
- if prediction_client
- else constants.API_BASE_PATH
- )
+ region = location_override or self.location
+ region = region.lower()
+
+ utils.validate_region(region)
+
+ service_base_path = api_base_path_override or (
+ constants.PREDICTION_API_BASE_PATH
+ if prediction_client
+ else constants.API_BASE_PATH
+ )
+
+ api_endpoint = (
+ f"{region}-{service_base_path}"
+ if not api_path_override
+ else api_path_override
+ )
 
- api_endpoint = (
- f"{region}-{service_base_path}"
- if not api_path_override
- else api_path_override
- )
  return client_options.ClientOptions(api_endpoint=api_endpoint)
 
  def common_location_path(

tests/unit/aiplatform/test_initializer.py

@@ -17,19 +17,19 @@
 
 import importlib
 import os
-import pytest
+from typing import Optional
 from unittest import mock
 from unittest.mock import patch
 
+import pytest
+
 import google.auth
 from google.auth import credentials
-
 from google.cloud.aiplatform import initializer
 from google.cloud.aiplatform.metadata.metadata import _experiment_tracker
 from google.cloud.aiplatform.constants import base as constants
 from google.cloud.aiplatform import utils
 from google.cloud.aiplatform.utils import resource_manager_utils
-
 from google.cloud.aiplatform.compat.services import (
  model_service_client,
 )
@@ -307,30 +307,64 @@ def test_create_client_appended_user_agent(self):
  assert " " + appended_user_agent[0] in user_agent
  assert " " + appended_user_agent[1] in user_agent
 
+ def test_set_api_endpoint(self):
+ initializer.global_config.init(
+ project=_TEST_PROJECT,
+ location=_TEST_LOCATION,
+ api_endpoint="test.googleapis.com",
+ )
+
+ assert initializer.global_config.api_endpoint == "test.googleapis.com"
+
+ def test_not_set_api_endpoint(self):
+ initializer.global_config.init(
+ project=_TEST_PROJECT,
+ location=_TEST_LOCATION,
+ )
+
+ assert initializer.global_config.api_endpoint is None
+
  @pytest.mark.parametrize(
- "init_location, location_override, expected_endpoint",
+ "init_location, location_override, api_endpoint, expected_endpoint",
  [
- ("us-central1", None, "us-central1-aiplatform.googleapis.com"),
+ ("us-central1", None, None, "us-central1-aiplatform.googleapis.com"),
  (
  "us-central1",
  "europe-west4",
+ None,
  "europe-west4-aiplatform.googleapis.com",
  ),
- ("asia-east1", None, "asia-east1-aiplatform.googleapis.com"),
+ ("asia-east1", None, None, "asia-east1-aiplatform.googleapis.com"),
  (
  "asia-southeast1",
  "australia-southeast1",
+ None,
  "australia-southeast1-aiplatform.googleapis.com",
  ),
+ (
+ "asia-east1",
+ None,
+ "us-central1-aiplatform.googleapis.com",
+ "us-central1-aiplatform.googleapis.com",
+ ),
+ (
+ "us-central1",
+ None,
+ "test.aiplatform.googleapis.com",
+ "test.aiplatform.googleapis.com",
+ ),
  ],
  )
  def test_get_client_options(
  self,
  init_location: str,
- location_override: str,
+ location_override: Optional[str],
+ api_endpoint: Optional[str],
  expected_endpoint: str,
  ):
- initializer.global_config.init(location=init_location)
+ initializer.global_config.init(
+ location=init_location, api_endpoint=api_endpoint
+ )
 
  assert (
  initializer.global_config.get_client_options(