Skip to content

Commit adb9049

Browse files
jaycee-licopybara-github
jaycee-li
authored and
copybara-github
committed
chore: system tests for experiment models
PiperOrigin-RevId: 507377021
1 parent dc46b23 commit adb9049

2 files changed

Lines changed: 411 additions & 1 deletion

File tree

‎tests/system/aiplatform/test_experiment_model.py‎

Lines changed: 367 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,367 @@
1+
# -*- coding: utf-8 -*-
2+
3+
# Copyright 2023 Google LLC
4+
#
5+
# Licensed under the Apache License, Version 2.0 (the "License");
6+
# you may not use this file except in compliance with the License.
7+
# You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing, software
12+
# distributed under the License is distributed on an "AS IS" BASIS,
13+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
# See the License for the specific language governing permissions and
15+
# limitations under the License.
16+
#
17+
18+
import random
19+
20+
from google.cloud import aiplatform
21+
from google.cloud.aiplatform import models
22+
from google.cloud.aiplatform.constants import prediction
23+
from tests.system.aiplatform import e2e_base
24+
import numpy as np
25+
import pytest
26+
import sklearn
27+
from sklearn.linear_model import LinearRegression
28+
import tensorflow as tf
29+
import xgboost as xgb
30+
31+
32+
CONTAINER_MAP = prediction._SERVING_CONTAINER_URI_MAP[
33+
e2e_base._LOCATION.split("-", 1)[0]
34+
]
35+
36+
37+
@pytest.mark.usefixtures(
38+
"prepare_staging_bucket", "delete_staging_bucket", "tear_down_resources"
39+
)
40+
class TestExperimentModel(e2e_base.TestEndToEnd):
41+
42+
_temp_prefix = "test-vertex-sdk-e2e-experiment-model"
43+
registered_models_cpu = []
44+
registered_models_gpu = []
45+
46+
def test_sklearn_model(self, shared_state):
47+
aiplatform.init(
48+
project=e2e_base._PROJECT,
49+
location=e2e_base._LOCATION,
50+
staging_bucket=f"gs://{shared_state['staging_bucket_name']}",
51+
)
52+
53+
train_x = np.array([[1, 1], [1, 2], [2, 2], [2, 3]])
54+
train_y = np.dot(train_x, np.array([1, 2])) + 3
55+
model = LinearRegression()
56+
model.fit(train_x, train_y)
57+
58+
# Test save sklearn model
59+
aiplatform.save_model(model, "sk-model")
60+
61+
# Test get ExperimentModel with aritfact id
62+
model_artifact = aiplatform.get_experiment_model("sk-model")
63+
assert model_artifact.uri.endswith("sklearn-model")
64+
65+
shared_state["resources"] = [model_artifact]
66+
67+
# Test get model info from ExperimentModel
68+
model_info = model_artifact.get_model_info()
69+
assert model_info == {
70+
"model_class": "sklearn.linear_model._base.LinearRegression",
71+
"framework_name": "sklearn",
72+
"framework_version": sklearn.__version__,
73+
}
74+
75+
# Test load model and make prediction
76+
loaded_model = model_artifact.load_model()
77+
preds = loaded_model.predict(train_x)
78+
assert isinstance(preds, np.ndarray)
79+
80+
# Test register model
81+
# Check the highest pre-built container's version, if lower than the
82+
# ML framework version, use the highest version we have.
83+
version, container_uri = max(CONTAINER_MAP["sklearn"]["cpu"].items())
84+
if version >= sklearn.__version__:
85+
registered_model = model_artifact.register_model()
86+
else:
87+
registered_model = model_artifact.register_model(
88+
serving_container_image_uri=container_uri
89+
)
90+
assert registered_model.display_name.startswith("sklearn model")
91+
92+
self.registered_models_cpu.append(registered_model)
93+
shared_state["resources"].append(registered_model)
94+
95+
def test_xgboost_booster_with_custom_uri(self, shared_state):
96+
aiplatform.init(
97+
project=e2e_base._PROJECT,
98+
location=e2e_base._LOCATION,
99+
staging_bucket=f"gs://{shared_state['staging_bucket_name']}",
100+
)
101+
102+
train_x = np.array([[1, 1], [1, 2], [2, 2], [2, 3]])
103+
train_y = np.array([1, 1, 0, 0])
104+
dtrain = xgb.DMatrix(data=train_x, label=train_y)
105+
booster = xgb.train(
106+
params={"num_parallel_tree": 4, "subsample": 0.5, "num_class": 2},
107+
dtrain=dtrain,
108+
)
109+
110+
# Test save xgboost booster model with custom uri
111+
uri = f"gs://{shared_state['staging_bucket_name']}/custom-uri"
112+
aiplatform.save_model(
113+
model=booster,
114+
artifact_id="xgb-booster",
115+
uri=uri,
116+
)
117+
118+
# Test get ExperimentModel with aritfact id
119+
model_artifact = aiplatform.get_experiment_model("xgb-booster")
120+
assert model_artifact.uri == uri
121+
122+
shared_state["resources"].append(model_artifact)
123+
124+
# Test get model info from ExperimentModel
125+
model_info = model_artifact.get_model_info()
126+
assert model_info == {
127+
"model_class": "xgboost.core.Booster",
128+
"framework_name": "xgboost",
129+
"framework_version": xgb.__version__,
130+
}
131+
132+
# Test load model and make prediction
133+
loaded_model = model_artifact.load_model()
134+
preds = loaded_model.predict(xgb.DMatrix(data=train_x))
135+
assert isinstance(preds, np.ndarray)
136+
137+
# Test register model
138+
# Check the highest pre-built container's version, if lower than the
139+
# ML framework version, use the highest version we have.
140+
version, container_uri = max(CONTAINER_MAP["xgboost"]["cpu"].items())
141+
if version >= xgb.__version__:
142+
registered_model = model_artifact.register_model()
143+
else:
144+
registered_model = model_artifact.register_model(
145+
serving_container_image_uri=container_uri
146+
)
147+
assert registered_model.display_name.startswith("xgboost model")
148+
149+
self.registered_models_cpu.append(registered_model)
150+
shared_state["resources"].append(registered_model)
151+
152+
def test_xgboost_xgbmodel_with_custom_names(self, shared_state):
153+
aiplatform.init(
154+
project=e2e_base._PROJECT,
155+
location=e2e_base._LOCATION,
156+
staging_bucket=f"gs://{shared_state['staging_bucket_name']}",
157+
)
158+
159+
train_x = np.array([[1, 1], [1, 2], [2, 2], [2, 3]])
160+
train_y = np.array([1, 1, 0, 0])
161+
xgb_model = xgb.XGBClassifier()
162+
xgb_model.fit(train_x, train_y)
163+
164+
# Test save xgboost xgbmodel with custom display_name
165+
aiplatform.save_model(
166+
model=xgb_model,
167+
artifact_id="xgboost-xgbmodel",
168+
display_name="custom-experiment-model-name",
169+
)
170+
171+
# Test get ExperimentModel with aritfact id
172+
model_artifact = aiplatform.get_experiment_model("xgboost-xgbmodel")
173+
assert model_artifact.uri.endswith("xgboost-model")
174+
assert model_artifact.display_name == "custom-experiment-model-name"
175+
176+
shared_state["resources"].append(model_artifact)
177+
178+
# Test get model info from ExperimentModel
179+
model_info = model_artifact.get_model_info()
180+
assert model_info == {
181+
"model_class": "xgboost.sklearn.XGBClassifier",
182+
"framework_name": "xgboost",
183+
"framework_version": xgb.__version__,
184+
}
185+
186+
# Test load model and make prediction
187+
loaded_model = model_artifact.load_model()
188+
preds = loaded_model.predict(train_x)
189+
assert isinstance(preds, np.ndarray)
190+
191+
# Test register model with custom display name
192+
# Check the highest pre-built container's version, if lower than the
193+
# ML framework version, use the highest version we have.
194+
version, container_uri = max(CONTAINER_MAP["xgboost"]["cpu"].items())
195+
if version >= xgb.__version__:
196+
registered_model = model_artifact.register_model(
197+
display_name="custom-registered-model-name",
198+
)
199+
else:
200+
registered_model = model_artifact.register_model(
201+
serving_container_image_uri=container_uri,
202+
display_name="custom-registered-model-name",
203+
)
204+
assert registered_model.display_name == "custom-registered-model-name"
205+
206+
self.registered_models_cpu.append(registered_model)
207+
shared_state["resources"].append(registered_model)
208+
209+
def test_tensorflow_keras_model_with_input_example(self, shared_state):
210+
aiplatform.init(
211+
project=e2e_base._PROJECT,
212+
location=e2e_base._LOCATION,
213+
staging_bucket=f"gs://{shared_state['staging_bucket_name']}",
214+
)
215+
216+
train_x = np.random.random((100, 2))
217+
train_y = np.random.random((100, 1))
218+
model = tf.keras.Sequential(
219+
[tf.keras.layers.Dense(5, input_shape=(2,)), tf.keras.layers.Softmax()]
220+
)
221+
model.compile(optimizer="adam", loss="mean_squared_error")
222+
model.fit(train_x, train_y)
223+
224+
# Test save tf.keras model with input example
225+
aiplatform.save_model(
226+
model=model,
227+
artifact_id="keras-model",
228+
input_example=train_x,
229+
)
230+
231+
# Test get ExperimentModel with aritfact id
232+
model_artifact = aiplatform.get_experiment_model("keras-model")
233+
assert model_artifact.uri.endswith("tensorflow-model")
234+
235+
shared_state["resources"].append(model_artifact)
236+
237+
# Test get model info from ExperimentModel
238+
model_info = model_artifact.get_model_info()
239+
assert model_info == {
240+
"model_class": "tensorflow.keras.Model",
241+
"framework_name": "tensorflow",
242+
"framework_version": tf.__version__,
243+
"input_example": {
244+
"data": train_x[:5].tolist(),
245+
"type": "numpy.ndarray",
246+
},
247+
}
248+
249+
# Test load model and make prediction
250+
loaded_model = model_artifact.load_model()
251+
preds = loaded_model.predict(train_x)
252+
assert isinstance(preds, np.ndarray)
253+
254+
# Test register model
255+
# Check the highest pre-built container's version, if lower than the
256+
# ML framework version, use the highest version we have.
257+
version, container_uri = max(CONTAINER_MAP["tensorflow"]["cpu"].items())
258+
if version >= tf.__version__:
259+
registered_model = model_artifact.register_model()
260+
else:
261+
registered_model = model_artifact.register_model(
262+
serving_container_image_uri=container_uri
263+
)
264+
assert registered_model.display_name.startswith("tensorflow model")
265+
266+
self.registered_models_cpu.append(registered_model)
267+
shared_state["resources"].append(registered_model)
268+
269+
def test_tensorflow_module_with_gpu_container(self, shared_state):
270+
aiplatform.init(
271+
project=e2e_base._PROJECT,
272+
location=e2e_base._LOCATION,
273+
staging_bucket=f"gs://{shared_state['staging_bucket_name']}",
274+
)
275+
276+
class Adder(tf.Module):
277+
@tf.function(
278+
input_signature=[
279+
tf.TensorSpec(
280+
shape=[
281+
2,
282+
],
283+
dtype=tf.float32,
284+
)
285+
]
286+
)
287+
def add(self, x):
288+
return x + x
289+
290+
model = Adder()
291+
292+
# Test save tf.Module model
293+
aiplatform.save_model(model, "tf-module")
294+
295+
# Test get ExperimentModel with aritfact id
296+
model_artifact = aiplatform.get_experiment_model("tf-module")
297+
assert model_artifact.uri.endswith("tensorflow-model")
298+
299+
shared_state["resources"].append(model_artifact)
300+
301+
# Test get model info from ExperimentModel
302+
model_info = model_artifact.get_model_info()
303+
assert model_info == {
304+
"model_class": "tensorflow.Module",
305+
"framework_name": "tensorflow",
306+
"framework_version": tf.__version__,
307+
}
308+
309+
# Test load model and make prediction
310+
loaded_model = model_artifact.load_model()
311+
preds = loaded_model.add([1, 2])
312+
assert isinstance(preds, tf.Tensor)
313+
314+
# Test register model with gpu container
315+
# Check the highest pre-built container's version, if lower than the
316+
# ML framework version, use the highest version we have.
317+
version, container_uri = max(CONTAINER_MAP["tensorflow"]["gpu"].items())
318+
if version >= tf.__version__:
319+
registered_model = model_artifact.register_model(use_gpu=True)
320+
else:
321+
registered_model = model_artifact.register_model(
322+
serving_container_image_uri=container_uri,
323+
use_gpu=True,
324+
)
325+
assert registered_model.display_name.startswith("tensorflow model")
326+
327+
self.registered_models_gpu.append(registered_model)
328+
shared_state["resources"].append(registered_model)
329+
330+
def test_deploy_model_with_cpu_container(self, shared_state):
331+
aiplatform.init(
332+
project=e2e_base._PROJECT,
333+
location=e2e_base._LOCATION,
334+
staging_bucket=f"gs://{shared_state['staging_bucket_name']}",
335+
)
336+
337+
# It takes long time to deploy a model. To reduce the system test run
338+
# time, we randomly choose one registered model to test deployment.
339+
registered_model = random.choice(self.registered_models_cpu)
340+
341+
# Deploy the registered model
342+
endpoint = registered_model.deploy()
343+
344+
pred = endpoint.predict([[1, 2]])
345+
assert isinstance(pred, models.Prediction)
346+
shared_state["resources"].append(endpoint)
347+
348+
def test_deploy_model_with_gpu_container(self, shared_state):
349+
aiplatform.init(
350+
project=e2e_base._PROJECT,
351+
location=e2e_base._LOCATION,
352+
staging_bucket=f"gs://{shared_state['staging_bucket_name']}",
353+
)
354+
355+
# It takes long time to deploy a model. To reduce the system test run
356+
# time, we randomly choose one registered model to test deployment.
357+
registered_model = random.choice(self.registered_models_gpu)
358+
359+
# Deploy the registered model
360+
# Since we are using gpu, we need to specify accelerator_type and count
361+
endpoint = registered_model.deploy(
362+
accelerator_type="NVIDIA_TESLA_T4", accelerator_count=1, sync=False
363+
)
364+
365+
pred = endpoint.predict([[1, 2]])
366+
assert isinstance(pred, models.Prediction)
367+
shared_state["resources"].append(endpoint)

0 commit comments

Comments
 (0)