feature: Vertex AI Prediction Custom Prediction Routine (#1535)

* feature: Vertex AI Prediction Custom Prediction Routine

* chore: Changed SDK requirement for CPR  to main pypi package

Author: kthytang

.kokoro/presubmit/system.cfg

@@ -11,3 +11,6 @@ env_vars: {
     key: "PYTEST_ADDOPTS"
     value: "-n=auto --dist=loadscope"
 }
+
+# Kokoro VM timeout of 5 hours for system tests
+timeout_mins: 300

google/cloud/aiplatform/base.py

@@ -506,6 +506,7 @@ def _instantiate_client(
         cls,
         location: Optional[str] = None,
         credentials: Optional[auth_credentials.Credentials] = None,
+        appended_user_agent: Optional[List[str]] = None,
     ) -> utils.VertexAiServiceClientWithOverride:
         """Helper method to instantiate service client for resource noun.
 
@@ -514,6 +515,9 @@ def _instantiate_client(
             credentials (google.auth.credentials.Credentials):
                 Optional custom credentials to use when accessing interacting with
                 resource noun.
+            appended_user_agent (List[str]):
+                Optional. User agent appended in the client info. If more than one,
+                it will be separated by spaces.
         Returns:
             client (utils.VertexAiServiceClientWithOverride):
                 Initialized service client for this service noun with optional overrides.
@@ -522,6 +526,7 @@ def _instantiate_client(
             client_class=cls.client_class,
             credentials=credentials,
             location_override=location,
+            appended_user_agent=appended_user_agent,
         )
 
     @classmethod

google/cloud/aiplatform/constants/prediction.py

@@ -22,6 +22,12 @@
     r"(?P<framework>[\w]+)\-(?P<accelerator>[\w]+)\.(?P<version>[\d-]+):latest"
 )
 
+CONTAINER_URI_REGEX = (
+    r"^(us|europe|asia)-docker.pkg.dev/"
+    r"vertex-ai/prediction/"
+    r"(tf|sklearn|xgboost).+$"
+)
+
 SKLEARN = "sklearn"
 TF = "tf"
 TF2 = "tf2"
@@ -145,3 +151,30 @@
 _SERVING_CONTAINER_DOCUMENTATION_URL = (
     "https://cloud.google.com/vertex-ai/docs/predictions/pre-built-containers"
 )
+
+# Variables set by Vertex AI. For more details, please refer to
+# https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#aip-variables
+DEFAULT_AIP_HTTP_PORT = 8080
+AIP_HTTP_PORT = "AIP_HTTP_PORT"
+AIP_HEALTH_ROUTE = "AIP_HEALTH_ROUTE"
+AIP_PREDICT_ROUTE = "AIP_PREDICT_ROUTE"
+AIP_STORAGE_URI = "AIP_STORAGE_URI"
+
+# Default values for Prediction local experience.
+DEFAULT_LOCAL_PREDICT_ROUTE = "/predict"
+DEFAULT_LOCAL_HEALTH_ROUTE = "/health"
+DEFAULT_LOCAL_RUN_GPU_CAPABILITIES = [["utility", "compute"]]
+DEFAULT_LOCAL_RUN_GPU_COUNT = -1
+
+CUSTOM_PREDICTION_ROUTINES = "custom-prediction-routines"
+
+# Headers' related constants for the handler usage.
+CONTENT_TYPE_HEADER_REGEX = re.compile("^[Cc]ontent-?[Tt]ype$")
+ACCEPT_HEADER_REGEX = re.compile("^[Aa]ccept$")
+ANY_ACCEPT_TYPE = "*/*"
+DEFAULT_ACCEPT_VALUE = "application/json"
+
+# Model filenames.
+MODEL_FILENAME_BST = "model.bst"
+MODEL_FILENAME_JOBLIB = "model.joblib"
+MODEL_FILENAME_PKL = "model.pkl"

google/cloud/aiplatform/docker_utils/__init__.py

@@ -0,0 +1,16 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2022 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#