Add opt-in lineage columns, stamped by the prediction writer

3rdparty/notices/inventory.json

@@ -515,6 +515,7 @@
         "rust:cpp-demangle@0.4.5",
         "rust:cpufeatures@0.2.17",
         "rust:cpufeatures@0.3.0",
+        "rust:crc32fast@1.5.0",
         "rust:crossbeam-deque@0.8.6",
         "rust:crossbeam-epoch@0.9.21",
         "rust:crossbeam-utils@0.8.21",
@@ -3642,7 +3643,8 @@
     },
     {
       "artifact_ids": [
-        "diskgraph-extension"
+        "diskgraph-extension",
+        "diskgraph-ingest"
       ],
       "copyright": [
         "Copyright (c) 2018 Sam Rijs, Alex Crichton and contributors"

examples/mlops/README.md

@@ -0,0 +1,61 @@
+# Scheduled kumo-engine jobs
+
+Two plain Python job scripts that use the normal `kumo.engine` API inside
+`kumo.job()`, and the wiring that runs them on a schedule:
+
+- `fit_job.py` builds a snapshot (exported under `KUMO_ARTIFACT_ROOT` when
+  set), trains the users-churn model, publishes it and moves the `champion`
+  alias to it.
+- `score_job.py` loads `champion`, scores a fresh snapshot (or the exported
+  one named by `KUMO_SNAPSHOT`) with `ANCHOR_TIMESTAMP` on every row, and
+  writes the predictions. It passes a 0.5 threshold unless the model carries
+  scoring defaults.
+- `jobs.py` holds what both share: the graph (users, items, orders), the
+  connectors and the registry choice, all read from environment variables
+  (listed in its docstring). 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 |
+
+## Run locally
+
+With `users.parquet`, `items.parquet` and `orders.parquet` in a directory:
+
+```bash
+cd examples/mlops
+export KUMO_SOURCE=file KUMO_DATA_ROOT=<data dir> KUMO_ARTIFACT_ROOT=<artifact dir>
+python fit_job.py --job-id local --job-run-id run-1
+KUMO_MODEL_PATH=runs/run-1/model python score_job.py --job-id local --job-run-id run-2
+```
+
+Without `KUMO_REGISTRY` the fit publishes the model to
+`<artifact dir>/runs/<fit run>/model` (its summary's `outputs.model.uri`; the
+`version` there is null, as no registry assigned one),
+and the score job loads the path `KUMO_MODEL_PATH` names under
+`KUMO_ARTIFACT_ROOT`; scheduled score runs have their own run ids, so set it to
+the fit run you want to score with. The query needs
+several 30-day windows of order history; for a small sample set `KUMO_QUERY`
+to a shorter window and `KUMO_NUM_EXPERIMENTS=1`.
+
+## What every scheduler sees
+
+- **Exit codes.** `0` ok, `1` failure, `143` SIGTERM (timeout or cancel).
+  The summary's `error.kind` says whether a failure is `transient` (worth a
+  retry), `permanent` or `unknown`.
+- **Output.** Each script prints one JSON line on stdout when it ends:
+  `status`, `exit_code`, `scheduler`, `bundle_id`, `snapshot_id`, `run_id`,
+  `outputs` (full URIs of the exported snapshot and published model, or of
+  the prediction destinations), `warnings` and `error`. Logs go to stderr.
+- **Job ids.** `--job-id`/`--job-run-id` (or `KUMO_JOB_ID`/`KUMO_JOB_RUN_ID`)
+  are recorded as `scheduler: {platform, job_id, job_run_id}` in the snapshot
+  and training-table manifests and the bundle's `provenance.json`, and as
+  `kumo.scheduler`/`kumo.job_id`/`kumo.job_run_id` on a tracked fit run.
+  Without them the engine detects SageMaker and Snowflake ML Jobs itself.
+
+## Limitations
+
+`score_job.py` does not yet tag output rows with lineage columns or record a
+tracked prediction run; both arrive with the prediction-run PRs.

examples/mlops/databricks/README.md

@@ -0,0 +1,85 @@
+# 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
+`--job-id {{job.id}} --job-run-id {{job.run_id}}`; the cluster's
+`spark_env_vars` configure `jobs.py`.
+
+- **`kumo_train`** runs `fit_job.py` on the `users`, `items` and `orders`
+  tables in `<catalog>.<schema>`. `KUMO_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. The
+  snapshot is exported to the volume.
+- **`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. Set
+  `KUMO_SNAPSHOT=snapshots/<id>` to score an exported 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__`). Databricks sets no job-id
+  environment variable, so ids come from `{{job.id}}`/`{{job.run_id}}`.
+- **Dedicated access mode** (`SINGLE_USER`): standard compute runs as a
+  low-privilege user and has no GPUs. `KUMO_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
+
+The bundle was deployed and both jobs were run on a single-node cluster against
+synthetic `users`, `items` and `orders` tables: both exit 0 with the JSON
+summary and the Databricks job and run ids, the fit registers a new
+`users_churn` version and moves `champion`, and the score job loads it and
+writes the predictions table.
+
+## 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,136 @@
+# 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:
+              KUMO_SOURCE: databricks
+              KUMO_DATABRICKS_CATALOG: ${var.catalog}
+              KUMO_DATABRICKS_SCHEMA: ${var.schema}
+              KUMO_DATABRICKS_WAREHOUSE_ID: ${var.warehouse_id}
+              KUMO_ARTIFACT_ROOT: /Volumes/${var.catalog}/${var.schema}/${var.volume}/kumo
+              MLFLOW_TRACKING_URI: databricks
+              MLFLOW_REGISTRY_URI: databricks-uc
+              MLFLOW_EXPERIMENT_NAME: ${var.experiment}
+              KUMO_REGISTRY: mlflow
+              KUMO_SCRATCH_DIR: /local_disk0/kumo
+              KUMO_MODEL_NAME: ${var.catalog}.${var.schema}.users_churn
+      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:
+              - --job-id
+              - "{{job.id}}"
+              - --job-run-id
+              - "{{job.run_id}}"
+
+    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:
+              KUMO_SOURCE: databricks
+              KUMO_DATABRICKS_CATALOG: ${var.catalog}
+              KUMO_DATABRICKS_SCHEMA: ${var.schema}
+              KUMO_DATABRICKS_WAREHOUSE_ID: ${var.warehouse_id}
+              KUMO_ARTIFACT_ROOT: /Volumes/${var.catalog}/${var.schema}/${var.volume}/kumo
+              MLFLOW_TRACKING_URI: databricks
+              MLFLOW_REGISTRY_URI: databricks-uc
+              MLFLOW_EXPERIMENT_NAME: ${var.experiment}
+              KUMO_REGISTRY: mlflow
+              KUMO_SCRATCH_DIR: /local_disk0/kumo
+              KUMO_MODEL_NAME: ${var.catalog}.${var.schema}.users_churn
+              KUMO_PREDICTIONS_TABLE: users_churn_predictions
+      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:
+              - --job-id
+              - "{{job.id}}"
+              - --job-run-id
+              - "{{job.run_id}}"
+
+targets:
+  dev:
+    default: true
+    workspace:
+      root_path: /Workspace/Shared/kumo-engine/.bundle/${bundle.name}/${bundle.target}

examples/mlops/fit_job.py

@@ -0,0 +1,59 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Train and publish the users-churn model: ``python fit_job.py [--job-id ID
+--job-run-id ID]``; see README.md."""
+
+import os
+import sys
+from dataclasses import asdict
+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
+
+# 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 main(argv: Optional[Sequence[str]] = None) -> None:
+    job_id, job_run_id = jobs.start(argv)
+    with kumo.job(job_id=job_id, job_run_id=job_run_id) as run:
+        connector = jobs.source()
+        graph = jobs.build_graph(connector)
+        snapshot = graph.snapshot()
+        run.record(snapshot_id=snapshot.id)
+        if os.environ.get('KUMO_ARTIFACT_ROOT'):
+            run.record(outputs={
+                'snapshot':
+                jobs.artifact_uri(snapshot.export(jobs.artifacts(), 'snapshots'))
+            })
+
+        query = kumo.PredictiveQuery(graph=graph,
+                                     query=os.environ.get('KUMO_QUERY', QUERY))
+        table = query.generate_training_table(snapshot)
+        plan = query.suggest_model_plan(snapshot)
+        if os.environ.get('KUMO_NUM_EXPERIMENTS'):
+            plan.training_job.num_experiments = int(
+                os.environ['KUMO_NUM_EXPERIMENTS'])
+        tracking, export_config = jobs.model_destination(connector)
+        result = kumo.Trainer(plan).fit(table, snapshot=snapshot,
+                                        tracking=tracking,
+                                        export_config=export_config)
+
+        registered = getattr(tracking, 'registered', None)
+        published = registered or result.published_model
+        model = published and asdict(published)
+        if model and not os.environ.get('KUMO_REGISTRY'):
+            model['version'] = None
+        run.record(bundle_id=result.bundle_id,
+                   run_id=getattr(tracking, 'run_id', None),
+                   outputs={'model': model})
+
+
+if __name__ == '__main__':
+    main()