feat: allow the prediction endpoint to be overridden (#461)

Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly:
- [ ] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/python-aiplatform/issues/new/choose) before writing your code!  That way we can discuss the change, evaluate designs, and agree on the general idea
- [ ] Ensure the tests and linter pass
- [ ] Code coverage does not decrease (if any source code was changed)
- [ ] Appropriate docs were updated (if necessary)

Fixes #<issue_number_goes_here> 🦕

Author: geraint0923

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 = (

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,
         }