Snowflake Tasks Recipe (Preview)

Register two Snowflake ML Jobs and two Tasks that train and score a Kumo Engine model with MLOps turned on. When you finish this page, a train task runs every Sunday at 02:00 UTC: it fits the model on your Snowflake tables, stores it under a Snowflake stage, records the fit in Snowflake Experiment Tracking, and publishes a new version of USERS_CHURN in the Snowflake Model Registry with the champion alias. A score task runs daily at 03:00 UTC: it loads the version that holds champion, scores your users with trainer.predict() as in Score in Batch, and appends the scores to the table USERS_CHURN_PREDICTIONS. By default every train run moves champion, so a new version is used without review; see Review Before Promoting.

This recipe is in preview. It uses ML Job definitions (MLJobDefinition), a preview Snowflake API, and the statement they generate to run a job from a Task is not documented by Snowflake; both can change without notice. Nothing in this recipe has been registered or run yet.

The jobs and tasks come from examples/mlops/snowflake in the Kumo Engine repository. Each job runs one of the job scripts in examples/mlops, fit_job.py or score_job.py, on a Snowpark Container Services compute pool, on an example graph of users, items, and orders. They run the same kumo.MLOps workflow as Turn On MLOps, with a stage as the artifact root and Snowflake as the tracker and registry. To use your own graph, edit build_graph, QUERY, and TEXT_COLUMNS in jobs.py as described in Adapt the Template. For text columns, pass the wheel with [snowflake,text] as the --requirement.

The examples use the database ML and the schema MODELS. The source tables, stages, tasks, model, experiment, and predictions table all live in one schema; replace ML.MODELS with yours throughout.

Before You Start

You need:

Three tables in one schema: USERS with primary key USER_ID, ITEMS with primary key ITEM_ID, and ORDERS with primary key ORDER_ID, time column ORDER_TIME, and the foreign keys USER_ID and ITEM_ID. The default query needs a few months of orders.
A warehouse, which the tasks and the jobs’ queries use.
A GPU compute pool for the jobs, created with CREATE COMPUTE POOL.
An external access integration with a network rule that allows PyPI, from which each job installs the Kumo Engine wheel’s dependencies.
The kumo-engine wheel, from PyPI or the Kumo Engine releases page. Turn On MLOps shows how to check that a build includes MLOps.
A local Python environment with that wheel installed with the snowflake extra, and a connections.toml entry for the role that will own the tasks.

The role that owns the tasks needs:

For	Privileges
Compute	`USAGE` on the database, schema, warehouse, and compute pool; `CREATE SERVICE` on the schema; `EXECUTE TASK` on the account
Tasks and outputs	`CREATE TASK`, `CREATE STAGE`, `CREATE TABLE`, `CREATE MODEL`, and `CREATE EXPERIMENT` on the schema
Stages	Read and write access to the payload stage and the artifact root stage
Packages	`USAGE` on the external access integration
Data and model	`SELECT` on the source tables and `INSERT` on the predictions table. The first train run creates the model, and the role must own it to move its alias

Create the Stages

Create a stage for the job payloads and a stage for the artifact root. The artifact root stage needs server-side encryption:

CREATE STAGE ML.MODELS.KUMO_JOBS;
CREATE STAGE ML.MODELS.KUMO ENCRYPTION = (TYPE = 'SNOWFLAKE_SSE');

SHOW STAGES IN SCHEMA ML.MODELS lists both.

Add the Wheel to the Job Payload

Clone the Kumo Engine repository. Each job uploads the examples/mlops directory as its payload and installs its requirements inside it, so copy the kumo-engine wheel into a wheels directory there:

git clone https://github.com/NVIDIA/kumo-engine.git
cd kumo-engine
mkdir -p examples/mlops/wheels
cp <wheel dir>/<wheel>.whl examples/mlops/wheels/

The job refers to the wheel by its path in the payload, ./wheels/<wheel>.whl. A stage path does not work there, because pip cannot read from a stage.

Register the Jobs and Tasks

Run register_jobs.py with your connection, compute pool, stages, location, and warehouse. Replace the placeholders; <eai> is your external access integration:

python examples/mlops/snowflake/register_jobs.py --connection <connection> \
  --compute-pool <compute-pool> --stage ML.MODELS.KUMO_JOBS \
  --database ML --schema MODELS --warehouse <warehouse> \
  --artifact-root @ML.MODELS.KUMO/kumo \
  --requirement './wheels/<wheel>.whl[snowflake]' \
  --external-access-integration <eai>

The command registers the ML Job definitions KUMO_TRAIN and KUMO_SCORE, creates the tasks ML.MODELS.KUMO_TRAIN_TASK and ML.MODELS.KUMO_SCORE_TASK, and prints each task’s state:

ML.MODELS.KUMO_TRAIN_TASK: suspended (ALTER TASK ML.MODELS.KUMO_TRAIN_TASK RESUME)
ML.MODELS.KUMO_SCORE_TASK: suspended (ALTER TASK ML.MODELS.KUMO_SCORE_TASK RESUME)

--model-name (default USERS_CHURN, also used as the experiment name), --table-prefix (default empty, a prefix of the table names such as DEMO_), and --on-error (default raise) are passed to both jobs, and --predictions-table (default USERS_CHURN_PREDICTIONS) to the score job. A task is suspended after three failures in a row.

Run the Tasks Once

New tasks are suspended. Run the train task once, then the score task, and check each before you turn on the schedules:

EXECUTE TASK ML.MODELS.KUMO_TRAIN_TASK;

Check the task’s runs:

SELECT NAME, STATE, ERROR_MESSAGE, COMPLETED_TIME
FROM TABLE(ML.INFORMATION_SCHEMA.TASK_HISTORY(TASK_NAME => 'KUMO_TRAIN_TASK'))
ORDER BY SCHEDULED_TIME DESC;

STATE shows whether the task ran, but a failed job might not fail its task, so also check the job and its outputs. List the recent jobs with their name and status, using the same connection, then print a job’s log. Replace <job name> with a name from the list:

from snowflake.ml.jobs import get_job, list_jobs
from snowflake.snowpark import Session
session = Session.builder.config("connection_name", "<connection>").create()
print(list_jobs(database="ML", schema="MODELS", session=session))
print(get_job("<job name>", session=session).get_logs())

A finished job’s status is DONE, and the train job’s log ends with model: models/<model id>. A successful train run also leaves:

A new version of ML.MODELS.USERS_CHURN that holds champion, listed by SHOW VERSIONS IN MODEL ML.MODELS.USERS_CHURN.
A fit run in the experiment USERS_CHURN, in Snowsight under AI & ML, then Experiments.
The snapshot, training table, and model under @ML.MODELS.KUMO/kumo, listed by LIST @ML.MODELS.KUMO/kumo.

Then run the score task:

EXECUTE TASK ML.MODELS.KUMO_SCORE_TASK;

A successful score run appends one row per user to ML.MODELS.USERS_CHURN_PREDICTIONS, with the prediction run’s id in PREDICTION_ID and the model id in BUNDLE_ID, and stores the prediction run under the artifact root.

Turn On the Schedules

Resume both tasks:

ALTER TASK ML.MODELS.KUMO_TRAIN_TASK RESUME;
ALTER TASK ML.MODELS.KUMO_SCORE_TASK RESUME;

SHOW TASKS IN SCHEMA ML.MODELS shows both tasks as started. The train task runs Sundays at 02:00 UTC and the score task daily at 03:00 UTC.

Running register_jobs.py again replaces both jobs and tasks, and resumes each task that was started.

Review Before Promoting

Every train run moves champion to its new version, so the next score run uses it. To review each version first, make the train job move candidate instead: in register_jobs.py, change job_args=(flags if kind == 'train' else ... to job_args=([*flags, '--alias', 'candidate'] if kind == 'train' else ..., then run register_jobs.py again. Move champion yourself as described in Promote a Version.

Next Steps

Compare a new version with the champion before you move the alias.

MLOps Reference

Look up what each stored artifact records and where it is stored.