Split Test AI/ML in Snowflake

Cortex

Test different models with SQL

This example demonstrates how you can run a SQL query that produces an output from Cortex’s COMPLETE method, testing different model versions for various customers. Winning Variant is used to determine which model should be used per customer ID. Later, we would tie business KPIs (i.e., revenue, churn, returns, etc) back to each variant to see which model should be deployed. We have the following declared for this example:

A 50/50 test between control (llama3.2-1b model) and treatment (claude-3-5-sonnet)
The Winning Variant experiment ID is cortex-test
The ID of your user in this example is abc123

-- Get an assignment from Winning Variant for the experiment
WITH assignment AS (
    SELECT create_assignment('abc123', 'cortex-test') AS variant
)

-- If the assigned variant for this user is 'TREATMENT',
-- use 'claude-3-5-sonnet', otherwise use 'llama3.2-1b' (CONTROL)
SELECT
  CASE
    WHEN variant = 'TREATMENT'
    THEN SNOWFLAKE.CORTEX.COMPLETE('claude-3-5-sonnet', 'Is a hot dog a sandwich?')
    ELSE SNOWFLAKE.CORTEX.COMPLETE('llama3.2-1b', 'Is a hot dog a sandwich?')
  END AS result
FROM assignment;

ML Models

This following examples demonstrate how you would use Winning Variant to run predictions for a set of customers using two model versions from the Snowflake Model Registry. The experiment has the following definition:

A 50/50 test between control (v1 of our model) and treatment (v2)
Winning Variant experiment ID of ml-test

customer_test_data in these examples contains the following fields:

customer_id: ID of the customer
age: how long the customer has been a customer
usage: some arbitrary integer that defines how much the customer has used the product

The sample model customer_churn_model has been trained to predict churn given an age and usage for a customer. The “Python (Simple)” example below uses the Winning Variant Python SDK.

-- Make/get assignments for all customers IDs
WITH assignments AS (
    SELECT
		customer_id,
		create_assignment(customer_id, 'ml-test') AS variant
	FROM customer_test_data
),

-- Pull model v1 (our CONTROL variant)
model_v1 AS MODEL customer_churn_model VERSION v1,

-- Pull model v2 (our TREATMENT variant)
model_v2 AS MODEL customer_churn_model VERSION v2

-- If the assigned variant for the customer is 'TREATMENT',
-- use v2, otherwise use v1 (CONTROL)
SELECT
    c.customer_id,
    a.variant as experiment_variant,
    CASE
        WHEN a.variant = 'TREATMENT'
        THEN model_v2!predict(age, usage)
        ELSE model_v1!predict(age, usage)
    END AS prediction
FROM customer_test_data c
    LEFT JOIN assignments a ON c.customer_id = a.customer_id;

import snowflake.snowpark as snowpark
from snowflake.snowpark.functions import call_udf, col
from snowflake.ml.registry import Registry

# Load the model from the registry
registry = Registry(session)
model = registry.get_model("customer_churn_model")

# Get customer data to run predictions on
df_customers = session.sql('SELECT * FROM customer_test_data')

# Use the Winning Variant UDF to get/make assignments for customers
customer_list = df_customers.with_column("experiment_variant", call_udf("create_assignment", col("customer_id"), "ml-test")).collect()

# Split customers into experiment variants
control_rows = [row for row in customer_list if row["EXPERIMENT_VARIANT"] == "CONTROL"]
treatment_rows = [row for row in customer_list if row["EXPERIMENT_VARIANT"] == "TREATMENT"]

# Convert the row lists back into DataFrames
df_control = session.create_dataframe(control_rows)
df_treatment = session.create_dataframe(treatment_rows)
    
# Run inference
v1_predictions = model.version('v1').run(df_control, function_name='predict')
v2_predictions = model.version('v2').run(df_treatment, function_name='predict')
    
# Show predictions
v1_predictions.show()
v2_predictions.show()

from snowflake.ml.registry import Registry
from winningvariant import WinningVariantClient

# Initialize the Winning Variant SDK
wv = WinningVariantClient(session)

# Load the model from the registry
registry = Registry(session)
model = registry.get_model("customer_churn_model")

# Get the assignment for the user we want to run a prediction for
assignment = wv.create_assignment(subject_id="user_123", experiment_id="my-exp")

if assignment == "v2":
  prediction = model.version('v2').run(...)
elif assignment == "v1":
  prediction = model.version('v1').run(...)
else:
  # Something went wrong. Fall back to v1.
  prediction = model.version('v1').run(...)

# Show prediction
prediction.show()

​Cortex

​Test different models with SQL

​ML Models

Cortex

Test different models with SQL

ML Models