From b104d63abcb4d9850a4dc1f1dce824cc291565cb Mon Sep 17 00:00:00 2001 From: Vinny Senthil Date: Fri, 29 Oct 2021 15:38:02 -0700 Subject: [PATCH 1/8] Initial container URI builder --- google/cloud/aiplatform/__init__.py | 2 + google/cloud/aiplatform/constants/__init__.py | 18 +++ .../{constants.py => constants/base.py} | 0 .../cloud/aiplatform/constants/prediction.py | 138 ++++++++++++++++++ google/cloud/aiplatform/helpers/__init__.py | 19 ++- .../helpers/container_uri_builders.py | 136 +++++++++++++++++ 6 files changed, 312 insertions(+), 1 deletion(-) create mode 100644 google/cloud/aiplatform/constants/__init__.py rename google/cloud/aiplatform/{constants.py => constants/base.py} (100%) create mode 100644 google/cloud/aiplatform/constants/prediction.py create mode 100644 google/cloud/aiplatform/helpers/container_uri_builders.py diff --git a/google/cloud/aiplatform/__init__.py b/google/cloud/aiplatform/__init__.py index 17c6952f35..60982f74c0 100644 --- a/google/cloud/aiplatform/__init__.py +++ b/google/cloud/aiplatform/__init__.py @@ -47,6 +47,7 @@ AutoMLTextTrainingJob, AutoMLVideoTrainingJob, ) +from google.cloud.aiplatform import helpers """ Usage: @@ -67,6 +68,7 @@ "explain", "gapic", "init", + "helpers", "hyperparameter_tuning", "log_params", "log_metrics", diff --git a/google/cloud/aiplatform/constants/__init__.py b/google/cloud/aiplatform/constants/__init__.py new file mode 100644 index 0000000000..e83de5cefb --- /dev/null +++ b/google/cloud/aiplatform/constants/__init__.py @@ -0,0 +1,18 @@ +# Copyright 2021 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 +# +# https://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. + +from google.cloud.aiplatform.constants.base import * +from google.cloud.aiplatform.constants.prediction import * + +__all__ = () diff --git a/google/cloud/aiplatform/constants.py b/google/cloud/aiplatform/constants/base.py similarity index 100% rename from google/cloud/aiplatform/constants.py rename to google/cloud/aiplatform/constants/base.py diff --git a/google/cloud/aiplatform/constants/prediction.py b/google/cloud/aiplatform/constants/prediction.py new file mode 100644 index 0000000000..5e9cc95189 --- /dev/null +++ b/google/cloud/aiplatform/constants/prediction.py @@ -0,0 +1,138 @@ +# Copyright 2021 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 +# +# https://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. + +import re + + +# [region]-docker.pkg.dev/vertex-ai/prediction/[framework]-[accelerator].[version]:latest +CONTAINER_URI_PATTERN = re.compile( + r"(?P[\w]+)\-docker\.pkg\.dev\/vertex\-ai\/prediction\/" + r"(?P[\w]+)\-(?P[\w]+)\.(?P[\d-]+):latest" +) + +SKLEARN = "sklearn" +TF = "tf" +TF2 = "tf2" +XGBOOST = "xgboost" + +XGBOOST_CONTAINER_URIS = [ + "us-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.1-4:latest", + "europe-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.1-4:latest", + "asia-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.1-4:latest", + "us-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.1-3:latest", + "europe-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.1-3:latest", + "asia-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.1-3:latest", + "us-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.1-2:latest", + "europe-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.1-2:latest", + "asia-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.1-2:latest", + "us-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.1-1:latest", + "europe-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.1-1:latest", + "asia-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.1-1:latest", + "us-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.0-90:latest", + "europe-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.0-90:latest", + "asia-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.0-90:latest", + "us-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.0-82:latest", + "europe-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.0-82:latest", + "asia-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.0-82:latest", +] + +SKLEARN_CONTAINER_URIS = [ + "us-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.0-24:latest", + "europe-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.0-24:latest", + "asia-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.0-24:latest", + "us-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.0-23:latest", + "europe-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.0-23:latest", + "asia-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.0-23:latest", + "us-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.0-22:latest", + "europe-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.0-22:latest", + "asia-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.0-22:latest", + "us-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.0-20:latest", + "europe-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.0-20:latest", + "asia-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.0-20:latest", +] + +TF_CONTAINER_URIS = [ + "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-6:latest", + "europe-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-6:latest", + "asia-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-6:latest", + "us-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-6:latest", + "europe-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-6:latest", + "asia-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-6:latest", + "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-5:latest", + "europe-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-5:latest", + "asia-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-5:latest", + "us-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-5:latest", + "europe-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-5:latest", + "asia-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-5:latest", + "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-4:latest", + "europe-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-4:latest", + "asia-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-4:latest", + "us-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-4:latest", + "europe-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-4:latest", + "asia-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-4:latest", + "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-3:latest", + "europe-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-3:latest", + "asia-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-3:latest", + "us-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-3:latest", + "europe-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-3:latest", + "asia-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-3:latest", + "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-2:latest", + "europe-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-2:latest", + "asia-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-2:latest", + "us-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-2:latest", + "europe-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-2:latest", + "asia-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-2:latest", + "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-1:latest", + "europe-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-1:latest", + "asia-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-1:latest", + "us-docker.pkg.dev/vertex-ai/prediction/tf-cpu.1-15:latest", + "europe-docker.pkg.dev/vertex-ai/prediction/tf-cpu.1-15:latest", + "asia-docker.pkg.dev/vertex-ai/prediction/tf-cpu.1-15:latest", + "us-docker.pkg.dev/vertex-ai/prediction/tf-gpu.1-15:latest", + "europe-docker.pkg.dev/vertex-ai/prediction/tf-gpu.1-15:latest", + "asia-docker.pkg.dev/vertex-ai/prediction/tf-gpu.1-15:latest", +] + +# Pattern: +# {user_provided_framework: framework_name_in_uri} +_FRAMEWORK_TO_URI_REF = { + SKLEARN: SKLEARN, + "scikitlearn": SKLEARN, + "scikit-learn": SKLEARN, + TF: TF, + "tensorflow": TF, + XGBOOST: XGBOOST, + "xgb": XGBOOST, +} + +# Pattern: +# {framework: (accelerator, no_accelerator)} +_ACCELERATOR_TO_URI_REF = { + # If `None`, accelerator is not supported + SKLEARN: (None, "cpu"), + TF: ("gpu", "cpu"), + TF2: ("gpu", "cpu"), + XGBOOST: (None, "cpu"), +} + +SERVING_CONTAINER_URIS = ( + SKLEARN_CONTAINER_URIS + TF_CONTAINER_URIS + XGBOOST_CONTAINER_URIS +) + +# All Vertex AI Prediction first-party prediction containers as a string +_SERVING_CONTAINER_URIS_STR = " ".join(SERVING_CONTAINER_URIS) + +_SERVING_CONTAINER_DOCUMENTATION_URL = ( + "https://cloud.google.com/vertex-ai/docs/predictions/pre-built-containers" +) diff --git a/google/cloud/aiplatform/helpers/__init__.py b/google/cloud/aiplatform/helpers/__init__.py index 3f031f2bb4..93092661bb 100644 --- a/google/cloud/aiplatform/helpers/__init__.py +++ b/google/cloud/aiplatform/helpers/__init__.py @@ -1,3 +1,20 @@ +# Copyright 2021 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. + from google.cloud.aiplatform.helpers import value_converter +from google.cloud.aiplatform.helpers.container_uri_builders import ( + get_prediction_container_uri, +) -__all__ = (value_converter,) +__all__ = ("get_prediction_container_uri" "value_converter",) diff --git a/google/cloud/aiplatform/helpers/container_uri_builders.py b/google/cloud/aiplatform/helpers/container_uri_builders.py new file mode 100644 index 0000000000..29df554b82 --- /dev/null +++ b/google/cloud/aiplatform/helpers/container_uri_builders.py @@ -0,0 +1,136 @@ +# Copyright 2021 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. + +import re + +from google.cloud.aiplatform.constants import prediction +from google.cloud.aiplatform.initializer import global_config + + +def get_prediction_container_uri( + framework: str, + framework_version: str, + region: str = None, + with_accelerator: bool = False, +) -> str: + """ + Get a Vertex AI pre-built prediction Docker container URI for + a given framework, version, region, and accelerator use. + + Example usage: + ``` + uri = aiplatform.helpers.get_prediction_container_uri( + framework="tensorflow", + framework_version="2.6", + with_accelerator=True + ) + + model = aiplatform.Model.upload( + display_name="boston_housing_", + artifact_uri="gs://my-bucket/my-model/", + serving_container_image_uri=uri + ) + ``` + + Args: + framework (str): + The ML framework of the pre-built container. For example, + "tensorflow", "xgboost", or "sklearn" + framework_version (str): + The version of the specified ML framework as a string. + region (str): + A Vertex AI region or multi-region. Used to select the correct + Artifact Registry multi-region repository and reduce latency. + Must start with "us", "asia" or "europe". If not set, defaults + to location set by `aiplatform.init()`. + with_accelerator (bool): + If set to `True`, return container URI that supports GPU usage. + Default is `False`. + + Returns: + uri (str): + A Vertex AI prediction container URI + + Raises: + ValueError: If containers for provided framework are unavailable, + the container does not support accelerators, or is not available + in the specified version or region. + """ + + DOCS_URI_MESSAGE = ( + f"See {prediction._SERVING_CONTAINER_DOCUMENTATION_URL} " + "for complete list of supported containers" + ) + alt_versions = [] # Alternative framework versions if provided does not exist + + # Validate provided framework + try: + framework = framework.lower() + framework = prediction._FRAMEWORK_TO_URI_REF[framework] + except KeyError: + raise ValueError( + f"No containers found for framework `{framework}`. {DOCS_URI_MESSAGE}" + ) + + version = framework_version.replace(".", "-") + + # Tensorflow 2.x is has a different framework name in URI + if framework == prediction.TF and version.startswith("2"): + framework = prediction.TF2 + + accelerator = prediction._ACCELERATOR_TO_URI_REF[framework][ + 0 if with_accelerator else 1 + ] + + if accelerator is None: + raise ValueError( + f"{framework} containers do not support accelerators. " + f"Please set `with_accelerator` to False. {DOCS_URI_MESSAGE}" + ) + + # If region not provided, use initializer location + region = region or global_config.location + region = region.split("-", 1)[0] + + for uri in re.finditer( + prediction.CONTAINER_URI_PATTERN, prediction._SERVING_CONTAINER_URIS_STR + ): + match = uri.groups() + if framework == match[1] and region == match[0] and accelerator == match[2]: + if version != match[3]: + # If URI matches all but version, add to alternative suggestion + alt_versions.append(match[3].replace("-", ".")) + else: + return uri.group() + + ALTERNATE_VERSIONS_MESSAGE = ( + f"Supported versions for {framework} include {', '.join(alt_versions)}. " + if alt_versions + else "" + ) + + raise ValueError( + f"No serving container for {framework} {framework_version} " + f"{'with accelerator ' if with_accelerator else ''}found. " + f"{ALTERNATE_VERSIONS_MESSAGE}{DOCS_URI_MESSAGE}" + ) + + +def get_training_container_uri( + framework: str, + framework_version: str, + region: str = None, + with_accelerator: bool = False, +): + raise NotImplementedError() From dde683467d85c69009b45692fed33e6244b34b6e Mon Sep 17 00:00:00 2001 From: Vinny Senthil Date: Tue, 16 Nov 2021 12:06:47 -0800 Subject: [PATCH 2/8] Address minor requested changes --- google/cloud/aiplatform/helpers/__init__.py | 4 ++-- .../helpers/container_uri_builders.py | 22 ++++++------------- 2 files changed, 9 insertions(+), 17 deletions(-) diff --git a/google/cloud/aiplatform/helpers/__init__.py b/google/cloud/aiplatform/helpers/__init__.py index 93092661bb..16a758dd0d 100644 --- a/google/cloud/aiplatform/helpers/__init__.py +++ b/google/cloud/aiplatform/helpers/__init__.py @@ -14,7 +14,7 @@ from google.cloud.aiplatform.helpers import value_converter from google.cloud.aiplatform.helpers.container_uri_builders import ( - get_prediction_container_uri, + get_prebuilt_prediction_container_uri, ) -__all__ = ("get_prediction_container_uri" "value_converter",) +__all__ = ("get_prebuilt_prediction_container_uri" "value_converter",) diff --git a/google/cloud/aiplatform/helpers/container_uri_builders.py b/google/cloud/aiplatform/helpers/container_uri_builders.py index 29df554b82..b97ce2f9be 100644 --- a/google/cloud/aiplatform/helpers/container_uri_builders.py +++ b/google/cloud/aiplatform/helpers/container_uri_builders.py @@ -13,15 +13,16 @@ # limitations under the License. import re +from typing import Optional from google.cloud.aiplatform.constants import prediction from google.cloud.aiplatform.initializer import global_config -def get_prediction_container_uri( +def get_prebuilt_prediction_container_uri( framework: str, framework_version: str, - region: str = None, + region: Optional[str] = None, with_accelerator: bool = False, ) -> str: """ @@ -45,17 +46,17 @@ def get_prediction_container_uri( Args: framework (str): - The ML framework of the pre-built container. For example, + Required. The ML framework of the pre-built container. For example, "tensorflow", "xgboost", or "sklearn" framework_version (str): - The version of the specified ML framework as a string. + Required. The version of the specified ML framework as a string. region (str): - A Vertex AI region or multi-region. Used to select the correct + Optional. AI region or multi-region. Used to select the correct Artifact Registry multi-region repository and reduce latency. Must start with "us", "asia" or "europe". If not set, defaults to location set by `aiplatform.init()`. with_accelerator (bool): - If set to `True`, return container URI that supports GPU usage. + Optional. If set to `True`, return container URI that supports GPU usage. Default is `False`. Returns: @@ -125,12 +126,3 @@ def get_prediction_container_uri( f"{'with accelerator ' if with_accelerator else ''}found. " f"{ALTERNATE_VERSIONS_MESSAGE}{DOCS_URI_MESSAGE}" ) - - -def get_training_container_uri( - framework: str, - framework_version: str, - region: str = None, - with_accelerator: bool = False, -): - raise NotImplementedError() From d34ba3d9e36ff59e7d7fc8870b8d705276db0e73 Mon Sep 17 00:00:00 2001 From: Vinny Senthil Date: Wed, 17 Nov 2021 13:46:56 -0800 Subject: [PATCH 3/8] Move enhanced_library from helpers to utils --- google/cloud/aiplatform/gapic/schema/__init__.py | 2 +- google/cloud/aiplatform/helpers/__init__.py | 3 +-- .../aiplatform/utils/enhanced_library/__init__.py | 13 +++++++++++++ .../enhanced_library}/_decorators.py | 2 +- .../enhanced_library}/value_converter.py | 0 tests/unit/enhanced_library/test_value_converter.py | 2 +- 6 files changed, 17 insertions(+), 5 deletions(-) create mode 100644 google/cloud/aiplatform/utils/enhanced_library/__init__.py rename google/cloud/aiplatform/{helpers => utils/enhanced_library}/_decorators.py (97%) rename google/cloud/aiplatform/{helpers => utils/enhanced_library}/value_converter.py (100%) diff --git a/google/cloud/aiplatform/gapic/schema/__init__.py b/google/cloud/aiplatform/gapic/schema/__init__.py index e726749c77..5d31a70f1f 100644 --- a/google/cloud/aiplatform/gapic/schema/__init__.py +++ b/google/cloud/aiplatform/gapic/schema/__init__.py @@ -14,7 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from google.cloud.aiplatform.helpers import _decorators +from google.cloud.aiplatform.utils.enhanced_library import _decorators from google.cloud.aiplatform.v1.schema import predict from google.cloud.aiplatform.v1.schema import trainingjob from google.cloud.aiplatform.v1beta1.schema import predict as predict_v1beta1 diff --git a/google/cloud/aiplatform/helpers/__init__.py b/google/cloud/aiplatform/helpers/__init__.py index 16a758dd0d..e3802910ed 100644 --- a/google/cloud/aiplatform/helpers/__init__.py +++ b/google/cloud/aiplatform/helpers/__init__.py @@ -12,9 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -from google.cloud.aiplatform.helpers import value_converter from google.cloud.aiplatform.helpers.container_uri_builders import ( get_prebuilt_prediction_container_uri, ) -__all__ = ("get_prebuilt_prediction_container_uri" "value_converter",) +__all__ = "get_prebuilt_prediction_container_uri" diff --git a/google/cloud/aiplatform/utils/enhanced_library/__init__.py b/google/cloud/aiplatform/utils/enhanced_library/__init__.py new file mode 100644 index 0000000000..c6334245ae --- /dev/null +++ b/google/cloud/aiplatform/utils/enhanced_library/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2020 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. diff --git a/google/cloud/aiplatform/helpers/_decorators.py b/google/cloud/aiplatform/utils/enhanced_library/_decorators.py similarity index 97% rename from google/cloud/aiplatform/helpers/_decorators.py rename to google/cloud/aiplatform/utils/enhanced_library/_decorators.py index 95aac31c4f..43e395393b 100644 --- a/google/cloud/aiplatform/helpers/_decorators.py +++ b/google/cloud/aiplatform/utils/enhanced_library/_decorators.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from __future__ import absolute_import -from google.cloud.aiplatform.helpers import value_converter +from google.cloud.aiplatform.utils.enhanced_library import value_converter from proto.marshal import Marshal from proto.marshal.rules.struct import ValueRule diff --git a/google/cloud/aiplatform/helpers/value_converter.py b/google/cloud/aiplatform/utils/enhanced_library/value_converter.py similarity index 100% rename from google/cloud/aiplatform/helpers/value_converter.py rename to google/cloud/aiplatform/utils/enhanced_library/value_converter.py diff --git a/tests/unit/enhanced_library/test_value_converter.py b/tests/unit/enhanced_library/test_value_converter.py index 4cbaeed7cf..4bb24372e2 100644 --- a/tests/unit/enhanced_library/test_value_converter.py +++ b/tests/unit/enhanced_library/test_value_converter.py @@ -13,7 +13,7 @@ # limitations under the License. from __future__ import absolute_import -from google.cloud.aiplatform.helpers import value_converter +from google.cloud.aiplatform.utils.enhanced_library import value_converter from google.protobuf import json_format from google.protobuf.struct_pb2 import Value import proto From 9c8742ad65e89c94a86c554fb81f4eedaf74c02d Mon Sep 17 00:00:00 2001 From: Vinny Senthil Date: Wed, 17 Nov 2021 21:04:45 -0800 Subject: [PATCH 4/8] Update logic of uri helper to use programmatic map --- .../cloud/aiplatform/constants/prediction.py | 39 +++----- google/cloud/aiplatform/helpers/__init__.py | 6 +- .../helpers/container_uri_builders.py | 99 ++++++++----------- 3 files changed, 61 insertions(+), 83 deletions(-) diff --git a/google/cloud/aiplatform/constants/prediction.py b/google/cloud/aiplatform/constants/prediction.py index 5e9cc95189..11843b5c85 100644 --- a/google/cloud/aiplatform/constants/prediction.py +++ b/google/cloud/aiplatform/constants/prediction.py @@ -14,6 +14,7 @@ import re +from collections import defaultdict # [region]-docker.pkg.dev/vertex-ai/prediction/[framework]-[accelerator].[version]:latest CONTAINER_URI_PATTERN = re.compile( @@ -104,34 +105,24 @@ "asia-docker.pkg.dev/vertex-ai/prediction/tf-gpu.1-15:latest", ] -# Pattern: -# {user_provided_framework: framework_name_in_uri} -_FRAMEWORK_TO_URI_REF = { - SKLEARN: SKLEARN, - "scikitlearn": SKLEARN, - "scikit-learn": SKLEARN, - TF: TF, - "tensorflow": TF, - XGBOOST: XGBOOST, - "xgb": XGBOOST, -} - -# Pattern: -# {framework: (accelerator, no_accelerator)} -_ACCELERATOR_TO_URI_REF = { - # If `None`, accelerator is not supported - SKLEARN: (None, "cpu"), - TF: ("gpu", "cpu"), - TF2: ("gpu", "cpu"), - XGBOOST: (None, "cpu"), -} - SERVING_CONTAINER_URIS = ( SKLEARN_CONTAINER_URIS + TF_CONTAINER_URIS + XGBOOST_CONTAINER_URIS ) -# All Vertex AI Prediction first-party prediction containers as a string -_SERVING_CONTAINER_URIS_STR = " ".join(SERVING_CONTAINER_URIS) +# Map of all first-party prediction containers +d = defaultdict(lambda: defaultdict(lambda: defaultdict(lambda: defaultdict(str)))) + +for container_uri in SERVING_CONTAINER_URIS: + m = CONTAINER_URI_PATTERN.match(container_uri) + region, framework, accelerator, version = m[1], m[2], m[3], m[4] + version = version.replace("-", ".") + + if framework == TF2: # Store both `tf`, `tf2` as `tf` + framework = TF + + d[region][framework][accelerator][version] = container_uri + +_SERVING_CONTAINER_URI_MAP = d _SERVING_CONTAINER_DOCUMENTATION_URL = ( "https://cloud.google.com/vertex-ai/docs/predictions/pre-built-containers" diff --git a/google/cloud/aiplatform/helpers/__init__.py b/google/cloud/aiplatform/helpers/__init__.py index e3802910ed..e5fa8f665d 100644 --- a/google/cloud/aiplatform/helpers/__init__.py +++ b/google/cloud/aiplatform/helpers/__init__.py @@ -12,8 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -from google.cloud.aiplatform.helpers.container_uri_builders import ( - get_prebuilt_prediction_container_uri, +from google.cloud.aiplatform.helpers import container_uri_builders + +get_prebuilt_prediction_container_uri = ( + container_uri_builders.get_prebuilt_prediction_container_uri ) __all__ = "get_prebuilt_prediction_container_uri" diff --git a/google/cloud/aiplatform/helpers/container_uri_builders.py b/google/cloud/aiplatform/helpers/container_uri_builders.py index b97ce2f9be..f68cd4e331 100644 --- a/google/cloud/aiplatform/helpers/container_uri_builders.py +++ b/google/cloud/aiplatform/helpers/container_uri_builders.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import re from typing import Optional from google.cloud.aiplatform.constants import prediction @@ -23,7 +22,7 @@ def get_prebuilt_prediction_container_uri( framework: str, framework_version: str, region: Optional[str] = None, - with_accelerator: bool = False, + accelerator: Optional[str] = "cpu", ) -> str: """ Get a Vertex AI pre-built prediction Docker container URI for @@ -31,10 +30,10 @@ def get_prebuilt_prediction_container_uri( Example usage: ``` - uri = aiplatform.helpers.get_prediction_container_uri( + uri = aiplatform.helpers.get_prebuilt_prediction_container_uri( framework="tensorflow", framework_version="2.6", - with_accelerator=True + accelerator="gpu" ) model = aiplatform.Model.upload( @@ -47,82 +46,68 @@ def get_prebuilt_prediction_container_uri( Args: framework (str): Required. The ML framework of the pre-built container. For example, - "tensorflow", "xgboost", or "sklearn" + `"tensorflow"`, `"xgboost"`, or `"sklearn"` framework_version (str): Required. The version of the specified ML framework as a string. region (str): Optional. AI region or multi-region. Used to select the correct Artifact Registry multi-region repository and reduce latency. - Must start with "us", "asia" or "europe". If not set, defaults - to location set by `aiplatform.init()`. - with_accelerator (bool): - Optional. If set to `True`, return container URI that supports GPU usage. - Default is `False`. + Must start with `"us"`, `"asia"` or `"europe"`. + Default is location set by `aiplatform.init()`. + accelerator (str): + Optional. The type of accelerator support provided by container. For + example: `"cpu"` or `"gpu"` + Default is `"cpu"`. Returns: uri (str): A Vertex AI prediction container URI Raises: - ValueError: If containers for provided framework are unavailable, - the container does not support accelerators, or is not available - in the specified version or region. + ValueError: If containers for provided framework are unavailable or the + container does not support the specified version, accelerator, or region. """ - + URI_MAP = prediction._SERVING_CONTAINER_URI_MAP DOCS_URI_MESSAGE = ( f"See {prediction._SERVING_CONTAINER_DOCUMENTATION_URL} " "for complete list of supported containers" ) - alt_versions = [] # Alternative framework versions if provided does not exist - # Validate provided framework - try: - framework = framework.lower() - framework = prediction._FRAMEWORK_TO_URI_REF[framework] - except KeyError: + # If region not provided, use initializer location + region = region or global_config.location + region = region.split("-", 1)[0] + framework = framework.lower() + + if not URI_MAP[region]: + del URI_MAP[region] raise ValueError( - f"No containers found for framework `{framework}`. {DOCS_URI_MESSAGE}" + f"Unsupported container region `{region}`, supported regions are " + f"{', '.join(URI_MAP.keys())}. " + f"{DOCS_URI_MESSAGE}" ) - version = framework_version.replace(".", "-") - - # Tensorflow 2.x is has a different framework name in URI - if framework == prediction.TF and version.startswith("2"): - framework = prediction.TF2 - - accelerator = prediction._ACCELERATOR_TO_URI_REF[framework][ - 0 if with_accelerator else 1 - ] + if not URI_MAP[region][framework]: + del URI_MAP[region][framework] + raise ValueError( + f"No containers found for framework `{framework}`. Supported frameworks are " + f"{', '.join(URI_MAP[region].keys())} {DOCS_URI_MESSAGE}" + ) - if accelerator is None: + if not URI_MAP[region][framework][accelerator]: + del URI_MAP[region][framework][accelerator] raise ValueError( - f"{framework} containers do not support accelerators. " - f"Please set `with_accelerator` to False. {DOCS_URI_MESSAGE}" + f"{framework} containers do not support `{accelerator}` accelerator. Supported accelerators " + f"are {', '.join(URI_MAP[region][framework].keys())}. {DOCS_URI_MESSAGE}" ) - # If region not provided, use initializer location - region = region or global_config.location - region = region.split("-", 1)[0] + final_uri = URI_MAP[region][framework][accelerator][framework_version] - for uri in re.finditer( - prediction.CONTAINER_URI_PATTERN, prediction._SERVING_CONTAINER_URIS_STR - ): - match = uri.groups() - if framework == match[1] and region == match[0] and accelerator == match[2]: - if version != match[3]: - # If URI matches all but version, add to alternative suggestion - alt_versions.append(match[3].replace("-", ".")) - else: - return uri.group() - - ALTERNATE_VERSIONS_MESSAGE = ( - f"Supported versions for {framework} include {', '.join(alt_versions)}. " - if alt_versions - else "" - ) + if not final_uri: + del URI_MAP[region][framework][accelerator][framework_version] + raise ValueError( + f"No serving container for `{framework}` version `{framework_version}` " + f"with accelerator `{accelerator}` found. Supported versions " + f"include {', '.join(URI_MAP[region][framework][accelerator].keys())}. {DOCS_URI_MESSAGE}" + ) - raise ValueError( - f"No serving container for {framework} {framework_version} " - f"{'with accelerator ' if with_accelerator else ''}found. " - f"{ALTERNATE_VERSIONS_MESSAGE}{DOCS_URI_MESSAGE}" - ) + return final_uri From 26947053eff7a2d683b7bf650493ac99da40eb6d Mon Sep 17 00:00:00 2001 From: Vinny Senthil Date: Wed, 17 Nov 2021 21:16:06 -0800 Subject: [PATCH 5/8] Split aiplatform.constants into multiple modules --- google/cloud/aiplatform/constants/__init__.py | 6 +++--- google/cloud/aiplatform/initializer.py | 2 +- google/cloud/aiplatform/jobs.py | 2 +- google/cloud/aiplatform/tensorboard/uploader_main.py | 3 ++- google/cloud/aiplatform/training_jobs.py | 2 +- google/cloud/aiplatform/utils/__init__.py | 2 +- tests/unit/aiplatform/test_datasets.py | 3 ++- tests/unit/aiplatform/test_initializer.py | 2 +- 8 files changed, 12 insertions(+), 10 deletions(-) diff --git a/google/cloud/aiplatform/constants/__init__.py b/google/cloud/aiplatform/constants/__init__.py index e83de5cefb..95f437a335 100644 --- a/google/cloud/aiplatform/constants/__init__.py +++ b/google/cloud/aiplatform/constants/__init__.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from google.cloud.aiplatform.constants.base import * -from google.cloud.aiplatform.constants.prediction import * +from google.cloud.aiplatform.constants import base +from google.cloud.aiplatform.constants import prediction -__all__ = () +__all__ = ("base", "prediction") diff --git a/google/cloud/aiplatform/initializer.py b/google/cloud/aiplatform/initializer.py index ea1a51c8a7..00f6b19b40 100644 --- a/google/cloud/aiplatform/initializer.py +++ b/google/cloud/aiplatform/initializer.py @@ -29,7 +29,7 @@ from google.auth.exceptions import GoogleAuthError from google.cloud.aiplatform import compat -from google.cloud.aiplatform import constants +from google.cloud.aiplatform.constants import base as constants from google.cloud.aiplatform import utils from google.cloud.aiplatform.metadata import metadata diff --git a/google/cloud/aiplatform/jobs.py b/google/cloud/aiplatform/jobs.py index c1ba9739cc..94aed4d70f 100644 --- a/google/cloud/aiplatform/jobs.py +++ b/google/cloud/aiplatform/jobs.py @@ -44,7 +44,7 @@ machine_resources as gca_machine_resources_compat, study as gca_study_compat, ) -from google.cloud.aiplatform import constants +from google.cloud.aiplatform.constants import base as constants from google.cloud.aiplatform import initializer from google.cloud.aiplatform import hyperparameter_tuning from google.cloud.aiplatform import utils diff --git a/google/cloud/aiplatform/tensorboard/uploader_main.py b/google/cloud/aiplatform/tensorboard/uploader_main.py index 7c868fcfb7..163d6c46e8 100644 --- a/google/cloud/aiplatform/tensorboard/uploader_main.py +++ b/google/cloud/aiplatform/tensorboard/uploader_main.py @@ -30,6 +30,7 @@ from google.cloud import storage from google.cloud import aiplatform +from google.cloud.aiplatform.constants import base as constants from google.cloud.aiplatform.tensorboard import uploader from google.cloud.aiplatform.utils import TensorboardClientWithOverride @@ -89,7 +90,7 @@ def main(argv): if len(argv) > 1: raise app.UsageError("Too many command-line arguments.") - aiplatform.constants.API_BASE_PATH = FLAGS.api_uri + constants.API_BASE_PATH = FLAGS.api_uri m = re.match( "projects/(.*)/locations/(.*)/tensorboards/.*", FLAGS.tensorboard_resource_name ) diff --git a/google/cloud/aiplatform/training_jobs.py b/google/cloud/aiplatform/training_jobs.py index 679bb277ab..4afd4920db 100644 --- a/google/cloud/aiplatform/training_jobs.py +++ b/google/cloud/aiplatform/training_jobs.py @@ -23,7 +23,7 @@ from google.auth import credentials as auth_credentials from google.cloud.aiplatform import base -from google.cloud.aiplatform import constants +from google.cloud.aiplatform.constants import base as constants from google.cloud.aiplatform import datasets from google.cloud.aiplatform import initializer from google.cloud.aiplatform import models diff --git a/google/cloud/aiplatform/utils/__init__.py b/google/cloud/aiplatform/utils/__init__.py index 56baa42a9d..7844dd5bae 100644 --- a/google/cloud/aiplatform/utils/__init__.py +++ b/google/cloud/aiplatform/utils/__init__.py @@ -30,7 +30,7 @@ from google.cloud import storage from google.cloud.aiplatform import compat -from google.cloud.aiplatform import constants +from google.cloud.aiplatform.constants import base as constants from google.cloud.aiplatform import initializer from google.cloud.aiplatform.compat.services import ( diff --git a/tests/unit/aiplatform/test_datasets.py b/tests/unit/aiplatform/test_datasets.py index f80beac240..71ca5907ab 100644 --- a/tests/unit/aiplatform/test_datasets.py +++ b/tests/unit/aiplatform/test_datasets.py @@ -30,6 +30,7 @@ from google.cloud import aiplatform from google.cloud.aiplatform import base from google.cloud.aiplatform import compat +from google.cloud.aiplatform.constants import base as constants from google.cloud.aiplatform import datasets from google.cloud.aiplatform import initializer from google.cloud.aiplatform import schema @@ -514,7 +515,7 @@ def test_init_dataset_with_alt_location(self, get_dataset_tabular_gcs_mock): assert ( ds.api_client._clients[compat.DEFAULT_VERSION]._client_options.api_endpoint - == f"{_TEST_LOCATION}-{aiplatform.constants.API_BASE_PATH}" + == f"{_TEST_LOCATION}-{constants.API_BASE_PATH}" ) assert _TEST_ALT_LOCATION != _TEST_LOCATION diff --git a/tests/unit/aiplatform/test_initializer.py b/tests/unit/aiplatform/test_initializer.py index 7e65c99b4c..f4043a5eba 100644 --- a/tests/unit/aiplatform/test_initializer.py +++ b/tests/unit/aiplatform/test_initializer.py @@ -26,7 +26,7 @@ from google.cloud.aiplatform import initializer from google.cloud.aiplatform.metadata.metadata import metadata_service -from google.cloud.aiplatform import constants +from google.cloud.aiplatform.constants import base as constants from google.cloud.aiplatform import utils from google.cloud.aiplatform_v1.services.model_service import ( From e0ceb800ede17aa3474a379abcc7446ab1bc79d1 Mon Sep 17 00:00:00 2001 From: Vinny Senthil Date: Mon, 22 Nov 2021 10:46:51 -0800 Subject: [PATCH 6/8] Add tests for 1P pred container URI helper, fixes --- .../cloud/aiplatform/constants/prediction.py | 4 +- .../helpers/container_uri_builders.py | 4 +- tests/unit/aiplatform/test_helpers.py | 149 ++++++++++++++++++ 3 files changed, 153 insertions(+), 4 deletions(-) create mode 100644 tests/unit/aiplatform/test_helpers.py diff --git a/google/cloud/aiplatform/constants/prediction.py b/google/cloud/aiplatform/constants/prediction.py index 11843b5c85..56ba7c93f1 100644 --- a/google/cloud/aiplatform/constants/prediction.py +++ b/google/cloud/aiplatform/constants/prediction.py @@ -117,8 +117,8 @@ region, framework, accelerator, version = m[1], m[2], m[3], m[4] version = version.replace("-", ".") - if framework == TF2: # Store both `tf`, `tf2` as `tf` - framework = TF + if framework in (TF2, TF): # Store both `tf`, `tf2` as `tensorflow` + framework = "tensorflow" d[region][framework][accelerator][version] = container_uri diff --git a/google/cloud/aiplatform/helpers/container_uri_builders.py b/google/cloud/aiplatform/helpers/container_uri_builders.py index f68cd4e331..bdfba9bbef 100644 --- a/google/cloud/aiplatform/helpers/container_uri_builders.py +++ b/google/cloud/aiplatform/helpers/container_uri_builders.py @@ -15,7 +15,7 @@ from typing import Optional from google.cloud.aiplatform.constants import prediction -from google.cloud.aiplatform.initializer import global_config +from google.cloud.aiplatform import initializer def get_prebuilt_prediction_container_uri( @@ -74,7 +74,7 @@ def get_prebuilt_prediction_container_uri( ) # If region not provided, use initializer location - region = region or global_config.location + region = region or initializer.global_config.location region = region.split("-", 1)[0] framework = framework.lower() diff --git a/tests/unit/aiplatform/test_helpers.py b/tests/unit/aiplatform/test_helpers.py new file mode 100644 index 0000000000..4d9afe6fa5 --- /dev/null +++ b/tests/unit/aiplatform/test_helpers.py @@ -0,0 +1,149 @@ +# -*- coding: utf-8 -*- + +# Copyright 2021 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. +# + +import importlib +import pytest + +from typing import Sequence + +from google.cloud import aiplatform +from google.cloud.aiplatform import helpers +from google.cloud.aiplatform import initializer + + +class TestContainerUriHelpers: + def setup_method(self): + importlib.reload(initializer) + importlib.reload(aiplatform) + + def teardown_method(self): + initializer.global_pool.shutdown(wait=True) + + def _build_predict_uri_kwargs(self, args: Sequence[str]) -> dict: + """ + Takes list of values for all method parameters and return dict of kwargs, + dropping keywords that were set as None. + """ + func = helpers.get_prebuilt_prediction_container_uri + arg_names = func.__code__.co_varnames[: func.__code__.co_argcount] + return {k: v for k, v in dict(zip(arg_names, args)).items() if v is not None} + + @pytest.mark.parametrize( + "args, expected_uri", + [ + ( + ("tensorflow", "2.6", None, None), + "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-6:latest", + ), + ( + ("tensorflow", "1.15", "europe-west4", None), + "europe-docker.pkg.dev/vertex-ai/prediction/tf-cpu.1-15:latest", + ), + ( + ("tensorflow", "2.2", None, "gpu"), + "us-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-2:latest", + ), + ( + ("sklearn", "0.24", "asia", "cpu"), + "asia-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.0-24:latest", + ), + ( + ("sklearn", "0.20", None, None), + "us-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.0-20:latest", + ), + ( + ("xgboost", "1.3", None, None), + "us-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.1-3:latest", + ), + ( + ("xgboost", "0.90", "europe", None), + "europe-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.0-90:latest", + ), + ], + ) + def test_correct_prediction_uri_args(self, args, expected_uri): + uri = helpers.get_prebuilt_prediction_container_uri( + **self._build_predict_uri_kwargs(args) + ) + + assert uri == expected_uri + + def test_correct_prediction_uri_args_with_init_location(self): + """ + Ensure that aiplatform.init location is used when region + is not provided + """ + uri = aiplatform.helpers.get_prebuilt_prediction_container_uri( + "tensorflow", "2.6" + ) + # SDK default location is us-central1 + assert uri.startswith("us-docker.pkg.dev") + + aiplatform.init(location="asia-northeast3") + uri = aiplatform.helpers.get_prebuilt_prediction_container_uri( + "tensorflow", "2.6" + ) + assert uri.startswith("asia-docker.pkg.dev") + + aiplatform.init(location="europe-west2") + uri = aiplatform.helpers.get_prebuilt_prediction_container_uri( + "xgboost", "0.90" + ) + assert uri.startswith("europe-docker.pkg.dev") + + @pytest.mark.parametrize( + "args, expected_error_msg", + [ + ( + ("pytorch", "1.10", None, None), + "No containers found for framework `pytorch`. Supported frameworks are", + ), + ( + ("tensorflow", "9.15", None, None), + ( + "No serving container for `tensorflow` version `9.15` with accelerator " + "`cpu` found. Supported versions include" + ), + ), + ( + # Make sure region error supercedes version error + ("tensorflow", "9.15", "pluto", None), + "Unsupported container region `pluto`, supported regions are ", + ), + ( + ("tensorflow", "2.2", "narnia", None), + "Unsupported container region `narnia`, supported regions are ", + ), + ( + ("sklearn", "0.24", "asia", "gpu"), + "sklearn containers do not support `gpu` accelerator. Supported accelerators are cpu.", + ), + ( + # Make sure framework error supercedes accelerator error + ("onnx", "1.9", None, "gpu"), + "No containers found for framework `onnx`. Supported frameworks are", + ), + ], + ) + def test_invalid_prediction_uri_args(self, args, expected_error_msg): + + with pytest.raises(ValueError) as err: + helpers.get_prebuilt_prediction_container_uri( + **self._build_predict_uri_kwargs(args) + ) + + assert err.match(expected_error_msg) From 123aa2db6c256d51afe8356375bfa8ec351141fa Mon Sep 17 00:00:00 2001 From: Vinny Senthil Date: Tue, 23 Nov 2021 14:48:54 -0800 Subject: [PATCH 7/8] Address final requested changes, update OWNERS --- .github/CODEOWNERS | 3 +++ .../aiplatform/helpers/container_uri_builders.py | 14 +++++--------- .../aiplatform/utils/enhanced_library/__init__.py | 2 +- 3 files changed, 9 insertions(+), 10 deletions(-) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index b1ed9d116a..bb192279f6 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -26,3 +26,6 @@ CHANGELOG.md @googleapis/cloud-aiplatform-model-builder-sdk @googleapis/cdpe-cloudai README.rst @googleapis/cloud-aiplatform-model-builder-sdk @googleapis/cdpe-cloudai setup.py @googleapis/cloud-aiplatform-model-builder-sdk @googleapis/cdpe-cloudai + +# Vertex AI product team-specific ownership +/google/cloud/aiplatform/constants/prediction.py @googleapis/vertex-prediction-team diff --git a/google/cloud/aiplatform/helpers/container_uri_builders.py b/google/cloud/aiplatform/helpers/container_uri_builders.py index bdfba9bbef..6b49d3e230 100644 --- a/google/cloud/aiplatform/helpers/container_uri_builders.py +++ b/google/cloud/aiplatform/helpers/container_uri_builders.py @@ -22,7 +22,7 @@ def get_prebuilt_prediction_container_uri( framework: str, framework_version: str, region: Optional[str] = None, - accelerator: Optional[str] = "cpu", + accelerator: str = "cpu", ) -> str: """ Get a Vertex AI pre-built prediction Docker container URI for @@ -78,32 +78,28 @@ def get_prebuilt_prediction_container_uri( region = region.split("-", 1)[0] framework = framework.lower() - if not URI_MAP[region]: - del URI_MAP[region] + if not URI_MAP.get(region): raise ValueError( f"Unsupported container region `{region}`, supported regions are " f"{', '.join(URI_MAP.keys())}. " f"{DOCS_URI_MESSAGE}" ) - if not URI_MAP[region][framework]: - del URI_MAP[region][framework] + if not URI_MAP[region].get(framework): raise ValueError( f"No containers found for framework `{framework}`. Supported frameworks are " f"{', '.join(URI_MAP[region].keys())} {DOCS_URI_MESSAGE}" ) - if not URI_MAP[region][framework][accelerator]: - del URI_MAP[region][framework][accelerator] + if not URI_MAP[region][framework].get(accelerator): raise ValueError( f"{framework} containers do not support `{accelerator}` accelerator. Supported accelerators " f"are {', '.join(URI_MAP[region][framework].keys())}. {DOCS_URI_MESSAGE}" ) - final_uri = URI_MAP[region][framework][accelerator][framework_version] + final_uri = URI_MAP[region][framework][accelerator].get(framework_version) if not final_uri: - del URI_MAP[region][framework][accelerator][framework_version] raise ValueError( f"No serving container for `{framework}` version `{framework_version}` " f"with accelerator `{accelerator}` found. Supported versions " diff --git a/google/cloud/aiplatform/utils/enhanced_library/__init__.py b/google/cloud/aiplatform/utils/enhanced_library/__init__.py index c6334245ae..7e1ec16ec8 100644 --- a/google/cloud/aiplatform/utils/enhanced_library/__init__.py +++ b/google/cloud/aiplatform/utils/enhanced_library/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2020 Google LLC +# Copyright 2021 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. From d7a3ae904259a3d3274fdbd0da657ed98d79f12e Mon Sep 17 00:00:00 2001 From: Vinny Senthil Date: Wed, 24 Nov 2021 11:45:30 -0800 Subject: [PATCH 8/8] Add newest prediction containers --- google/cloud/aiplatform/constants/prediction.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/google/cloud/aiplatform/constants/prediction.py b/google/cloud/aiplatform/constants/prediction.py index 56ba7c93f1..682978ea58 100644 --- a/google/cloud/aiplatform/constants/prediction.py +++ b/google/cloud/aiplatform/constants/prediction.py @@ -49,6 +49,9 @@ ] SKLEARN_CONTAINER_URIS = [ + "us-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.1-0:latest", + "europe-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.1-0:latest", + "asia-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.1-0:latest", "us-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.0-24:latest", "europe-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.0-24:latest", "asia-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.0-24:latest", @@ -64,6 +67,12 @@ ] TF_CONTAINER_URIS = [ + "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-7:latest", + "europe-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-7:latest", + "asia-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-7:latest", + "us-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-7:latest", + "europe-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-7:latest", + "asia-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-7:latest", "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-6:latest", "europe-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-6:latest", "asia-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-6:latest",