From 1ea72bbdd7be5e76286e75a5e4f540ee16b25617 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 24 Oct 2026 18:38:32 +0000 Subject: [PATCH] Add Databricks, SageMaker and Snowflake job templates for kumo-engine Teams want to run kumo-engine fits and scoring on the scheduler they already have, but there is no runnable example of wiring the engine, MLOps settings and a registry into a scheduled job. examples/mlops holds two plain Python job scripts (fit_job.py, score_job.py) and a shared jobs.py. They take CLI flags (--artifact-root, --source, --registry, --model-name, --model-path, --predictions-table, ...), build kumo.MLOps from them and run inside `with kumo.mlops(cfg)`, so snapshots, training tables and models are stored under the artifact root (models//). With a registry the fit registers the model and moves `champion`; scoring loads it, or --model-path from the root. Platform wiring: a Databricks Asset Bundle, a SageMaker Pipeline and (experimental) Snowflake ML Jobs + Tasks, each passing the flags. SnowflakeConnector.session() exposes the connector's Snowpark session for SnowflakeTracking. Verified with test_job_templates.py (local fit then score) and test_snowflake_connector.py. Co-Authored-By: Claude Opus 5.5 --- examples/mlops/README.md | 57 ++++++ examples/mlops/databricks/README.md | 87 +++++++++ examples/mlops/databricks/databricks.yml | 145 +++++++++++++++ examples/mlops/fit_job.py | 36 ++++ examples/mlops/jobs.py | 155 ++++++++++++++++ examples/mlops/sagemaker/README.md | 81 +++++++++ examples/mlops/sagemaker/pipeline.py | 100 ++++++++++ examples/mlops/score_job.py | 40 ++++ examples/mlops/snowflake/README.md | 88 +++++++++ examples/mlops/snowflake/register_jobs.py | 106 +++++++++++ .../kumo/engine/connector/snowflake.py | 8 +- .../kumo-engine/test/test_job_templates.py | 172 ++++++++++++++++++ .../test/test_snowflake_connector.py | 15 ++ 13 files changed, 1089 insertions(+), 1 deletion(-) create mode 100644 examples/mlops/README.md create mode 100644 examples/mlops/databricks/README.md create mode 100644 examples/mlops/databricks/databricks.yml create mode 100644 examples/mlops/fit_job.py create mode 100644 examples/mlops/jobs.py create mode 100644 examples/mlops/sagemaker/README.md create mode 100644 examples/mlops/sagemaker/pipeline.py create mode 100644 examples/mlops/score_job.py create mode 100644 examples/mlops/snowflake/README.md create mode 100644 examples/mlops/snowflake/register_jobs.py create mode 100644 packages/kumo-engine/test/test_job_templates.py diff --git a/examples/mlops/README.md b/examples/mlops/README.md new file mode 100644 index 000000000..1079eccf6 --- /dev/null +++ b/examples/mlops/README.md @@ -0,0 +1,57 @@ +# Scheduled kumo-engine jobs + +Two plain Python job scripts that run the normal `kumo.engine` API inside +`with kumo.mlops(kumo.MLOps(...)):`, and the wiring that runs them on a +schedule. Both take their settings as command-line flags. + +- `fit_job.py` builds a snapshot, generates the users-churn training table and + trains the model. Inside the block each artifact is stored under the + artifact root (`snapshots/…`, `training_tables//`, `models//`). + With a registry, the model is registered and the `champion` alias moves to + it. It prints `model: models/`. +- `score_job.py` loads `champion` (or `--model-path`), scores a fresh snapshot + (or the stored one `--snapshot` names) and writes `--predictions-table` when + given. It passes a 0.5 threshold unless the model carries scoring defaults. +- `jobs.py` holds what both share: the flags, the graph (users, items, orders), + the connectors and the `kumo.MLOps` settings. Replace it with your own module. + +| Directory | Platform | Registry | Schedule | +|---|---|---|---| +| [databricks](databricks/README.md) | Databricks Jobs (Asset Bundle) | Unity Catalog through MLflow | train on demand, score daily | +| [sagemaker](sagemaker/README.md) | SageMaker Pipelines | SageMaker managed MLflow | fit step, then score step | +| [snowflake](snowflake/README.md) (experimental, preview APIs) | Snowflake ML Jobs + Tasks | Snowflake Model Registry, Snowflake Experiment Tracking | train weekly, score daily | + +## Flags + +| Flag | Meaning | +|---|---| +| `--artifact-root` | `s3://`, `gs://`, `file://` or a local or `/Volumes` path; required | +| `--source` | `databricks`, `snowflake`, or a directory holding `users`, `items` and `orders` Parquet files | +| `--catalog`, `--schema`, `--warehouse` | Databricks catalog, schema and SQL warehouse id, or Snowflake database, schema and warehouse | +| `--registry` | `mlflow` (`MLOps(tracking=MlflowTracking(), registry=ModelRegistryLocation(...))`) or `snowflake` (`MLOps(tracking=SnowflakeTracking(...), registry=ModelOutputConfig(...))`); omit for no registry | +| `--model-name` | registered model name, with `--registry` | +| `--model-path` | score only, without a registry: `models/` under the artifact root | +| `--snapshot` | score only: a stored snapshot under the artifact root | +| `--predictions-table` | score only: a table of the source connector to write | +| `--query`, `--num-experiments` | fit only: the predictive query and the number of experiments | +| `--scratch-dir` | local disk for snapshots, tables and bundles | + +## Run locally + +With `users.parquet`, `items.parquet` and `orders.parquet` in a directory: + +```bash +cd examples/mlops +python fit_job.py --source --artifact-root +python score_job.py --source --artifact-root \ + --model-path models/ +``` + +The query needs several 30-day windows of order history; for a small sample +pass a shorter `--query` and `--num-experiments 1`. + +## What every scheduler sees + +The scripts are ordinary Python: they exit 0 on success and non-zero with a +traceback on failure, and log to stderr. Databricks Jobs, SageMaker Pipelines +and Snowflake Tasks keep their own run ids, status and logs. diff --git a/examples/mlops/databricks/README.md b/examples/mlops/databricks/README.md new file mode 100644 index 000000000..0df30c2f9 --- /dev/null +++ b/examples/mlops/databricks/README.md @@ -0,0 +1,87 @@ +# Databricks Jobs + +`databricks.yml` is an Asset Bundle with two jobs on a single-node CPU job +cluster. Each task is a `spark_python_task` running a job script with its +flags as `parameters`; the cluster's `spark_env_vars` hold only MLflow's own +settings. + +- **`kumo_train`** runs `fit_job.py` on the `users`, `items` and `orders` + tables in `.` inside `kumo.mlops(...)` with the artifact + root `/Volumes////kumo`: the snapshot, training + table and model are stored under it. `--registry mlflow` logs the fit to + `MLFLOW_EXPERIMENT_NAME`, registers the model as a new version of + `..users_churn` and moves the `champion` alias to it. +- **`kumo_score`** runs `score_job.py` daily at 03:00 UTC (created paused). + It loads `models:/..users_churn@champion`, builds a fresh + snapshot and writes `users_churn_predictions` in the connector's schema. Add + `--snapshot snapshots//` to score a stored snapshot + instead. + +## Before deploying + +1. Put the kumo-engine wheel and a `requirements.txt` for its dependencies + (kumo-core, kumo-api, kumo-pql, diskgraph, kumo-databricks, mlflow 3) on a + volume and set `kumo_engine_wheel` and `requirements`. Start the + requirements file with `--prefer-binary`: DBR 16.4's setuptools cannot + build the SQLAlchemy sdist MLflow pulls. Each task installs the + requirements, then the wheel. + Use a standard (non-ML) DBR runtime: the ML runtime's preinstalled + torch, torchvision and pyOpenSSL conflict with kumo-engine's pins. + Databricks GPU nodes require a GPU ML runtime, so use a CPU node + type with local NVMe disks (for example m5d.xlarge). +2. Replace the `<...>` placeholders in `databricks.yml` (or pass + `--var catalog=...`). The parent folder of the MLflow experiment + path must already exist in the workspace. +3. Adjust the `dev` target (workspace `root_path`) or add your own, then + `databricks bundle validate` and `databricks bundle deploy -t `. The bundle syncs this + whole `examples/mlops` directory (`sync.paths: [..]`). + +## Privileges for the run-as principal + +- `USE CATALOG`, `USE SCHEMA` and `SELECT` on the three tables; `CREATE TABLE` + and `MODIFY` on the schema for predictions. +- `READ VOLUME`/`WRITE VOLUME` on ``, and `CREATE VOLUME` on the + schema (or a pre-created `.._kumo_staging` volume) for + staged table writes. +- `CAN USE` on the SQL warehouse; edit permission on the MLflow experiment. +- `CREATE MODEL` on the schema and `EXECUTE` on the model; moving an alias + needs rights on the model itself. + +## Notes + +- **Script, not `python_wheel_task`.** Databricks overrides a custom + `PYTHONPATH` at driver startup, so a wheel task cannot import `jobs.py`; the + scripts put their own directory on `sys.path` (from `sys.argv[0]`: a + `spark_python_task` runs the file without `__file__`). +- **Dedicated access mode** (`SINGLE_USER`): standard compute runs as a + low-privilege user and has no GPUs. `--scratch-dir` is + `/local_disk0/kumo`. +- **MLflow.** `MLFLOW_TRACKING_URI=databricks` and + `MLFLOW_REGISTRY_URI=databricks-uc` are set explicitly; + `MLFLOW_EXPERIMENT_NAME` is required in jobs. +- **One registry format.** A `DatabricksConnector` publish and an MLflow + pyfunc model are different formats; use one per registered model. +- **Job clusters.** Both jobs repeat the cluster block; YAML anchors would + remove the copy but are untested with `databricks bundle validate`. + +## Verified + +An earlier revision of this bundle (configured through environment variables, +same engine calls) was deployed and both jobs were run on a single-node cluster +against synthetic `users`, `items` and `orders` tables: the fit registered a +new `users_churn` version and moved `champion`, and the score job loaded it and +wrote the predictions table. This revision is checked locally by +`test_job_templates.py`. + +## References + +- Task types and parameters: + https://docs.databricks.com/aws/en/dev-tools/bundles/job-task-types +- Custom `PYTHONPATH` overridden: + https://kb.databricks.com/en_US/clusters/cannot-set-a-custom-pythonpath +- Dynamic value references: + https://docs.databricks.com/aws/en/jobs/dynamic-value-references +- Standard compute limitations: + https://docs.databricks.com/aws/en/compute/standard-limitations +- Unity Catalog models: + https://docs.databricks.com/aws/en/machine-learning/manage-model-lifecycle/ diff --git a/examples/mlops/databricks/databricks.yml b/examples/mlops/databricks/databricks.yml new file mode 100644 index 000000000..df4dd4039 --- /dev/null +++ b/examples/mlops/databricks/databricks.yml @@ -0,0 +1,145 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Train and daily score jobs for kumo-engine; see README.md. + +bundle: + name: kumo-engine-users-churn + +sync: + paths: + - .. + +variables: + catalog: + description: Unity Catalog catalog holding the users, items and orders tables + default: + schema: + description: schema holding the tables, the registered model and the predictions + default: + volume: + description: Unity Catalog volume for exported models and snapshots + default: + warehouse_id: + description: SQL warehouse that reads the source tables and writes predictions + default: + kumo_engine_wheel: + description: kumo-engine wheel on a volume or workspace path + default: /Volumes////wheels/.whl + requirements: + description: requirements.txt for kumo-engine's dependencies, starting with --prefer-binary + default: /Volumes////wheels/requirements.txt + spark_version: + description: a standard (non-ML) runtime, for example 16.4.x-scala2.12 + default: + node_type_id: + description: a CPU node type with local NVMe disks, for example m5d.xlarge + default: + experiment: + description: MLflow experiment (workspace path) for train and score runs + default: /Shared/kumo-engine/users-churn + +resources: + jobs: + kumo_train: + name: kumo-engine-train + max_concurrent_runs: 1 + job_clusters: + - job_cluster_key: kumo + new_cluster: + spark_version: ${var.spark_version} + node_type_id: ${var.node_type_id} + num_workers: 0 + data_security_mode: SINGLE_USER + spark_conf: + spark.databricks.cluster.profile: singleNode + spark.master: local[*] + custom_tags: + ResourceClass: SingleNode + spark_env_vars: + MLFLOW_TRACKING_URI: databricks + MLFLOW_REGISTRY_URI: databricks-uc + MLFLOW_EXPERIMENT_NAME: ${var.experiment} + tasks: + - task_key: fit + job_cluster_key: kumo + libraries: + - requirements: ${var.requirements} + - whl: ${var.kumo_engine_wheel} + spark_python_task: + python_file: ${workspace.file_path}/fit_job.py + parameters: + - --artifact-root + - /Volumes/${var.catalog}/${var.schema}/${var.volume}/kumo + - --source + - databricks + - --catalog + - ${var.catalog} + - --schema + - ${var.schema} + - --warehouse + - ${var.warehouse_id} + - --registry + - mlflow + - --model-name + - ${var.catalog}.${var.schema}.users_churn + - --scratch-dir + - /local_disk0/kumo + + kumo_score: + name: kumo-engine-score + max_concurrent_runs: 1 + schedule: + quartz_cron_expression: 0 0 3 * * ? + timezone_id: UTC + pause_status: PAUSED + job_clusters: + - job_cluster_key: kumo + new_cluster: + spark_version: ${var.spark_version} + node_type_id: ${var.node_type_id} + num_workers: 0 + data_security_mode: SINGLE_USER + spark_conf: + spark.databricks.cluster.profile: singleNode + spark.master: local[*] + custom_tags: + ResourceClass: SingleNode + spark_env_vars: + MLFLOW_TRACKING_URI: databricks + MLFLOW_REGISTRY_URI: databricks-uc + MLFLOW_EXPERIMENT_NAME: ${var.experiment} + tasks: + - task_key: score + job_cluster_key: kumo + libraries: + - requirements: ${var.requirements} + - whl: ${var.kumo_engine_wheel} + spark_python_task: + python_file: ${workspace.file_path}/score_job.py + parameters: + - --artifact-root + - /Volumes/${var.catalog}/${var.schema}/${var.volume}/kumo + - --source + - databricks + - --catalog + - ${var.catalog} + - --schema + - ${var.schema} + - --warehouse + - ${var.warehouse_id} + - --registry + - mlflow + - --model-name + - ${var.catalog}.${var.schema}.users_churn + - --scratch-dir + - /local_disk0/kumo + - --predictions-table + - users_churn_predictions + +targets: + dev: + default: true + workspace: + root_path: /Workspace/Shared/kumo-engine/.bundle/${bundle.name}/${bundle.target} diff --git a/examples/mlops/fit_job.py b/examples/mlops/fit_job.py new file mode 100644 index 000000000..f56f7b7ac --- /dev/null +++ b/examples/mlops/fit_job.py @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Train the users-churn model: ``python fit_job.py --artifact-root ROOT +--source SOURCE [...]``; see README.md.""" + +import os +import sys +from typing import Optional, Sequence + +sys.path.insert(0, os.path.dirname(os.path.abspath( + globals().get('__file__') or sys.argv[0]))) + +import jobs # noqa: E402 + +import kumo.engine as kumo # noqa: E402 + + +def main(argv: Optional[Sequence[str]] = None) -> kumo.TrainingResult: + args = jobs.parse_args(argv, score=False) + connector = jobs.source(args) + graph = jobs.build_graph(connector) + with kumo.mlops(jobs.mlops(args, connector)): + snapshot = graph.snapshot() + query = kumo.PredictiveQuery(graph=graph, query=args.query) + table = query.generate_training_table(snapshot) + plan = query.suggest_model_plan(snapshot) + if args.num_experiments: + plan.training_job.num_experiments = args.num_experiments + result = kumo.Trainer(plan).fit(table, snapshot=snapshot) + print(f'model: models/{result.id}') + return result + + +if __name__ == '__main__': + main() diff --git a/examples/mlops/jobs.py b/examples/mlops/jobs.py new file mode 100644 index 000000000..09af9d596 --- /dev/null +++ b/examples/mlops/jobs.py @@ -0,0 +1,155 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The flags, graph, connectors and MLOps settings both job scripts share.""" + +import argparse +from typing import Optional, Sequence + +from kumoapi.jobs import MetadataField, PredictionArtifactType +from kumoapi.typing import Stype + +import kumo.engine as kumo +from kumo.engine import ( + Edge, + Graph, + ModelOutputConfig, + ModelRegistryLocation, + PredictionOutputConfig, + Table, + Trainer, +) +from kumo.engine.connector import ( + Connector, + DatabricksConnector, + FileConnector, + S3Connector, + SnowflakeConnector, + connector_for_root, +) + +CHAMPION = 'champion' +# Needs several 30-day windows of order history; shorten it for small data. +QUERY = 'PREDICT COUNT(orders.*, 0, 30, days) > 0 FOR EACH users.USER_ID' + + +def parse_args(argv: Optional[Sequence[str]], *, + score: bool) -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument( + '--artifact-root', required=True, + help='s3://, gs://, file:// or a local or /Volumes path') + parser.add_argument( + '--source', required=True, + help='databricks, snowflake, or a directory holding users, items ' + 'and orders parquet files') + parser.add_argument('--catalog', + help='Databricks catalog or Snowflake database') + parser.add_argument('--schema') + parser.add_argument( + '--warehouse', help='Databricks SQL warehouse id or Snowflake ' + 'warehouse') + parser.add_argument('--registry', choices=('mlflow', 'snowflake')) + parser.add_argument('--model-name') + parser.add_argument('--scratch-dir') + if score: + parser.add_argument( + '--model-path', + help='models/ under the artifact root, without a ' + 'registry') + parser.add_argument('--snapshot', + help='a stored snapshot under the artifact root') + parser.add_argument('--predictions-table') + else: + parser.add_argument('--query', default=QUERY) + parser.add_argument('--num-experiments', type=int) + args = parser.parse_args(argv) + if args.registry and not args.model_name: + parser.error('--registry needs --model-name') + if score and not args.registry and not args.model_path: + parser.error('without --registry, pass --model-path') + if args.scratch_dir: + kumo.configure(scratch_dir=args.scratch_dir) + return args + + +def source(args: argparse.Namespace) -> Connector: + if args.source == 'databricks': + return DatabricksConnector.from_notebook(catalog=args.catalog, + schema_name=args.schema, + warehouse_id=args.warehouse) + if args.source == 'snowflake': + return SnowflakeConnector.from_spcs_env(warehouse=args.warehouse, + database=args.catalog, + schema_name=args.schema) + return connector_for_root(args.source) + + +def mlops(args: argparse.Namespace, connector: Connector) -> kumo.MLOps: + if args.registry == 'mlflow': + from kumo.engine.integrations.mlflow import MlflowTracking + + return kumo.MLOps( + args.artifact_root, tracking=MlflowTracking(), + registry=ModelRegistryLocation(args.model_name, alias=CHAMPION)) + if args.registry == 'snowflake': + from kumo.engine.integrations.snowflake import SnowflakeTracking + + assert isinstance(connector, SnowflakeConnector) + return kumo.MLOps( + args.artifact_root, + tracking=SnowflakeTracking(connector.session(), + experiment_name=args.model_name), + registry=ModelOutputConfig( + connector, location=ModelRegistryLocation( + args.model_name, alias=CHAMPION))) + return kumo.MLOps(args.artifact_root) + + +def build_graph(connector: Connector) -> Graph: + + def table(alias: str, **kwargs: str) -> Table: + name = { + S3Connector: f'{alias}.parquet', + FileConnector: f'{alias}.parquet', + SnowflakeConnector: alias.upper(), + }.get(type(connector), alias) + return Table.from_source_table(connector[name], **kwargs) + + orders = table('orders', primary_key='ORDER_ID', + time_column='ORDER_TIME').with_column( + 'USER_ID', + stype=Stype.ID).with_column('ITEM_ID', stype=Stype.ID) + return Graph( + tables={ + 'users': table('users', primary_key='USER_ID'), + 'items': table('items', primary_key='ITEM_ID'), + 'orders': orders, + }, edges=[ + Edge(src_table='orders', fkey='USER_ID', dst_table='users'), + Edge(src_table='orders', fkey='ITEM_ID', dst_table='items'), + ]) + + +def load_model(args: argparse.Namespace, connector: Connector) -> Trainer: + if args.registry == 'mlflow': + from kumo.engine.integrations.mlflow import load_model + + return load_model(f'models:/{args.model_name}@{CHAMPION}') + if args.registry == 'snowflake': + return Trainer.load_from_connector( + connector, + location=ModelRegistryLocation(args.model_name, alias=CHAMPION)) + return Trainer.load_from_connector(connector_for_root(args.artifact_root), + path=args.model_path) + + +def prediction_output( + args: argparse.Namespace, + connector: Connector) -> Optional[PredictionOutputConfig]: + if not args.predictions_table: + return None + return PredictionOutputConfig( + connector, artifact_type=PredictionArtifactType.PREDICTIONS, + table_name=args.predictions_table, + extra_fields=[MetadataField.ANCHOR_TIMESTAMP]) diff --git a/examples/mlops/sagemaker/README.md b/examples/mlops/sagemaker/README.md new file mode 100644 index 000000000..d54412398 --- /dev/null +++ b/examples/mlops/sagemaker/README.md @@ -0,0 +1,81 @@ +# SageMaker Pipelines + +`pipeline.py` defines and upserts a two-step pipeline: + +1. **`Fit`** (training job) runs `python fit_job.py` on `users.parquet`, + `items.parquet` and `orders.parquet` under `--data-root`, inside + `kumo.mlops(...)` with `--artifact-root`: the snapshot, training table and + model are stored under it. `--registry mlflow` logs the fit to the managed + MLflow tracking server, registers the model as a new version of + `users-churn` and moves the `champion` alias to it. +2. **`Predict`** (processing job, after `Fit`) runs `python score_job.py`, + loading `models:/users-churn@champion` and scoring a fresh snapshot. + +To score on a schedule without retraining, put `Predict` in its own pipeline +with a `PipelineSchedule` trigger (`pipeline.put_triggers`). + +## Before running + +1. Build an ECR image with `kumo-engine[mlflow,s3]`, its dependency wheels + and `sagemaker-mlflow` (plus the CUDA stack for GPU instance types). +2. Upload the code channel: `aws s3 sync examples/mlops s3:///kumo/code/`. +3. With the SageMaker Python SDK v2 (`pip install 'sagemaker<3'`): + ```bash + python examples/mlops/sagemaker/pipeline.py --role \ + --image .dkr.ecr..amazonaws.com/: \ + --code s3:///kumo/code/ --data-root s3:///kumo/data/ \ + --artifact-root s3:///kumo/artifacts/ \ + --tracking-arn \ + --instance-type --start + ``` + SDK v3 replaced `Estimator` with `ModelTrainer` and moved the workflow + classes (`sagemaker.mlops.workflow.*`); this script is not ported. + +## Privileges + +- **Caller:** `sagemaker:CreatePipeline`, `UpdatePipeline`, + `StartPipelineExecution`, and `iam:PassRole` on ``. +- **Execution role** (`--role` also runs the steps): `sagemaker:CreateTrainingJob`, + `DescribeTrainingJob`, `StopTrainingJob`, `CreateProcessingJob`, + `DescribeProcessingJob`, `StopProcessingJob` and `AddTags`; `iam:PassRole` + on itself; S3 read on the code and data prefixes and read/write on + the artifact prefix; ECR pull; CloudWatch Logs; the `sagemaker-mlflow:*` + actions for tracking and the registry (`CreateRun`, `LogBatch`, + `CreateRegisteredModel`, `CreateModelVersion`, `SetRegisteredModelAlias`, + `GetModelVersionByAlias`, ...). + +## Notes + +- **Flags.** Both steps pass the job scripts' flags as plain strings + (`container_arguments` for the training step, `arguments` for the + processing step); only MLflow's own settings are environment variables. +- **MLflow.** `MLFLOW_TRACKING_URI` is the tracking server or MLflow App ARN. + Managed MLflow servers run MLflow 3.0; the engine needs a newer client + (3.11 or later), which works against them, so do not pin the client to the + server's version. Aliases are letters, digits and underscores, so + SageMaker's reserved `SageMakerLifecycle-` aliases cannot be set. +- **Scratch.** The fit step uses `/tmp/kumo` (on NVMe instance types + SageMaker puts `/tmp` on the instance store; on EBS-only types raise + `volume_size`); the score step uses `/opt/ml/processing/scratch`. +- **Retries.** Step `retry_policies` key on SageMaker failure reasons, never + on the exit code; a non-zero exit is not retried. + +## Verified + +An earlier revision of this pipeline (configured through environment +variables, same engine calls) was upserted and run on `ml.m5.xlarge` with a CPU image, the +tracking server ARN as `MLFLOW_TRACKING_URI` and 180 days of synthetic orders +(the 30-day query needs about 150 days of history): both steps exit 0, the +fit registers a new `users-churn` version and moves `champion`, and the score +step loads it. An MLflow App ARN as tracking URI is unverified. This revision +is checked locally by `test_job_templates.py`. + +## References + +- Pipeline steps: https://docs.aws.amazon.com/sagemaker/latest/dg/build-and-manage-steps-types.html +- Retry policies: https://docs.aws.amazon.com/sagemaker/latest/dg/pipelines-retry-policy.html +- Instance store and `/tmp`: + https://docs.aws.amazon.com/sagemaker/latest/dg/model-train-storage-tips-considerations.html +- Managed MLflow and IAM: + https://docs.aws.amazon.com/sagemaker/latest/dg/mlflow-track-experiments.html, + https://docs.aws.amazon.com/sagemaker/latest/dg/mlflow-create-tracking-server-iam.html diff --git a/examples/mlops/sagemaker/pipeline.py b/examples/mlops/sagemaker/pipeline.py new file mode 100644 index 000000000..88c04abc9 --- /dev/null +++ b/examples/mlops/sagemaker/pipeline.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Upsert a SageMaker Pipeline (SDK v2) that fits and scores with kumo-engine; +see README.md.""" + +import argparse +from typing import Any, Optional, Sequence + +TRAIN_CODE = '/opt/ml/input/data/code' +PROCESS_CODE = '/opt/ml/processing/input/code' + + +def build_pipeline( + *, + name: str, + role: str, + image: str, + code: str, + data_root: str, + artifact_root: str, + tracking_arn: str, + experiment: str, + instance_type: str, +) -> Any: + from sagemaker.estimator import Estimator + from sagemaker.inputs import TrainingInput + from sagemaker.processing import ProcessingInput, Processor + from sagemaker.workflow.pipeline import Pipeline + from sagemaker.workflow.pipeline_context import PipelineSession + from sagemaker.workflow.steps import ProcessingStep, TrainingStep + + session = PipelineSession() + env = { + 'MLFLOW_TRACKING_URI': tracking_arn, + 'MLFLOW_EXPERIMENT_NAME': experiment, + } + flags = [ + '--artifact-root', artifact_root, '--source', data_root, + '--registry', 'mlflow', '--model-name', 'users-churn' + ] + + estimator = Estimator( + image_uri=image, + role=role, + instance_count=1, + instance_type=instance_type, + container_entry_point=['python', f'{TRAIN_CODE}/fit_job.py'], + container_arguments=[*flags, '--scratch-dir', '/tmp/kumo'], + environment=env, + output_path=f'{artifact_root.rstrip("/")}/sagemaker-output', + sagemaker_session=session, + ) + fit = TrainingStep( + name='Fit', + step_args=estimator.fit(inputs={'code': TrainingInput(code)}), + ) + + processor = Processor( + role=role, + image_uri=image, + instance_count=1, + instance_type=instance_type, + entrypoint=['python', f'{PROCESS_CODE}/score_job.py'], + env=env, + sagemaker_session=session, + ) + predict = ProcessingStep( + name='Predict', + step_args=processor.run( + inputs=[ProcessingInput(source=code, destination=PROCESS_CODE)], + arguments=[ + *flags, '--scratch-dir', '/opt/ml/processing/scratch' + ], + ), + depends_on=[fit], + ) + return Pipeline(name=name, steps=[fit, predict], + sagemaker_session=session) + + +def main(argv: Optional[Sequence[str]] = None) -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--name', default='kumo-engine-users-churn') + for flag in ('--role', '--image', '--code', '--data-root', + '--artifact-root', '--tracking-arn'): + parser.add_argument(flag, required=True) + parser.add_argument('--experiment', default='kumo-engine-users-churn') + parser.add_argument('--instance-type', required=True) + parser.add_argument('--start', action='store_true') + args = vars(parser.parse_args(argv)) + start = args.pop('start') + pipeline = build_pipeline(**args) + pipeline.upsert(role_arn=args['role']) + if start: + print(pipeline.start().arn) + + +if __name__ == '__main__': + main() diff --git a/examples/mlops/score_job.py b/examples/mlops/score_job.py new file mode 100644 index 000000000..859a78ec5 --- /dev/null +++ b/examples/mlops/score_job.py @@ -0,0 +1,40 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Score with the champion model, or the one ``--model-path`` names: +``python score_job.py --artifact-root ROOT --source SOURCE [...]``; see +README.md.""" + +import os +import sys +from typing import Optional, Sequence + +sys.path.insert(0, os.path.dirname(os.path.abspath( + globals().get('__file__') or sys.argv[0]))) + +import jobs # noqa: E402 + +import kumo.engine as kumo # noqa: E402 +from kumo.engine.connector import connector_for_root # noqa: E402 + + +def main(argv: Optional[Sequence[str]] = None) -> kumo.PredictionResult: + args = jobs.parse_args(argv, score=True) + connector = jobs.source(args) + graph = jobs.build_graph(connector) + with kumo.mlops(jobs.mlops(args, connector)): + trainer = jobs.load_model(args, connector) + snapshot = (kumo.GraphSnapshot.load_from_connector( + connector_for_root(args.artifact_root), path=args.snapshot, + graph=graph) if args.snapshot else graph.snapshot()) + defaults = trainer.result.scoring_defaults + result = trainer.predict( + snapshot, output_config=jobs.prediction_output(args, connector), + binary_classification_threshold=None + if defaults.binary_classification_threshold is not None else 0.5) + print(f'predictions: {result.num_rows} rows') + return result + + +if __name__ == '__main__': + main() diff --git a/examples/mlops/snowflake/README.md b/examples/mlops/snowflake/README.md new file mode 100644 index 000000000..e2a3de931 --- /dev/null +++ b/examples/mlops/snowflake/README.md @@ -0,0 +1,88 @@ +# Snowflake ML Jobs and Tasks (experimental) + +**Experimental.** This template is built on preview Snowflake APIs +(`MLJobDefinition` is in preview and the `SYSTEM$EXECUTE_ML_JOB` statement it +generates is undocumented) and has not been run end to end; see +[Unverified](#unverified). They may change without notice. + +`register_jobs.py` registers two ML Job definitions whose payload is this +`examples/mlops` directory, and one Task per definition that runs it. + +- **`KUMO_TRAIN`** (`KUMO_TRAIN_TASK`, Sundays 02:00 UTC) runs `fit_job.py` + on `USERS`, `ITEMS` and `ORDERS` in `.`, publishes a new + version of model `USERS_CHURN` in that schema's Model Registry and moves the + `champion` alias to it (`--registry snowflake`, which builds + `kumo.MLOps(registry=ModelOutputConfig(connector, location=...))`). The + snapshot, training table and model are stored under `--artifact-root`. The + fit is recorded in + Snowflake Experiment Tracking through the connector's own Snowpark session + (`SnowflakeConnector.session()`, the container's OAuth token). +- **`KUMO_SCORE`** (`KUMO_SCORE_TASK`, daily 03:00 UTC) runs `score_job.py`, + loading the version `champion` points at and writing + `USERS_CHURN_PREDICTIONS`. + +To review a model before it is scored, drop the alias from the registry +location in `jobs.mlops` and move `champion` yourself with +`kumo.engine.set_model_alias`. +If moving the alias fails, the fit still publishes, warns, and scoring keeps +the previous `champion`. + +Each Task runs its job synchronously with the statement +`MLJobDefinition.to_sql(job_args=, use_async=False)` generates. `USER_TASK_TIMEOUT_MS` is 24 h for train and 6 h for score +(default 1 h); a job that outlives it, or the warehouse's +`STATEMENT_TIMEOUT_IN_SECONDS`, is cancelled. + +## Before running + +1. Make `kumo-engine[snowflake]` and its dependency wheels installable in the + job container (package index via an external access integration, or a + custom image). +2. Create the GPU compute pool `` and the stage ``. The + artifact root must outlive the job container: an `s3://` or `gs://` URI + the external access integration reaches. +3. Run: + ```bash + python examples/mlops/snowflake/register_jobs.py --connection \ + --compute-pool --stage --database \ + --schema --warehouse \ + --artifact-root s3:///kumo/ \ + --requirement 'kumo-engine[snowflake]' --external-access-integration + ``` +4. Tasks are created suspended: `ALTER TASK RESUME`, or + `EXECUTE TASK ` to run one now. + +## Privileges for the role that owns the tasks + +- `USAGE` on the database, schema, warehouse and compute pool; + `CREATE SERVICE`, `CREATE TASK`, `CREATE TABLE` and `CREATE MODEL` on the + schema; `EXECUTE TASK` on the account. +- Read and write on the stage; `USAGE` on the external access integration (or + `SNOWFLAKE.PYPI_REPOSITORY_USER`). +- `SELECT` on the source tables, `INSERT` for predictions, and ownership of + the model to move its alias. + +## Notes + +- The ML Jobs launcher runs the job script from the payload directory, so + `jobs.py` imports; a non-zero exit fails the job (its `mljob_result` is + still `None`, so judge a run by its status). +- The Task's synchronous `CALL` holds the task warehouse while the job runs; + a serverless task avoids that. +- `--scratch-dir` is `/tmp/kumo`; add a volume through `spec_overrides` if the + graph needs more disk. Tasks retry only via `TASK_AUTO_RETRY_ATTEMPTS`. + +## Unverified + +Nothing has been registered or run. `SYSTEM$EXECUTE_ML_JOB` (what `to_sql` +generates) is not a documented SQL API and `MLJobDefinition` is in preview. +Unverified: reaching the reader's presigned URLs without an external access +integration, a failed job failing the Task, and creating experiments from a job's OAuth session. + +## References + +- `MLJobDefinition`: + https://docs.snowflake.com/en/developer-guide/snowpark-ml/reference/latest/api/jobs/snowflake.ml.jobs.MLJobDefinition +- ML Jobs: https://docs.snowflake.com/en/developer-guide/snowflake-ml/ml-jobs/overview +- `CREATE TASK`: https://docs.snowflake.com/en/sql-reference/sql/create-task +- Model versions: + https://docs.snowflake.com/en/developer-guide/snowflake-ml/model-registry/model-management diff --git a/examples/mlops/snowflake/register_jobs.py b/examples/mlops/snowflake/register_jobs.py new file mode 100644 index 000000000..5ea485e33 --- /dev/null +++ b/examples/mlops/snowflake/register_jobs.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Experimental: register fit_job.py/score_job.py as Snowflake ML Jobs with +one Task each (snowflake-ml-python 2.6, preview API); see README.md.""" + +import argparse +from pathlib import Path +from typing import Any, Dict, List, Optional, Sequence + +PAYLOAD = Path(__file__).resolve().parent.parent +JOBS = { + 'train': { + 'entrypoint': 'fit_job.py', + 'schedule': 'USING CRON 0 2 * * 0 UTC', + 'timeout_ms': 24 * 3600 * 1000, + 'args': [], + }, + 'score': { + 'entrypoint': 'score_job.py', + 'schedule': 'USING CRON 0 3 * * * UTC', + 'timeout_ms': 6 * 3600 * 1000, + 'args': ['--predictions-table', 'USERS_CHURN_PREDICTIONS'], + }, +} + + +def task_sql(*, task_name: str, warehouse: str, schedule: str, + timeout_ms: int, execute_ml_job_sql: str) -> str: + return (f"CREATE OR REPLACE TASK {task_name}\n" + f" WAREHOUSE = {warehouse}\n" + f" SCHEDULE = '{schedule}'\n" + f" USER_TASK_TIMEOUT_MS = {timeout_ms}\n" + f" SUSPEND_TASK_AFTER_NUM_FAILURES = 3\n" + f"AS\n{execute_ml_job_sql}") + + +def register(session: Any, *, compute_pool: str, stage: str, database: str, + schema: str, warehouse: str, artifact_root: str, + requirements: List[str], + external_access_integrations: List[str]) -> Dict[str, str]: + from snowflake.ml.jobs import MLJobDefinition + + flags = [ + '--artifact-root', artifact_root, '--source', 'snowflake', + '--catalog', database, '--schema', schema, '--warehouse', warehouse, + '--registry', 'snowflake', '--model-name', 'USERS_CHURN', + '--scratch-dir', '/tmp/kumo' + ] + tasks = {} + for kind, job in JOBS.items(): + definition = MLJobDefinition.register( + str(PAYLOAD), + compute_pool, + stage, + session=session, + name=f'KUMO_{kind.upper()}', + entrypoint=str(job['entrypoint']), + pip_requirements=requirements, + external_access_integrations=external_access_integrations, + query_warehouse=warehouse, + database=database, + schema=schema, + overwrite=True, + ) + name = f'{database}.{schema}.KUMO_{kind.upper()}_TASK' + session.sql( + task_sql( + task_name=name, warehouse=warehouse, + schedule=str(job['schedule']), + timeout_ms=int(job['timeout_ms']), + execute_ml_job_sql=definition.to_sql( + job_args=[*flags, *job['args']], use_async=False), + )).collect() + tasks[kind] = name + return tasks + + +def main(argv: Optional[Sequence[str]] = None) -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--connection', help='connections.toml entry') + for flag in ('--compute-pool', '--stage', '--database', '--schema', + '--warehouse', '--artifact-root'): + parser.add_argument(flag, required=True) + parser.add_argument('--requirement', action='append', required=True) + parser.add_argument('--external-access-integration', action='append', + default=[]) + args = parser.parse_args(argv) + + from snowflake.snowpark import Session + + builder = Session.builder + if args.connection is not None: + builder = builder.config('connection_name', args.connection) + tasks = register( + builder.create(), compute_pool=args.compute_pool, stage=args.stage, + database=args.database, schema=args.schema, warehouse=args.warehouse, + artifact_root=args.artifact_root, + requirements=args.requirement, + external_access_integrations=args.external_access_integration) + for kind, name in tasks.items(): + print(f'{kind}: {name} (suspended; ALTER TASK {name} RESUME)') + + +if __name__ == '__main__': + main() diff --git a/packages/kumo-engine/kumo/engine/connector/snowflake.py b/packages/kumo-engine/kumo/engine/connector/snowflake.py index 0e75a9947..bdfc97387 100644 --- a/packages/kumo-engine/kumo/engine/connector/snowflake.py +++ b/packages/kumo-engine/kumo/engine/connector/snowflake.py @@ -11,7 +11,7 @@ import os from functools import cached_property from pathlib import Path -from typing import Dict, Optional +from typing import Any, Dict, Optional from kumoapi.data_source import ( DataSourceType, @@ -100,6 +100,12 @@ def sink_backend(self) -> ArtifactSink: "snowflake").SnowflakeSinkBackend return SnowflakeSinkBackend(identity=self.identity) + def session(self) -> Any: + r"""The Snowpark session this connector reads through: its identity, + pointed at its warehouse, database and schema, kept alive.""" + with self.source_backend._session(self) as session: + return session + def source_table_for(self, table_name: str) -> SnowflakeSourceTable: return SnowflakeSourceTable( snowflake_connector_id=self.id, diff --git a/packages/kumo-engine/test/test_job_templates.py b/packages/kumo-engine/test/test_job_templates.py new file mode 100644 index 000000000..7606902f4 --- /dev/null +++ b/packages/kumo-engine/test/test_job_templates.py @@ -0,0 +1,172 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import importlib.util +import py_compile +import sys +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest +from kumoapi.model_plan import ModelPlan + +from kumo.engine.pquery import PredictiveQuery + + +def _examples() -> Path: + for parent in Path(__file__).resolve().parents: + candidate = parent / 'examples' / 'mlops' + if (candidate / 'jobs.py').is_file(): + return candidate + pytest.skip('examples/mlops is not in this checkout') + + +def _module(name: str): + spec = importlib.util.spec_from_file_location( + f'template_{name}', _examples() / f'{name}.py') + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.mark.parametrize('name', [ + 'jobs.py', 'fit_job.py', 'score_job.py', 'sagemaker/pipeline.py', + 'snowflake/register_jobs.py' +]) +def test_python_compiles(name, tmp_path): + py_compile.compile(str(_examples() / name), cfile=str(tmp_path / 'out'), + doraise=True) + + +def test_databricks_tasks_pass_flags_to_the_job_scripts(): + yaml = pytest.importorskip('yaml') + bundle = yaml.safe_load( + (_examples() / 'databricks' / 'databricks.yml').read_text()) + tasks = [ + task['spark_python_task'] for job in bundle['resources']['jobs'].values() + for task in job['tasks'] + ] + assert [task['python_file'].rsplit('/', 1)[-1] + for task in tasks] == ['fit_job.py', 'score_job.py'] + fit, score = (dict(zip(task['parameters'][::2], task['parameters'][1::2])) + for task in tasks) + assert fit['--registry'] == score['--registry'] == 'mlflow' + assert fit['--artifact-root'] == score['--artifact-root'] + assert '--predictions-table' in score + + +def test_snowflake_tasks_run_the_jobs_with_flags(monkeypatch): + calls = [] + + class Definition: + @staticmethod + def register(*args, **kwargs): + return Definition() + + def to_sql(self, job_args, use_async): + calls.append(job_args) + return 'CALL JOB' + + class Session: + def sql(self, text): + calls.append(text) + return self + + def collect(self): + pass + + jobs_module = type(sys)('snowflake.ml.jobs') + jobs_module.MLJobDefinition = Definition + monkeypatch.setitem(sys.modules, 'snowflake.ml.jobs', jobs_module) + tasks = _module('snowflake/register_jobs').register( + Session(), compute_pool='P', stage='@S', database='DB', schema='SC', + warehouse='WH', artifact_root='s3://bucket/kumo', requirements=[], + external_access_integrations=[]) + assert tasks == {'train': 'DB.SC.KUMO_TRAIN_TASK', + 'score': 'DB.SC.KUMO_SCORE_TASK'} + fit_args, fit_sql, score_args, score_sql = calls + assert fit_args[:2] == ['--artifact-root', 's3://bucket/kumo'] + assert '--predictions-table' in score_args + assert '--predictions-table' not in fit_args + assert fit_sql.startswith('CREATE OR REPLACE TASK DB.SC.KUMO_TRAIN_TASK') + assert score_sql.endswith('AS\nCALL JOB') + + +def test_sagemaker_and_snowflake_name_the_job_scripts(): + pipeline = (_examples() / 'sagemaker' / 'pipeline.py').read_text() + register = (_examples() / 'snowflake' / 'register_jobs.py').read_text() + for script in ('fit_job.py', 'score_job.py'): + assert (_examples() / script).is_file() + assert f"/{script}'" in pipeline + assert f"'entrypoint': '{script}'" in register + + +def _data(root: Path) -> None: + root.mkdir() + rng = np.random.default_rng(0) + pd.DataFrame({'USER_ID': np.arange(50)}).to_parquet(root / 'users.parquet') + pd.DataFrame({'ITEM_ID': np.arange(20)}).to_parquet(root / 'items.parquet') + pd.DataFrame({ + 'ORDER_ID': np.arange(3000), + 'USER_ID': rng.integers(0, 50, 3000), + 'ITEM_ID': rng.integers(0, 20, 3000), + 'ORDER_TIME': pd.Timestamp('2024-01-01') + + pd.to_timedelta(rng.integers(0, 120, 3000), unit='D'), + }).to_parquet(root / 'orders.parquet') + + +def test_fit_then_score_locally(tmp_path, monkeypatch, capsys): + _data(tmp_path / 'data') + plan = ModelPlan() + plan.training_job.num_experiments = 1 + plan.training_job.disable_explain = True + plan.optimization.max_epochs = 1 + plan.optimization.max_steps_per_epoch = 2 + monkeypatch.setattr(PredictiveQuery, 'suggest_model_plan', + lambda self, snapshot: plan) + monkeypatch.syspath_prepend(str(_examples())) + monkeypatch.delitem(sys.modules, 'jobs', raising=False) + root = tmp_path / 'artifacts' + flags = [ + '--source', + str(tmp_path / 'data'), '--artifact-root', + str(root), '--scratch-dir', + str(tmp_path / 'scratch') + ] + + fit = _module('fit_job').main([ + *flags, '--query', 'PREDICT COUNT(orders.*, 0, 7, days) > 0 ' + 'FOR EACH users.USER_ID' + ]) + assert capsys.readouterr().out.splitlines()[-1] == f'model: models/{fit.id}' + assert (root / 'models' / fit.id / '_provenance.json').is_file() + assert any((root / 'snapshots').iterdir()) + assert any((root / 'training_tables').iterdir()) + + score = _module('score_job').main( + [*flags, '--model-path', f'models/{fit.id}']) + assert score.num_rows == 50 + + +def test_sagemaker_pipeline_needs_an_instance_type(capsys): + pipeline = _module('sagemaker/pipeline') + args = [ + '--role', 'r', '--image', 'i', '--code', 'c', '--data-root', 'd', + '--artifact-root', 'a', '--tracking-arn', 't' + ] + with pytest.raises(SystemExit): + pipeline.main(args) + assert '--instance-type' in capsys.readouterr().err + + +@pytest.mark.parametrize('name', ['fit_job', 'score_job']) +def test_scripts_import_without_dunder_file(name, monkeypatch): + path = _examples() / f'{name}.py' + monkeypatch.setattr(sys, 'argv', [str(path)]) + monkeypatch.setattr(sys, 'path', list(sys.path)) + monkeypatch.delitem(sys.modules, 'jobs', raising=False) + namespace = {'__name__': 'spark_python_task'} + exec(compile(path.read_text(), str(path), 'exec'), namespace) + assert callable(namespace['main']) diff --git a/packages/kumo-engine/test/test_snowflake_connector.py b/packages/kumo-engine/test/test_snowflake_connector.py index 1e4fcc687..b2b9254c8 100644 --- a/packages/kumo-engine/test/test_snowflake_connector.py +++ b/packages/kumo-engine/test/test_snowflake_connector.py @@ -108,3 +108,18 @@ def test_a_missing_target_field_with_no_override_is_reported( with pytest.raises(ValueError, match="names no warehouse"): SnowflakeConnector.from_spcs_env() + + +def test_session_is_the_readers_session(monkeypatch): + from kumo_snowflake import session as sessions + + connector = SnowflakeConnector(warehouse='WH', database='DB', + schema_name='S') + calls = [] + monkeypatch.setattr( + sessions._sessions, 'get', + lambda principal, target, fresh=False: calls.append(target) or 'live') + monkeypatch.setattr(type(connector.source_backend), '_principal', + lambda self, target: 'principal') + assert connector.session() == 'live' + assert calls == [connector]