Separating Environments for an ML Platform on Snowflake

Sharing ML Models Across Accounts Is Now Possible

Snowflake's Direct Share now supports ML models.

(Reference: Snowflake Model Registry - Sharing models)

This landed without me noticing. A good reminder that you really do have to keep up with the documentation.

This feature significantly widens the set of options available when you use Snowflake as an ML platform, so I want to lay them out.

Checking What Is Now Possible

First, let me go through what Direct Share actually lets you do.

I prepared a model using Snowflake's example helper. Model accuracy is irrelevant here, so the model itself is thrown together.

As I explain later, sharing behaves differently depending on whether you run inference in a warehouse or on SPCS, so the sample code below builds both.

Preparing a prediction model in the provider account

Environment:

• Snowflake Notebook
• Container Runtime v2.6

from snowflake.ml.feature_store.examples.example_helper import ExampleHelper
from snowflake.ml.feature_store import (
    FeatureStore,
    FeatureView,
    Entity,
    CreationMode,
    FeatureViewStatus,
)
from snowflake.ml.registry import Registry
from snowflake.ml.model.target_platform import TargetPlatform

import xgboost as xgb
from sklearn.model_selection import train_test_split
import pandas as pd

from snowflake.snowpark.context import get_active_session
session = get_active_session()

example_helper = ExampleHelper(session, session.get_current_database(), 'PUBLIC')


source_tables = example_helper.load_example('new_york_taxi_features')

fs = FeatureStore(
    session=session, 
    database=session.get_current_database(), 
    name='PUBLIC', 
    default_warehouse=session.get_current_warehouse(),
    creation_mode=CreationMode.CREATE_IF_NOT_EXIST,
)


for fv in example_helper.load_draft_feature_views():
    fs.register_feature_view(
        feature_view=fv,
        version='1.0'
    )

entity_key_names = ','.join(my_entity.join_keys)
spine_df = session.sql(f"SELECT {entity_key_names} FROM {source_tables[0]}").sample(n=1000)

training_fv = fs.get_feature_view(target_feature_view, '1.0')

training_data_df = fs.generate_training_set(
    spine_df=spine_df,
    features=[training_fv]
)


df = training_data_df.to_pandas()
feature_cols = ['PASSENGER_COUNT', 'TRIP_DISTANCE', 'TIP_AMOUNT', 'TOLLS_AMOUNT', 'PICKUP_LOCATION_ID', 'DROPOFF_LOCATION_ID']
target_col = 'FARE_AMOUNT'

X = df[feature_cols]
y = df[target_col]

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

model = xgb.XGBRegressor(n_estimators=100, max_depth=4, learning_rate=0.1, random_state=42)
model.fit(X_train, y_train)

reg = Registry(session=session, database_name='ML_SHARE_TEST', schema_name='PUBLIC')
reg.log_model(
    model=model,
    model_name='taxi_fare_xgboost',
    version_name='v1',
    sample_input_data=X_train[:10],
    conda_dependencies=['xgboost'],
)

reg.log_model(
    model=model,
    model_name='taxi_fare_xgboost',
    version_name='v2_warehouse',
    sample_input_data=X_train[:10],
    conda_dependencies=['xgboost'],
    target_platforms=[TargetPlatform.WAREHOUSE],
)

Version v1 runs on SPCS, and v2_warehouse runs in a warehouse.

1. Provider Side: Create the Share Object

Create the share on the provider side.

USE ROLE ACCOUNTADMIN;
CREATE SHARE ML_MODEL_SHARE
    SECURE_OBJECT_ONLY = FALSE;

2. Grant Privileges to the Share

There are two ways to do this. One is to grant the model privileges directly to the share; the other is to grant them to a database role and then hand that role to the share.

The database role approach makes it easier for the consumer to reproduce the provider's role structure, so that is what I use here.

CREATE DATABASE ROLE DB_ROLE_SHARE;

-- Grant schema USAGE to the database role
GRANT USAGE ON SCHEMA ML_SHARE_TEST.PUBLIC TO DATABASE ROLE DB_ROLE_SHARE;

-- Grant the model privilege to the database role
GRANT USAGE ON MODEL ML_SHARE_TEST.PUBLIC.TAXI_FARE_XGBOOST TO DATABASE ROLE DB_ROLE_SHARE;

-- Put the database role into the share
GRANT USAGE ON DATABASE ML_SHARE_TEST TO SHARE ML_MODEL_SHARE;
GRANT DATABASE ROLE DB_ROLE_SHARE TO SHARE ML_MODEL_SHARE;

-- Add the consumer account to the share
ALTER SHARE ML_MODEL_SHARE ADD ACCOUNTS = <consumer_account_locator>;

You can view the share you created under External Sharing in the Data Sharing tab of Snowsight.

The database role is indeed included in the share.

3. Consumer Side: Create a Database from the Share

Now we move to the consumer account and create a shared database from the share.

CREATE DATABASE SHARED_ML_DB FROM SHARE <provider_account>.ML_MODEL_SHARE;
GRANT DATABASE ROLE SHARED_ML_DB.DB_ROLE_SHARE TO ROLE <custom_role>;

Here the shared database role is inherited by an appropriate custom role on the consumer side.

The model is now visible in the consumer's database explorer.

4. Consumer Side: Run Inference

You can run inference with the shared model.

from snowflake.ml.registry import Registry

reg = Registry(session, database_name='SHARED_ML_DB', schema_name='PUBLIC')
model = reg.get_model('TAXI_FARE_XGBOOST')
mv = model.version('V2_WAREHOUSE')

input_df = session.table('SHARED_ML_DB.PUBLIC."TAXI_TRIP_FEATURES$1.0"').select(
    'PASSENGER_COUNT', 'TRIP_DISTANCE', 'TIP_AMOUNT', 
    'TOLLS_AMOUNT', 'PICKUP_LOCATION_ID', 'DROPOFF_LOCATION_ID'
)

result_wh = mv.run(input_df, function_name='PREDICT')
result_wh.show(10)

The code is no different from working with a normal model. As long as the consumer provides the data and the compute, the same inference code that runs on the provider side runs here too.

Note 1: The Difference Between Privileges

There are two privileges you can hand to a share from a model: USAGE and READ. The easiest way to think about which one you need is in terms of the target_platform the model was created with.

target_platform	What you want to do	Required privilege
Warehouse	Inference via mv.run	USAGE
SPCS	Creating an inference service, inference on SPCS	READ

If you grant only USAGE on an SPCS model, the model's artifact files cannot be read and therefore cannot be loaded onto SPCS. That rules out both creating an inference service and running inference on SPCS with something like run_batch.

Also, with either privilege, you cannot copy a shared model into a model of your own.

-- This does not work
CREATE MODEL CONSUMER_DB.PUBLIC.SHARE_TEST_MODEL
FROM MODEL SHARED_ML_DB.PUBLIC.TAXI_FARE_XGBOOST VERSION V1;

Note 2: Replication

Model objects support not only sharing but also replication. Replicating one materializes the model on the consumer side. The materialized object is a replica of the source object, and its contents cannot be modified.

Designing Environments Around This Feature

From here I want to look at the possibilities that open up now that models can be shared.

Problems You Run Into with MLOps

One of the big advantages of practising MLOps or LLMOps on Snowflake is that you can build the platform directly on top of your data, with nothing in between. ML and AI only work when the underlying data is trustworthy, so being able to develop an ML platform as one more capability of the data platform is a significant benefit.

That said, building your environment around the data platform sometimes imposes constraints on the ML platform. Separate development and production accounts are a good example. When validation and production live in different Snowflake accounts, problems like these come up:

• Even after confirming a model's accuracy in the validation environment, you have to retrain it in production, which means the accuracy of the model you actually operate is unknown
• You end up repeating experiments in the production environment, risking an outage from human error
• Even for identical processing, the features you can build differ between the validation and production environments, which makes validation meaningless

How much each of these matters depends on the business problem you are applying ML to and on the values of your team and company. So rather than simply chasing best practices, you need to think hard about which design actually fits what you want to do.

This article presents a few options that look reasonable, but each has its trade-offs and none of them is strictly superior. I hope it gives you a foothold for reaching the MLOps setup that is right for you.

Comparing the Concrete Options

It is easier to organise the patterns if you focus on three points:

1. Where training happens
2. Where accuracy validation happens
3. How a model is promoted to production

Pattern Overview

	A. Single environment	B. Training in production	C. Training in development
Training location	Production	Production	Development
Validation location	Production	Development	Development
Promotion method	Alias	Alias	Replication + alias

Let me go through each pattern with a diagram.

A. Single Environment

The simplest setup, where training and validation of the ML model are completed entirely within the production account.

A trunk-based branching strategy tends to keep development fast here, so I think this is the form to start with for a PoC or a small project.

By small project I mean the blast radius of the MLOps work, not the size of the data. For instance, when the goal is simply to build a model and produce predictions (when ML sits at the very end of the workflow), a few errors or some bad data may be tolerable. When most of the blast radius is under the control of the people doing MLOps, this approach has a lot going for it.

Pros: the model you validated is the model you operate. No data movement is involved, so you can start with nothing more than your regular processing pipeline.

Cons: ML model developers need access to the production environment, and as the scale grows, privilege management tends to become the bottleneck.

B. Training in Production

A setup where a model trained in the production account is exposed to the development account via Direct Share, and only validation work happens on the development side.

It is easier to get around the privilege constraints of production while still giving ML model developers a reasonable degree of freedom, but deploying a model to production involves more steps, so development slows down somewhat. On the other hand, it combines well with a wide range of branching strategies, and I think it adapts to almost any project.

Pros: because training happens in production, the validated model and the model running in production are the same object. Even when the training data contains sensitive information, you can train without exposing that data to the development environment. You can also validate safely in the development environment without running commands directly against production or triggering production pipelines.

Cons: since training runs in production, the branching strategy tends to get complicated, or you end up with several approval phases. You have to weigh how much developer velocity the data scientists need against how much clutter you are willing to accept in production, and find a development flow that fits.

C. Training in Development

An approach where a model trained in the development environment is replicated into production and operated there. Replication breaks the dependency between the environments and lets you operate the model as an immutable object.

Trial and error on the model happens in the development environment, which allows for flexible development. On the other hand, you have to guarantee the credibility of a model built in development and keep it consistent with objects outside the model itself, which makes the environment harder to build.

Some companies have operational policies that forbid putting objects created in a development environment into production, so it is important to sketch out a workable shape first. That demands a high level of engineering.

Pros: high freedom for trial and error and a fast iteration loop for developers. As with B, the model running in production and the model used for validation are the same.

Cons: if some data has to be masked in the development environment, you may not be able to develop the model properly. You also have to tighten the privilege design in the development environment as well, which tends to come back as operational and maintenance cost.

If your data platform is already running solidly and you are adding an ML environment on top of it, this may be an easy option to take.

Closing

In this article I organised the options for designing an MLOps environment around Snowflake's model sharing and replication features.

Each of the three patterns has its own pros and cons, and no single one is always the right answer. The right setup shifts with the scale and phase of the project and with your organisation's security policy.

Start with a simple setup, then revisit the architecture in stages as the project grows and the data platform changes. Snowflake gives you features with exactly that kind of flexibility.