Vertex AI Feature Store

Credit: Google Cloud Icon Library

Introduction

Before Vertex AI Feature Store

It is known that in the last half century, the world of science and technology has made incomparably great progress, and the fantasy that existed a century ago is already real. A century ago, who would have thought that people, thousands of kilometers away, could communicate with each other by video call. Today, people already think about transmitting the scent of food or perfume thousands of kilometers away in just a few seconds, and I wouldn’t be surprised if in a few years it will be possible. However, let’s not deviate from our main topic and ask օurselves a question: what is the reason and the real motivation of scientific and technical progress?

Credit: Internet (Not mine)

I found two answers from my personal analysis: natural curiosity and the desire to facilitate and improve human life. The focus of this article is on the basis of the second answer.

The average person created 1.7 MB of data every second in 2022. There are 5 billion people on the internet today. Do the math, and you will know that there is a lot of data created daily!

The total size of data held by the big four members of the Nasdaq (namely Google, Amazon, Microsoft, and Facebook) alone is already 1.2 million terabytes of data. Specialists, like data and machine learning engineers, work with that data, make predictions, find solutions to daily problems, and make people’s life better. As data grows, it becomes increasingly difficult to work with. Consequently, more flexible tools are created to facilitate people’s work and make it easier.

Vertex AI’s recent product, Feature Store, has an important role in scientific and technical progress. The article will start with an overview and discuss benefits of the product. It’ll also cover the architecture of Vertex AI Feature Store, concluding with a practical example.

What is Vertex AI Feature Store?

Let’s talk about an imaginary company called Armcomp. Since deploying their first ML models using Auto ML and BigQuery ML, the Armcomp team has added feature engineering to their skillset. However, the team is now divided into multiple sub teams recreating the same features, which leads to redundant work. Each time a team needs to add a new feature to a model, they are dependent on the Ops team. Team members also spend too much time trying to manage the lifecycle of a feature starting with training and ending with deployment and serving.

Credit: Internet (Not mine)

The Armcomp team’s software engineer needs to understand how to add ML to their apps via online prediction endpoints without having to understand how to procure or generate the features required by the models at prediction time.

Credit: Internet (Not mine)

Armcomp’s DevOps resource is new to the Armcomp team and needs to understand how to monitor and maintain feature serving and management infrastructure.

Here are the key challenges that often come up for people working in ML features:

1. Features are hard to share and reuse.
2. Reliably serving in production with low latency is a challenge.
3. Inadvertent skew in feature values between training and serving is common.

These three key challenges are all considered feature management pain points. So how can the expanding machine learning team of Armcomp overcome these challenges? This is where Vertex AI Feature Store comes in to solve these feature management problems.

Before using the Feature Store, Armcomp computed feature values and saved them in various locations such as tables in BigQuery or hedge files in Cloud Storage. They also built and managed separate solutions for storage and consumption of feature values. Vertex AI Feature Store provides a centralized repository for organizing, storing, and serving machine learning features. By using a central feature store, the Armcomp team can efficiently share, discover, and reuse ML features at scale, letting the team increase the speed at which they can develop and deploy new management learning applications. Feature Store is a fully managed solution. It manages and scales the underlying infrastructure for us, which means that our data scientists can focus on the feature computation logic instead of worrying about the challenges of deploying features into production.

Benefits of Vertex AI Feature Store

1. Teams can share and reuse ML features across use cases because Feature Store has a centralized feature repository with easy APIs to create, modify, search, and discover features, fetch them for training and serving, and also managing permissions.
2. Teams can serve ML features at scale with low latency because operational overhead is handled by Feature Store.
3. A common problem in machine learning is termed training-serving skew, which means the data we train with may be skewed or different from the one that is served in production. Feature Store alleviates training-serving skew as it enables its users to compute feature values once in a single repository then reuse it for both training and serving. Drift and other quality issues can be tracked and monitored as well.
4. Feature Store also allows for feature ingestion in large batches or in real time as data streams in. With Vertex AI Feature Store, the team can store features with batch and stream import APIs and register the feature to its feature registry.

Architecture of Vertex AI Feature Store

There are three main concepts in the Feature Store.

Featurestore -> EntityType -> Feature

A Feature Store is a top-level container for our features and their values. When we set up a Feature Store, permitted users can add and share their features without additional engineering support. Users can define features and then ingest or import feature values from various data sources.

To move further on, let’s compare Feature Store architecture with the Relational Database architecture.

Credit: GCP documentation

An entity type in a Feature Store is like the primary key column in a relational database. It is a collection of semantically related features.

An entity is an instance of an entity type. It is like a value of the primary key column in RDBMS. For example, movie_01 and movie_02 are entities of the movie’s entity type. In a Feature Store, each entity must have a unique ID like every value of the primary key column in RDBMS and must be of type STRING.

A feature in a Feature Store is like a non primary column in RDBMS. It is a measurable property or attribute of an entity type. For example, the movie entity type has features such as average_rating and title that track various properties of movies. Features are associated with entity types. They must be distinct within a given entity type, but they don’t need to be globally unique. For example, if we use title for two different entity types, Vertex AI Feature Store interprets title as two different features. When reading feature values, we provide the feature and its entity type as part of the request.

A feature value is like a value of a non primary key column in a relational database. Vertex AI Feature Store captures feature values for a feature at a specific point in time. In other words, we can have multiple values for a given entity and feature. For example, the movie_01 entity can have multiple feature values for the average_rating feature. The value can be 4.4 at one time and 4.8 at some later time. Vertex AI Feature Store associates a tuple identifier with each feature value (entity_id, feature_id, timestamp), which Vertex AI Feature Store uses to look up values at serving time.

Timestamp is a separate concept in Vertex AI Feature Store that does not have its alternative in a relational database. For batch ingestions, Vertex AI Feature Store requires user-provided timestamps for the ingested feature values. We can specify a particular timestamp for each value or specify the same timestamp for all values:

• If the timestamps for feature values are different, specify the timestamps in a column in your source data. Each row must have its own timestamp indicating when the feature value was generated. In your ingestion request, you specify the column name to identify the timestamp column.
• If the timestamp for all feature values is the same, you can specify it as a parameter in your ingestion request. You can also specify the timestamp in a column in your source data, where each row has the same timestamp.

What about the feature ingestion and the feature serving? Those are the 2 other important concepts in the feature store and somehow are opposite actions.

Feature ingestion

The process of importing feature values computed by our feature engineering jobs into a featurestore. Before we can ingest data, the corresponding entity type and features must be defined in the featurestore. Vertex AI Feature Store offers batch and streaming ingestion, letting us add feature values in bulk or in real time.

Feature serving

Get Arman Malkhasyan’s stories in your inbox

Join Medium for free to get updates from this writer.

Remember me for faster sign in

The process of exporting stored feature values for training or inference. Vertex AI Feature Store offers two methods for serving features: batch and online. Batch serving is for high throughput and serving large volumes of data for offline processing. Online serving is for low-latency data retrieval of small batches of data for real-time processing.

Demo

Now, let’s dive into the practical world and go through the concepts that I introduced above. Let’s create an automated pipeline which

1. Creates a Feature Store, an entity type, a feature
2. Encodes and ingests 100 images from Cloud Storage into the feature store (feature ingestion)
3. Reads a randomly selected image from the feature store decodes and transfers back to the cloud storage (feature serving)
4. Deletes Feature store

What should we have on hand?

• Local or Virtual machine with any OS (in this demo we will use Ubuntu)
• Python 3.6+
• gcloud CLI installed

All codes are publicly available on my GitHub repo

Credit: mine

The UI of Vertex AI Feature Store is quite simple. We can create a feature store directly from UI, but for our application we will use Python script.

1. Specify GCP credentials

Our application will run on a VM environment with Ubuntu OS. So we should specify GCP credentials to run the application. For that purpose we should create a Service Account taking count the principle of least privileges. We should grant the following roles:

• Vertex AI Administrator
• Storage Object Admin
• BigQuery Data Editor (When ingesting from DataFrame to feature store a temporary BigQuery dataset is automatically created and deleted in the same GCP project, which is used as the intermediary storage for ingesting feature values from dataframe to featurestore.)

After we create a Service Account in GCP we should generate a new key in json format, download it and save it in the VM where the application would run. Then we can specify GCP credentials in Python environment.

import os

# Specify GCP credentials
def specifyGoogleCredentials():
   key_path = os.path.dirname(__file__) + '/key.json'
   os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = key_path

2. Read from GCS

As we will ingest data from DataFrame to GCS let’s read 100 images from Cloud Storage, encode them and store them in a pandas dataframe.

import pandas as pd
from datetime import datetime, timezone

from google.cloud import storage

def readFromGCS(SOURCE_BUCKET: str):
# Read images from GCS bucket and convert them from jpeg to string
# Then define DataFrame in put all the information there

   image_id = []
   image_byte = []
   update_time = []

   client = storage.Client()
   bucket = client.get_bucket(SOURCE_BUCKET)
   blobs = list(bucket.list_blobs())
   i = 1
   for blob in blobs:
       image_id.append('image-' + str(i))
      
       img_byte = blob.download_as_bytes()
       image_byte.append(img_byte)
      
       feature_time_str = datetime.now().isoformat(sep=" ", timespec="milliseconds")
       feature_time = datetime.strptime(feature_time_str, "%Y-%m-%d %H:%M:%S.%f")
       update_time.append(feature_time)
      
       i += 1
      
   imageDF = pd.DataFrame(list(zip(image_id, image_byte, update_time)), columns =['image_id', 'image_byte', 'update_time']) 
  
   print('===================================================================')
   print(f'==========> 1. Image dataframe is created with {len(imageDF)} rows')
   print('===================================================================')
   return imageDF

Let’s pay attention to how we encode images. We just downloaded them as bytes.

There are also other data types:

• BOOL
• BOOL_ARRAY
• DOUBLE
• DOUBLE_ARRAY
• INT64
• INT64_ARRAY
• STRING
• STRING_ARRAY
• BYTES

Why do we choose bytes? Initially there have been two options that easily work with jpeg format. They are STRING and BYTES. The download_as_string() method is deprecated in favor of using download_as_bytes() to download the contents of the blob as a bytes object.

Dataframe consists of 3 columns:

• image_id
• image_byte
• update_time

The source, whether it is a DataFrame, csv or avro file in cloud storage, or BigQuery, must include at least 3 columns. The first one is always entity type, the second one is a feature (we can have more than one feature) and the last column is update_time which includes user-provided timestamps for the ingested feature values. We can specify a particular timestamp for each value or specify the same timestamp for all values.

3. Create a feature store, an entity type, a feature and ingest data

import pandas as pd

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

def FeatureStore(PROJECT_ID: str, REGION: str, FEATURESTORE_ID: str, ONLINE_STORE_FIXED_NODE_COUNT: int, source: pd.DataFrame):
   # Create Featurestore
   fs = Featurestore.create(
       featurestore_id=FEATURESTORE_ID,
       online_store_fixed_node_count=ONLINE_STORE_FIXED_NODE_COUNT,
       project=PROJECT_ID,
       location=REGION,
       sync=True,
   )

   print('============================================================')
   print('==========> 2.', FEATURESTORE_ID, 'feature store is created.')
   print('============================================================')

   # Create Images Entity Type
   images_entity_type = fs.create_entity_type(
       entity_type_id="images",
       description="Images entity",
   )

   print('=============================================')
   print('==========> 3. images entity type is created.')
   print('=============================================')

   image_byte = images_entity_type.create_feature(
       feature_id="image_byte",
       value_type="BYTES",
       description="image converted from jpeg to byte",
   )

   print('=============================================')
   print('==========> 4. image_byte feature is created.')
   print('=============================================')

   # Import to Feature Store from dataframe
   IMAGES_FEATURES_IDS = [feature.name for feature in images_entity_type.list_features()]
   IMAGES_FEATURE_TIME = "update_time"
   IMAGES_ENTITY_ID_FIELD = "image_id"
           
   # DATAFRAME
   images_entity_type.ingest_from_df(
       feature_ids = IMAGES_FEATURES_IDS,
       feature_time = IMAGES_FEATURE_TIME,
       df_source = source,
       entity_id_field = IMAGES_ENTITY_ID_FIELD,
   )

   print('===============================================')
   print('==========> 5. Batch ingestion process is done.')
   print('===============================================')

FeatureStore function takes care of creating Feature Store, Entity type called images and a single feature called image_byte which includes encoded images data. Afterwards, the function organizes the data ingestion process from DataFrame to feature store. When ingesting from DataFrame to feature store a temporary BigQuery dataset (intermediary storage for ingesting feature values from DataFrame to featurestore.) is automatically created for transit after the latter is deleted in the same GCP project.

4. Online serving

from google.cloud import aiplatform

def onlineServing(PROJECT_ID: str, REGION: str, FEATURESTORE_ID: str, entity_ids: str):

   aiplatform.init(project=PROJECT_ID, location=REGION)
   my_entity_type = aiplatform.featurestore.EntityType(
           entity_type_name='images', featurestore_id=FEATURESTORE_ID
       )
   my_dataframe = my_entity_type.read(entity_ids=entity_ids, feature_ids='image_byte')

   print('========================================================')
   print(f'==========> 6. Computer randomly selected {entity_ids}.')
   print('========================================================')

   return my_dataframe

Online serving lets us serve feature values for small batches of entities at low latency. For each request, you can only serve feature values from a single entity type. Vertex AI Feature Store returns only the latest non-null value of each feature.

In the onlineServing function a randomly selected entity_ids is read from feature store to dataframe.

5. Migrate from feature store to Cloud Storage

import pandas as pd
from google.cloud import storage

def backToBucket(df: pd.DataFrame, DESTINATION_BUCKET: str):
   imgByte = df['image_byte'][0]
   imgName = 'img' + df['entity_id'][0][6:] + '.jpg'

   f = open(imgName, 'wb')
   f.write(imgByte)
   f.close()

   client = storage.Client()
   bucket = client.get_bucket(DESTINATION_BUCKET)
   blob = bucket.blob(imgName)
   blob.upload_from_filename(imgName)

print('=============================================================================')
print(f'==========> 7. Randomly selected image is converted from bytes to jpg and uploaded into {DESTINATION_BUCKET}.')
   print('=============================================================================')
print('==============================')
print('==========> Just go and check!')
print('==============================')

Here, a randomly selected image is decoded and migrated back to cloud storage.

6. Delete feature store

The last action that our application does is to delete a feature store just to save money.

from google.cloud import aiplatform

def deletFS(PROJECT_ID: str, REGION: str, FEATURESTORE_ID: str):

   aiplatform.init(project=PROJECT_ID, location=REGION)
   fs = aiplatform.featurestore.Featurestore(featurestore_name=FEATURESTORE_ID)
   fs.delete(sync=True, force=True)

   print('===========================================================')
   print('==========> 8.', FEATURESTORE_ID, 'feature store is deleted')
   print('===========================================================')

7. Executing program above

First of all, let’s gather all the necessary information.

PROJECT_ID = ...
REGION = ...
SOURCE_BUCKET = ...
DESTINATION_BUCKET = ...
FS = ...
ONLINE_STORE_FIXED_NODE_COUNT = ...
FEATURESTORE_ID = ...

Now, let’s run the application.

import random

# GOOGLE  CREDENTIALS
gcpCredentials.specifyGoogleCredentials()

# READ  IMAGES  FROM  GCS,  CREATE  DATAFRAME
df = readFromGCS(SOURCE_BUCKET)

# CREATE  FEATURE  STORE, MANAGE BATCH INGESTION
FeatureStore(PROJECT_ID, REGION, FEATURESTORE_ID, ONLINE_STORE_FIXED_NODE_COUNT, df)

# ONLINE SERVING, READ A RANDOMLY SELECTED IMAGE
entity_ids = 'image-' + str(random.randint(1, 100))
testDF = onlineServing(PROJECT_ID, REGION, FEATURESTORE_ID, entity_ids)

# DECODE A RANDOMLY SELECTED IMAGE AND PUT IT BACK INTO BUCKET
backToBucket(testDF, DESTINATION_BUCKET)

# DELETE FEATURE STORE
deletFS(PROJECT_ID, REGION, FEATURESTORE_ID)

This is how you can develop a pipeline which does Feature ingestion and serving. Enjoy the process!

Connect with me over LinkedIn.