Record prediction runs and monitor score drift

docs/development/testing.md

@@ -23,6 +23,14 @@ Run from the repository root unless a component guide says otherwise.
 Python package test suites have different roots and fixtures; inspect the
 package's configuration before invoking pytest over multiple packages.
 
+kumo-engine and kumo-core tests load the DiskGraph native module
+(`diskgraph/_c.abi3.so`) and its `bin/ingest` binary from the installed
+`diskgraph` package. Both must be built from the same tree as the Python code
+under test: a build from an older checkout fails with missing attributes or
+arguments (for example `RuntimeIndex.iter(projections=...)` or
+`compute_col_stats_for_parquet(percentiles=...)`). After pulling DiskGraph
+changes, rebuild with `make install-diskgraph`, which builds and installs both.
+
 For connector changes, start with the [public test command](public-tests.md).
 It builds its prerequisites and runs without cloud credentials. The
 [test matrix](../connector-test-matrix.md) lists coverage and CI limits.

examples/mlops/README.md

@@ -0,0 +1,69 @@
+# 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>/`).
+  The model carries its `--threshold` (0.5 by default) as scoring defaults.
+  With a registry, the model is registered and the `--alias` alias
+  (`champion` by default) moves to it. It prints `model: models/<bundle id>`.
+- `score_job.py` loads the version `--alias` points at (or `--model-path`)
+  through `Trainer.load_from_connector(cfg, ...)`, scores a fresh snapshot
+  (or the stored one `--snapshot` names) and writes `--predictions-table` when
+  given. Inside the block the run is stored under `predictions/<id>/` with its
+  `_provenance.json`.
+- `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;
+  list text columns in `TEXT_COLUMNS` (e.g. `('users.BIO',)`), since inference
+  reads short text as categorical. With `--on-error raise`, a score job whose
+  graph differs from the one the model was trained on fails before scoring.
+
+| 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://`, a local or `/Volumes` path, or with `--source snowflake` an `@DB.SCHEMA.STAGE/...` 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 |
+| `--table-prefix` | prefix of the `users`, `items` and `orders` table or file names, e.g. `demo_` for `demo_users`; default none |
+| `--registry` | `mlflow` (`MLOps(tracking=MlflowTracking(), registry=ModelRegistryLocation(...))`) or `snowflake` (`MLOps(tracking=SnowflakeTracking(connector, ...), registry=ModelRegistryLocation(...))`); omit for no registry |
+| `--model-name` | registered model name, with `--registry`: `users-churn` on SageMaker (letters, digits and hyphens only), `<catalog>.<schema>.users_churn` in Unity Catalog, `USERS_CHURN` in Snowflake |
+| `--alias` | with `--registry`: the alias the fit moves to its new version and the score job loads; default `champion`. With MLflow, register under `candidate` and promote yourself to gate scoring |
+| `--model-path` | score only, without a registry: `models/<bundle id>` under the artifact root |
+| `--snapshot` | score only: the id of a snapshot stored under the artifact root |
+| `--predictions-table` | score only: a table of the source connector to write |
+| `--query`, `--num-experiments`, `--threshold` | fit only: the predictive query, the number of experiments and the classification threshold the model scores with |
+| `--scratch-dir` | local disk for snapshots, tables and bundles |
+| `--on-error` | `raise` (default): a failed tracking call or registration fails the job after its result is stored (`kumo.MLOps(on_error="raise")`); `warn`: the job warns and exits 0 |
+
+## 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. A failed registration or tracking
+call is a failure too (`--on-error raise`), so a fit that stored its model but
+did not register it does not leave the score job on the previous version
+unnoticed. Databricks Jobs, SageMaker Pipelines
+and Snowflake Tasks keep their own run ids, status and logs.

examples/mlops/databricks/README.md

@@ -0,0 +1,103 @@
+# Databricks Jobs
+
+`databricks.yml` is an Asset Bundle with two jobs on a single-node job
+cluster, a GPU node with the GPU ML runtime by default. 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 (each prefixed with `table_prefix`) 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>.<model_name>` (default `users_churn`) and moves the
+  `alias` variable's alias (default `candidate`) to it; move `champion` to it
+  once it is better (see
+  [Promote a model](../../../packages/kumo-engine/docs/mlops.md#promote-a-model)).
+  Set `alias` to `champion` to serve every fit without a gate.
+- **`kumo_score`** runs `score_job.py` daily at 03:00 UTC (created paused).
+  It loads the version `score_alias` (default `champion`) points at in
+  `<catalog>.<schema>.<model_name>`, builds a fresh snapshot and appends to
+  `<predictions_table>` (default `users_churn_predictions`) in the
+  connector's schema. Add `--snapshot <snapshot id>` 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-skinny)
+   on a volume and set `kumo_engine_wheel` and `requirements`. Each task
+   installs the requirements, then the wheel. Start the requirements file with
+   `--prefer-binary`, so pip installs wheels rather than building sdists.
+   - **Runtime:** Databricks Runtime 17.3 LTS ML or later: the GPU flavour
+     (`17.3.x-gpu-ml-scala2.13`) on a GPU node such as `g5.xlarge`, or the
+     CPU flavour (`17.3.x-cpu-ml-scala2.13`) on a node type with local NVMe
+     disks such as `m5d.xlarge`. These runtimes already carry a torch that
+     kumo-core accepts, so pip keeps the runtime's torch, torchvision,
+     MLflow and cryptography; leave torch out of the requirements file.
+     kumo-engine accepts pyarrow 16 through 24, so the runtime's pyarrow
+     (19 on 17.3) is kept too. On 16.4 LTS ML, whose torch 2.6 is
+     older than kumo-core needs, pip replaces torch and breaks the
+     preinstalled torchvision.
+2. Replace the `<...>` placeholders in `databricks.yml` (or pass
+   `--var catalog=...`). `model_name`, `predictions_table`, `table_prefix`,
+   `experiment` and `on_error` have defaults; set them to run beside existing
+   objects in the same schema. The parent folder of the MLflow experiment
+   path must already exist in the workspace.
+3. Adjust the `dev` target (workspace `root_path`, under your user folder by
+   default) 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
+
+- **Strict mode.** Both jobs pass `--on-error raise` (variable `on_error`):
+  if MLflow tracking or Unity Catalog registration fails, the task fails after
+  the model or predictions are stored, and `champion` stays where it was.
+
+- **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 access mode 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.
+- **Job clusters.** Both jobs repeat the cluster block; YAML anchors would
+  remove the copy but are untested with `databricks bundle validate`.
+
+## Verified
+
+This bundle was deployed and both jobs were run on a single-node cluster
+against synthetic `users`, `items` and `orders` tables, on a GPU ML runtime
+on a `g5.xlarge` (A10G; the fit trained on the GPU): the fit registered a
+new `users_churn` version and moved `champion`, and the score job loaded it and
+wrote the predictions table. `test_job_templates.py` checks the scripts
+locally.
+
+## 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,173 @@
+# 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; see README.md
+    default: /Volumes/<catalog>/<schema>/<volume>/wheels/requirements.txt
+  spark_version:
+    description: Databricks Runtime 17.3 LTS ML or later, for example 17.3.x-gpu-ml-scala2.13 (or 17.3.x-cpu-ml-scala2.13 on a CPU node)
+    default: <dbr_version>
+  node_type_id:
+    description: a GPU node type, for example g5.xlarge (or a CPU node type with local NVMe disks, such as m5d.xlarge)
+    default: <node_type>
+  experiment:
+    description: MLflow experiment (workspace path) for train and score runs
+    default: /Shared/kumo-engine/users-churn
+  model_name:
+    description: registered model, created as <catalog>.<schema>.<model_name>
+    default: users_churn
+  predictions_table:
+    description: table in <catalog>.<schema> the score job appends to
+    default: users_churn_predictions
+  table_prefix:
+    description: prefix of the users, items and orders table names, e.g. demo_ for demo_users
+    default: ""
+  on_error:
+    description: raise fails a job whose tracking or registration fails; warn lets it succeed
+    default: raise
+  alias:
+    description: alias kumo_train moves to its new version; promote it to score_alias yourself
+    default: candidate
+  score_alias:
+    description: alias kumo_score loads
+    default: champion
+
+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}.${var.model_name}
+              - --table-prefix=${var.table_prefix}
+              - --on-error
+              - ${var.on_error}
+              - --alias
+              - ${var.alias}
+              - --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}.${var.model_name}
+              - --table-prefix=${var.table_prefix}
+              - --on-error
+              - ${var.on_error}
+              - --alias
+              - ${var.score_alias}
+              - --scratch-dir
+              - /local_disk0/kumo
+              - --predictions-table
+              - ${var.predictions_table}
+
+targets:
+  dev:
+    default: true
+    workspace:
+      root_path: /Workspace/Users/${workspace.current_user.userName}/.bundle/${bundle.name}/${bundle.target}