Skip to content

Commit 91dd3c0

Browse files
vinnysenthilvinnysenthil
authored
feat: Add prediction container URI builder method (googleapis#805)
* Initial container URI builder * Address minor requested changes * Move enhanced_library from helpers to utils * Update logic of uri helper to use programmatic map * Split aiplatform.constants into multiple modules * Add tests for 1P pred container URI helper, fixes * Address final requested changes, update OWNERS * Add newest prediction containers
1 parent 004bf5f commit 91dd3c0

20 files changed

Lines changed: 464 additions & 12 deletions

File tree

‎.github/CODEOWNERS‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,3 +26,6 @@
2626
CHANGELOG.md @googleapis/cloud-aiplatform-model-builder-sdk @googleapis/cdpe-cloudai
2727
README.rst @googleapis/cloud-aiplatform-model-builder-sdk @googleapis/cdpe-cloudai
2828
setup.py @googleapis/cloud-aiplatform-model-builder-sdk @googleapis/cdpe-cloudai
29+
30+
# Vertex AI product team-specific ownership
31+
/google/cloud/aiplatform/constants/prediction.py @googleapis/vertex-prediction-team

‎google/cloud/aiplatform/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@
5353
AutoMLTextTrainingJob,
5454
AutoMLVideoTrainingJob,
5555
)
56+
from google.cloud.aiplatform import helpers
5657

5758
"""
5859
Usage:
@@ -73,6 +74,7 @@
7374
"explain",
7475
"gapic",
7576
"init",
77+
"helpers",
7678
"hyperparameter_tuning",
7779
"log_params",
7880
"log_metrics",

‎google/cloud/aiplatform/constants/__init__.py‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
# Copyright 2021 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from google.cloud.aiplatform.constants import base
16+
from google.cloud.aiplatform.constants import prediction
17+
18+
__all__ = ("base", "prediction")

google/cloud/aiplatform/constants.py renamed to google/cloud/aiplatform/constants/base.py

File renamed without changes.

‎google/cloud/aiplatform/constants/prediction.py‎

Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,138 @@
1+
# Copyright 2021 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import re
16+
17+
from collections import defaultdict
18+
19+
# [region]-docker.pkg.dev/vertex-ai/prediction/[framework]-[accelerator].[version]:latest
20+
CONTAINER_URI_PATTERN = re.compile(
21+
r"(?P<region>[\w]+)\-docker\.pkg\.dev\/vertex\-ai\/prediction\/"
22+
r"(?P<framework>[\w]+)\-(?P<accelerator>[\w]+)\.(?P<version>[\d-]+):latest"
23+
)
24+
25+
SKLEARN = "sklearn"
26+
TF = "tf"
27+
TF2 = "tf2"
28+
XGBOOST = "xgboost"
29+
30+
XGBOOST_CONTAINER_URIS = [
31+
"us-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.1-4:latest",
32+
"europe-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.1-4:latest",
33+
"asia-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.1-4:latest",
34+
"us-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.1-3:latest",
35+
"europe-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.1-3:latest",
36+
"asia-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.1-3:latest",
37+
"us-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.1-2:latest",
38+
"europe-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.1-2:latest",
39+
"asia-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.1-2:latest",
40+
"us-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.1-1:latest",
41+
"europe-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.1-1:latest",
42+
"asia-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.1-1:latest",
43+
"us-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.0-90:latest",
44+
"europe-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.0-90:latest",
45+
"asia-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.0-90:latest",
46+
"us-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.0-82:latest",
47+
"europe-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.0-82:latest",
48+
"asia-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.0-82:latest",
49+
]
50+
51+
SKLEARN_CONTAINER_URIS = [
52+
"us-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.1-0:latest",
53+
"europe-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.1-0:latest",
54+
"asia-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.1-0:latest",
55+
"us-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.0-24:latest",
56+
"europe-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.0-24:latest",
57+
"asia-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.0-24:latest",
58+
"us-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.0-23:latest",
59+
"europe-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.0-23:latest",
60+
"asia-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.0-23:latest",
61+
"us-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.0-22:latest",
62+
"europe-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.0-22:latest",
63+
"asia-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.0-22:latest",
64+
"us-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.0-20:latest",
65+
"europe-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.0-20:latest",
66+
"asia-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.0-20:latest",
67+
]
68+
69+
TF_CONTAINER_URIS = [
70+
"us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-7:latest",
71+
"europe-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-7:latest",
72+
"asia-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-7:latest",
73+
"us-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-7:latest",
74+
"europe-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-7:latest",
75+
"asia-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-7:latest",
76+
"us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-6:latest",
77+
"europe-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-6:latest",
78+
"asia-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-6:latest",
79+
"us-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-6:latest",
80+
"europe-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-6:latest",
81+
"asia-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-6:latest",
82+
"us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-5:latest",
83+
"europe-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-5:latest",
84+
"asia-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-5:latest",
85+
"us-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-5:latest",
86+
"europe-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-5:latest",
87+
"asia-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-5:latest",
88+
"us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-4:latest",
89+
"europe-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-4:latest",
90+
"asia-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-4:latest",
91+
"us-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-4:latest",
92+
"europe-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-4:latest",
93+
"asia-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-4:latest",
94+
"us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-3:latest",
95+
"europe-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-3:latest",
96+
"asia-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-3:latest",
97+
"us-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-3:latest",
98+
"europe-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-3:latest",
99+
"asia-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-3:latest",
100+
"us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-2:latest",
101+
"europe-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-2:latest",
102+
"asia-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-2:latest",
103+
"us-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-2:latest",
104+
"europe-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-2:latest",
105+
"asia-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-2:latest",
106+
"us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-1:latest",
107+
"europe-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-1:latest",
108+
"asia-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-1:latest",
109+
"us-docker.pkg.dev/vertex-ai/prediction/tf-cpu.1-15:latest",
110+
"europe-docker.pkg.dev/vertex-ai/prediction/tf-cpu.1-15:latest",
111+
"asia-docker.pkg.dev/vertex-ai/prediction/tf-cpu.1-15:latest",
112+
"us-docker.pkg.dev/vertex-ai/prediction/tf-gpu.1-15:latest",
113+
"europe-docker.pkg.dev/vertex-ai/prediction/tf-gpu.1-15:latest",
114+
"asia-docker.pkg.dev/vertex-ai/prediction/tf-gpu.1-15:latest",
115+
]
116+
117+
SERVING_CONTAINER_URIS = (
118+
SKLEARN_CONTAINER_URIS + TF_CONTAINER_URIS + XGBOOST_CONTAINER_URIS
119+
)
120+
121+
# Map of all first-party prediction containers
122+
d = defaultdict(lambda: defaultdict(lambda: defaultdict(lambda: defaultdict(str))))
123+
124+
for container_uri in SERVING_CONTAINER_URIS:
125+
m = CONTAINER_URI_PATTERN.match(container_uri)
126+
region, framework, accelerator, version = m[1], m[2], m[3], m[4]
127+
version = version.replace("-", ".")
128+
129+
if framework in (TF2, TF): # Store both `tf`, `tf2` as `tensorflow`
130+
framework = "tensorflow"
131+
132+
d[region][framework][accelerator][version] = container_uri
133+
134+
_SERVING_CONTAINER_URI_MAP = d
135+
136+
_SERVING_CONTAINER_DOCUMENTATION_URL = (
137+
"https://cloud.google.com/vertex-ai/docs/predictions/pre-built-containers"
138+
)

‎google/cloud/aiplatform/gapic/schema/__init__.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
# See the License for the specific language governing permissions and
1515
# limitations under the License.
1616

17-
from google.cloud.aiplatform.helpers import _decorators
17+
from google.cloud.aiplatform.utils.enhanced_library import _decorators
1818
from google.cloud.aiplatform.v1.schema import predict
1919
from google.cloud.aiplatform.v1.schema import trainingjob
2020
from google.cloud.aiplatform.v1beta1.schema import predict as predict_v1beta1

‎google/cloud/aiplatform/helpers/__init__.py‎

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,21 @@
1-
from google.cloud.aiplatform.helpers import value_converter
1+
# Copyright 2021 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
214

3-
__all__ = (value_converter,)
15+
from google.cloud.aiplatform.helpers import container_uri_builders
16+
17+
get_prebuilt_prediction_container_uri = (
18+
container_uri_builders.get_prebuilt_prediction_container_uri
19+
)
20+
21+
__all__ = "get_prebuilt_prediction_container_uri"

‎google/cloud/aiplatform/helpers/container_uri_builders.py‎

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
# Copyright 2021 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from typing import Optional
16+
17+
from google.cloud.aiplatform.constants import prediction
18+
from google.cloud.aiplatform import initializer
19+
20+
21+
def get_prebuilt_prediction_container_uri(
22+
framework: str,
23+
framework_version: str,
24+
region: Optional[str] = None,
25+
accelerator: str = "cpu",
26+
) -> str:
27+
"""
28+
Get a Vertex AI pre-built prediction Docker container URI for
29+
a given framework, version, region, and accelerator use.
30+
31+
Example usage:
32+
```
33+
uri = aiplatform.helpers.get_prebuilt_prediction_container_uri(
34+
framework="tensorflow",
35+
framework_version="2.6",
36+
accelerator="gpu"
37+
)
38+
39+
model = aiplatform.Model.upload(
40+
display_name="boston_housing_",
41+
artifact_uri="gs://my-bucket/my-model/",
42+
serving_container_image_uri=uri
43+
)
44+
```
45+
46+
Args:
47+
framework (str):
48+
Required. The ML framework of the pre-built container. For example,
49+
`"tensorflow"`, `"xgboost"`, or `"sklearn"`
50+
framework_version (str):
51+
Required. The version of the specified ML framework as a string.
52+
region (str):
53+
Optional. AI region or multi-region. Used to select the correct
54+
Artifact Registry multi-region repository and reduce latency.
55+
Must start with `"us"`, `"asia"` or `"europe"`.
56+
Default is location set by `aiplatform.init()`.
57+
accelerator (str):
58+
Optional. The type of accelerator support provided by container. For
59+
example: `"cpu"` or `"gpu"`
60+
Default is `"cpu"`.
61+
62+
Returns:
63+
uri (str):
64+
A Vertex AI prediction container URI
65+
66+
Raises:
67+
ValueError: If containers for provided framework are unavailable or the
68+
container does not support the specified version, accelerator, or region.
69+
"""
70+
URI_MAP = prediction._SERVING_CONTAINER_URI_MAP
71+
DOCS_URI_MESSAGE = (
72+
f"See {prediction._SERVING_CONTAINER_DOCUMENTATION_URL} "
73+
"for complete list of supported containers"
74+
)
75+
76+
# If region not provided, use initializer location
77+
region = region or initializer.global_config.location
78+
region = region.split("-", 1)[0]
79+
framework = framework.lower()
80+
81+
if not URI_MAP.get(region):
82+
raise ValueError(
83+
f"Unsupported container region `{region}`, supported regions are "
84+
f"{', '.join(URI_MAP.keys())}. "
85+
f"{DOCS_URI_MESSAGE}"
86+
)
87+
88+
if not URI_MAP[region].get(framework):
89+
raise ValueError(
90+
f"No containers found for framework `{framework}`. Supported frameworks are "
91+
f"{', '.join(URI_MAP[region].keys())} {DOCS_URI_MESSAGE}"
92+
)
93+
94+
if not URI_MAP[region][framework].get(accelerator):
95+
raise ValueError(
96+
f"{framework} containers do not support `{accelerator}` accelerator. Supported accelerators "
97+
f"are {', '.join(URI_MAP[region][framework].keys())}. {DOCS_URI_MESSAGE}"
98+
)
99+
100+
final_uri = URI_MAP[region][framework][accelerator].get(framework_version)
101+
102+
if not final_uri:
103+
raise ValueError(
104+
f"No serving container for `{framework}` version `{framework_version}` "
105+
f"with accelerator `{accelerator}` found. Supported versions "
106+
f"include {', '.join(URI_MAP[region][framework][accelerator].keys())}. {DOCS_URI_MESSAGE}"
107+
)
108+
109+
return final_uri

‎google/cloud/aiplatform/initializer.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@
2929
from google.auth.exceptions import GoogleAuthError
3030

3131
from google.cloud.aiplatform import compat
32-
from google.cloud.aiplatform import constants
32+
from google.cloud.aiplatform.constants import base as constants
3333
from google.cloud.aiplatform import utils
3434
from google.cloud.aiplatform.metadata import metadata
3535

‎google/cloud/aiplatform/jobs.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@
4242
machine_resources as gca_machine_resources_compat,
4343
study as gca_study_compat,
4444
)
45-
from google.cloud.aiplatform import constants
45+
from google.cloud.aiplatform.constants import base as constants
4646
from google.cloud.aiplatform import initializer
4747
from google.cloud.aiplatform import hyperparameter_tuning
4848
from google.cloud.aiplatform import utils

0 commit comments

Comments
 (0)