Orchestrating ML Pipelines with Dagster and Snowflake ML

The challenge with ML applications is rarely a single training job; it is the operational lifecycle around it. Moving reliably from experimentation to evaluation, approval, deployment, and iteration requires an orchestrator, and Dagster has become a popular choice for teams building production ML workflows.

In this post, we walk through an end-to-end pipeline that trains a model on Snowpark Container Services compute pools, evaluates it against deployment thresholds, registers it in the Snowflake Model Registry, and deploys a live inference endpoint — all orchestrated by Dagster and executed on Snowflake.

For teams standardizing ML workflows on Snowflake, Dagster offers a compelling orchestration layer when pipelines need more than scheduling alone. It brings parameterized, typed steps; visual lineage and observability; conditional deployment gates; and step-level retries or re-execution after failure.

For ML teams already using Dagster, Snowflake ML Jobs provide a straightforward way to extend existing data pipelines into end-to-end ML workflows by running training, evaluation, and deployment steps on Snowflake-managed compute.

Importantly, the pattern described in this post does not require Dagster to run inside Snowflake. Teams can keep Dagster in their existing infrastructure and still use Snowflake ML Jobs as the execution layer, which is often the most practical model for organizations already invested in external orchestrators.

What We’re Building

The goal is to provide teams with a reusable and configurable framework for running ML workloads at scale on Snowflake using Dagster. The project structure makes it trivial to onboard new use-cases: create a folder under pipelines/, drop in a config, a job definition, and a training script, and the framework handles Snowflake connectivity, stage management, model registry, and deployment. Below, the walkthrough covers building an example pipeline and deploying it to Snowflake.

snowflake-dagster-ml/
├── definitions.py              # Dagster entry point
├── pipelines/
│   ├── __init__.py
│   ├── connection.py           # Shared Snowflake session helper
│   └── ml_use_case_name/
│       ├── config.py           # Infra + model config
│       ├── jobs.py             # Dagster ops + job
│       └── scripts/
│           └── train.py        # Training script (runs on SPCS)
├── deploy/
│   ├── Dockerfile
│   ├── build_and_push.sh
│   └── deploy.sql
└── pyproject.toml

Each pipeline has a **config.py** that centralizes infrastructure and model settings in one place — compute pool, stage name, model name, accuracy threshold, etc. This keeps the job logic clean and makes it easy to adjust settings without touching pipeline code:

# pipelines/iris/config.py
DATABASE = "ML_DEMO"
SCHEMA = "DAGSTER_ML"

COMPUTE_POOL = "COCO_ML_COMPUTE_POOL"
RUNTIME_ENVIRONMENT = "2.6.0"
STAGE_NAME = f"{DATABASE}.{SCHEMA}.PAYLOAD_STAGE"
MODEL_NAME = "IRIS_RF_DAGSTER"
MODEL_ARTIFACT = "iris_rf_model.pkl"
ACCURACY_THRESHOLD = 0.90
INFERENCE_SERVICE_NAME = "IRIS_RF_SERVICE"

ML Example using Dagster on Snowflake:

The following steps build a standard classification pipeline that trains a RandomForest model using the Iris dataset, evaluates it against a quality threshold, registers it to the Snowflake Model Registry, and deploys it as a live inference endpoint. The same pipeline runs identically in two environments: locally during development (via dagster dev) and in production on SPCS (deployed as a containerized service).

Each step is implemented as a Dagster op — the right abstraction here because each ML step is an imperative action with side effects (submitting a job to a compute pool, writing a model to the registry, creating a service). Each op can be re-executed independently if it fails, making it easy to resume a pipeline from a broken step without re-training from scratch.

Snowflake Connection & Authentication

A shared connection.py module handles authentication transparently across environments. In local development, it uses the SNOWFLAKE_CONNECTION_NAME environment variable to resolve credentials from ~/.snowflake/connections.toml.

In SPCS, it detects the DAGSTER_SPCS_MODE environment variable and switches to token-based OAuth using the service credential file at /snowflake/session/token—eliminating the need to manage secrets or bake credentials into the container image.

# pipelines/connection.py
import os
from snowflake.snowpark import Session

def create_session(database: str, schema: str) -> Session:
    if os.getenv("DAGSTER_SPCS_MODE"):
        token = open("/snowflake/session/token").read()
        account = os.environ["SNOWFLAKE_ACCOUNT"]
        return Session.builder.configs(
            {
                "account": account,
                "host": os.environ.get(
                    "SNOWFLAKE_HOST", f"{account}.snowflakecomputing.com"
                ),
                "authenticator": "oauth",
                "token": token,
                "database": database,
                "schema": schema,
                "warehouse": os.environ.get("SNOWFLAKE_WAREHOUSE", "COMPUTE_WH"),
            }
        ).create()
    else:
        return Session.builder.configs(
            {
                "connection_name": os.environ["SNOWFLAKE_CONNECTION_NAME"],
                "database": database,
                "schema": schema,
            }
        ).create()

Every op calls create_session() without worrying about which environment it's in. The same code, the same pipeline, zero auth changes between dev and prod.

How Compute Works

Rather than training locally or managing a separate Kubernetes cluster, each op in this pipeline submits work to Snowflake ML Jobs via submit_file(). This uploads a Python training script to a Snowflake stage and executes it on a SPCS compute pool — a managed container environment with ML libraries (scikit-learn, XGBoost, PyTorch, etc.) pre-installed. The Dagster process itself only orchestrates: it submits the job, polls for completion, reads results from stage, and decides what to do next.

This means:

— Training runs on Snowflake compute (CPU or GPU), close to the data, with no data egress — Dagster handles the workflow logic: sequencing, gating, error handling, observability — Artifacts (model files, metrics) persist on Snowflake stages between steps — The Dagster UI provides a visual timeline of every run, with logs from each ML Job

Note: When Dagster orchestrates Snowflake ML Jobs externally, only orchestration metadata leaves Snowflake — job status, run IDs, and completion signals — while your training data and model artifacts stay entirely within Snowflake’s compute boundary.

Prerequisites

• Snowflake account with Snowpark Container Services enabled
• Python 3.11+
• snowflake-ml-python==1.44.0
• snowflake-snowpark-python==1.52.0
• dagster==1.13.11

Step 1: The Training Script

This is the code that actually runs on Snowflake’s compute pool. It’s a plain Python script — no Dagster dependency, no Snowflake connection needed inside the container.

# pipelines/iris/scripts/train.py
import argparse, json, os, pickle
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, classification_report

def main():

    # Define configurable hyperparameters as CLI arguments
    parser = argparse.ArgumentParser()
    parser.add_argument("--n-estimators", type=int, default=100)
    parser.add_argument("--test-size", type=float, default=0.2)
    parser.add_argument("--random-state", type=int, default=42)
    parser.add_argument("--model-output", type=str,
                        default="/mnt/job_stage/artifacts/iris_rf_model.pkl")
    args = parser.parse_args()

    # Load data and split into train/test sets
    iris = load_iris()
    X_train, X_test, y_train, y_test = train_test_split(
        iris.data, iris.target, test_size=args.test_size,
        random_state=args.random_state
    )

    # Train the model
    model = RandomForestClassifier(
        n_estimators=args.n_estimators, random_state=args.random_state
    )
    model.fit(X_train, y_train)
    y_pred = model.predict(X_test)
    accuracy = accuracy_score(y_test, y_pred)

    # Save model artifact to stage-mounted path (persists after job completes)
    os.makedirs(os.path.dirname(args.model_output), exist_ok=True)
    with open(args.model_output, "wb") as f:
        pickle.dump(model, f)

    # Save metrics as JSON for downstream evaluation gate to read
    metrics_path = os.path.join(os.path.dirname(args.model_output), "metrics.json")
    with open(metrics_path, "w") as f:
        json.dump({"accuracy": accuracy, "n_train": len(X_train),
                   "n_test": len(X_test)}, f)
    print(f"Accuracy: {accuracy:.4f}")

if __name__ == "__main__":
    main()

Key points: — The script writes to /mnt/job_stage/artifacts/ — a stage-mounted path that persists after the job completes - Metrics are saved as JSON so downstream ops can read them - All hyperparameters are CLI arguments, making them configurable from Dagster

Step 2: The Dagster Pipeline

Here’s where orchestration happens. Each op connects to Snowflake, submits work, and passes results downstream.

# pipelines/iris/jobs.py
from dagster import op, job, Config, OpExecutionContext
from snowflake.ml.jobs import submit_file
from snowflake.ml.registry import Registry
from ..connection import create_session
from . import config as cfg

SCRIPTS_DIR = Path(__file__).parent / "scripts"
# Dagster Config class - these appear as editable fields in the UI launchpad
class TrainingJobConfig(Config):
    n_estimators: int = 100
    test_size: float = 0.2
    random_state: int = 42

@op
def ml_job_training_snowflake(context: OpExecutionContext,
                               config: TrainingJobConfig) -> str:

    # Create a Snowpark session (handles both local and SPCS auth automatically)
    session = create_session(cfg.DATABASE, cfg.SCHEMA)

    # Submit the training script to run on Snowflake's managed compute pool
    ml_job = submit_file(
        file_path=str(SCRIPTS_DIR / "train.py"),
        compute_pool=cfg.COMPUTE_POOL,
        stage_name=cfg.STAGE_NAME,
        args=["--n-estimators", str(config.n_estimators),
              "--test-size", str(config.test_size),
              "--random-state", str(config.random_state),
              "--model-output", "/mnt/job_stage/artifacts/iris_rf_model.pkl"],
        session=session,
        runtime_environment="2.6.0", 
    )
    context.log.info(f"Submitted ML Job: {ml_job.id}")
    ml_job.wait()  # Block until the job finishes

    # Fail the Dagster step if the ML Job didn't succeed
    if ml_job.status != "DONE":
        raise RuntimeError(f"ML Job failed: {ml_job.status}")
    session.close()
    return ml_job.id  # Pass job ID downstream so other ops can find artifacts

The submit_file call uploads the training script to a Snowflake stage and runs it on the compute pool. No Docker image needed for training code — Snowflake provides the runtime environment with scikit-learn, XGBoost, PyTorch, etc. pre-installed.

Step 3: Conditional Deployment Gate

This is what separates a toy pipeline from a production one. If the model doesn’t meet the accuracy threshold, it never gets register and deployed:

# Configurable threshold — override per-run in the Dagster UI launchpad
class EvalConfig(Config):
    accuracy_threshold: float = 0.90

@op
def evaluate_model(context: OpExecutionContext, config: EvalConfig,
                   job_id: str) -> str:
    session = create_session(cfg.DATABASE, cfg.SCHEMA)

    # Download metrics.json that the training script wrote to stage
    stage_prefix = _get_stage_prefix(job_id)
    tmp_dir = tempfile.mkdtemp()
    session.sql(f"GET {stage_prefix}/metrics.json file://{tmp_dir}").collect()
    session.close()

    # Parse the metrics
    with open(os.path.join(tmp_dir, "metrics.json")) as f:
        metrics = json.load(f)
    accuracy = metrics["accuracy"]
    context.log.info(f"Accuracy: {accuracy:.4f} (threshold: {config.accuracy_threshold})")

    # Gate: if accuracy is below threshold, raise an error to block downstream ops
    if accuracy < config.accuracy_threshold:
        raise RuntimeError(
            f"Model accuracy {accuracy:.4f} below threshold. Deployment skipped."
        )
    return job_id  # Pass through to registration step

In the Dagster UI, accuracy_threshold can be overridden per-run — useful for experimenting with stricter gates before promoting to production.

Step 4: Model Registration + Deployment

Once the model passes the gate, register it and spin up an inference service:

@op
def register_model(context: OpExecutionContext, job_id: str) -> str:
    session = create_session(cfg.DATABASE, cfg.SCHEMA)

    # Download the trained model pickle from the ML Job's stage artifacts
    # ... (GET from stage, pickle.load) ...
    # Register the model to Snowflake Model Registry with auto-versioning
    registry = Registry(session=session, database_name=cfg.DATABASE,
                       schema_name=cfg.SCHEMA)
    mv = registry.log_model(model=model, model_name="IRIS_RF_DAGSTER",
                            sample_input_data=sample_input)
    return mv.version_name  # e.g., "POLITE_GECKO_3"

@op
def deploy_inference_service(context: OpExecutionContext,
                             version_name: str) -> str:
    session = create_session(cfg.DATABASE, cfg.SCHEMA)

    # Retrieve the registered model version
    registry = Registry(session=session, database_name=cfg.DATABASE,
                       schema_name=cfg.SCHEMA)
    mv = registry.get_model("IRIS_RF_DAGSTER").version(version_name)

    # Deploy as a live REST inference endpoint on SPCS
    mv.create_service(
        service_name="IRIS_RF_SERVICE",
        service_compute_pool=cfg.COMPUTE_POOL,
        image_build_compute_pool=cfg.COMPUTE_POOL,
        ingress_enabled=True,  # Makes endpoint publicly accessible
    )
    return "IRIS_RF_SERVICE"

After this step, a live REST endpoint is serving predictions on Snowflake infrastructure.

Step 5: Putting It All Together

# Define the DAG — each op's output feeds into the next op's input
@job
def iris_training_job():
    job_id = ml_job_training_snowflake()            # Step 1: Train on SPCS
    evaluated_job_id = evaluate_model(job_id)       # Step 2: Gate on accuracy
    version_name = register_model(evaluated_job_id) # Step 3: Register to Model Registry
    deploy_inference_service(version_name)          # Step 4: Deploy inference endpoint

Running Dagster Server Locally

python3.11 -m venv .venv && source .venv/bin/activate
pip install -e ".[dev]"

DAGSTER_HOME=.dagster_home SNOWFLAKE_CONNECTION_NAME=<your_connection> \
  dagster dev -m definitions --port 3000

Open http://localhost:3000, navigate to the iris job, and hit Launch Run. Each op executes in sequence, with logs streaming from the ML Job.

Deploying Dagster to SPCS

Running locally is great for development. For production, Dagster itself can be deployed to Snowflake.

Dockerfile:

FROM python:3.11-slim
WORKDIR /app
COPY pyproject.toml definitions.py ./
COPY pipelines/ pipelines/
RUN pip install --no-cache-dir . snowflake-ml-python snowflake-snowpark-python
ENV DAGSTER_HOME=/dagster_home
RUN mkdir -p $DAGSTER_HOME
EXPOSE 3000
CMD ["dagster", "dev", "-m", "definitions", "-h", "0.0.0.0", "-p", "3000"]

Build and push:

snow spcs image-registry login --connection <your_connection>
docker buildx build --platform linux/amd64 --load -t <registry_url>/dagster-ml:latest .
docker push <registry_url>/dagster-ml:latest

Set up network access (required for the Dagster container to call back to Snowflake APIs):

-- Create a network rule allowing outbound traffic to Snowflake
CREATE OR REPLACE NETWORK RULE ML_DEMO.DAGSTER_ML.SNOWFLAKE_EGRESS_RULE
  MODE = EGRESS
  TYPE = HOST_PORT
  VALUE_LIST = ('<your_account>.snowflakecomputing.com:443',
                'ocsp.snowflakecomputing.com:80', 'ocsp.digicert.com:80');

-- Create an External Access Integration referencing the rule
CREATE OR REPLACE EXTERNAL ACCESS INTEGRATION DAGSTER_SNOWFLAKE_EAI
  ALLOWED_NETWORK_RULES = (ML_DEMO.DAGSTER_ML.SNOWFLAKE_EGRESS_RULE)
  ENABLED = TRUE;

Deploy the service:

CREATE SERVICE ML_DEMO.DAGSTER_ML.DAGSTER_SERVICE
  IN COMPUTE POOL COCO_ML_COMPUTE_POOL
  FROM SPECIFICATION $$
spec:
  containers:
    - name: dagster
      image: /ml_demo/dagster_ml/dagster_repo/dagster-ml:latest
      env:
        DAGSTER_HOME: /dagster_home
        DAGSTER_SPCS_MODE: "true"
  endpoints:
    - name: dagster-ui
      port: 3000
      public: true
  volumes:
    - name: dagster-storage
      source: local
$$
EXTERNAL_ACCESS_INTEGRATIONS = (DAGSTER_SNOWFLAKE_EAI)
MIN_INSTANCES = 1
MAX_INSTANCES = 1;

The EXTERNAL_ACCESS_INTEGRATIONS is critical — without it, the container cannot make outbound network calls to Snowflake APIs for submitting ML Jobs and registering models.

Once deployed, the Dagster UI is accessible at a .snowflakecomputing.app URL, authenticated via Snowflake SSO.

Example Dagster ML Pipeline deployed on SPCS

Fast Iteration: Code Sync Without Docker Rebuilds

The Docker image includes only Dagster and its dependencies. The pipeline source code — including training scripts, ops, and configs — can be deployed to a Snowflake stage through your CI/CD workflow, or run locally, using the following file sync command:

SNOWFLAKE_CONNECTION_NAME=<your_connection> python scripts/sync_code.py

This uploads all .py files to @CODE_STAGE. The submit_file calls always upload the latest training scripts from the local (or staged) code, so ML Jobs automatically run the newest version. Only rebuild the Docker image when pip dependencies change.

Key Takeaways

1. Separation of concerns: Dagster handles orchestration, Snowflake handles compute. Training scripts are pure Python with no infrastructure code.
2. Conditional deployment gates: Don’t deploy bad models. The evaluation op is a configurable quality gate that blocks downstream steps.
3. No Kubernetes: Snowflake ML Jobs provide managed containers with ML libraries pre-installed. Submit a .py file and get results.
4. Fast iteration: Sync code to stage in seconds. Only rebuild Docker when dependencies change.
5. Production-ready: Model versioning via Snowflake Model Registry, live inference endpoints via SPCS, all observable in the Dagster UI.
6. The full source code is available on GitHub.

References

• Snowflake ML Jobs — Run Python scripts on managed SPCS compute pools
• Snowflake Model Registry — Version, manage, and deploy ML models
• Snowpark Container Services (SPCS) — Run containerized workloads on Snowflake
• Model Inference Services — Deploy models as REST endpoints
• Dagster Documentation — Orchestration framework for data and ML pipelines