Skip to content

Commit e3ea683

Browse files
lclc19lclc19
authored
feat: Add data plane code snippets for feature store service (googleapis#713)
* feat: Add data plane code snippets for feature store service
1 parent 8e06ced commit e3ea683

13 files changed

Lines changed: 583 additions & 8 deletions

‎samples/snippets/conftest.py‎

Lines changed: 40 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
from uuid import uuid4
1717

1818
from google.cloud import aiplatform, aiplatform_v1beta1
19+
from google.cloud import bigquery
1920
from google.cloud import storage
2021
import pytest
2122

@@ -91,6 +92,14 @@ def featurestore_client():
9192
yield featurestore_client
9293

9394

95+
@pytest.fixture
96+
def bigquery_client():
97+
bigquery_client = bigquery.Client(
98+
project=os.getenv("BUILD_SPECIFIC_GCLOUD_PROJECT")
99+
)
100+
yield bigquery_client
101+
102+
94103
# Shared setup/teardown.
95104
@pytest.fixture()
96105
def teardown_batch_prediction_job(shared_state, job_client):
@@ -213,16 +222,24 @@ def teardown_dataset(shared_state, dataset_client):
213222
def teardown_featurestore(shared_state, featurestore_client):
214223
yield
215224

216-
# Delete the created featurestore
217-
featurestore_client.delete_featurestore(name=shared_state["featurestore_name"])
225+
# Force delete the created featurestore
226+
force_delete_featurestore_request = {
227+
"name": shared_state["featurestore_name"],
228+
"force": True,
229+
}
230+
featurestore_client.delete_featurestore(request=force_delete_featurestore_request)
218231

219232

220233
@pytest.fixture()
221234
def teardown_entity_type(shared_state, featurestore_client):
222235
yield
223236

224-
# Delete the created entity type
225-
featurestore_client.delete_entity_type(name=shared_state["entity_type_name"])
237+
# Force delete the created entity type
238+
force_delete_entity_type_request = {
239+
"name": shared_state["entity_type_name"],
240+
"force": True,
241+
}
242+
featurestore_client.delete_entity_type(request=force_delete_entity_type_request)
226243

227244

228245
@pytest.fixture()
@@ -233,6 +250,25 @@ def teardown_feature(shared_state, featurestore_client):
233250
featurestore_client.delete_feature(name=shared_state["feature_name"])
234251

235252

253+
@pytest.fixture()
254+
def teardown_features(shared_state, featurestore_client):
255+
yield
256+
257+
# Delete the created features
258+
for feature_name in shared_state["feature_names"]:
259+
featurestore_client.delete_feature(name=feature_name)
260+
261+
262+
@pytest.fixture()
263+
def teardown_batch_read_feature_values(shared_state, bigquery_client):
264+
yield
265+
266+
# Delete the created dataset
267+
bigquery_client.delete_dataset(
268+
shared_state["destination_data_set"], delete_contents=True, not_found_ok=True
269+
)
270+
271+
236272
@pytest.fixture()
237273
def create_endpoint(shared_state, endpoint_client):
238274
def create(project, location, test_name="temp_deploy_model_test"):

‎samples/snippets/feature_store_service/batch_create_features_sample.py‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,74 @@
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+
# Create features in bulk for an existing type.
16+
# See https://cloud.google.com/vertex-ai/docs/featurestore/setup before running
17+
# the code snippet
18+
19+
# [START aiplatform_batch_create_features_sample]
20+
from google.cloud import aiplatform_v1beta1 as aiplatform
21+
22+
23+
def batch_create_features_sample(
24+
project: str,
25+
featurestore_id: str,
26+
entity_type_id: str,
27+
location: str = "us-central1",
28+
api_endpoint: str = "us-central1-aiplatform.googleapis.com",
29+
timeout: int = 300,
30+
):
31+
# The AI Platform services require regional API endpoints, which need to be
32+
# in the same region or multi-region overlap with the Feature Store location.
33+
client_options = {"api_endpoint": api_endpoint}
34+
# Initialize client that will be used to create and send requests.
35+
# This client only needs to be created once, and can be reused for multiple requests.
36+
client = aiplatform.FeaturestoreServiceClient(client_options=client_options)
37+
parent = f"projects/{project}/locations/{location}/featurestores/{featurestore_id}/entityTypes/{entity_type_id}"
38+
age_feature = aiplatform.Feature(
39+
value_type=aiplatform.Feature.ValueType.INT64, description="User age",
40+
)
41+
age_feature_request = aiplatform.CreateFeatureRequest(
42+
feature=age_feature, feature_id="age"
43+
)
44+
45+
gender_feature = aiplatform.Feature(
46+
value_type=aiplatform.Feature.ValueType.STRING, description="User gender"
47+
)
48+
gender_feature_request = aiplatform.CreateFeatureRequest(
49+
feature=gender_feature, feature_id="gender"
50+
)
51+
52+
liked_genres_feature = aiplatform.Feature(
53+
value_type=aiplatform.Feature.ValueType.STRING_ARRAY,
54+
description="An array of genres that this user liked",
55+
)
56+
liked_genres_feature_request = aiplatform.CreateFeatureRequest(
57+
feature=liked_genres_feature, feature_id="liked_genres"
58+
)
59+
60+
requests = [
61+
age_feature_request,
62+
gender_feature_request,
63+
liked_genres_feature_request,
64+
]
65+
batch_create_features_request = aiplatform.BatchCreateFeaturesRequest(
66+
parent=parent, requests=requests
67+
)
68+
lro_response = client.batch_create_features(request=batch_create_features_request)
69+
print("Long running operation:", lro_response.operation.name)
70+
batch_create_features_response = lro_response.result(timeout=timeout)
71+
print("batch_create_features_response:", batch_create_features_response)
72+
73+
74+
# [END aiplatform_batch_create_features_sample]

‎samples/snippets/feature_store_service/batch_create_features_sample_test.py‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
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+
import os
16+
from uuid import uuid4
17+
18+
import batch_create_features_sample
19+
import create_entity_type_sample
20+
21+
import pytest
22+
23+
import helpers
24+
25+
PROJECT_ID = os.getenv("BUILD_SPECIFIC_GCLOUD_PROJECT")
26+
27+
28+
@pytest.fixture(scope="function", autouse=True)
29+
def teardown(teardown_entity_type):
30+
yield
31+
32+
33+
def setup_temp_entity_type(featurestore_id, entity_type_id, capsys):
34+
create_entity_type_sample.create_entity_type_sample(
35+
project=PROJECT_ID,
36+
featurestore_id=featurestore_id,
37+
entity_type_id=entity_type_id,
38+
)
39+
out, _ = capsys.readouterr()
40+
assert "create_entity_type_response" in out
41+
return helpers.get_featurestore_resource_name(out)
42+
43+
44+
def test_ucaip_generated_batch_create_features_sample_vision(capsys, shared_state):
45+
featurestore_id = "perm_sample_featurestore"
46+
entity_type_id = f"users_{uuid4()}".replace("-", "_")[:60]
47+
entity_type_name = setup_temp_entity_type(featurestore_id, entity_type_id, capsys)
48+
location = "us-central1"
49+
batch_create_features_sample.batch_create_features_sample(
50+
project=PROJECT_ID,
51+
featurestore_id=featurestore_id,
52+
entity_type_id=entity_type_id,
53+
location=location,
54+
)
55+
out, _ = capsys.readouterr()
56+
assert "batch_create_features_response" in out
57+
58+
shared_state["entity_type_name"] = entity_type_name

‎samples/snippets/feature_store_service/batch_read_feature_values_sample.py‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
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+
# Batch read feature values from a featurestore, as determined by your read
16+
# instances list file, to export data.
17+
# See https://cloud.google.com/vertex-ai/docs/featurestore/setup before running
18+
# the code snippet
19+
20+
# [START aiplatform_batch_read_feature_values_sample]
21+
from google.cloud import aiplatform_v1beta1 as aiplatform
22+
23+
24+
def batch_read_feature_values_sample(
25+
project: str,
26+
featurestore_id: str,
27+
input_csv_file: str,
28+
destination_table_uri: str,
29+
location: str = "us-central1",
30+
api_endpoint: str = "us-central1-aiplatform.googleapis.com",
31+
timeout: int = 300,
32+
):
33+
# The AI Platform services require regional API endpoints, which need to be
34+
# in the same region or multi-region overlap with the Feature Store location.
35+
client_options = {"api_endpoint": api_endpoint}
36+
# Initialize client that will be used to create and send requests.
37+
# This client only needs to be created once, and can be reused for multiple requests.
38+
client = aiplatform.FeaturestoreServiceClient(client_options=client_options)
39+
featurestore = (
40+
f"projects/{project}/locations/{location}/featurestores/{featurestore_id}"
41+
)
42+
csv_read_instances = aiplatform.CsvSource(
43+
gcs_source=aiplatform.GcsSource(uris=[input_csv_file])
44+
)
45+
destination = aiplatform.FeatureValueDestination(
46+
bigquery_destination=aiplatform.BigQueryDestination(
47+
# Output to BigQuery table created earlier
48+
output_uri=destination_table_uri
49+
)
50+
)
51+
52+
users_feature_selector = aiplatform.FeatureSelector(
53+
id_matcher=aiplatform.IdMatcher(ids=["age", "gender", "liked_genres"])
54+
)
55+
users_entity_type_spec = aiplatform.BatchReadFeatureValuesRequest.EntityTypeSpec(
56+
# Read the 'age', 'gender' and 'liked_genres' features from the 'perm_users' entity
57+
entity_type_id="perm_users",
58+
feature_selector=users_feature_selector,
59+
)
60+
61+
movies_feature_selector = aiplatform.FeatureSelector(
62+
id_matcher=aiplatform.IdMatcher(ids=["*"])
63+
)
64+
movies_entity_type_spec = aiplatform.BatchReadFeatureValuesRequest.EntityTypeSpec(
65+
# Read the all features from the 'perm_movies' entity
66+
entity_type_id="perm_movies",
67+
feature_selector=movies_feature_selector,
68+
)
69+
70+
entity_type_specs = [users_entity_type_spec, movies_entity_type_spec]
71+
# Batch serving request from CSV
72+
batch_read_feature_values_request = aiplatform.BatchReadFeatureValuesRequest(
73+
featurestore=featurestore,
74+
csv_read_instances=csv_read_instances,
75+
destination=destination,
76+
entity_type_specs=entity_type_specs,
77+
)
78+
lro_response = client.batch_read_feature_values(
79+
request=batch_read_feature_values_request
80+
)
81+
print("Long running operation:", lro_response.operation.name)
82+
batch_read_feature_values_response = lro_response.result(timeout=timeout)
83+
print("batch_read_feature_values_response:", batch_read_feature_values_response)
84+
85+
86+
# [END aiplatform_batch_read_feature_values_sample]

‎samples/snippets/feature_store_service/batch_read_feature_values_sample_test.py‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
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 datetime import datetime
16+
import os
17+
18+
import batch_read_feature_values_sample
19+
from google.cloud import bigquery
20+
21+
import pytest
22+
23+
PROJECT_ID = os.getenv("BUILD_SPECIFIC_GCLOUD_PROJECT")
24+
LOCATION = "us-central1"
25+
INPUT_CSV_FILE = "gs://cloud-samples-data-us-central1/vertex-ai/feature-store/datasets/movie_prediction_perm.csv"
26+
27+
28+
@pytest.fixture(scope="function", autouse=True)
29+
def teardown(teardown_batch_read_feature_values):
30+
yield
31+
32+
33+
def setup_test():
34+
# Output dataset
35+
destination_data_set = "movie_predictions_" + datetime.now().strftime(
36+
"%Y%m%d%H%M%S"
37+
)
38+
# Output table. Make sure that the table does NOT already exist, the
39+
# BatchReadFeatureValues API cannot overwrite an existing table.
40+
destination_table_name = "training_data"
41+
DESTINATION_PATTERN = "bq://{project}.{dataset}.{table}"
42+
destination_table_uri = DESTINATION_PATTERN.format(
43+
project=PROJECT_ID, dataset=destination_data_set, table=destination_table_name
44+
)
45+
# Create dataset
46+
bq_client = bigquery.Client(project=PROJECT_ID)
47+
dataset_id = "{}.{}".format(bq_client.project, destination_data_set)
48+
dataset = bigquery.Dataset(dataset_id)
49+
dataset.location = LOCATION
50+
dataset = bq_client.create_dataset(dataset)
51+
print("Created dataset {}.{}".format(bq_client.project, dataset.dataset_id))
52+
return destination_data_set, destination_table_uri
53+
54+
55+
def test_ucaip_generated_batch_read_feature_values_sample_vision(capsys, shared_state):
56+
destination_data_set, destination_table_uri = setup_test()
57+
featurestore_id = "perm_sample_featurestore"
58+
59+
batch_read_feature_values_sample.batch_read_feature_values_sample(
60+
project=PROJECT_ID,
61+
featurestore_id=featurestore_id,
62+
input_csv_file=INPUT_CSV_FILE,
63+
destination_table_uri=destination_table_uri,
64+
)
65+
out, _ = capsys.readouterr()
66+
assert "batch_read_feature_values_response" in out
67+
with capsys.disabled():
68+
print(out)
69+
70+
shared_state["destination_data_set"] = destination_data_set

‎samples/snippets/feature_store_service/create_entity_type_sample.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,10 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15+
# Create an entity type so that you can create its related features.
16+
# See https://cloud.google.com/vertex-ai/docs/featurestore/setup before running
17+
# the code snippet
18+
1519
# [START aiplatform_create_entity_type_sample]
1620
from google.cloud import aiplatform_v1beta1 as aiplatform
1721

@@ -25,7 +29,8 @@ def create_entity_type_sample(
2529
api_endpoint: str = "us-central1-aiplatform.googleapis.com",
2630
timeout: int = 300,
2731
):
28-
# The AI Platform services require regional API endpoints.
32+
# The AI Platform services require regional API endpoints, which need to be
33+
# in the same region or multi-region overlap with the Feature Store location.
2934
client_options = {"api_endpoint": api_endpoint}
3035
# Initialize client that will be used to create and send requests.
3136
# This client only needs to be created once, and can be reused for multiple requests.

0 commit comments

Comments
 (0)