• Required snowflake-ml-python version: >=1.6.1
• Last updated on: 8/26/2024

End to end ML with Feature Store and Model Registry

This notebook demonstrates an end-to-end ML experiment cycle including feature creation, training data generation, model training and inference. The workflow touches on key Snowflake ML features including Snowflake Feature Store, Dataset, ML Lineage, Snowpark ML Modeling and Snowflake Model Registry.

Note: there may be a delay in the availability of the newest snowflake-ml-python package in the Snowflake Conda channel. To install the latest snowflake-ml-python package which includes all of necessary components used in this notebook, please follow the install instructions here.

Table of contents

• Set up test environment
  • Connect to Snowflake
  • Select your example
• Create features with Feature Store
  • Initialize Feature Store
  • Register entities and feature views
  • Examine features in Snowflake UI
• Generate Training Data
• Train model with Snowpark ML
• Log models in Model Registry
  • Examine model in Snowflake UI
• Query lineage (Preview Feature)
• Predict with model
  • Predict with local model
  • Predict with Model Registry
• Clean up notebook

Set up test environment

Connect to Snowflake

Let's start with setting up our test environment. We will create a session and a schema. The schema FS_DEMO_SCHEMA will be used as the Feature Store. It will be cleaned up at the end of the demo. You need to fill the connection_parameters with your Snowflake connection information. Follow this guide for more details about how to connect to Snowflake.

Select your example

We have prepared some examples that you can find in our open source repo. Each example contains the source dataset, feature view and entity definitions which will be used in this demo. ExampleHelper (included in snowflake-ml-python) will setup everything with simple APIs and you don't have to worry about the details.

load_example() will load the source data into Snowflake tables. In the example below, we are using the “new_york_taxi_features” example. You can replace this with any example listed above. Execution of the cell below may take some time depending on the size of the dataset.

Create features with Feature Store

Initialize Feature Store

Let's first create a feature store client. With CREATE_IF_NOT_EXIST mode, it will try to create a new Feature Store schema and all necessary feature store metadata if it doesn't exist already. It is required for the first time to set up a Feature Store. Afterwards, you can use FAIL_IF_NOT_EXIST mode to connect to an existing Feature Store.

Note that the database being used must already exist. Feature Store will NOT try to create the database even in CREATE_IF_NOT_EXIST mode.

Register entities and feature views

Next we register new entities and feature views in Feature Store. Entities will be the join keys used to generate training data. Feature Views contains all the features you need for your model training and inference. We have entities and feature views for this example defined in our open source repo. We will load the definitions with load_entities() and load_draft_feature_views() for simplicity.

We can examine all features in a feature view.

Examine features in Snowflake UI

Now you should be able to see registered entities and feature views in Snowflake UI.

Generate Training Data

After our feature pipelines are fully setup, we can use them to generate Snowflake Dataset and later do model training. Generating training data is easy since materialized FeatureViews already carry most of the metadata like join keys, timestamp for point-in-time lookup, etc. We just need to provide the spine data (it's called spine because it is the list of entity IDs that we are essentially enriching by joining features with it).

generate_dataset() returns a Snowflake Dataset object, which is best for distributed training with deep learning frameworks like TensorFlow or Pytorch which requires fine-grained file-level access. It creates a new Dataset object (which is versioned and immutable) in Snowflake which materializes the data in Parquet files. If you train models with classic ML libraries like Snowpark ML or scikit-learn, you can use generate_training_set() which returns a classic Snowflake table. The Cell below demonstrates generate_dataset().

Retrieve some metadata columns that are essential when generating training data.

Create a spine dataframe that's sampled from source table.

Generate dataset object from spine dataframe and feature views.

Convert dataset to a snowpark dataframe and examine all the features in it.

Train model with Snowpark ML

Now let's train a simple random forest model, and evaluate the prediction accuracy. When you call fit() on a DataFrame that is created from a Dataset, the linkage between the trained model and dataset is automatically wired up. Later, you can easily retrieve the training dataset from this model, or you can query the lineage about the dataset and model. This is work-in-progress and will be available soon in an upcoming release.

Log model in Model Registry

After the model is trained, we can save the model into Model Registry so we can manage the model, its metadata including metrics, versions, and use it later for inference. Also, ML lineage is built automatically between the model, dataset and feature views.

Log model into Model Registry.

Examine model in Snowflake UI

Now you should be able to see the model in Snowflake UI.

Query lineage (Preview Feature)

We can now query the lineage from an object. You can call lineage() on any object and it returns a set of objects that it has dependency with.

Predict with model

Finally we are almost ready for prediction! For this, we can look up the latest feature values from Feature Store for the specific data records that we are running prediction on. One of the key benefits of using the Feature Store is that it provides a way to automatically serve up the right feature values during prediction with point-in-time correct feature values. load_feature_views_from_dataset() gets the same feature views used in training, then retrieve_feature_values() lookups the latest feature values.

[Optional 1] predict with local model

Now we can predict with a local model and the feature values retrieved from feature store.

[Option 2] Predict with Model Registry

We can also retrieve the model from model registry and run predictions on the model using latest feature values.

Clean up notebook

This cell will drop the schemas have been created at beginning of this notebook, and also drop all objects live in the schemas including source data tables, feature views, datasets, and models.