Building an ML pipeline with modular code and Snowflake ML Jobs

1.From notebook cells to a chained pipeline

Snowflake splits an ML workflow into four broad stages: data exploration and preparation, data engineering, model development, and model deployment. Early work is usually iterative, done in a notebook or a local IDE. Once a model proves its value, the goal changes to operationalization. The pipeline gets hardened and automated, so that every change to code, data or model is built, tested and deployed the same way each time.

The documentation recommends structuring things for this from the start. Parameterize inputs (tables, stages, hyperparameters) and give each step its own cell or function. Before you operationalize, turn each major step (data preparation, feature engineering, training, evaluation) into a function with clear inputs and outputs. Move every configuration value out of the code, so the same pipeline can run in dev and in prod. Also write one entrypoint script that runs the whole pipeline locally, which you can use for debugging.

    # Load configuration
    config = load_config(args.config, args.env)

    # Execute pipeline stages
    raw_data = load_raw_data(config.data.source_table)
    validate_data_quality(raw_data, config.data.quality_checks)
    features = create_features(raw_data, config.features.transformations)
    model = train_model(features, config.model.hyperparameters)
    metrics = evaluate_model(model, features, config.model.eval_metrics)
    register_model(model, metrics, config.model.registry_name)

This script is the chain the exam guide describes, written out in order. If validate_data_quality fails, training never starts. If a model has not been evaluated, it never reaches register_model. The --env argument (default dev) picks the configuration, so promoting the pipeline to production means changing an argument, not editing code. The example project layout puts this entrypoint in scripts/ next to a dag.py, with stage code under src/ml_pipeline/ split into data/, features/, models/ and inference/.

The documentation names a Snowflake tool for each lifecycle stage. The rest of this lesson uses the ones in the table below.

Sources1

2.Running pipeline steps as ML Jobs from the Python SDK

Training usually needs more compute than a warehouse provides. Snowflake ML Jobs run ML workloads inside Snowflake's container runtimes on compute pools, including GPU and high-memory CPU instances, and you can submit them from any development environment. The SDK is the snowflake-ml-python package, and ML Jobs need version 1.26.0 or later. The documentation also lists integration with orchestration tools such as Apache Airflow, so an external scheduler can call the same API.

Most code written in Snowflake Notebooks works in ML Jobs without changes. One exception: some distributed ML APIs exist only inside the Container Runtime, so you must import them inside the job payload, not at the top of your local script.

Every ML Jobs call needs a Snowpark Session. The documentation builds one with Session.builder.getOrCreate(), which reads a valid ~/.snowflake/config.toml. Once the session exists, you can pass it explicitly or let the API infer it from context.

There are two ways to submit work. Function dispatch uses the @remote decorator. It serializes the function and its dependencies, uploads them to a stage, and runs them in a Container Runtime. File dispatch submits whole files or project directories, which suits multi-module projects, scripts that take command-line arguments, and existing projects that were not written for Snowflake. File dispatch has three methods:

from snowflake.ml.jobs import submit_file

# Run a single file
job1 = submit_file(
  "train.py",
  "MY_COMPUTE_POOL",
  stage_name="payload_stage",
  args=["--data-table", "my_training_data"],
  session=session,
)

from snowflake.ml.jobs import submit_directory

# Run from a directory
job2 = submit_directory(
  "./ml_project/",
  "MY_COMPUTE_POOL",
   ? ="train.py",
  stage_name="payload_stage",
  session=session,
)

Sources21

3.Getting results back and pinning the runtime

A pipeline can only chain steps if each step's output is available to the next. Every submission method returns an MLJob object that you use to manage and monitor the run. Inside the job, a Snowpark Session is already available, so the payload can read tables and stages and write results back. For function dispatch, you can declare a session: Session parameter and the session is injected automatically. That only works for function payloads. Session.builder.getOrCreate() works for every payload type.

How you return a value depends on the dispatch style. A @remote function just returns it. A file-based job assigns it to the special __return__ variable, for example __return__ = main(). On the client, MLJob.result() blocks until the job reaches a terminal state. It then returns the payload's value, or raises an exception if the run failed. If the payload returned nothing, a successful job gives None.

from snowflake.ml.jobs import remote

@remote("MY_COMPUTE_POOL", stage_name="payload_stage", session=session, runtime_environment="2.3.0")
def train_model(data_table: str):
  # Provide your ML code here, including imports and function calls
  ...

If you leave out runtime_environment, the job uses the latest Container Runtime available on the compute pool. That is fine for experiments, but a production pipeline should pin a full version string so runs are reproducible. For package or compliance needs, you can register a custom image and reference it as runtime_environment="cre@<name>". Compute pools are sized with MIN_NODES and MAX_NODES. The default pool uses the CPU_X64_S instance family, with a minimum of 1 node and a maximum of 25.

Sources2

Exam traps

Each one states something that sounds right. Open it to see what is actually true.

1. 1.ML Jobs can only be launched from a Snowflake Notebook or worksheet.Why is that wrong?

   You can submit ML Jobs from any development environment with snowflake-ml-python, including local IDEs and external orchestrators such as Airflow.

   Covered in Running pipeline steps as ML Jobs from the Python SDK

2. 2.A file-based ML Job returns whatever its main() function returns, the same way a @remote function does.Why is that wrong?

   A file payload has to assign its result to the special __return__ variable. Otherwise result() gives None on success.

   Covered in Getting results back and pinning the runtime