Skip to content

Commit bc9e2cf

Browse files
vertex-sdk-botcopybara-github
vertex-sdk-bot
authored and
copybara-github
committed
docs(samples): Feature Store: Streaming ingestion code sample and test
PiperOrigin-RevId: 494255355
1 parent 99313e0 commit bc9e2cf

4 files changed

Lines changed: 143 additions & 22 deletions

File tree

‎samples/model-builder/conftest.py‎

Lines changed: 39 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -114,7 +114,7 @@ def mock_create_image_dataset(mock_image_dataset):
114114
@pytest.fixture
115115
def mock_create_tabular_dataset(mock_tabular_dataset):
116116
with patch.object(
117-
aiplatform.TabularDataset, "create"
117+
aiplatform.TabularDataset, "create"
118118
) as mock_create_tabular_dataset:
119119
mock_create_tabular_dataset.return_value = mock_tabular_dataset
120120
yield mock_create_tabular_dataset
@@ -123,7 +123,7 @@ def mock_create_tabular_dataset(mock_tabular_dataset):
123123
@pytest.fixture
124124
def mock_create_time_series_dataset(mock_time_series_dataset):
125125
with patch.object(
126-
aiplatform.TimeSeriesDataset, "create"
126+
aiplatform.TimeSeriesDataset, "create"
127127
) as mock_create_time_series_dataset:
128128
mock_create_time_series_dataset.return_value = mock_time_series_dataset
129129
yield mock_create_time_series_dataset
@@ -251,7 +251,9 @@ def mock_run_automl_forecasting_training_job(mock_forecasting_training_job):
251251

252252
@pytest.fixture
253253
def mock_get_automl_forecasting_seq2seq_training_job(mock_forecasting_training_job):
254-
with patch.object(aiplatform, "SequenceToSequencePlusForecastingTrainingJob") as mock:
254+
with patch.object(
255+
aiplatform, "SequenceToSequencePlusForecastingTrainingJob"
256+
) as mock:
255257
mock.return_value = mock_forecasting_training_job
256258
yield mock
257259

@@ -445,6 +447,7 @@ def mock_endpoint_explain(mock_endpoint):
445447
mock_get_endpoint.return_value = mock_endpoint
446448
yield mock_endpoint_explain
447449

450+
448451
# ----------------------------------------------------------------------------
449452
# Hyperparameter Tuning Job Fixtures
450453
# ----------------------------------------------------------------------------
@@ -471,7 +474,9 @@ def mock_run_hyperparameter_tuning_job(mock_hyperparameter_tuning_job):
471474

472475
@pytest.fixture
473476
def mock_hyperparameter_tuning_job_get(mock_hyperparameter_tuning_job):
474-
with patch.object(aiplatform.HyperparameterTuningJob, "get") as mock_hyperparameter_tuning_job_get:
477+
with patch.object(
478+
aiplatform.HyperparameterTuningJob, "get"
479+
) as mock_hyperparameter_tuning_job_get:
475480
mock_hyperparameter_tuning_job_get.return_value = mock_hyperparameter_tuning_job
476481
yield mock_hyperparameter_tuning_job_get
477482

@@ -530,7 +535,7 @@ def mock_get_entity_type(mock_entity_type):
530535
@pytest.fixture
531536
def mock_create_featurestore(mock_featurestore):
532537
with patch.object(
533-
aiplatform.featurestore.Featurestore, "create"
538+
aiplatform.featurestore.Featurestore, "create"
534539
) as mock_create_featurestore:
535540
mock_create_featurestore.return_value = mock_featurestore
536541
yield mock_create_featurestore
@@ -539,7 +544,7 @@ def mock_create_featurestore(mock_featurestore):
539544
@pytest.fixture
540545
def mock_create_entity_type(mock_entity_type):
541546
with patch.object(
542-
aiplatform.featurestore.EntityType, "create"
547+
aiplatform.featurestore.EntityType, "create"
543548
) as mock_create_entity_type:
544549
mock_create_entity_type.return_value = mock_entity_type
545550
yield mock_create_entity_type
@@ -567,7 +572,7 @@ def mock_batch_serve_to_bq(mock_featurestore):
567572
@pytest.fixture
568573
def mock_batch_create_features(mock_entity_type):
569574
with patch.object(
570-
mock_entity_type, "batch_create_features"
575+
mock_entity_type, "batch_create_features"
571576
) as mock_batch_create_features:
572577
yield mock_batch_create_features
573578

@@ -581,11 +586,19 @@ def mock_read_feature_values(mock_entity_type):
581586
@pytest.fixture
582587
def mock_import_feature_values(mock_entity_type):
583588
with patch.object(
584-
mock_entity_type, "ingest_from_gcs"
589+
mock_entity_type, "ingest_from_gcs"
585590
) as mock_import_feature_values:
586591
yield mock_import_feature_values
587592

588593

594+
@pytest.fixture
595+
def mock_write_feature_values(mock_entity_type):
596+
with patch.object(
597+
mock_entity_type.preview, "write_feature_values"
598+
) as mock_write_feature_values:
599+
yield mock_write_feature_values
600+
601+
589602
"""
590603
----------------------------------------------------------------------------
591604
Experiment Tracking Fixtures
@@ -724,7 +737,7 @@ def mock_context_list(mock_context):
724737
@pytest.fixture
725738
def mock_create_schema_base_context(mock_context):
726739
with patch.object(
727-
aiplatform.metadata.schema.base_context.BaseContextSchema, "create"
740+
aiplatform.metadata.schema.base_context.BaseContextSchema, "create"
728741
) as mock_create_schema_base_context:
729742
mock_create_schema_base_context.return_value = mock_context
730743
yield mock_create_schema_base_context
@@ -782,7 +795,7 @@ def mock_create_artifact(mock_artifact):
782795
@pytest.fixture
783796
def mock_create_schema_base_artifact(mock_artifact):
784797
with patch.object(
785-
aiplatform.metadata.schema.base_artifact.BaseArtifactSchema, "create"
798+
aiplatform.metadata.schema.base_artifact.BaseArtifactSchema, "create"
786799
) as mock_create_schema_base_artifact:
787800
mock_create_schema_base_artifact.return_value = mock_artifact
788801
yield mock_create_schema_base_artifact
@@ -791,7 +804,7 @@ def mock_create_schema_base_artifact(mock_artifact):
791804
@pytest.fixture
792805
def mock_create_schema_base_execution(mock_execution):
793806
with patch.object(
794-
aiplatform.metadata.schema.base_execution.BaseExecutionSchema, "create"
807+
aiplatform.metadata.schema.base_execution.BaseExecutionSchema, "create"
795808
) as mock_create_schema_base_execution:
796809
mock_create_schema_base_execution.return_value = mock_execution
797810
yield mock_create_schema_base_execution
@@ -837,7 +850,7 @@ def mock_log_metrics():
837850
@pytest.fixture
838851
def mock_log_time_series_metrics():
839852
with patch.object(
840-
aiplatform, "log_time_series_metrics"
853+
aiplatform, "log_time_series_metrics"
841854
) as mock_log_time_series_metrics:
842855
mock_log_time_series_metrics.return_value = None
843856
yield mock_log_time_series_metrics
@@ -909,24 +922,24 @@ def mock_get_params(mock_params, mock_experiment_run):
909922
@pytest.fixture
910923
def mock_get_time_series_metrics(mock_time_series_metrics, mock_experiment_run):
911924
with patch.object(
912-
mock_experiment_run, "get_time_series_data_frame"
925+
mock_experiment_run, "get_time_series_data_frame"
913926
) as mock_get_time_series_metrics:
914927
mock_get_time_series_metrics.return_value = mock_time_series_metrics
915928
yield mock_get_time_series_metrics
916929

917930

918931
@pytest.fixture
919932
def mock_get_classification_metrics(mock_classification_metrics, mock_experiment_run):
920-
with patch.object(mock_experiment_run, "get_classification_metrics") as mock_get_classification_metrics:
933+
with patch.object(
934+
mock_experiment_run, "get_classification_metrics"
935+
) as mock_get_classification_metrics:
921936
mock_get_classification_metrics.return_value = mock_classification_metrics
922937
yield mock_get_classification_metrics
923938

924939

925940
@pytest.fixture
926941
def mock_get_artifacts(mock_artifacts, mock_experiment_run):
927-
with patch.object(
928-
mock_experiment_run, "get_artifacts"
929-
) as mock_get_artifacts:
942+
with patch.object(mock_experiment_run, "get_artifacts") as mock_get_artifacts:
930943
mock_get_artifacts.return_value = mock_artifacts
931944
yield mock_get_artifacts
932945

@@ -966,7 +979,9 @@ def mock_get_model(mock_model_registry):
966979

967980
@pytest.fixture
968981
def mock_get_model_version_info(mock_model_registry):
969-
with patch.object(mock_model_registry, "get_version_info") as mock_get_model_version_info:
982+
with patch.object(
983+
mock_model_registry, "get_version_info"
984+
) as mock_get_model_version_info:
970985
mock_get_model_version_info.return_value = mock_version_info
971986
yield mock_get_model_version_info
972987

@@ -987,13 +1002,17 @@ def mock_delete_version(mock_model_registry):
9871002

9881003
@pytest.fixture
9891004
def mock_add_version_aliases(mock_model_registry):
990-
with patch.object(mock_model_registry, "add_version_aliases") as mock_add_version_aliases:
1005+
with patch.object(
1006+
mock_model_registry, "add_version_aliases"
1007+
) as mock_add_version_aliases:
9911008
mock_add_version_aliases.return_value = None
9921009
yield mock_add_version_aliases
9931010

9941011

9951012
@pytest.fixture
9961013
def mock_remove_version_aliases(mock_model_registry):
997-
with patch.object(mock_model_registry, "remove_version_aliases") as mock_remove_version_aliases:
1014+
with patch.object(
1015+
mock_model_registry, "remove_version_aliases"
1016+
) as mock_remove_version_aliases:
9981017
mock_remove_version_aliases.return_value = None
9991018
yield mock_remove_version_aliases

‎samples/model-builder/test_constants.py‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -213,6 +213,18 @@
213213
ENTITY_TYPE_ID = "users"
214214
ENTITY_IDS = ["alice", "bob"]
215215
ENTITY_TYPE_NAME = f"projects/{PROJECT}/locations/{LOCATION}/featurestores/{FEATURESTORE_ID}/entityTypes/{ENTITY_TYPE_ID}"
216+
ENTITY_INSTANCES = {
217+
"movie_01": {
218+
"title": "The Shawshank Redemption",
219+
"average_rating": 4.7,
220+
"genre": "Drama",
221+
},
222+
"movie_02": {
223+
"title": "Everything Everywhere All At Once",
224+
"average_rating": 4.4,
225+
"genre": "Adventure",
226+
},
227+
}
216228
FEATURE_ID = "liked_genres"
217229
FEATURE_IDS = ["age", "gender", "liked_genres"]
218230
FEATURE_NAME = f"projects/{PROJECT}/locations/{LOCATION}/featurestores/{FEATURESTORE_ID}/entityTypes/{ENTITY_TYPE_ID}/features/{FEATURE_ID}"
@@ -290,10 +302,10 @@
290302
# Hyperparameter tuning job
291303
HYPERPARAMETER_TUNING_JOB_DISPLAY_NAME = "hpt_job"
292304
HYPERPARAMETER_TUNING_JOB_ID = "4447046521673744384"
293-
HYPERPARAMETER_TUNING_JOB_METRIC_SPEC = {'loss': 'minimize'}
305+
HYPERPARAMETER_TUNING_JOB_METRIC_SPEC = {"loss": "minimize"}
294306
HYPERPARAMETER_TUNING_JOB_MAX_TRIAL_COUNT = 128
295307
HYPERPARAMETER_TUNING_JOB_PARALLEL_TRIAL_COUNT = 8
296-
HYPERPARAMETER_TUNING_JOB_LABELS = {'my_key': 'my_value'}
308+
HYPERPARAMETER_TUNING_JOB_LABELS = {"my_key": "my_value"}
297309

298310
# Custom job
299311
CUSTOM_JOB_DISPLAY_NAME = "custom_job"

‎samples/model-builder/write_feature_values_sample.py‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
# Copyright 2022 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+
# Serve feature values from a single entity for a particular entity type.
16+
# See https://cloud.google.com/vertex-ai/docs/featurestore/setup before running
17+
# the code snippet
18+
19+
# [START aiplatform_write_feature_values_sample]
20+
from google.cloud import aiplatform
21+
22+
23+
def write_feature_values_sample(
24+
project: str, location: str, entity_type_id: str, featurestore_id: str
25+
):
26+
27+
aiplatform.init(project=project, location=location)
28+
29+
my_entity_type = aiplatform.featurestore.EntityType(
30+
entity_type_name=entity_type_id, featurestore_id=featurestore_id
31+
)
32+
33+
my_data = {
34+
"movie_01": {
35+
"title": "The Shawshank Redemption",
36+
"average_rating": 4.7,
37+
"genre": "Drama",
38+
},
39+
"movie_02": {
40+
"title": "Everything Everywhere All At Once",
41+
"average_rating": 4.4,
42+
"genre": "Adventure",
43+
},
44+
}
45+
46+
my_entity_type.preview.write_feature_values(instances=my_data)
47+
48+
49+
# [END aiplatform_write_feature_values_sample]

‎samples/model-builder/write_feature_values_sample_test.py‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
# Copyright 2022 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 test_constants as constants
16+
import write_feature_values_sample
17+
18+
19+
def test_write_feature_values_sample(
20+
mock_sdk_init, mock_get_entity_type, mock_write_feature_values
21+
):
22+
23+
write_feature_values_sample.write_feature_values_sample(
24+
project=constants.PROJECT,
25+
location=constants.LOCATION,
26+
entity_type_id=constants.ENTITY_TYPE_ID,
27+
featurestore_id=constants.FEATURESTORE_ID,
28+
)
29+
30+
mock_sdk_init.assert_called_once_with(
31+
project=constants.PROJECT, location=constants.LOCATION
32+
)
33+
34+
mock_get_entity_type.assert_called_once_with(
35+
entity_type_name=constants.ENTITY_TYPE_ID,
36+
featurestore_id=constants.FEATURESTORE_ID,
37+
)
38+
39+
mock_write_feature_values.assert_called_once_with(
40+
instances=constants.ENTITY_INSTANCES
41+
)

0 commit comments

Comments
 (0)