model_garden/notebooks/community/feature_store/sdk-feature-store.ipynb

In [ ]:

# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Run in Colab

View on GitHub

Overview

This Colab introduces Vertex AI Feature Store, a managed cloud service for machine learning engineers and data scientists to store, serve, manage and share machine learning features at a large scale.

This Colab assumes that you understand basic Google Cloud concepts such as Project, Storage and Vertex AI. Some machine learning knowledge is also helpful but not required.

Dataset

This Colab uses a movie recommendation dataset as an example throughout all the sessions. The task is to train a model to predict if a user is going to watch a movie and serve this model online.

Objective

In this notebook, you will learn how to:

* Create featurestore, entity type, and feature resources. 
* Import your features into Vertex AI Feature Store.
* Serve online prediction requests using the imported features.
* Access imported features in offline jobs, such as training jobs.

Costs

This tutorial uses billable components of Google Cloud:

• Vertex AI
• Cloud Storage
• Cloud BigQuery

Learn about Vertex AI pricing and Cloud Storage pricing, and use the Pricing Calculator to generate a cost estimate based on your projected usage.

Set up your local development environment

If you are using Colab or Google Cloud Notebooks, your environment already meets all the requirements to run this notebook. You can skip this step.

Otherwise, make sure your environment meets this notebook's requirements. You need the following:

• The Google Cloud SDK
• Git
• Python 3
• virtualenv
• Jupyter notebook running in a virtual environment with Python 3

The Google Cloud guide to setting up a Python development environment and the Jupyter installation guide provide detailed instructions for meeting these requirements. The following steps provide a condensed set of instructions:

1. Install and initialize the Cloud SDK.

2. Install Python 3.

3. Install virtualenv and create a virtual environment that uses Python 3. Activate the virtual environment.

4. To install Jupyter, run pip install jupyter on the command-line in a terminal shell.

5. To launch Jupyter, run jupyter notebook on the command-line in a terminal shell.

6. Open this notebook in the Jupyter Notebook dashboard.

Before you begin

Install additional packages

For this Colab, you need the Vertex SDK for Python.

In [ ]:

import os

# The Google Cloud Notebook product has specific requirements
IS_GOOGLE_CLOUD_NOTEBOOK = os.path.exists("/opt/deeplearning/metadata/env_version")

# Google Cloud Notebook requires dependencies to be installed with '--user'
USER_FLAG = ""
if IS_GOOGLE_CLOUD_NOTEBOOK:
    USER_FLAG = "--user"

In [ ]:

! pip install {USER_FLAG} --upgrade google-cloud-aiplatform

Restart the kernel

After you install the SDK, you need to restart the notebook kernel so it can find the packages. You can restart kernel from Kernel -> Restart Kernel, or running the following:

In [ ]:

# Automatically restart kernel after installs
import os

if not os.getenv("IS_TESTING"):
    # Automatically restart kernel after installs
    import IPython

    app = IPython.Application.instance()
    app.kernel.do_shutdown(True)

Set up your Google Cloud project

The following steps are required, regardless of your notebook environment.

1. Select or create a Google Cloud project. When you first create an account, you get a $300 free credit towards your compute/storage costs.

2. Make sure that billing is enabled for your project.

3. Enable the Vertex AI API and Compute Engine API.

4. If you are running this notebook locally, you will need to install the Cloud SDK.

5. Enter your project ID in the cell below. Then run the cell to make sure the Cloud SDK uses the right project for all the commands in this notebook.

Note: Jupyter runs lines prefixed with ! as shell commands, and it interpolates Python variables prefixed with $ into these commands.

Set your project ID

If you don't know your project ID, you may be able to get your project ID using gcloud.

In [ ]:

import os

PROJECT_ID = ""

# Get your Google Cloud project ID from gcloud
if not os.getenv("IS_TESTING"):
    shell_output = !gcloud config list --format 'value(core.project)' 2>/dev/null
    PROJECT_ID = shell_output[0]
    print("Project ID: ", PROJECT_ID)

Otherwise, set your project ID here.

In [ ]:

if PROJECT_ID == "" or PROJECT_ID is None:
    PROJECT_ID = "[your-project-id]"  # @param {type:"string"}
print("Project ID: ", PROJECT_ID)

Authenticate your Google Cloud account

If you are using Google Cloud Notebooks, your environment is already authenticated. Skip this step.

If you are using Colab, run the cell below and follow the instructions when prompted to authenticate your account via oAuth.

Otherwise, follow these steps:

1. In the Cloud Console, go to the Create service account key page.

2. Click Create service account.

3. In the Service account name field, enter a name, and click Create.

4. In the Grant this service account access to project section, click the Role drop-down list. Type "Vertex AI" into the filter box, and select Vertex AI Administrator. Type "Storage Object Admin" into the filter box, and select Storage Object Admin.

5. Click Create. A JSON file that contains your key downloads to your local environment.

6. Enter the path to your service account key as the GOOGLE_APPLICATION_CREDENTIALS variable in the cell below and run the cell.

In [ ]:

import os
import sys

# If you are running this notebook in Colab, run this cell and follow the
# instructions to authenticate your GCP account. This provides access to your
# Cloud Storage bucket and lets you submit training jobs and prediction
# requests.

# The Google Cloud Notebook product has specific requirements
IS_GOOGLE_CLOUD_NOTEBOOK = os.path.exists("/opt/deeplearning/metadata/env_version")

# If on Google Cloud Notebooks, then don't execute this code
if not IS_GOOGLE_CLOUD_NOTEBOOK:
    if "google.colab" in sys.modules:
        from google.colab import auth as google_auth

        google_auth.authenticate_user()

    # If you are running this notebook locally, replace the string below with the
    # path to your service account key and run this cell to authenticate your GCP
    # account.
    elif not os.getenv("IS_TESTING"):
        %env GOOGLE_APPLICATION_CREDENTIALS ''

Import libraries and define constants

In [ ]:

from google.cloud import aiplatform
from google.cloud.aiplatform import Feature, Featurestore

REGION = "[your-region]"  # @param {type:"string"}
FEATURESTORE_ID = "movie_prediction"
INPUT_CSV_FILE = "gs://cloud-samples-data-us-central1/vertex-ai/feature-store/datasets/movie_prediction.csv"
ONLINE_STORE_FIXED_NODE_COUNT = 1

aiplatform.init(project=PROJECT_ID, location=REGION)

Terminology and Concept

Featurestore Data model

Vertex AI Feature Store organizes data with the following 3 important hierarchical concepts:

Featurestore -> Entity type -> Feature

• Featurestore: the place to store your features
• Entity type: under a Featurestore, an Entity type describes an object to be modeled, real one or virtual one.
• Feature: under an Entity type, a Feature describes an attribute of the Entity type

In the movie prediction example, you will create a featurestore called movie_prediction. This store has 2 entity types: users and movies. The users entity type has the age, gender, and liked_genres features. The movies entity type has the titles, genres, and average rating features.

Create Featurestore and Define Schemas

Create Featurestore

The method to create a Featurestore returns a long-running operation (LRO). An LRO starts an asynchronous job. LROs are returned for other API methods too, such as updating or deleting a featurestore. Running the code cell will create a featurestore and print the process log.

In [ ]:

fs = Featurestore.create(
    featurestore_id=FEATURESTORE_ID,
    online_store_fixed_node_count=ONLINE_STORE_FIXED_NODE_COUNT,
    project=PROJECT_ID,
    location=REGION,
    sync=True,
)

Use the function call below to retrieve a Featurestore and check that it has been created.

In [ ]:

fs = Featurestore(
    featurestore_name=FEATURESTORE_ID,
    project=PROJECT_ID,
    location=REGION,
)
print(fs.gca_resource)

Create Entity Type

Entity types can be created within the Featurestore class. Below, create the Users entity type and Movies entity type. A process log will be printed out.

In [ ]:

# Create users entity type
users_entity_type = fs.create_entity_type(
    entity_type_id="users",
    description="Users entity",
)

In [ ]:

# Create movies entity type
movies_entity_type = fs.create_entity_type(
    entity_type_id="movies",
    description="Movies entity",
)

To retrieve an entity type or check that it has been created use the get_entity_type or list_entity_types methods on the Featurestore object.

In [ ]:

users_entity_type = fs.get_entity_type(entity_type_id="users")
movies_entity_type = fs.get_entity_type(entity_type_id="movies")
print(users_entity_type)
print(movies_entity_type)

In [ ]:

fs.list_entity_types()

Create Feature

Features can be created within each entity type. Add defining features to the Users entity type and Movies entity type by using the following methods.

In [ ]:

# to create features one at a time use
users_feature_age = users_entity_type.create_feature(
    feature_id="age",
    value_type="INT64",
    description="User age",
)

users_feature_gender = users_entity_type.create_feature(
    feature_id="gender",
    value_type="STRING",
    description="User gender",
)

users_feature_liked_genres = users_entity_type.create_feature(
    feature_id="liked_genres",
    value_type="STRING_ARRAY",
    description="An array of genres this user liked",
)

Use the list_features method to list all the features of a given entity type.

In [ ]:

users_entity_type.list_features()

In [ ]:

movies_feature_configs = {
    "title": {
        "value_type": "STRING",
        "description": "The title of the movie",
    },
    "genres": {
        "value_type": "STRING",
        "description": "The genre of the movie",
    },
    "average_rating": {
        "value_type": "DOUBLE",
        "description": "The average rating for the movie, range is [1.0-5.0]",
    },
}

In [ ]:

movie_features = movies_entity_type.batch_create_features(
    feature_configs=movies_feature_configs,
)

Search created features

While the list_features method allows you to easily view all features of a single entity type, the search method in the Feature class searches across all featurestores and entity types in a given location (such as us-central1), and returns a list of features. This can help you discover features that were created by someone else.

You can query based on feature properties including feature ID, entity type ID, and feature description. You can also limit results by filtering on a specific featurestore, feature value type, and/or labels. Some search examples are shown below.

Search for all features within a featurestore with the code snippet below.

In [ ]:

my_features = Feature.search(query="featurestore_id={}".format(FEATURESTORE_ID))
my_features

Now, narrow down the search to features that are of type DOUBLE.

In [ ]:

double_features = Feature.search(
    query="value_type=DOUBLE AND featurestore_id={}".format(FEATURESTORE_ID)
)
double_features[0].gca_resource

Or, limit the search results to features with specific keywords in their ID and type.

In [ ]:

title_features = Feature.search(
    query="feature_id:title AND value_type=STRING AND featurestore_id={}".format(
        FEATURESTORE_ID
    )
)
title_features[0].gca_resource

Import Feature Values

You need to import feature values before you can use them for online/offline serving. In this step, you will learn how to import feature values by ingesting the values from GCS (Google Cloud Storage). We can also import feature values from BigQuery or a Pandas dataframe.

Source Data Format and Layout

BigQuery table/Avro/CSV are supported as input data types. No matter what format you are using, each imported entity must have an ID; also, each entity can optionally have a timestamp, specifying when the feature values are generated. This Colab uses Avro as an input, located at this public bucket. The Avro schemas are as follows:

For the Users entity:

schema = {
  "type": "record",
  "name": "User",
  "fields": [
      {
       "name":"user_id",
       "type":["null","string"]
      },
      {
       "name":"age",
       "type":["null","long"]
      },
      {
       "name":"gender",
       "type":["null","string"]
      },
      {
       "name":"liked_genres",
       "type":{"type":"array","items":"string"}
      },
      {
       "name":"update_time",
       "type":["null",{"type":"long","logicalType":"timestamp-micros"}]
      },
  ]
 }

For the Movies entity:

schema = {
 "type": "record",
 "name": "Movie",
 "fields": [
     {
      "name":"movie_id",
      "type":["null","string"]
     },
     {
      "name":"average_rating",
      "type":["null","double"]
     },
     {
      "name":"title",
      "type":["null","string"]
     },
     {
      "name":"genres",
      "type":["null","string"]
     },
     {
      "name":"update_time",
      "type":["null",{"type":"long","logicalType":"timestamp-micros"}]
     },
 ]
}

Import feature values for Users entity type

When importing, specify the following in your request:

• IDs of the features to import
• Data source URI
• Data source format: BigQuery Table/Avro/CSV

In [ ]:

USERS_FEATURES_IDS = [feature.name for feature in users_entity_type.list_features()]
USERS_FEATURE_TIME = "update_time"
USERS_ENTITY_ID_FIELD = "user_id"
USERS_GCS_SOURCE_URI = (
    "gs://cloud-samples-data-us-central1/vertex-ai/feature-store/datasets/users.avro"
)
GCS_SOURCE_TYPE = "avro"
WORKER_COUNT = 1
print(USERS_FEATURES_IDS)

In [ ]:

users_entity_type.ingest_from_gcs(
    feature_ids=USERS_FEATURES_IDS,
    feature_time=USERS_FEATURE_TIME,
    entity_id_field=USERS_ENTITY_ID_FIELD,
    gcs_source_uris=USERS_GCS_SOURCE_URI,
    gcs_source_type=GCS_SOURCE_TYPE,
    worker_count=WORKER_COUNT,
    sync=False,
)

Import feature values for Movies entity type

Similarly, import feature values for the Movies entity type into the featurestore.

In [ ]:

MOVIES_FEATURES_IDS = [feature.name for feature in movies_entity_type.list_features()]
MOVIES_FEATURE_TIME = "update_time"
MOVIES_ENTITY_ID_FIELD = "movie_id"
MOVIES_GCS_SOURCE_URI = (
    "gs://cloud-samples-data-us-central1/vertex-ai/feature-store/datasets/movies.avro"
)
GCS_SOURCE_TYPE = "avro"
WORKER_COUNT = 1
print(MOVIES_FEATURES_IDS)

In [ ]:

movies_entity_type.ingest_from_gcs(
    feature_ids=MOVIES_FEATURES_IDS,
    feature_time=MOVIES_FEATURE_TIME,
    entity_id_field=MOVIES_ENTITY_ID_FIELD,
    gcs_source_uris=MOVIES_GCS_SOURCE_URI,
    gcs_source_type=GCS_SOURCE_TYPE,
    worker_count=WORKER_COUNT,
    sync=False,
)

Online serving

Online serving lets you serve feature values for small batches of entities. It's designed for latency-sensitive service, such as online model prediction. For example, for a movie service, you might want to quickly show movies that the current user would most likely watch.

Read one entity per request

With the Python SDK, it is easy to read feature values of one entity. By default, the SDK will return the latest value of each feature, meaning the feature values with the most recent timestamp.

To read feature values, specify the entity type ID and features to read. By default all the features of an entity type will be selected. The response will output and display the selected entity type ID and the selected feature values as a Pandas dataframe.

In [ ]:

users_entity_type.read(entity_ids="bob")

In [ ]:

movies_entity_type.read(entity_ids="movie_01", feature_ids="title")

Read multiple entities per request

To read feature values from multiple entities, specify the different entity type IDs. By default all the features of an entity type will be selected. Note that fetching only a small number of entities is recommended when using this SDK due to its latency-sensitive nature.

In [ ]:

users_entity_type.read(entity_ids=["bob", "alice"])

In [ ]:

movies_entity_type.read(
    entity_ids=["movie_02", "movie_03", "movie_04"], feature_ids=["title, genres"]
)

Now that you have learned how to fetch imported feature values for online serving, the next step is learning how to use imported feature values for offline use cases.

Batch Serving

Batch Serving is used to fetch a large batch of feature values for high-throughput, and is typically used for training a model or batch prediction. In this section, you will learn how to prepare for training examples by using the Featurestore's batch serve function.

Use case

The task is to prepare a training dataset to train a model, which predicts if a given user will watch a given movie. To achieve this, you need 2 sets of input:

• Features: you already imported into the featurestore.
• Labels: the groud-truth data recorded that user X has watched movie Y.

To be more specific, the ground-truth observation is described in Table 1 and the desired training dataset is described in Table 2. Each row in Table 2 is a result of joining the imported feature values from Vertex AI Feature Store according to the entity IDs and timestamps in Table 1. In this example, the age, gender and liked_genres features from users and the titles, genres and average_rating features from movies are chosen to train the model. Note that only positive examples are shown in these 2 tables, i.e., you can imagine there is a label column whose values are all True.

batch_serve_to_bq takes Table 1 as input, joins all required feature values from the featurestore, and returns Table 2 for training.

Table 1. Ground-truth data

users	movies	timestamp
alice	Cinema Paradiso	2019-11-01T00:00:00Z
bob	The Shining	2019-11-15T18:09:43Z
...	...	...

Table 2. Expected training data generated by using batch serve

timestamp	entity_type_users	age	gender	liked_genres	entity_type_movies	title	genre	average_rating
2019-11-01T00:00:00Z	bob	35	M	[Action, Crime]	movie_02	The Shining	Horror	4.8
2019-11-01T00:00:00Z	alice	55	F	[Drama, Comedy]	movie_03	Cinema Paradiso	Romance	4.5
...	...	...	...	...	...	...	...	...

Why timestamp?

Note that there is a timestamp column in Table 2. This indicates the time when the ground-truth was observed. This is to avoid data inconsistency.

For example, the 2nd row of Table 2 indicates that user alice watched movie Cinema Paradiso on 2019-11-01T00:00:00Z. The featurestore keeps feature values for all timestamps but fetches feature values only at the given timestamp during batch serving. On that day, Alice might have been 54 years old, but now Alice might be 56; featurestore returns age=54 as Alice's age, instead of age=56, because that is the value of the feature at the observation time. Similarly, other features might be time-variant as well, such as liked_genres.

Create BigQuery dataset for output

You need a BigQuery dataset to host the output data in us-central1. Input the name of the dataset you want to create and specify the name of the table you want to store the output created later. These will be used in the next section.

Make sure that the table name does NOT already exist.

In [ ]:

from datetime import datetime

from google.cloud import bigquery

In [ ]:

# Output dataset
DESTINATION_DATA_SET = "movie_predictions"  # @param {type:"string"}
TIMESTAMP = datetime.now().strftime("%Y%m%d%H%M%S")
DESTINATION_DATA_SET = "{prefix}_{timestamp}".format(
    prefix=DESTINATION_DATA_SET, timestamp=TIMESTAMP
)

# Output table. Make sure that the table does NOT already exist; the BatchReadFeatureValues API cannot overwrite an existing table
DESTINATION_TABLE_NAME = "training_data"  # @param {type:"string"}

DESTINATION_PATTERN = "bq://{project}.{dataset}.{table}"
DESTINATION_TABLE_URI = DESTINATION_PATTERN.format(
    project=PROJECT_ID, dataset=DESTINATION_DATA_SET, table=DESTINATION_TABLE_NAME
)

In [ ]:

# Create dataset
client = bigquery.Client(project=PROJECT_ID)
dataset_id = "{}.{}".format(client.project, DESTINATION_DATA_SET)
dataset = bigquery.Dataset(dataset_id)
dataset.location = REGION
dataset = client.create_dataset(dataset)
print("Created dataset {}.{}".format(client.project, dataset.dataset_id))

Batch Read Feature Values

Assemble the request which specify the following info:

• Where is the label data, i.e., Table 1.
• Which features are read, i.e., the column names in Table 2.

The output is stored in the BigQuery table.

In [ ]:

SERVING_FEATURE_IDS = {
    # to choose all the features use 'entity_type_id: ['*']'
    "users": ["age", "gender", "liked_genres"],
    "movies": ["title", "average_rating", "genres"],
}

In [ ]:

fs.batch_serve_to_bq(
    bq_destination_output_uri=DESTINATION_TABLE_URI,
    serving_feature_ids=SERVING_FEATURE_IDS,
    read_instances_uri=INPUT_CSV_FILE,
)

After the LRO finishes, you should be able to see the result in the BigQuery console, as a new table under the BigQuery dataset created earlier.

Cleaning up

To clean up all Google Cloud resources used in this project, you can delete the Google Cloud project you used for the tutorial.

You can also keep the project but delete the featurestore and the BigQuery dataset by running the code below:

In [ ]:

# Delete Featurestore
fs.delete(force=True)

In [ ]:

# Delete BigQuery dataset
client = bigquery.Client(project=PROJECT_ID)
client.delete_dataset(
    DESTINATION_DATA_SET, delete_contents=True, not_found_ok=True
)  # Make an API request.

print("Deleted dataset '{}'.".format(DESTINATION_DATA_SET))