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. Environment: Version Preparing a prediction model in the provider account 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], ) 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 = ; 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 .ML_MODEL_SHARE; GRANT DATABASE ROLE SHARED_ML_DB.DB_ROLE_SHARE TO 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: - Where training happens - Where accuracy validation happens - 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 m