feat: allow the prediction endpoint to be overridden

googleapis · Jun 10, 2021 · d265924 · d265924
1 parent 8cfd611
commit d265924
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 3 deletions.
diff --git a/google/cloud/aiplatform/constants.py b/google/cloud/aiplatform/constants.py
@@ -33,6 +33,7 @@
 }
 
 API_BASE_PATH = "aiplatform.googleapis.com"
+PREDICTION_API_BASE_PATH = API_BASE_PATH
 
 # Batch Prediction
 BATCH_PREDICTION_INPUT_STORAGE_FORMATS = (

diff --git a/google/cloud/aiplatform/initializer.py b/google/cloud/aiplatform/initializer.py
@@ -194,7 +194,7 @@ def encryption_spec_key_name(self) -> Optional[str]:
         return self._encryption_spec_key_name
 
     def get_client_options(
-        self, location_override: Optional[str] = None
+        self, location_override: Optional[str] = None, prediction_client: bool = False
     ) -> client_options.ClientOptions:
         """Creates GAPIC client_options using location and type.
 
@@ -203,6 +203,8 @@ def get_client_options(
                 Set this parameter to get client options for a location different from
                 location set by initializer. Must be a GCP region supported by AI
                 Platform (Unified).
+            prediction_client (str): Optional flag to use a prediction endpoint.
+
 
         Returns:
             clients_options (google.api_core.client_options.ClientOptions):
@@ -220,8 +222,14 @@ def get_client_options(
 
         utils.validate_region(region)
 
+        service_base_path = (
+            constants.PREDICTION_API_BASE_PATH
+            if prediction_client
+            else constants.API_BASE_PATH
+        )
+
         return client_options.ClientOptions(
-            api_endpoint=f"{region}-{constants.API_BASE_PATH}"
+            api_endpoint=f"{region}-{service_base_path}"
         )
 
     def common_location_path(
@@ -278,7 +286,8 @@ def create_client(
         kwargs = {
             "credentials": credentials or self.credentials,
             "client_options": self.get_client_options(
-                location_override=location_override
+                location_override=location_override,
+                prediction_client=prediction_client,
             ),
             "client_info": client_info,
         }