Add example MLOps job scripts for Databricks, SageMaker and Snowflake

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/<id>/`, `models/<bundle id>/`).
+  With a registry, the model is registered and the `champion` alias moves to
+  it. It prints `model: models/<bundle id>`.
+- `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/<bundle id>` 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 <data dir> --artifact-root <artifact dir>
+python score_job.py --source <data dir> --artifact-root <artifact dir> \
+  --model-path models/<bundle id>
+```
+
+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.

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 `<catalog>.<schema>` inside `kumo.mlops(...)` with the artifact
+  root `/Volumes/<catalog>/<schema>/<volume>/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
+  `<catalog>.<schema>.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:/<catalog>.<schema>.users_churn@champion`, builds a fresh
+  snapshot and writes `users_churn_predictions` in the connector's schema. Add
+  `--snapshot snapshots/<graph key>/<version>` 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 <target>`. 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 `<volume>`, and `CREATE VOLUME` on the
+  schema (or a pre-created `<catalog>.<schema>._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/

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: <catalog>
+  schema:
+    description: schema holding the tables, the registered model and the predictions
+    default: <schema>
+  volume:
+    description: Unity Catalog volume for exported models and snapshots
+    default: <volume>
+  warehouse_id:
+    description: SQL warehouse that reads the source tables and writes predictions
+    default: <warehouse_id>
+  kumo_engine_wheel:
+    description: kumo-engine wheel on a volume or workspace path
+    default: /Volumes/<catalog>/<schema>/<volume>/wheels/<kumo_engine_wheel>.whl
+  requirements:
+    description: requirements.txt for kumo-engine's dependencies, starting with --prefer-binary
+    default: /Volumes/<catalog>/<schema>/<volume>/wheels/requirements.txt
+  spark_version:
+    description: a standard (non-ML) runtime, for example 16.4.x-scala2.12
+    default: <dbr_version>
+  node_type_id:
+    description: a CPU node type with local NVMe disks, for example m5d.xlarge
+    default: <node_type>
+  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}

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()