Skip to content

Commit 603cdbe

Browse files
jaycee-licopybara-github
jaycee-li
authored and
copybara-github
committed
chore: Add bigframes system tests & clean up resources after tests succeed
PiperOrigin-RevId: 570235473
1 parent 639cf10 commit 603cdbe

6 files changed

Lines changed: 332 additions & 5 deletions

File tree

‎setup.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -136,6 +136,7 @@
136136
full_extra_require
137137
+ profiler_extra_require
138138
+ [
139+
"bigframes; python_version>='3.11'",
139140
"grpcio-testing",
140141
"ipython",
141142
"kfp",

‎tests/system/vertexai/test_bigframes_sklearn.py‎

Lines changed: 188 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,188 @@
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 os
19+
from unittest import mock
20+
21+
from google.cloud import aiplatform
22+
import vertexai
23+
from tests.system.aiplatform import e2e_base
24+
from vertexai.preview._workflow.executor import training
25+
from vertexai.preview._workflow.serialization_engine import (
26+
serializers,
27+
)
28+
import numpy as np
29+
30+
import pytest
31+
from sklearn.linear_model import LogisticRegression
32+
from sklearn.preprocessing import StandardScaler
33+
34+
import bigframes.pandas as bf
35+
from bigframes.ml.model_selection import train_test_split as bf_train_test_split
36+
37+
38+
bf.options.bigquery.location = "us" # Dataset is in 'us' not 'us-central1'
39+
bf.options.bigquery.project = e2e_base._PROJECT
40+
41+
42+
# Wrap classes
43+
StandardScaler = vertexai.preview.remote(StandardScaler)
44+
LogisticRegression = vertexai.preview.remote(LogisticRegression)
45+
46+
47+
@mock.patch.object(
48+
training,
49+
"VERTEX_AI_DEPENDENCY_PATH",
50+
"google-cloud-aiplatform[preview] @ git+https://github.com/googleapis/"
51+
f"python-aiplatform.git@{os.environ['KOKORO_GIT_COMMIT']}"
52+
if os.environ.get("KOKORO_GIT_COMMIT")
53+
else "google-cloud-aiplatform[preview] @ git+https://github.com/googleapis/python-aiplatform.git@main",
54+
)
55+
@mock.patch.object(
56+
training,
57+
"VERTEX_AI_DEPENDENCY_PATH_AUTOLOGGING",
58+
"google-cloud-aiplatform[preview,autologging] @ git+https://github.com/googleapis/"
59+
f"python-aiplatform.git@{os.environ['KOKORO_GIT_COMMIT']}"
60+
if os.environ.get("KOKORO_GIT_COMMIT")
61+
else "google-cloud-aiplatform[preview,autologging] @ git+https://github.com/googleapis/python-aiplatform.git@main",
62+
)
63+
@pytest.mark.usefixtures(
64+
"prepare_staging_bucket", "delete_staging_bucket", "tear_down_resources"
65+
)
66+
class TestRemoteExecutionBigframesSklearn(e2e_base.TestEndToEnd):
67+
68+
_temp_prefix = "temp-vertexai-remote-execution"
69+
70+
def test_remote_execution_sklearn(self, shared_state):
71+
# Initialize vertexai
72+
vertexai.init(
73+
project=e2e_base._PROJECT,
74+
location=e2e_base._LOCATION,
75+
staging_bucket=f"gs://{shared_state['staging_bucket_name']}",
76+
)
77+
78+
# Prepare dataset
79+
df = bf.read_gbq("bigquery-public-data.ml_datasets.iris")
80+
species_categories = {
81+
"versicolor": 0,
82+
"virginica": 1,
83+
"setosa": 2,
84+
}
85+
df["species"] = df["species"].map(species_categories)
86+
index_col = "index"
87+
df.index.name = index_col
88+
feature_columns = df[
89+
["sepal_length", "sepal_width", "petal_length", "petal_width"]
90+
]
91+
label_columns = df[["species"]]
92+
train_X, test_X, train_y, _ = bf_train_test_split(
93+
feature_columns, label_columns, test_size=0.2
94+
)
95+
96+
# Remote fit_transform on bf train dataset
97+
vertexai.preview.init(remote=True)
98+
transformer = StandardScaler()
99+
transformer.fit_transform.vertex.remote_config.display_name = (
100+
self._make_display_name("bigframes-fit-transform")
101+
)
102+
X_train = transformer.fit_transform(train_X)
103+
104+
# Assert the right serializer is being used
105+
remote_job = aiplatform.CustomJob.list(
106+
filter=f'display_name="{transformer.fit_transform.vertex.remote_config.display_name}"'
107+
)[0]
108+
base_path = remote_job.job_spec.base_output_directory.output_uri_prefix
109+
110+
input_estimator_metadata = serializers._get_metadata(
111+
os.path.join(base_path, "input/input_estimator")
112+
)
113+
assert input_estimator_metadata["serializer"] == "SklearnEstimatorSerializer"
114+
115+
output_estimator_metadata = serializers._get_metadata(
116+
os.path.join(base_path, "output/output_estimator")
117+
)
118+
assert output_estimator_metadata["serializer"] == "SklearnEstimatorSerializer"
119+
120+
train_x_metadata = serializers._get_metadata(os.path.join(base_path, "input/X"))
121+
assert train_x_metadata["serializer"] == "BigframeSerializer"
122+
assert train_x_metadata["framework"] == "sklearn"
123+
124+
shared_state["resources"] = [remote_job]
125+
126+
assert type(X_train) is np.ndarray
127+
assert X_train.shape == (120, 4)
128+
129+
# Remote transform on bf test dataset
130+
transformer.transform.vertex.remote_config.display_name = (
131+
self._make_display_name("bigframes-transform")
132+
)
133+
X_test = transformer.transform(test_X)
134+
135+
# Assert the right serializer is being used
136+
remote_job = aiplatform.CustomJob.list(
137+
filter=f'display_name="{transformer.transform.vertex.remote_config.display_name}"'
138+
)[0]
139+
base_path = remote_job.job_spec.base_output_directory.output_uri_prefix
140+
141+
input_estimator_metadata = serializers._get_metadata(
142+
os.path.join(base_path, "input/input_estimator")
143+
)
144+
assert input_estimator_metadata["serializer"] == "SklearnEstimatorSerializer"
145+
146+
test_x_metadata = serializers._get_metadata(os.path.join(base_path, "input/X"))
147+
assert test_x_metadata["serializer"] == "BigframeSerializer"
148+
assert train_x_metadata["framework"] == "sklearn"
149+
150+
shared_state["resources"].append(remote_job)
151+
152+
assert type(X_test) is np.ndarray
153+
assert X_test.shape == (30, 4)
154+
155+
# Remote training on sklearn
156+
vertexai.preview.init(remote=True)
157+
158+
model = LogisticRegression(warm_start=True)
159+
model.fit.vertex.remote_config.display_name = self._make_display_name(
160+
"bigframes-sklearn-training"
161+
)
162+
model.fit(train_X, train_y)
163+
164+
# Assert the right serializer is being used
165+
remote_job = aiplatform.CustomJob.list(
166+
filter=f'display_name="{model.fit.vertex.remote_config.display_name}"'
167+
)[0]
168+
base_path = remote_job.job_spec.base_output_directory.output_uri_prefix
169+
170+
input_estimator_metadata = serializers._get_metadata(
171+
os.path.join(base_path, "input/input_estimator")
172+
)
173+
assert input_estimator_metadata["serializer"] == "SklearnEstimatorSerializer"
174+
175+
output_estimator_metadata = serializers._get_metadata(
176+
os.path.join(base_path, "output/output_estimator")
177+
)
178+
assert output_estimator_metadata["serializer"] == "SklearnEstimatorSerializer"
179+
180+
train_x_metadata = serializers._get_metadata(os.path.join(base_path, "input/X"))
181+
assert train_x_metadata["serializer"] == "BigframeSerializer"
182+
assert train_x_metadata["framework"] == "sklearn"
183+
184+
train_y_metadata = serializers._get_metadata(os.path.join(base_path, "input/y"))
185+
assert train_y_metadata["serializer"] == "BigframeSerializer"
186+
assert train_y_metadata["framework"] == "sklearn"
187+
188+
shared_state["resources"].append(remote_job)

‎tests/system/vertexai/test_bigframes_tensorflow.py‎

Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,125 @@
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 os
19+
from unittest import mock
20+
21+
from google.cloud import aiplatform
22+
import vertexai
23+
from tests.system.aiplatform import e2e_base
24+
from vertexai.preview._workflow.executor import training
25+
from vertexai.preview._workflow.serialization_engine import (
26+
serializers,
27+
)
28+
import pytest
29+
from tensorflow import keras
30+
31+
import bigframes.pandas as bf
32+
from bigframes.ml.model_selection import train_test_split as bf_train_test_split
33+
34+
35+
bf.options.bigquery.location = "us" # Dataset is in 'us' not 'us-central1'
36+
bf.options.bigquery.project = e2e_base._PROJECT
37+
38+
39+
# Wrap classes
40+
keras.Sequential = vertexai.preview.remote(keras.Sequential)
41+
42+
43+
@mock.patch.object(
44+
training,
45+
"VERTEX_AI_DEPENDENCY_PATH",
46+
"google-cloud-aiplatform[preview] @ git+https://github.com/googleapis/"
47+
f"python-aiplatform.git@{os.environ['KOKORO_GIT_COMMIT']}"
48+
if os.environ.get("KOKORO_GIT_COMMIT")
49+
else "google-cloud-aiplatform[preview] @ git+https://github.com/googleapis/python-aiplatform.git@main",
50+
)
51+
@mock.patch.object(
52+
training,
53+
"VERTEX_AI_DEPENDENCY_PATH_AUTOLOGGING",
54+
"google-cloud-aiplatform[preview,autologging] @ git+https://github.com/googleapis/"
55+
f"python-aiplatform.git@{os.environ['KOKORO_GIT_COMMIT']}"
56+
if os.environ.get("KOKORO_GIT_COMMIT")
57+
else "google-cloud-aiplatform[preview,autologging] @ git+https://github.com/googleapis/python-aiplatform.git@main",
58+
)
59+
@pytest.mark.usefixtures(
60+
"prepare_staging_bucket", "delete_staging_bucket", "tear_down_resources"
61+
)
62+
class TestRemoteExecutionBigframesTensorflow(e2e_base.TestEndToEnd):
63+
64+
_temp_prefix = "temp-vertexai-remote-execution"
65+
66+
def test_remote_execution_keras(self, shared_state):
67+
# Initialize vertexai
68+
vertexai.init(
69+
project=e2e_base._PROJECT,
70+
location=e2e_base._LOCATION,
71+
staging_bucket=f"gs://{shared_state['staging_bucket_name']}",
72+
)
73+
74+
# Prepare dataset
75+
df = bf.read_gbq("bigquery-public-data.ml_datasets.iris")
76+
77+
species_categories = {
78+
"versicolor": 0,
79+
"virginica": 1,
80+
"setosa": 2,
81+
}
82+
df["target"] = df["species"].map(species_categories)
83+
df = df.drop(columns=["species"])
84+
85+
train, _ = bf_train_test_split(df, test_size=0.2)
86+
87+
# Remote GPU training on Keras
88+
vertexai.preview.init(remote=True)
89+
90+
model = keras.Sequential(
91+
[keras.layers.Dense(5, input_shape=(4,)), keras.layers.Softmax()]
92+
)
93+
model.compile(optimizer="adam", loss="mean_squared_error")
94+
model.fit.vertex.set_config(
95+
enable_cuda=True,
96+
display_name=self._make_display_name("bigframes-keras-training"),
97+
)
98+
model.fit.vertex.remote_config.custom_commands = [
99+
"pip install tensorflow-io==0.32.0"
100+
]
101+
102+
# Train model on Vertex
103+
model.fit(train, epochs=10)
104+
105+
# Assert the right serializer is being used
106+
remote_job = aiplatform.CustomJob.list(
107+
filter=f'display_name="{model.fit.vertex.remote_config.display_name}"'
108+
)[0]
109+
base_path = remote_job.job_spec.base_output_directory.output_uri_prefix
110+
111+
input_estimator_metadata = serializers._get_metadata(
112+
os.path.join(base_path, "input/input_estimator")
113+
)
114+
assert input_estimator_metadata["serializer"] == "KerasModelSerializer"
115+
116+
output_estimator_metadata = serializers._get_metadata(
117+
os.path.join(base_path, "output/output_estimator")
118+
)
119+
assert output_estimator_metadata["serializer"] == "KerasModelSerializer"
120+
121+
train_x_metadata = serializers._get_metadata(os.path.join(base_path, "input/x"))
122+
assert train_x_metadata["serializer"] == "BigframeSerializer"
123+
assert train_x_metadata["framework"] == "tensorflow"
124+
125+
shared_state["resources"] = [remote_job]

‎tests/system/vertexai/test_pytorch.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -169,6 +169,12 @@ def predict(self, X):
169169
)
170170
model.predict(X_test)
171171

172+
# Add prediction job to teardown resource
173+
remote_job = aiplatform.CustomJob.list(
174+
filter=f'display_name="{model.predict.vertex.remote_config.display_name}"'
175+
)[0]
176+
shared_state["resources"].append(remote_job)
177+
172178
# Register trained model
173179
registered_model = vertexai.preview.register(model)
174180
shared_state["resources"].append(registered_model)

‎tests/system/vertexai/test_sklearn.py‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -122,11 +122,6 @@ def test_remote_execution_sklearn(self, shared_state):
122122
)
123123
assert input_estimator_metadata["serializer"] == "SklearnEstimatorSerializer"
124124

125-
output_estimator_metadata = serializers._get_metadata(
126-
os.path.join(base_path, "output/output_estimator")
127-
)
128-
assert output_estimator_metadata["serializer"] == "SklearnEstimatorSerializer"
129-
130125
shared_state["resources"].append(remote_job)
131126

132127
# Local transform on retrain data
@@ -169,6 +164,12 @@ def test_remote_execution_sklearn(self, shared_state):
169164
)
170165
model.predict(X_test)
171166

167+
# Add prediction job to teardown resource
168+
remote_job = aiplatform.CustomJob.list(
169+
filter=f'display_name="{model.predict.vertex.remote_config.display_name}"'
170+
)[0]
171+
shared_state["resources"].append(remote_job)
172+
172173
# Register trained model
173174
registered_model = vertexai.preview.register(model)
174175
shared_state["resources"].append(registered_model)

‎tests/system/vertexai/test_tensorflow.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,12 @@ def test_remote_execution_keras(self, shared_state):
132132
)
133133
model.predict(tf_remote_prediction_test_data)
134134

135+
# Add prediction job to teardown resource
136+
remote_job = aiplatform.CustomJob.list(
137+
filter=f'display_name="{model.predict.vertex.remote_config.display_name}"'
138+
)[0]
139+
shared_state["resources"].append(remote_job)
140+
135141
# Register trained model
136142
registered_model = vertexai.preview.register(model)
137143
shared_state["resources"].append(registered_model)

0 commit comments

Comments
 (0)