[Vedaant-J]

Series: kumo-engine MLOps series: every snapshot, training table, model and prediction run carries a self-contained _provenance.json; with kumo.mlops(...) stores each artifact under one artifact root, tracks fits and prediction runs in MLflow or Snowflake and registers models; volumes and stages can hold files; prediction runs are monitored for drift and realized performance.
Stack: #53 → #112 → #50 → #60 → #64; #114 on #112; #115 on #64
Merge order: one linear stack; merge bottom to top (#53 first); #114 sits on #112 and merges any time after it; #115 (draft) sits on #64 and merges after it
This PR: Adds records and monitoring for prediction runs: a manifest with output score statistics for every Trainer.predict, a holdout score baseline in every model manifest, Trainer.evaluate and Trainer.score_drift for realized performance, score drift and input-feature drift (from the snapshots' build-time column statistics), a record of the entities scoring skipped, and tracking of prediction and monitoring runs in MLflow or Snowflake.

Summary

This PR adds records and monitoring for prediction runs: a manifest with output score statistics for every Trainer.predict, a holdout score baseline in every model manifest, Trainer.evaluate and Trainer.score_drift for realized performance, score drift and input-feature drift (from the snapshots' build-time column statistics), a record of the entities scoring skipped, and tracking of prediction and monitoring runs in MLflow or Snowflake. It also ships job templates (examples/mlops) that run fit and scoring on Databricks Jobs, SageMaker Pipelines and Snowflake Tasks.

It is the top of the stack and builds on two PRs below it: #50 (snapshot, training-table and model manifests, kumo.MLOps and the artifact root) and #60 (the Tracking protocol, the MLflow and Snowflake trackers, model registration). No PR builds on this one.

Problem. Teams running batch scoring in production need to answer questions like these:

• Which model, snapshot and source table versions produced the scores in this table?
• Was this run scored with a threshold the caller passed, the model's default, or the engine's?
• This week's scores moved. Compared with what: the model's holdout, or last week's run? By how much?
• We scored three weeks ago and the outcomes are in now. How good were those scores?
• Can a separate scheduled job do this from the exported predictions, after the scoring job's scratch is gone?
• How do we feed Databricks data quality monitoring or a Snowflake Model Monitor from these runs?
• Where are the scoring runs in MLflow or Snowsight, next to the training runs of the same model?

Today none of these can be answered with kumo-engine without hand-written SQL that re-implements the query.

Cause. A prediction run leaves only the scorer's selection and attempt files in local scratch, and an export directory holds only Parquet parts and the export markers:

<scratch>/trainers/<trainer>/predict_output/<id>/
  _selected_artifacts.json  _selections/  _attempts/   # scorer selection and attempts

<out>/preds/                                           # predict(output_config=...) or result.export(...)
  part-00000-00000.parquet  _KUMO_EXPORT  _SUCCESS

None of these record the model, the snapshot, the scoring options or where they came from, the destinations, timings, or any statistic of the scores. The model bundle has no baseline either: its holdout predictions stay in the trainer's scratch, outside the bundle, and metrics.json holds aggregate quality, not the score distribution. Labels can be generated only for training, and generation refuses outcome tables that could not train a model (one class, fewer classes than TOP K). The Tracking protocol (#60) records fits only.

trainer.fit(...)                -> TrainingResult (model manifest, #50; no score baseline)
trainer.predict(snapshot, ...)  -> PredictionResult (selection only)
# no run record, realized metrics, drift, monitor tables, or tracked scoring runs

Change.

<bundle>/
  _provenance.json                                     # gains "holdout"

<scratch>/trainers/<trainer>/predict_output/<id>/
  ...same files as above...
  _provenance.json                                     # new: PredictionManifest
  actuals/<snapshot id>/part-00000.parquet             # new: evaluated rows, written by evaluate()

<out>/preds/
  part-00000-00000.parquet  _KUMO_EXPORT  _SUCCESS
  _provenance.json                                     # new, written before _SUCCESS

<root>/predictions/<id>/predictions/                   # new, MLOps on (export kind "prediction")
  part-00000-00000.parquet  _KUMO_EXPORT  _SUCCESS  _provenance.json

trainer.predict(snapshot, ..., mlops=None)                          -> PredictionResult (+ .id .manifest .bundle_id .snapshot_id)
trainer.evaluate(run, snapshot=newer, mlops=None)                   -> EvaluationResult    # new
evaluation.export(PredictionOutputConfig(...))                      # new: evaluated rows to a table or directory
trainer.score_drift(run, previous=None, snapshot=None,
                    training_snapshot=None, mlops=None)             -> DriftReport         # new
PredictionResult.load_from_connector(connector, path=...)           -> PredictionResult    # new
training_result.evaluate(...), training_result.score_drift(...)     # pass through to the trainer

The examples below come from one real CPU run: a synthetic users (500) / orders graph and PREDICT COUNT(orders.*, 0, 7, days) > 0 FOR EACH users.ID. The model was fitted on a snapshot whose orders stop at 2024-03-31 (2,984 orders) inside kumo.mlops(MLOps(artifact_root=..., tracking=MlflowTracking(experiment_name="users-churn"))). The run was scored on that snapshot, exported to <out>/preds, loaded back with load_from_connector, and evaluated on a later snapshot with all 4,000 orders. A second run scored the later snapshot. Ids are shortened, percentiles elided, and paths replaced. The model was trained for a few steps, so the numbers are real but say nothing about model quality.

1. Holdout baseline in the model manifest

After fit, _provenance.json (ModelManifest, #50) gains holdout: statistics of the scores the published model gave its own holdout rows.

"holdout": {
  "truncated": false,
  "rows": 500,
  "label_positive_rate": 0.336,
  "columns": {
    "True_PROB": {"count": 500, "max": 0.35833, "mean": 0.33387, "min": 0.30085,
                  "percentiles": [0.30789, 0.30958, "…", 0.33654, "…", 0.34891, 0.35260]}
  }
}

• Which model: always the published model's own holdout, the fit or refit_trainval model named by fit_mode. A refit_full model has no holdout of its own and gets no holdout key; another model's scores are never recorded in its place. Link prediction writes no holdout frame, so it has none either.
• truncated: true when holdout rows were dropped to fit the holdout's 1 GiB in-memory cap; null when that could not be read.
• rows: the number of holdout predictions (entities for multiclass and multilabel).
• columns: one summary per score column. count counts finite values only (null, NaN and infinite scores are not counted). mean, min, max and percentiles (99 values, q01..q99, approximate, from DiskGraph's t-digest) are all over finite values. Which column depends on the task:
  • binary: True_PROB, plus label_positive_rate, the share of labelled holdout rows with a positive label;
  • regression and forecasting: TARGET_PRED and TARGET;
  • multiclass, multilabel and ranking: top_score, each prediction's highest score. A prediction is a run of consecutive rows with equal non-class columns (entity, and anchor when present), with its rows in any score order.
• In Python: result.manifest.holdout (or trainer.result.manifest.holdout), a HoldoutStats (truncated, rows, columns, label_positive_rate). holdout is in ModelManifest.EMBED_OMITS, so manifests that embed the model leave it out. A model derived with with_scoring_defaults scores identically and keeps its parent's holdout.

2. The prediction manifest

Every Trainer.predict writes _provenance.json (PredictionManifest, public as kumo.engine.PredictionManifest) into the run directory. Exports to Parquet directories get a copy next to the parts (the leading underscore keeps Spark and Hive readers off it). The manifest of the first run:

{
  "kind": "prediction", "schema_version": 1,
  "id": "5a549ffa…-predict-da72a70e…",
  "created_at": "2026-10-26T14:50:15+00:00",
  "environment": {"kumo_engine": "…", "build_commit": "8ed78eee", "diskgraph": "0.1.0", "diskgraph_build_commit": "26cc3192490d", "torch": "2.12.0+cu130", "python": "3.14.7"},
  "parents": {
    "model": {
      "kind": "model", "id": "8c0c3459…", "fit_mode": "refit_trainval",
      "fit": {"started_at": "2026-10-26T14:47:58+00:00", "finished_at": "2026-10-26T14:50:11+00:00", "seconds": 133.489},
      "parents": {"training_table": {"kind": "training_table", "id": "tt-aba48cfa…",
        "query": "PREDICT COUNT(orders.*, 0, 7, days) > 0 FOR EACH users.ID",
        "task_type": "binary_classification", "split_seed": null, "generation_plan": {"…": "…"},
        "parents": {"snapshot": {"kind": "snapshot", "id": "9b04d531…@83364c76319c05b3", "…": "…"}}}}},
    "snapshot": {
      "kind": "snapshot", "id": "9b04d531…@83364c76319c05b3", "record_sources": true,
      "parents": {"source_tables": {
        "orders": {"connector": "LOCAL_FILE", "location": "<data>/orders.parquet", "version": "1793022475", "version_kind": "timestamp"},
        "users":  {"connector": "LOCAL_FILE", "location": "<data>/users.parquet",  "version": "1793022475", "version_kind": "timestamp"}}}}},
  "prediction_table": {"query_sha256": "69dab700…", "source": "pql", "anchors": ["2024-03-30T00:00:00+00:00"], "rows": 500,
                       "requested_rows": 500, "scored_rows": 500, "skipped_rows": 0, "skipped_sample": []},
  "scoring": {"binary_classification_threshold": 0.5, "threshold_source": "argument",
              "num_classes_to_return": null, "count_source": "none", "num_workers": 1},
  "outputs": [{"artifact": "predictions", "destination": {"connector": "LOCAL_FILE", "location": "<out>/preds"}}],
  "status": "succeeded",
  "timings": {"started_at": "2026-10-26T14:50:11+00:00", "finished_at": "2026-10-26T14:50:15+00:00"},
  "stats": {"predictions": {"rows": 500, "predicted_positive_rate": 0.0,
    "columns": {"True_PROB": {"count": 500, "max": 0.35972, "mean": 0.33446, "min": 0.30316,
                              "percentiles": [0.30790, 0.30984, "…", 0.35029]}}}}
}

• id: the prediction id, PredictionResult.id; run directories and <root>/predictions/<id>/ are per id.
• parents embeds two manifests without their statistics and later-changing fields:
  • model: the model manifest without holdout and tracking, with its training table and that table's snapshot embedded. Its own warm-start or derived-from models are id stubs, so size stays bounded across retrains.
  • snapshot: the scoring snapshot's manifest, with each source table's connector, location and version. When it equals the training snapshot, both copies are kept.
  • A parent without a manifest (a bundle saved before this series, a snapshot loaded without one) is a stub {kind, id, reconstructed: true}. PredictionResult.bundle_id and .snapshot_id read the two ids from here.
• prediction_table: the query's SHA-256 (the text is in the embedded training table), source (pql when generated from the query, custom when built from a source table), the anchor times in ISO UTC, and the row count.
• Skipped entities. Supplied rows (entities=, a supplied prediction_table=, the MLflow pyfunc model) whose entity is not in the scoring snapshot used to be dropped with only a DiskGraph log line that spoke of training rows. They are still skipped, never refused, and now recorded:
  • prediction_table.requested_rows, scored_rows, skipped_rows (0 for a generated population, never missing), skipped_sample (up to 5 skipped entity ids per reason, read from the source rows DiskGraph names) and skip_reason: not in snapshot, or null entity id, null anchor time, not a canonical node in the snapshot (DiskGraph's other predict-path skip conditions), joined with counts when several apply;
  • PredictionResult.requested_rows, .scored_rows, .skipped_rows;
  • one UserWarning, Prediction skipped 4 of 14 requested entities (28.6%): not in snapshot, which replaces DiskGraph's log line on the predict path;
  • with tracking on, predictions.skipped_rows and predictions.skipped_fraction on the prediction run.
  • The pyfunc model returns only the scored rows; the warning carries its counts (MLflow's pyfunc predict returns a frame, and each call's scratch, manifest included, is removed when it returns).
• scoring: the effective options and where each came from: argument (passed by the caller), bundle (the model's scoring defaults), default (filled in by the engine), query (from the PQL), or none.
• outputs: the destinations of predict(output_config=...) as {connector, location}, with the same location rules as source tables (absolute path, full object-store URI, a volume or stage path as given, or a fully qualified warehouse table). A destination is recorded once its export completes, so a failing export leaves only the completed ones; null means the outputs could not be described. A later result.export(...) copies the files and the manifest as they are and records nothing new.
• stats: the row count per artifact and, for predictions, the same score summary as the holdout baseline, computed by DiskGraph over the output Parquet:
  • binary: True_PROB, plus predicted_positive_rate at the threshold;
  • regression and forecasting: TARGET_PRED;
  • multiclass, multilabel and link prediction: top_score, the highest score per entity and anchor.
• status: succeeded, or failed when the output_config export raised; the manifest is written either way and predict re-raises. timings: start and finish of the call.
• In Python, PredictionResult.manifest returns it (run.manifest.stats['predictions'].columns), with PredictionParents, PredictionTableRecord, RunScoring, PredictionOutput, OutputDestination, RunTimings and ArtifactStats in kumo.engine._manifest. PredictionManifest.EMBED_OMITS is outputs, stats, status, timings.

Order inside predict: score, build the manifest, run the output_config export (recording outputs), write the manifest, then with MLOps on store <root>/predictions/<id>/ (whose manifest includes those outputs), then track the run. The store is skipped if the export raised, and a failed store only warns.

3. Realized performance: Trainer.evaluate

evaluate regenerates the query's labels at the run's anchors (read from its manifest) on the given snapshot with DiskGraph training-table generation, then joins them with the run's prediction Parquet in DiskGraph. The new diskgraph.evaluate_prediction_parquet streams both sides by row group and hash-partitions them on ENTITY (and ANCHOR_TIMESTAMP when the run has several anchors). It joins at most 8 partitions of about 1M rows at a time, each spilled to Arrow IPC first, and accumulates the metrics; nothing is collected into pandas. Metrics use the fit-time test_* definitions:

• binary: auroc and ap over True_PROB, exact (torchmetrics definitions, ties counted half, scores compared at float32);
• regression: mae and rmse over TARGET_PRED;
• multiclass: acc of the top-scored CLASS per entity and anchor;
• link prediction: map@k and precision@k as PyG's LinkPredMAP / LinkPredPrecision, where k is the number of classes returned per entity and anchor, over predictions with at least one target (logged to MLflow as map_at_k and precision_at_k);
• multilabel and forecasting: no metrics, with a reason.

EvaluationResult of the first run on the later snapshot (dataclasses.asdict):

{
  "prediction_id": "5a549ffa…-predict-da72a70e…", "bundle_id": "8c0c3459…",
  "snapshot_id": "9b04d531…@e27d964a25bec319",
  "task_type": "binary_classification",
  "anchors": ["2024-03-30T00:00:00+00:00"],
  "metrics": {"ap": 0.36784, "auroc": 0.51919},
  "rows": 500, "labelled": 500, "labelled_fraction": 1.0,
  "matured": true, "bundle_matches": true, "reasons": {},
  "actuals": "<scratch>/evaluations/5a549ffa…-predict-da72a70e…/actuals/9b04d531…@e27d964a25bec319"
}

actuals holds 500 rows such as {"ENTITY": 1, "ANCHOR_TIMESTAMP": "2024-03-30 00:00:00", "TARGET": true, "EVALUATED_AT": "2026-10-26 14:50:22.363663571+00:00", "True_PROB": 0.33915, "TARGET_PRED": false}.

Nothing is refused; the fields say how far to trust a number:

• snapshot_id: the snapshot the labels were read from. prediction_id and bundle_id are the run's own.
• matured: False while the label window ends after the earliest of the target tables' latest event times, a conservative proxy for label completeness. An immature run still returns its metrics, over the outcomes observed so far, and warns (UserWarning) that its labels are not mature yet. A static query reads labels as of the snapshot: anchors is empty, matured is None, and the evaluated rows carry the run's prediction time as ANCHOR_TIMESTAMP.
• bundle_matches: whether the run's model is the trainer's. Evaluating a run with a different model is allowed and flagged.
• labelled_fraction: labelled predictions / run predictions. For link prediction, the share of entities with any future link.
• actuals: a Parquet directory of the evaluated rows, next to the run (under <scratch>/evaluations/<id>/ for a run loaded from an export): ENTITY, ANCHOR_TIMESTAMP, TARGET, the matched prediction (binary True_PROB and the predicted label TARGET_PRED; regression TARGET_PRED; multiclass the top CLASS and its SCORE; link prediction none), and EVALUATED_AT, one UTC timestamp for the whole call. None when no labels were joined.
• reasons: when a step cannot be computed (no anchors, no labels, unsupported task, a failed step), the affected field is None or empty and reasons says why.

EVALUATED_AT exists because warehouse monitors profile only rows whose timestamp falls in a recent lookback window, and a run's anchors are older than that by at least the label window. On a live Databricks Data Quality monitor over evaluated rows, keying on ANCHOR_TIMESTAMP (59 days old) left the metric tables empty; rows timestamped 1–5 days old produced 60 profile rows and 48 drift rows. Both warehouse recipes key on EVALUATED_AT.

EvaluationResult.export(PredictionOutputConfig(...)) copies the evaluated rows to the config's predictions destination, a table or a directory. A Parquet directory also gets _evaluation.json (every other field), and the new EvaluationResult.load_from_connector(connector, path=...) loads the evaluation back with its rows, so every artifact has one export(<Kind>OutputConfig) and one load_from_connector; docs/mlops.md tabulates them. Exporting a prediction run or an evaluation again to a directory that already holds it complete rewrites nothing (an OVERWRITE destination is no longer emptied first in that case; #50 skips the upload).

4. Score drift: Trainer.score_drift

score_drift compares a run's score summary with the model's holdout and, optionally, with a previous run, with no pass over the data. Each side stores min, q01..q99 and max: 101 monotone knots of a piecewise-linear CDF.

statistic	definition	error from the stored curves
psi	PSI over the baseline's deciles	—
ks	sup of |F_a − F_b|, at every knot of both sides (left and right limits, for point masses)	≤ 0.02 + 2ε (ε: the sketch's rank error)
wasserstein	∫₀¹ |Q_a − Q_b| dp, exact per segment between knots	≤ 0.01 × (sum of both score ranges) + sketch error

wasserstein is in score units: comparable over time for probabilities, scale-dependent for regression. With both snapshots passed, input_changes holds per-table row and per-edge count changes between the training and the scoring snapshot (reverse edges skipped).

DriftReport of the second run (scored on the later snapshot), against the holdout and against the first run:

{
  "prediction_id": "5a549ffa…-predict-2633f84b…", "bundle_id": "8c0c3459…",
  "score": "True_PROB",
  "psi": 0.26605, "ks": 0.18113, "wasserstein": 0.00385,
  "holdout_truncated": false, "bundle_matches": true,
  "previous_prediction_id": "5a549ffa…-predict-da72a70e…",
  "previous_psi": 0.17006, "previous_ks": 0.16781, "previous_wasserstein": 0.00323,
  "input_changes": {"tables.orders.rows": [2984, 4000], "edges.orders_USER_ID_to_users_ID": [2984, 4000]},
  "reasons": {}
}

Missing records (a run without a manifest, a model without holdout, snapshots not passed) leave the affected fields None with a reasons entry. holdout_truncated repeats the holdout's truncated flag.

Input features. DriftReport.features compares, for each table and column the model encodes (the bundle's encoders), the column statistics DiskGraph already wrote into each snapshot's col_stats.json when it was built: no new pass over data and no Rust change. The training snapshot is the one passed, or the one named in the model manifest's embedded training table, found in the local snapshot store; the scoring snapshot likewise from the run's manifest. From a toy users table whose AGE mean moved from 40 to 55 and whose PLAN mix moved from 60/30/10 to 10/30/60 (values rounded; training, scoring and relative omitted):

{"users": {
  "AGE":  {"type": "numerical", "standardized_mean_shift": 3.361,
           "delta": {"mean": 15.311, "std": 0.535, "min": 14.031, "max": 18.258, "q25": 14.942, "q50": 15.108, "q75": 15.849,
                     "skewness": -0.041, "kurtosis": -0.268, "null_rate": 0.0, "num_unique": -3.0}},
  "PLAN": {"type": "categorical", "psi": 1.522, "js": 0.248, "category_counts_capped": false,
           "delta": {"null_rate": 0.0, "num_unique": 0.0}}}}

• Numeric columns: training and scoring values, delta and relative (to the training value) for mean, std, min, max, q25, q50, q75, skewness, kurtosis, null rate and distinct count, plus standardized_mean_shift (mean change over the training std; None when that std is 0).
• Categorical and multi-categorical columns: null rate and distinct count, PSI and Jensen-Shannon divergence (base 2) between the stored category distributions (smoothed by 1e-6), and category_counts_capped when either side's counts were truncated.
• Timestamp columns: the shift of the earliest and latest time (min_secs, max_secs).
• A used column missing on a side, a snapshot not found or without statistics, or a bundle without encoders leaves features partial or empty with a reasons['features'] entry; nothing raises.
• These are whole-snapshot statistics: a recent or slow change in a large event table is diluted by its history.
• Tracking: the monitoring run logs the 10 columns with the largest |standardized shift| and the 10 with the largest PSI (features.shift.<table>.<column>, features.psi.<table>.<column>), the largest of each per table for up to 10 tables, at most 40 metrics with names reduced to letters, digits, _, - and .; the full comparison is the run's feature_drift.json artifact (MLflow).

5. Monitoring from a later job

PredictionResult.load_from_connector(connector, path=...) takes any connector whose sink holds files at the path (S3, GCS, local, or a volume or stage). It requires the export's _SUCCESS marker and fetches only its _provenance.json. The Parquet parts are downloaded in parallel when files is first read (by evaluate, for example), into a scratch directory removed when the result is garbage-collected, so score_drift from another job fetches nothing else.

trainer = Trainer.load_from_connector(store, path="models/<bundle id>")
run = PredictionResult.load_from_connector(store, path="predictions/2024-06-01/")
evaluation = trainer.evaluate(run, snapshot=newer_snapshot)
drift = trainer.score_drift(run)

6. Tracking prediction and monitoring runs

With kumo.MLOps(tracking=...) (a block or mlops=), each predict is recorded as one run in the tracker's experiment, and each evaluate and score_drift call as a new monitoring run. Both go through two optional Tracking hooks, log_prediction_run(record) and log_monitoring(record). They run on a tracking thread, only on global rank 0, for up to 30 s; errors are logged, never raised. The MLflow and Snowflake trackers implement both hooks as their own start(record) then finish(record, error), so params, metrics, tags, inputs and artifacts go through the same code as a fit run.

_tracking.prediction_record(manifest, error=None) is the one mapping from a prediction manifest to a RunRecord:

Part	Contents
name	the prediction id
params	the manifest's scoring options plus anchors; None options are left out
metrics	prediction_table.rows, <artifact>.rows, predictions.predicted_positive_rate (binary), and per score column count, mean, min, max and seven percentile levels q01, q05, q25, q50, q75, q95, q99 (the manifest keeps all 99)
inputs	the scoring snapshot's source tables; MLflow logs them as run inputs
tags	kumo.kind=prediction, kumo.run_id (the prediction id), kumo.bundle_id, and kumo.error when the export failed
artifact	_provenance.json, the full manifest

If the output_config export raises, the run is still recorded and ended as failed with kumo.error, and predict raises the export's error. A monitoring run is named evaluate-<prediction id> or drift-<prediction id>, tagged kumo.kind=monitoring, kumo.prediction_run=<prediction id> and kumo.bundle_id, and logs the metrics, the drift statistics and the input count changes (input.<key>.before / .after). A new run is used because Snowflake cannot reliably reopen a finished run. MLflow also tags the prediction run with kumo.monitored=<monitoring run id>.

The two runs above, read back with MlflowClient.search_runs from a local store (rounded):

{"run_name": "5a549ffa…-predict-2633f84b…", "status": "FINISHED",
 "tags": {"kumo.kind": "prediction", "kumo.run_id": "5a549ffa…-predict-2633f84b…", "kumo.bundle_id": "8c0c3459…",
          "kumo.monitored": "41a503bf…"},
 "params": {"binary_classification_threshold": "0.5", "threshold_source": "argument", "count_source": "none",
            "num_workers": "1", "anchors": "[\"2024-04-29T00:00:00+00:00\"]"},
 "metrics": {"prediction_table.rows": 500.0, "predictions.rows": 500.0, "predictions.predicted_positive_rate": 0.0,
             "predictions.True_PROB.count": 500.0, "predictions.True_PROB.mean": 0.33763,
             "predictions.True_PROB.min": 0.30943, "predictions.True_PROB.max": 0.35574,
             "predictions.True_PROB.q01": 0.31291, "predictions.True_PROB.q05": 0.32228, "predictions.True_PROB.q25": 0.33396,
             "predictions.True_PROB.q50": 0.33902, "predictions.True_PROB.q75": 0.34256,
             "predictions.True_PROB.q95": 0.34757, "predictions.True_PROB.q99": 0.34999}}

{"run_name": "monitor-5a549ffa…-predict-2633f84b…", "status": "FINISHED",
 "tags": {"kumo.kind": "monitoring", "kumo.run_id": "5a549ffa…-predict-2633f84b…-monitor-f8d5f68f…",
          "kumo.prediction_run": "5a549ffa…-predict-2633f84b…", "kumo.bundle_id": "8c0c3459…"},
 "params": {"prediction_id": "5a549ffa…-predict-2633f84b…", "bundle_id": "8c0c3459…", "score": "True_PROB",
            "previous_prediction_id": "5a549ffa…-predict-da72a70e…"},
 "metrics": {"psi": 0.26605, "ks": 0.18113, "wasserstein": 0.00385,
             "previous_psi": 0.17006, "previous_ks": 0.16781, "previous_wasserstein": 0.00323,
             "input.tables.orders.rows.before": 2984.0, "input.tables.orders.rows.after": 4000.0,
             "input.edges.orders_USER_ID_to_users_ID.before": 2984.0, "input.edges.orders_USER_ID_to_users_ID.after": 4000.0}}

The fit run in the same experiment carries the same kumo.bundle_id, so filtering on it returns the training run, every prediction run and every monitoring run of that model. kumo.kind separates them. Every run of a model has the same metric names, so predictions.True_PROB.q50 or predictions.predicted_positive_rate can be charted across runs.

• MLflow: one run per call in the configured experiment, nested under the caller's active run if there is one (the same rule as fits).
• Snowflake: one run named by the prediction id as an unquoted identifier (characters other than letters, digits and _ become _; e.g. EVALUATE_<ID>/DRIFT_<ID>). A quoted name broke snowflake-ml's artifact URI (Run "<id>" does not exist on a live account), so kumo.run_id keeps the exact id. kumo.* tags become params, as for fits. An existing run is never resumed: a taken name gets a suffix. If the process's ExperimentTracking instance is in use, the run is not recorded there, with a warning.
• run_name and tags apply to fit runs only. The tracker's run_id attribute is the most recent tracked run, fit or prediction.

7. Warehouse model monitors

docs/mlops.md ("Monitoring prediction runs") gives two recipes over the evaluated rows, exported to a warehouse table with evaluation.export(...). Neither platform monitors link prediction or ranking outputs.

• Databricks Data Quality Monitoring, inference log (databricks-sdk ≥ 0.68): export to a Unity Catalog Delta table, add a model_id column with the bundle id, then w.data_quality.create_monitor(Monitor(..., DataProfilingConfig(..., inference_log=InferenceLogConfig(problem_type, prediction_column="TARGET_PRED", label_column="TARGET", model_id_column="model_id", timestamp_column="EVALUATED_AT", granularities=[1 day])))) and w.data_quality.create_refresh(...). The inference log has no probability column, so True_PROB is profiled as an ordinary column.
• Snowflake Model Monitor: one CREATE MODEL MONITOR ... FUNCTION=predict SOURCE=<evaluated rows> BASELINE=<holdout evaluated rows> TIMESTAMP_COLUMN=EVALUATED_AT ID_COLUMNS=(ENTITY) PREDICTION_SCORE_COLUMNS=("True_PROB") ACTUAL_CLASS_COLUMNS=(TARGET), with the per-task column choices in the doc. The model must be a Snowflake Model Registry version (kumo.MLOps(registry=ModelOutputConfig(sf, location=...))). The baseline is holdout_df() with an EVALUATED_AT column added; the monitor needs TIMESTAMP_NTZ, so the source table declares EVALUATED_AT TIMESTAMP_NTZ (the export writes UTC). Verified live: on a binary classification registry version the monitor was created, refreshed and returned ROC_AUC, COUNT and drift metrics; True_PROB must be double-quoted (unquoted, Snowflake looks for TRUE_PROB).

8. Job templates

examples/
├── sdk-local/
└── mlops/                                   # new
    ├── README.md        # overview, flags, local run
    ├── fit_job.py       # inside kumo.mlops(): snapshot, training table, fit (stored, tracked, registered)
    ├── score_job.py     # load "champion" or --model-path, fresh or stored snapshot, trainer.predict
    ├── jobs.py          # flags, graph (users, items, orders), connectors, kumo.MLOps settings
    ├── databricks/  databricks.yml  README.md
    ├── sagemaker/   pipeline.py  README.md
    └── snowflake/   register_jobs.py  README.md

Platform	Train	Score	Model handoff
Databricks (Asset Bundle)	fit task: fit_job.py on a single-node job cluster (Databricks Runtime 17.3 LTS ML or later, GPU or CPU flavour; the runtime's torch, pyarrow (16 through 24 accepted) and MLflow are kept)	score task: daily score_job.py into a UC table (schedule created paused)	--registry mlflow: MLOps(tracking=MlflowTracking(), registry=ModelRegistryLocation("<catalog>.<schema>.users_churn", alias="champion")); scoring loads champion
SageMaker Pipelines (SDK v2)	TrainingStep	ProcessingStep	--registry mlflow against SageMaker managed MLflow
Snowflake ML Jobs + Tasks (experimental: preview APIs, not run)	weekly Task	daily Task into a table	--registry snowflake: MLOps(tracking=SnowflakeTracking(connector.session(), ...), registry=ModelOutputConfig(...))
Plain VM or workstation	python fit_job.py	python score_job.py	no registry: the fit prints model: models/<bundle id>; scoring loads --model-path models/<bundle id>

• The scripts use only the public API and run inside with kumo.mlops(cfg), so the snapshot, training table, model and each prediction run are stored under --artifact-root. Settings are flags (--artifact-root, --source, --catalog/--schema/--warehouse, --registry, --model-name, --model-path, --snapshot, --predictions-table, --table-prefix, --query, --num-experiments, --scratch-dir, --on-error); credentials come from the platform identity, and only MLflow's own settings (MLFLOW_TRACKING_URI, ...) stay environment variables.
• The scripts exit 0 on success and non-zero with a traceback on failure; the scheduler tracks its own runs. They build kumo.MLOps(on_error="raise") unless --on-error warn is passed, so a fit whose model is stored but not registered (or a run that is not tracked) fails the job instead of leaving the score step on the previous version. The Databricks bundle, the SageMaker pipeline and the Snowflake registration take the model name, the predictions table, a source-table prefix and on_error as variables or flags with defaults; the bundle's default root_path is under the deploying user's folder. Every score passes a 0.5 binary threshold unless the model carries scoring defaults, because the template query is binary.
• These templates were deployed and run live: both Databricks bundle jobs on a single-node cluster, and the SageMaker pipeline with both steps; each fit registered a new version and moved champion, and each score loaded it. The Databricks fit trained on a GPU (a GPU ML runtime on a g5.xlarge). test_job_templates.py checks the scripts locally. The Snowflake templates are not run.
• Databricks: each task is a spark_python_task whose parameters are the flags. A spark_python_task runs a script without __file__, so the scripts find their directory from sys.argv[0].

Supporting changes

• DiskGraph compute_col_stats_for_parquet (the existing Rust batch-prediction statistics call) gains keyword options whose defaults keep today's output byte-identical:

  • percentiles=True emits PERCENTILES (q01..q99) from each numerical column's t-digest.
  • summary_stats_by_class=False skips the per-class map, which dominates cost with millions of classes.
  • histograms=False skips the second read that builds histograms.
  • category_counts_limit=N keeps the top N CATEGORY_COUNTS and records CATEGORY_COUNTS_CAPPED.
  • max_per=[keys] computes statistics over each column's maximum within each run of equal key values. Runs are reduced while streaming each row group, and a run that crosses row groups or files is joined in order afterwards, so the result is exact. It needs numerical columns, histograms=False and summary_stats_by_class=False.

  Row-group partial statistics are merged in fixed chunks in file and row-group order, so the output does not depend on thread count or scheduling. The Rust compute_parquet_col_stats is replaced by compute_parquet_col_stats_with_options (ColStatsOptions::default() gives the same output).

• DiskGraph evaluate_prediction_parquet(prediction_paths, label_paths, spec_json) (new bp_evaluate.rs): the streaming join and metrics above. The spec's prediction_columns lists prediction columns copied next to each matched label (for grouped tasks, from the top-scored row), and constant_timestamps adds constant UTC timestamp columns (EVALUATED_AT, and a static query's prediction time as ANCHOR_TIMESTAMP). It reuses the projected row-group reader of the statistics call.

• Engine: kumo.engine._score_stats owns the score summaries for both the holdout and prediction runs; kumo.engine._monitoring owns EvaluationResult, DriftReport, evaluate, score_drift and monitoring-run records; kumo.engine.trainer._run_manifest owns building, writing, storing and tracking the prediction manifest. _tracking.log_run(tracking, hook, record) logs one standalone run through an optional hook. PredictionTable.source (pql or custom) is new.

• Artifact root: prediction joins the export kinds (kumo.MLOps(export=...)), so each run is stored under <root>/predictions/<id>/. GraphSnapshot._counts reads row and edge counts from DiskGraph's ingest_stats.json for input_changes.

• kumo-core: generate_train_table(..., validate_targets=True); False keeps outcome tables that could not train a model, so labels can be read. pquery._target_maximum(query, index_path, min|max) sits behind latest_temporal_timestamp; evaluate uses its minimum for matured.

• PredictiveQuery: one private _generate helper for training-table and label generation, and a private _label_parts context manager that keeps label parts on disk for the join.

• Docs: docs/mlops.md describes the holdout record (Models), prediction runs, prediction and monitoring tags (Experiment tracking), and a new "Monitoring prediction runs" section.

Table outputs, scheduled jobs and docs

• Table outputs append. PredictionOutputConfig.write_mode defaults to append when a destination is a warehouse table and to overwrite for file destinations; an explicit WriteMode is kept. A daily scoring job therefore keeps every day's rows instead of replacing the table. Table outputs (and evaluation.export to a table) carry PREDICTION_ID (the run id) and BUNDLE_ID (the model), so rows tie to a run and a model; the Databricks monitor recipe uses model_id_column="BUNDLE_ID" and no longer adds a column by hand. A table created before these columns existed takes one WriteMode.OVERWRITE (or an ALTER TABLE ... ADD COLUMNS) first.
• on_error="raise" (from Record fits in experiment tracking and publish models to registries #60) also covers prediction-run tracking, after the outputs and manifest are written, and monitoring-run tracking.
• Reported once (as in Record fits in experiment tracking and publish models to registries #60): in warn mode a prediction-run or monitoring-run tracking failure is one UserWarning with the full error, pointing at the predict, evaluate or score_drift call; with "raise", one RuntimeError and no warning before it. Both hooks run through the tracker's client and credentials (MlflowTracking(target=...), Record fits in experiment tracking and publish models to registries #60).
• Docs. docs/mlops.md opens with a quickstart that runs as is on CPU against a local MLflow store (install, snapshot, tracked and registered fit, promotion, batch prediction, drift, evaluation), and "Monitoring in a later job" is a recipe that finds the champion's prediction runs by tag, loads them and their snapshots from the artifact root (with the store per root: local, s3://, gs://, /Volumes/...) and passes both snapshots to score_drift. Both scripts were run against this tree.

File outputs, drift from the artifact root, SageMaker template

• File outputs per run. A file or object-store PredictionOutputConfig (S3, GCS, local) with the default write_mode now writes each run under <path>/<prediction id>/, so a daily job's second run no longer fails; WriteMode.OVERWRITE replaces <path> (on object stores a reader in between sees no _SUCCESS); an explicit APPEND to files stays unsupported. One story with tables: tables append with PREDICTION_ID/BUNDLE_ID, files get a per-run directory, OVERWRITE replaces.
• Checks before scoring. The destination check runs before any scorer starts (object-store sinks gained the pre-scoring check in kumo-core), so an occupied destination fails before the work is done; the error keeps the full URI (s3://…).
• Drift from the artifact root. With MLOps on, score_drift reads the training and scoring snapshots' col_stats.json and ingest_stats.json from <root>/snapshots/… when the snapshots are not in the local store, so a monitoring job on another machine gets feature and input drift without loading snapshots; otherwise the reasons entry names the argument to pass.
• GPU drivers. On a plain GPU VM or a SageMaker image (not on Databricks Runtime 17.3 LTS ML, whose torch is kept), the quickstart says to install a torch build matching the NVIDIA driver first (cu126: driver ≥ 525.60 via CUDA 12 minor-version compatibility; PyPI's default cu130: ≥ 580; SageMaker's default hosting image has 470), with a link to NVIDIA's compatibility guide. No dependency pins change.
• SageMaker template. Ships examples/mlops/sagemaker/Dockerfile (cu126 GPU, or --build-arg TORCH=cpu), separate --fit-instance-type / --score-instance-type, and --alias instead of a fixed champion; the README notes that the GPU Predict step is a Processing job needing a GPU processing quota, and that pipeline names may not contain -aws-. The new flags and this Dockerfile have not yet run on SageMaker.

Back-compat

This PR only records and exposes information. It adds no new checks or errors on any existing path.

Situation	Behaviour
New fit whose published model has a holdout	_provenance.json gains holdout. Every other bundle file is unchanged, and _provenance.json stays out of the bundle id.
Published model is refit_full, or a link-prediction fit	No holdout record.
Holdout statistics fail	A warning; the model manifest is written without holdout. The fit is unaffected.
Bundle written before this PR	Loads and scores as before; its manifest has no holdout, and a prediction run embeds it as is (or as an id stub if it has no manifest).
New prediction run, MLOps off	Manifest written in the run directory. Predictions and exported data files are byte-identical to before.
New prediction run, MLOps on	As above, plus <root>/predictions/<id>/ after the output_config export (skipped with export lacking prediction).
Export to a Parquet directory (file, object store, volume or stage sink)	The manifest is written before _SUCCESS. Data files are unchanged.
Export to a warehouse table, a single file or a CSV directory	No sidecar. The run's manifest still lists the destination.
Output statistics fail	A warning; the manifest is written with only row counts in stats.
Destinations cannot be described	A warning; outputs is null. Exports go ahead.
The output_config export raises	The error is raised as before. The manifest is still written with status: failed and only the completed exports, the run is tracked as failed, and nothing is stored in the root.
Manifest build or write fails	A warning; no manifest and no sidecar. The result and exports are as before.
Prediction manifest unreadable, or a newer schema_version	manifest reads as None, with a warning.
Manifest read by an older engine	Nothing in an older engine reads it; the leading underscore keeps data readers from listing it.
PredictionResult from a run before this PR	manifest, bundle_id and snapshot_id are None; id is the run id. export() works as before. PredictionResult(selection, artifact_type) still constructs positionally.
predict() with tracking on	One run per call if the tracker has log_prediction_run; user Tracking implementations with fit hooks only still satisfy the protocol and get no call. Fit runs are unchanged.
Tracker error or hang	One warning; the prediction or monitoring result is unaffected; the call waits at most 30 s. Non-zero ranks log nothing.
evaluate / score_drift (new)	Return a result. A step that fails leaves its field None (or metrics empty) with a reason. evaluate raises only on a trainer that was never fitted or loaded; an immature run returns its metrics with matured=False and a UserWarning.
EvaluationResult.export without evaluated rows	Raises ValueError with the reasons.
EvaluationResult.load_from_connector (new)	TypeError for a table destination; FileNotFoundError without _evaluation.json (e.g. an export to a table or a CSV) or _SUCCESS.
PredictionResult.load_from_connector (new)	TypeError for a connector whose sink does not hold files at the path; FileNotFoundError without _SUCCESS. A completed directory without a manifest loads, with no ids or anchors.
Existing callers of diskgraph.compute_col_stats_for_parquet	Same positional signature; the new options are keyword-only and off by default. Default output is byte-identical (a test pins its SHA-256).
generate_train_table callers	validate_targets defaults to True, today's behaviour.

Performance

Measured on CPU with 8 rayon threads, synthetic Parquet on local NVMe, three runs each in a fresh process; peak RSS is the increase over the process after imports. The box was shared, so times are upper bounds.

Score statistics (compute_col_stats_for_parquet with the options above):

Statistics	Rows	Time	Peak RSS increase
Holdout, binary (True_PROB, TARGET), 1 file	10M / 50M	0.40–0.47 s / 1.46–1.71 s	24–27 MB / 28–32 MB
Holdout, regression (TARGET_PRED, TARGET), 1 file	10M / 50M	0.51–0.60 s / 1.84–2.08 s	33–35 MB / 42–43 MB
Holdout, multiclass top_score (×10 classes, unsorted, predictions split across row groups)	10M / 50M	0.14–0.21 s / 0.52–0.60 s	36–37 MB / 44–48 MB
Prediction output, binary, 8 files	10M / 50M	0.42–0.52 s / 1.48–1.72 s	26–27 MB / 34–37 MB
Prediction output, link prediction top_score (×10, unsorted), 8 files	10M / 50M	0.12–0.18 s / 0.42–0.53 s	36 MB / 43–45 MB

Memory stays flat with row count (about one row group of the needed columns per thread). The holdout is capped at 1 GiB in memory (about 30M binary rows), so 50M is beyond what a fit writes. Inside a full Trainer.predict (A10G GPU scoring, export to a local Parquet directory), the statistics took 0.16–0.18% of 234–275 s for 10M binary rows and 0.16–0.24% of 68–74 s for 10M link-prediction rows; the rest of the manifest adds at most 0.01 s.

evaluate's join and metrics, against a pandas baseline (concat of label parts, merge, per-row Python ranking); one anchor, 90% of predictions labelled, link prediction with k = 10 and 1–5 targets per entity; label generation not included. Measured without the copied prediction and timestamp columns, which add one column copy per matched row.

Output	Rows	pandas	DiskGraph
Binary	10M	9.6 s, 2.3 GB	2.1 s, 0.8 GB
Binary	100M	96 s, 14.9 GB	27 s, 1.9 GB
Link prediction	10M (1M × 10)	34.9 s, 2.1 GB	1.2 s, 0.5 GB
Link prediction	100M (10M × 10)	343 s, 16.8 GB	16.1 s, 0.8 GB

Metrics agree between the two (AUROC/AP to 1e-8, MAP@10 and precision@10 to 1e-12). Memory is bounded by the partitions in flight plus 8 bytes per labelled row for the exact binary ranking. score_drift works on 101 knots per side from manifest metadata.

Known limitations

• matured is a proxy: it uses the earliest latest-event time across the query's target tables.
• Binary AUROC/AP hold one 8-byte key per labelled row for the exact global ranking (11 GB at 1.4B rows).
• Link-prediction evaluated rows carry no prediction columns (several classes per entity), and neither warehouse monitor supports link prediction or ranking.
• top_score treats each run of consecutive rows with equal entity and anchor as one prediction; the scorer writes each prediction's rows together.
• Input drift covers table row and edge counts only; per-feature distribution drift between the training and scoring snapshots is not computed.
• Without MLOps, a run exported only to a warehouse table keeps its manifest only in local scratch. On an ephemeral machine, turn MLOps on to keep a copy under the artifact root.
• MLflow and Snowflake time a tracked prediction run from when it is logged; the run's own start and finish are in the manifest artifact.
• The job scripts run end to end locally in test_job_templates.py; the platform wiring (Databricks bundle, SageMaker pipeline, Snowflake tasks) is not run in CI. The Snowflake template is experimental (preview APIs, never run), and SYSTEM$EXECUTE_ML_JOB, which MLJobDefinition.to_sql generates, is not a documented public API. Use a standard (non-ML) Databricks runtime: the ML runtime's preinstalled torch, torchvision and pyOpenSSL conflict with kumo-engine's pins.

Testing Done

• Live Snowflake account (synthetic users/orders): SnowflakeTracking(connector=...) recorded fit, prediction, drift and evaluation runs in Snowflake Experiment Tracking (after the run-name fix); the Model Monitor recipe above refreshed and returned metrics. test_snowflake_tracking and test_tracking_contract: 27 passed, 8 skipped; the experiment fake now refuses artifacts for a run whose name needs quoting.

Run at this PR's head, one file at a time, with the DiskGraph extension rebuilt from it (CPU):

• docs/development/testing.md now says engine and core tests need DiskGraph (_c.abi3.so and bin/ingest) built from the same tree, and that a stale build fails with missing projections= / percentiles= arguments; rebuild with make install-diskgraph.

• After the rebase onto main (Push down feature projections through DiskGraph and Kumo loaders #46, Accept pyarrow 16 through 24 in DiskGraph; explain custom S3 CA bundles below pyarrow 21 #124) and Record fits in experiment tracking and publish models to registries #60's credentials change, rerun at this head with DiskGraph rebuilt from it: test_mlflow_target 13 (2 skipped: sagemaker-mlflow not installed), test_mlflow_tracking 18, test_mlflow_model 18, test_snowflake_tracking 14 (8 skipped), test_tracking_contract 13, test_trainer_registry 37, test_monitoring 16, test_semantic_carry_forward 20, test_holdout_df 11, test_bundle_provenance 8.

Test file	Passed
test_custom_source	20
test_bundle_provenance	8
test_custom_sql	11
test_filestore_sink	6
test_fit_events	4
test_holdout_df	11
test_job_templates	11
test_mlflow_model	12
test_mlflow_tracking	11
test_mlops	4
test_monitoring	17
test_optional_imports	51
test_prediction_completion	14
test_prediction_manifest	11
test_prediction_workers	55
test_predictive_query	67
test_query_from_snapshot	8
test_score_stats	3
test_scoring_defaults	7
test_semantic_carry_forward	20
test_snapshot_export	36
test_snapshot_manifest	21
test_snowflake_connector	8
test_snowflake_tracking	13 (8 skipped: checks against an installed snowflake-ml)
test_tracking_contract	13
test_trainer_registry	37
test_training_explanations	16
test_training_table_export	18
test_warm_start	6

• kumo-core: 522 passed, 1 skipped. kumo-api: 12 passed. The changed tests in kumo-snowflake (32), kumo-databricks (56) and kumo-bigquery (7) pass.
• Rust: bp_evaluate tests (3 passed): binary AUROC equals a brute-force pair count in one and several spilled partitions, and each evaluated row carries its own entity's score; the AP step definition; regression with a constant timestamp, multiclass top class and score, link-prediction metrics. bp_col_stats tests cover default output unchanged, each option additive, percentile rank error, the category cap flag, max_per across row groups and files, and identical output on 1, 2, 5 and 16 threads. cargo clippy clean.
• What the engine tests check:
  • test_score_stats, test_holdout_df, test_bundle_provenance, test_scoring_defaults: top scores of unsorted multiclass/multilabel holdouts across row groups and files; non-finite scores not counted; the holdout record matches the holdout frame and is left out when embedded; refit_full has none; a statistics failure keeps the rest of the manifest; a derived model keeps its parent's record.
  • test_prediction_manifest: every manifest field of a run exported to two directories; sidecars equal the run's manifest; percentiles near numpy's; stub parents; the artifact-root store after the export; no sidecar for warehouse or CSV destinations; scoring sources; failures that only warn and leave exported bytes identical; a failing export that still writes and tracks a failed run; a tracked run whose frame equals an untracked run's.
  • test_mlflow_tracking: a link-prediction monitoring run logs map@10/precision@10 as map_at_10/precision_at_10.
  • test_monitoring (feature drift): a toy graph with a numeric mean shift (standardized shift > 2) and a categorical mix shift (PSI > 0.5) is detected, identical snapshots give 0, a missing column and absent statistics leave reasons, features are found through the manifests, and a 1,200-column report logs at most 40 feature metrics (under 100 in total) plus feature_drift.json.
  • test_query_from_snapshot, test_mlflow_model (skipped entities): 4 of 14 requested ids absent from the snapshot give counts 14/10/4, the four sample ids, not in snapshot, the warning text and the two tracking metrics; a generated population records 0; the pyfunc model warns skipped 1 of 2.
  • test_monitoring: a matured binary run's labels against a pandas recomputation, emitted scores equal to the run's, AUROC against scikit-learn, one EVALUATED_AT within the call; an immature older snapshot; a run with no labelled row (actuals is None, export raises); several anchors; regression and static queries; EvaluationResult.export; metrics equal kumo-core BinaryAUROC / BinaryAveragePrecision and PyG LinkPredMAP / LinkPredPrecision; PSI, KS within 0.02 and Wasserstein within 2% of exact on 100k-sample normals; drift against holdout, previous run and snapshot counts; tracked monitoring runs; a run loaded from its export (no parts fetched by score_drift, same results, parts removed with the result, no _SUCCESS refused).
  • test_tracking_contract, test_mlflow_tracking, test_snowflake_tracking: the prediction record's params, metrics, tags, inputs and artifact, fed to both trackers and read back; nesting under an active MLflow run; a monitoring run tags its prediction run kumo.monitored; Snowflake never resumes a run and skips an in-use instance.
  • test_job_templates: fit_job.py then score_job.py end to end on local Parquet, storing predictions/<id>/ with the fitted model's id; both scripts import without __file__; every example file compiles.
• Live: a Databricks Data Quality monitor on evaluated rows produced empty metric tables when keyed on ANCHOR_TIMESTAMP (59 days old), and 60 profile rows and 48 drift rows on rows timestamped 1–5 days old.
• The examples above were captured from a CPU run of this head on synthetic data.

Checklist

• No breaking changes.
• Docs updated (docs/mlops.md, examples/mlops).
• Unit tests added.

🤖 Generated with Claude Code

[Vedaant-J]

Change map

File	What changed	Why
docs/development/testing.md	engine and core tests need DiskGraph (_c.abi3.so and bin/ingest) built from the same tree; a stale build fails with missing projections=/percentiles=; rebuild with make install-diskgraph	the series needs the DiskGraph APIs from #46 and its own column statistics
examples/mlops/README.md	New: overview of the two job scripts, the platform table, every flag, how to run locally, and what a scheduler sees (exit codes, stderr); notes that a scored run is stored under predictions/<id>/; the --table-prefix and --on-error flags, and that a failed registration or tracking call fails the job	Entry point for teams wiring kumo-engine into a scheduler
examples/mlops/databricks/README.md	New: deploy and run the Asset Bundle; artifact root on a UC volume, MLflow and UC registry settings, --prefer-binary requirements, how the bundle was verified live; Databricks Runtime 17.3 LTS ML or later (GPU or CPU flavour), whose torch, torchvision, MLflow and pyarrow are kept, with no torch or pyOpenSSL pins; the bundle variables with defaults (model_name, predictions_table, table_prefix, on_error) and strict mode	Databricks setup that the bundle cannot express
examples/mlops/databricks/databricks.yml	New: Asset Bundle with a train job and a daily score job (schedule created paused), each a single-node spark_python_task whose parameters are the flags (--registry mlflow); variables model_name, predictions_table, table_prefix and on_error (default raise) with defaults; root_path under the deploying user's folder; spark_version/node_type_id describe the GPU ML runtime and GPU node default and the CPU alternative; runtime examples are 17.3 LTS ML	Run fit and scoring as Databricks Jobs with the model handed over by the champion alias
examples/mlops/fit_job.py	New: builds the graph, snapshot, training table and model inside with kumo.mlops(...); prints model: models/<bundle id>; passes --table-prefix to the graph	Scheduled training that stores, tracks and registers every artifact through the public API
examples/mlops/jobs.py	New: shared flags, source connectors (Databricks, Snowflake, Parquet directory), the users/items/orders graph, kumo.MLOps settings per --registry, model loading by alias or path, prediction output config; --alias replaces the fixed champion; --on-error {raise,warn} (default raise) passed to kumo.MLOps; --table-prefix for the source table or file names	One module both scripts share, to be replaced by the user's own
examples/mlops/sagemaker/Dockerfile	New: GPU (cu126) or CPU image with kumo-engine, the mlflow and s3 extras and sagemaker-mlflow, installed from a wheel directory; installs kumo-engine[mlflow,sagemaker]	The pipeline's image, previously described but not shipped
examples/mlops/sagemaker/README.md	New: build and run the pipeline against SageMaker managed MLflow; required arguments; how the pipeline was verified live; the mlflow-skinny client; the Dockerfile, separate fit and score instance types, --alias, the GPU processing quota and the -aws- naming note; the --model-name, --table-prefix and --on-error flags; SDK v2 reads the region from AWS_DEFAULT_REGION; image contents named as kumo-engine[mlflow,sagemaker]	SageMaker setup
examples/mlops/sagemaker/pipeline.py	New: SageMaker Pipelines definition with a TrainingStep (fit) and a ProcessingStep (score) passing the flags; --instance-type required; --fit-instance-type, --score-instance-type and --alias; --model-name, --table-prefix and --on-error (default raise) passed to both steps	Run fit and scoring as SageMaker pipeline steps
examples/mlops/score_job.py	New: loads champion or --model-path, scores a fresh or stored snapshot inside with kumo.mlops(...) with a 0.5 threshold unless the model has scoring defaults; prints the row count; passes --table-prefix to the graph	Scheduled scoring whose prediction run is stored under the artifact root and tracked
examples/mlops/snowflake/README.md	New: register the jobs and Tasks; experimental status, preview APIs, artifact root reachable through external access; the --model-name, --predictions-table, --table-prefix and --on-error flags	Snowflake setup and its caveats
examples/mlops/snowflake/register_jobs.py	New: registers train and score MLJobDefinitions and creates weekly and daily Tasks that execute each job's to_sql(job_args=<flags>); --model-name, --predictions-table, --table-prefix and --on-error (default raise)	Run fit and scoring as Snowflake ML Jobs on a schedule
packages/diskgraph/diskgraph-ingest/src/col_stats.rs	FileColStats::numerical_quantiles(col_idx, qs) reads quantiles from a numerical column's existing t-digest	Percentiles for score statistics without a second sketch
packages/diskgraph/diskgraph-pybind/python/diskgraph/__init__.py	Exports evaluate_prediction_parquet	Public Python binding of the evaluation join
packages/diskgraph/diskgraph-pybind/python/diskgraph/tt_artifacts.py	The prediction-table builder no longer logs skipped rows as training rows; its caller reports them	One accurate warning per prediction
packages/diskgraph/diskgraph-pybind/python/tests/test_col_stats_options.py	New: the statistics options driven from Python; default output pinned by SHA-256	Python-level contract and byte-identical default
packages/diskgraph/diskgraph-pybind/python/tests/test_tt_artifacts_row_warnings.py	The training builder still logs skipped rows; the prediction builder leaves it to the caller	Warning ownership
packages/diskgraph/diskgraph-pybind/src/bp_col_stats.rs	ColStatsOptions (percentiles, summary_stats_by_class, histograms, category_counts_limit, max_per) and compute_parquet_col_stats_with_options replacing compute_parquet_col_stats; max_per reduces runs of equal keys per row group and stitches runs across row groups and files; row-group partials merged in fixed chunks in work order; row_group_batches shared as pub(crate); tests	Holdout and prediction score summaries (percentiles, per-prediction top score) at Parquet scale, deterministic across thread counts, with defaults unchanged
packages/diskgraph/diskgraph-pybind/src/bp_evaluate.rs	New: EvaluateSpec and evaluate_prediction_parquet: streaming hash-partitioned join of prediction Parquet with label parts (Arrow IPC spill, at most 8 joins in flight), fit-time metrics (exact AUROC/AP, MAE/RMSE, top-1 accuracy, MAP@k/precision@k), evaluated rows with label columns, the matched prediction's columns and constant timestamp columns; tests	evaluate with memory bounded by partition size, and the rows warehouse monitors read
packages/diskgraph/diskgraph-pybind/src/lib.rs	compute_col_stats_for_parquet takes keyword-only percentiles, summary_stats_by_class, histograms, category_counts_limit, max_per (defaults keep today's output); new evaluate_prediction_parquet(prediction_paths, label_paths, spec_json) binding; module registration	Python access to the new options and the join
packages/kumo-core/src/kumo_core/object_store_sink.py	Object-store sinks check a prediction destination before scoring (occupied unless OVERWRITE); errors keep the full URI	A daily run fails before scoring, not after
packages/kumo-core/src/kumo_core/pquery.py	Private _target_maximum(query, index_path, min/max) behind latest_temporal_timestamp	evaluate's matured flag uses the minimum of the target tables' latest events
packages/kumo-core/src/kumo_core/prediction.py	GeneratedPredictionTable.row_failures carries DiskGraph's report for supplied rows	Skipped entities reach the engine
packages/kumo-core/src/kumo_core/training/table.py	generate_train_table(..., validate_targets=True); False skips classification-target validation and the multilabel TOP K check	Read outcome labels that could not train a model
packages/kumo-core/tests/test_pquery.py	_target_maximum with min is order independent	Covers the new helper
packages/kumo-core/tests/test_train_table.py	validate_targets=False keeps degenerate labels and their staging	Covers the new keyword
packages/kumo-databricks/tests/test_write.py	An append stamps the extra id columns on each file	Table outputs append with run ids
packages/kumo-engine/docs/mlops.md	Holdout record under Models; new Prediction runs section; prediction and monitoring tags and hooks under Experiment tracking; new Monitoring prediction runs section (evaluate, score drift, tracking, a later job, Databricks and Snowflake recipes keyed on EVALUATED_AT); prediction in the export kinds; the Snowflake monitor recipe notes the registered version is a carrier the monitor never calls; skipped entities under Prediction runs; input-feature drift under Score drift; the Snowflake monitor recipe notes the recorded task type; a quickstart, a monitoring-job recipe with the store per artifact root, table write modes and id columns; one output story for tables (append) and files (per-run directory), drift from the artifact root, GPU driver guidance; an immature evaluate returns its metrics with matured=False and a warning; the job templates use on_error="raise"; the quickstart skips the torch step on Databricks 17.3 ML, installs a wheel-directory build by path and names the s3/gcs extras; new "Exporting and loading artifacts" table (artifact, export call, load call; fit(export_config=) and predict(output_config=) as shorthand; reuse of a complete copy); EvaluationResult.load_from_connector; the later-job recipe stops cleanly when the champion has no prediction runs; feature drift compares a column by the stype the model encodes it with; Snowflake run names as identifiers; Model Monitor recipe quotes "True_PROB" and is marked verified live; the Snowflake monitor baseline also works from a model loaded with Trainer.load_from_connector	User docs for everything this PR adds
packages/kumo-engine/kumo/engine/__init__.py	Exports PredictionManifest, EvaluationResult, DriftReport	Public API
packages/kumo-engine/kumo/engine/_manifest.py	HoldoutStats and ModelManifest.holdout (omitted when embedded); PredictionManifest with PredictionParents, PredictionTableRecord, RunScoring, OutputDestination, PredictionOutput, RunTimings, ArtifactStats; KINDS gains prediction; PredictionTableRecord.requested_rows, scored_rows, skipped_rows, skipped_sample, skip_reason	Typed records of the holdout baseline and of a prediction run
packages/kumo-engine/kumo/engine/_mlops.py	Docstrings name the prediction kind and stored prediction runs	Predictions join the artifact root
packages/kumo-engine/kumo/engine/_monitoring.py	New: EvaluationResult (+ export), DriftReport, evaluate, score_drift, track; PSI, KS and Wasserstein from percentile curves; snapshot count changes; reasons instead of exceptions; EVALUATED_AT per call; actuals is None when no row has a label; the join spill directory is removed on failure too; DriftReport.features: per used column, deltas of the build-time col_stats.json of the training and scoring snapshots (numeric stats and standardized mean shift; categorical PSI, JS and capped flag; timestamp range), feature_metrics (bounded top-N summary) and the feature_drift.json artifact; monitoring runs are named evaluate-<id>/drift-<id>; tracking failures raise with on_error="raise"; evaluation table exports append with the id columns; score_drift reads both snapshots' col_stats.json and ingest_stats.json from the artifact root when they are not local, and the reason names the argument to pass otherwise; a monitoring-record failure warns once (warn_failure, not under on_error="raise") and the end-of-call report points at the caller; evaluate warns when the run is not mature; EvaluationResult.export writes _evaluation.json to a Parquet directory and keys the export on the evaluated files; EvaluationResult.load_from_connector loads it back with its rows	Realized performance and score drift for prediction runs, logged as monitoring runs
packages/kumo-engine/kumo/engine/_score_stats.py	New: SCORE_COLUMNS per task, parquet_stats, top_score_stats, positive_rate over DiskGraph compute_col_stats_for_parquet	One score summary shared by the holdout baseline and prediction runs
packages/kumo-engine/kumo/engine/_tracking.py	prediction_record (manifest to RunRecord), track_prediction, record_failure, log_run (one standalone run on a tracking thread, rank 0, up to 30 s, never raises); fit_inputs reads a prediction manifest's scoring snapshot; predictions.skipped_rows and predictions.skipped_fraction on the prediction run; track_prediction / log_run log the full error, traceback only at DEBUG	Prediction and monitoring runs share the fit tracking code and record the same names in every tracker
packages/kumo-engine/kumo/engine/connector/file.py	The local-file destination check accepts an occupied directory under OVERWRITE	OVERWRITE replaces a file destination
packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py	pyfunc predict passes the active MLOps settings to Trainer._predict, so a pyfunc call inside an MLOps block is still a tracked prediction run	_predict replaced predict for the input-order fix in #60
packages/kumo-engine/kumo/engine/integrations/mlflow.py	start names a run after its record when set; log_prediction_run; log_monitoring tags the monitored prediction run kumo.monitored; both hooks go through the tracker's client and credentials scope (#60)	MLflow implementation of the two hooks
packages/kumo-engine/kumo/engine/integrations/snowflake.py	start names a run after its record as an unquoted identifier (_run_identifier) when set; _new_run_name takes the name; log_prediction_run, log_monitoring; warnings name the run kind	Snowflake implementation of the two hooks without ever resuming a run; a quoted name breaks snowflake-ml's artifact URI (seen live)
packages/kumo-engine/kumo/engine/pquery/prediction_table.py	PredictionTable.source: pql or custom; PredictionTable.skipped_rows, skipped_sample, skip_reason, requested_rows from DiskGraph's row_failures, with one warning naming the count, fraction and reason	The prediction manifest records where the table came from
packages/kumo-engine/kumo/engine/pquery/predictive_query.py	Private _generate shared by training-table and label generation; private _label_parts context manager yielding label parts at given anchors, removed on exit	Label regeneration for evaluate without a second workflow
packages/kumo-engine/kumo/engine/trainer/_provenance.py	write_provenance records holdout via holdout_baseline (warns on failure); record_derivation keeps the parent's holdout	Holdout baseline in every model manifest
packages/kumo-engine/kumo/engine/trainer/_run_manifest.py	New: start (build the manifest with output statistics), finish (export, record completed outputs, write, store under <root>/predictions/<id>/, track), store, output_records, destination_record, output_stats; records the skipped-entity fields; prediction-run tracking failures raise with on_error="raise" after outputs are written; outputs recorded at their per-run location; the prediction-run tracking failure is reported at the caller of predict	Owns the prediction manifest's lifecycle inside predict
packages/kumo-engine/kumo/engine/trainer/output_config.py	write_mode defaults to append for tables and overwrite for file destinations; an explicit mode is kept; file destinations default to a per-run directory <path>/<prediction id>/; OVERWRITE replaces the path; OVERWRITE keeps a directory that already holds this export complete, so the sink reuses it	No silent replace of a scoring table
packages/kumo-engine/kumo/engine/trainer/trainer.py	predict(mlops=) with scoring sources; Trainer.evaluate, Trainer.score_drift; TrainingResult.evaluate / score_drift; PredictionResult.id, manifest, bundle_id, snapshot_id, load_from_connector (_Export downloads parts lazily); exports write the manifest sidecar into Parquet directories; PredictionResult.requested_rows, scored_rows, skipped_rows; table exports carry PREDICTION_ID and BUNDLE_ID; a loaded bundle keeps its holdout alongside its scoring defaults (#50)	Public entry points for prediction records and monitoring
packages/kumo-engine/test/test_bundle_provenance.py	Stub fit passes the holdout arguments; a statistics failure keeps the manifest	Failure path of the holdout record
packages/kumo-engine/test/test_filestore_sink.py	Expects per-run directories and the pre-scoring check	File outputs per run
packages/kumo-engine/test/test_holdout_df.py	Holdout record matches the holdout frame per fit mode, refit_full has none, truncation flag, embedded copies omit it, loaded bundles keep it; a loaded or fetched model keeps the holdout record	Holdout record contract
packages/kumo-engine/test/test_job_templates.py	New: example files compile; fit_job.py then score_job.py end to end on local Parquet with predictions/<id>/ stored; scripts import without __file__; the local run uses --table-prefix; --on-error defaults to raise; the SageMaker and Snowflake templates pass the model name, prefix and --on-error	The examples keep working
packages/kumo-engine/test/test_mlflow_model.py	The pyfunc model warns with the skipped count when an entity is not in the snapshot	Pyfunc path of skipped entities
packages/kumo-engine/test/test_mlflow_tracking.py	A monitoring run is new and tags its prediction run; a prediction run nests under an active run or lands top level; link-prediction monitoring metrics logged as map_at_10/precision_at_10	MLflow hook behaviour
packages/kumo-engine/test/test_mlops.py	export=True includes prediction	Export kinds contract
packages/kumo-engine/test/test_monitoring.py	New: evaluate (matured, several anchors, regression, static, export, EVALUATED_AT), metric definitions, PSI/KS/Wasserstein accuracy, drift, tracked monitoring, a run loaded from its export; a run with no labelled row has no actuals and refuses export; feature drift on a toy graph with a numeric mean shift and a categorical mix shift, a missing column and no statistics as reasons, the metric bound, and features found through the manifests; evaluation table export with id columns; monitoring run names; on_error for monitoring; drift statistics read from the artifact root in an empty scratch dir, and the no-root reason; an immature run returns its metrics and warns; an exported evaluation loads back equal with the same rows; re-exporting an evaluation or a prediction run to the same directory rewrites nothing	Monitoring contract
packages/kumo-engine/test/test_prediction_completion.py	Expects per-run directories for file outputs	File outputs per run
packages/kumo-engine/test/test_prediction_manifest.py	New: manifest fields, sidecars, statistics, stub parents, artifact-root store, scoring sources, failure paths (outputs compared as frames to float tolerance, so GPU scoring passes), failed exports, tracking; the skipped-entity fields of a generated population	Prediction manifest contract
packages/kumo-engine/test/test_prediction_output_columns.py	Expects per-run directories for file outputs	File outputs per run
packages/kumo-engine/test/test_prediction_output_preflight.py	An occupied destination fails before the scorer runs, with the full URI; OVERWRITE passes the check	Destination checks before scoring
packages/kumo-engine/test/test_prediction_table_output.py	New: two runs into one table keep both runs' rows by PREDICTION_ID with BUNDLE_ID; overwrite replaces them; file defaults	Table outputs
packages/kumo-engine/test/test_predictive_query.py	Generation fakes accept keyword arguments; the custom-table stub carries a row_failures report	generate_train_table now takes validate_targets
packages/kumo-engine/test/test_query_from_snapshot.py	Entities absent from the snapshot: counts, sample ids, reason, warning text, tracking metrics, and 0 for a generated population	Skipped entities
packages/kumo-engine/test/test_score_stats.py	New: top scores of unsorted holdouts across row groups and files; non-finite scores; binary and regression summaries	Score summary contract
packages/kumo-engine/test/test_scoring_defaults.py	A derived model keeps its parent's holdout	Derived-model baseline
packages/kumo-engine/test/test_semantic_carry_forward.py	Expects the PREDICTION_ID column in the sink call	Id columns on table outputs
packages/kumo-engine/test/test_snowflake_tracking.py	Prediction runs skip an in-use instance and are never resumed (names TRAINER_PREDICT_1, TRAINER_PREDICT_1_1); updated warning text	Snowflake hook behaviour
packages/kumo-engine/test/_snowflake_experiment_fake.py	FakeExperimentTracking.log_artifact refuses artifacts for a run whose name needs quoting, as Snowflake's artifact URI does	The fake now fails the way the live account did
packages/kumo-engine/test/test_tracking_contract.py	PREDICTION_MANIFEST fixture; prediction_record fields; both trackers record what the record says; no manifest or no hook records nothing; Adapter.run_name compares names in the tracker's form	Tracking record contract
packages/kumo-snowflake/tests/test_sink.py	An append loads files stamped with the extra id columns	Table outputs append with run ids
tests/public/Dockerfile.dockerignore	Admits kumo_core/artifact_export.py	The public connector image runs the table-append tests, whose sinks import it

Requirement: every prediction run must be traceable to its model, snapshot and source versions, carry score statistics comparable with the model's holdout, be measurable against labels that arrive later from any job, and appear in the experiment tracker next to the model's training runs. Existing behaviour: snapshot, training-table and model manifests, the artifact root and fit tracking exist (#50, #60), but a prediction run leaves only scorer files, bundles carry no score baseline, labels are generated only for training, and trackers record only fits. Smallest change: one DiskGraph statistics call with options and one DiskGraph join for scale; a typed prediction manifest written inside predict; a holdout record on the existing model manifest; evaluate and score_drift that read manifests and return results with reasons instead of raising; two optional tracker hooks that reuse the fit paths; job templates that use only the public API. Verification: the engine, kumo-core, connector and Rust suites listed in the PR body, a real CPU fit, predict, evaluate and score_drift with MLflow tracking on synthetic data, and a live Databricks Data Quality monitor on evaluated rows. Stopping condition: every predict writes a manifest (and stores and tracks it with MLOps on), every fitted model with a holdout records its baseline, and evaluate and score_drift return results from a run in-process or loaded from its export, with all existing calls unchanged.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

• Last reviewed commit: cfdd76f
• Latest commit, not reviewed: 4e16e8e

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

Walkthrough

The PR adds configurable Parquet column statistics and a Parquet prediction-evaluation API to DiskGraph. Kumo Engine adds holdout statistics to model provenance, prediction manifests and tracking, realized-performance evaluation, and score- and feature-drift reports. It also adds shared fit and score job scripts, deployment templates for Databricks, SageMaker, and Snowflake, and related documentation. Core query helpers add optional target-validation control and reusable label-generation support.

Sequence Diagram(s)

sequenceDiagram
  participant Trainer
  participant RunManifest
  participant DiskGraphStats
  participant PredictionExport
  participant Tracking
  Trainer->>RunManifest: start prediction manifest
  RunManifest->>DiskGraphStats: compute prediction score summaries
  DiskGraphStats-->>RunManifest: return statistics
  Trainer->>PredictionExport: export prediction outputs
  PredictionExport-->>RunManifest: return output records
  RunManifest->>Tracking: log prediction run

Loading

Priority: Normal

Merge Risk: Moderate · previous review cfdd7

Resolve the open model-selection, scheduling, and permissions concerns before merging. Skipped-entity diagnostics and the all-skipped documentation also need correction.

Pre-merge checks: 1 · 1 · 1

Failed checks: 1 warning · 1 inconclusive

Check name	Status	Explanation	Resolution
Pr Size		The pull request changes 53 files, based on the authoritative review diff. This exceeds the limit of 40 files.	Split the pull request into smaller, reviewable units.
Change Map Design Review		The required PR conversation comment beginning with <!-- kumo-pr-change-map --> is not available. The supplied context contains only a summary of that comment, not its Change map, rationale, verific…	Provide the author-maintained PR conversation comment beginning with <!-- kumo-pr-change-map -->, including the complete Change map and its evidence, so each proposed change can be reviewed for necessity, abstraction layer, behavior prese…

Passed checks: 1 passed

Check name	Status	Explanation
Public Resource Access		PASS. The PR adds external references in the Databricks, SageMaker, and Snowflake example READMEs, plus the public PyTorch package index. Anonymous GET requests returned HTTP 200 for all 15 URLs, with…

Full details: Public Resource Access

Explanation

PASS. The PR adds external references in the Databricks, SageMaker, and Snowflake example READMEs, plus the public PyTorch package index. Anonymous GET requests returned HTTP 200 for all 15 URLs, with expected documentation content. The added relative repository links target files present in the PR. No added or modified external reference was confirmed inaccessible.

Full details: Change Map Design Review

Explanation

The required PR conversation comment beginning with <!-- kumo-pr-change-map --> is not available. The supplied context contains only a summary of that comment, not its Change map, rationale, verification, or stopping condition. Repository inspection found the marker only in the template and review guidance, and the reviewed commit has no map in its message. I cannot assess the proposed design against the required comment.

Resolution

Provide the author-maintained PR conversation comment beginning with <!-- kumo-pr-change-map -->, including the complete Change map and its evidence, so each proposed change can be reviewed for necessity, abstraction layer, behavior preservation, verification, and stopping condition.

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 2

Note

Quiet mode is enabled, so only the most important comments were posted inline. Other review comments are grouped below.

Other comments: 5

packages/kumo-engine/kumo/engine/graph.py: 1

438-439: Minor · Stability & Availability · Quick win

Validate nested ingest-stat values before conversion.

If a persisted stats file contains null or a scalar for either field, dict(...) raises instead of using the empty-mapping fallback. Normalize non-dictionary values to {} before conversion.

Suggested fix

-            'rows': dict(stats.get('source_rows_by_table', {})),
-            'edges': dict(stats.get('edge_counts', {})),
+            'rows': (dict(stats.get('source_rows_by_table'))
+                     if isinstance(stats.get('source_rows_by_table'), dict)
+                     else {}),
+            'edges': (dict(stats.get('edge_counts'))
+                      if isinstance(stats.get('edge_counts'), dict) else {}),

packages/kumo-engine/kumo/engine/_tracking.py: 1

501-527: Minor · Stability & Availability · Quick win

Do not start a second standalone tracking thread while one is still running on the same tracker.

log_run stops waiting after FINISHED_SECONDS and leaves the worker thread running in the background. The next predict, evaluate or score_drift call in the same process starts a new worker on the same tracker instance. MlflowTracking.start calls _reset() and reassigns _parent, _client, _experiment_id and _guard. finish reads those same fields and resets them again.

Two workers on one MlflowTracking instance can therefore interleave:

• The metrics and tags of one run can go to the other run.
• A run can be terminated with the wrong status.
• One worker can close the other worker's HTTP-limit guard.

This can happen in a scheduled loop of predictions when the tracking server is slow. MLflow HTTP calls use a 30 s timeout with 2 retries, so one call can take longer than the 30 s wait. Predictions are not affected, but the recorded runs are incorrect.

Track the live worker per tracker. Skip with a warning, or wait, while that worker is still alive.

🔒️ Proposed fix

+_LIVE: 'weakref.WeakKeyDictionary[Any, threading.Thread]' = weakref.WeakKeyDictionary()
+_LIVE_LOCK = threading.Lock()
+
 def log_run(tracking: Any, hook: str, record: RunRecord) -> None:
     ...
     method = getattr(tracking, hook, None)
     if method is None or _rank() != 0:
         return
+    with _LIVE_LOCK:
+        busy = _LIVE.get(tracking)
+        if busy is not None and busy.is_alive():
+            logger.warning('Experiment tracking is still recording an '
+                           'earlier run; the %s run is not tracked',
+                           record.tags.get('kumo.kind'))
+            return
     ...
     worker = threading.Thread(target=log, name='kumo-tracking', daemon=True)
+    with _LIVE_LOCK:
+        _LIVE[tracking] = worker
     worker.start()

If a tracker class cannot hold weak references, key the registry on id(tracking).

packages/kumo-engine/kumo/engine/_monitoring.py: 1

179-187: Minor · Data Integrity & Integration · Quick win

Give each evaluation its own actuals directory.

_actuals_dir builds the directory from the run and the snapshot id only. Every evaluate(run, snapshot=s) call therefore writes to the same path. Line 187 deletes that path before the join. This causes two problems:

• An earlier EvaluationResult.actuals silently changes to the rows of the later call, with a different EVALUATED_AT. If the later join fails after the rmtree, the earlier result's directory is gone or partial, so export() on the earlier result writes the wrong rows or fails.
• EvaluationResult.export uses Path(self.actuals).name in export_id. Two different evaluations of the same run and snapshot therefore get the same export id.

Two concurrent evaluations of one run against one snapshot also write into the same directory at the same time.

Add a unique per-evaluation suffix. The existing uuid4 import is enough.

🐛 Proposed fix

-    return root / 'actuals' / str(snapshot.id).replace('/', '_')
+    return (root / 'actuals' / str(snapshot.id).replace('/', '_') /
+            uuid4().hex)

-            shutil.rmtree(actuals, ignore_errors=True)
+            actuals.parent.mkdir(parents=True, exist_ok=True)

In EvaluationResult.export, the existing export_id then becomes unique per evaluation:

-                         export_id=f'{self.prediction_id}-evaluation-'
-                         f'{Path(self.actuals).name}')
+                         export_id=f'{self.prediction_id}-evaluation-'
+                         f'{Path(self.actuals).name}')

spill is derived from actuals.parent. Keep it outside the new unique directory, or derive it from actuals itself.

packages/diskgraph/diskgraph-pybind/src/bp_evaluate.rs: 2

548-550: Minor · Data Integrity & Integration · Quick win

Clear reused actuals_dir in the public evaluator.

The monitoring workflow clears its run-scoped directory before calling diskgraph.evaluate_prediction_parquet, but the evaluator is also publicly exposed and accepts a caller-supplied actuals_dir. It creates that directory without removing existing part-*.parquet files. A reused directory can therefore retain stale rows when a later run has fewer partitions or no matches for a prior partition.

Use a fresh directory, remove existing partition files before writing, or write to a staging directory and rename it after completion.

---

547-547: Minor · Data Integrity & Integration · Quick win

Use a private spill subdirectory and clean it on every exit path.

evaluate_prediction_parquet writes directly under spec.spill_dir and removes that directory only after successful evaluation. Calls that share a directory can overwrite each other’s files, and direct callers can lose pre-existing files. Errors can leave spill files behind at the Rust API boundary.

The engine monitoring caller currently supplies a UUID-based directory and removes it in finally, so this does not create a collision for that caller. The Rust API should still isolate its own files.

🛡️ Suggested fix

     std::fs::create_dir_all(&spec.spill_dir)?;
+    let spill = tempfile::Builder::new()
+        .prefix("evaluate-")
+        .tempdir_in(&spec.spill_dir)?;
     if let Some(dir) = &spec.actuals_dir {
         std::fs::create_dir_all(dir)?;
     }
...
-                    spec.spill_dir.join(format!("{side}-{p:05}.arrow")),
+                    spill.path().join(format!("{side}-{p:05}.arrow")),
...
-    let _ = std::fs::remove_dir_all(&spec.spill_dir);
+    drop(spill);

Move tempfile from [dev-dependencies] to [dependencies], or use an equivalent drop guard.

🧹 Nitpick comments: 1

packages/diskgraph/diskgraph-pybind/src/bp_col_stats.rs: 1

3076-3091: Trivial · Maintainability & Code Quality · Quick win

The default-equivalence test compares one call with an identical call.

default and unset both call compute_parquet_col_stats_with_options with ColStatsOptions::default(). The assertion at Line 3091 therefore only proves that the function is deterministic. It does not prove the claim in the test name: that unset options produce the same output as the old code. The legacy entry point is gone, so a change to a default (for example histograms or summary_stats_by_class) would still pass this test.

Pin the default output explicitly. One option is to compare it with the legacy-output test helper (Lines 1755-1765). Another option is to assert that the default-only keys are present, such as HISTOGRAM and SUMMARY_STATS_BY_CLASS, and that no PERCENTILES or CATEGORY_COUNTS_CAPPED keys appear.

♻️ Proposed assertion on the default key set

-        let unset = compute_parquet_col_stats_with_options(
-            &paths,
-            &columns,
-            "v1",
-            MAX_CATEGORIES,
-            &ColStatsOptions::default(),
-        )
-        .unwrap();
-        assert_eq!((default.clone(), default_min.clone()), unset);
+        let default_keys = parse(&default);
+        assert!(default_keys.contains_key("SCORE__kumo.numerical__HISTOGRAM__False"));
+        assert!(default_keys.contains_key("SCORE__kumo.numerical__SUMMARY_STATS_BY_CLASS__False"));
+        assert!(!default_keys.keys().any(|k| k.contains("PERCENTILES") || k.contains("CATEGORY_COUNTS_CAPPED")));

---

Review info

Run configuration

• Source mode: Auto
• Repository: NVIDIA/kumo-engine
• Branch: vedaant/mlops-mlflow-tracking
• File: .coderabbit.yaml
• Commit: 19e40fe
• Layers: Repository YAML
• Review profile: QUIET
• Plan: Enterprise
• Run ID: 579be13d-96a4-4382-884d-69ff62bb9c51

Commits

Reviewing files that changed from the base of the PR and between cb72a69 and 41c548f.

Files selected for processing: 40

• examples/mlops/README.md
• examples/mlops/databricks/README.md
• examples/mlops/databricks/databricks.yml
• examples/mlops/fit_job.py
• examples/mlops/jobs.py
• examples/mlops/sagemaker/README.md
• examples/mlops/sagemaker/pipeline.py
• examples/mlops/score_job.py
• examples/mlops/snowflake/README.md
• examples/mlops/snowflake/register_jobs.py
• packages/diskgraph/diskgraph-pybind/python/diskgraph/__init__.py
• packages/diskgraph/diskgraph-pybind/src/bp_col_stats.rs
• packages/diskgraph/diskgraph-pybind/src/bp_evaluate.rs
• packages/diskgraph/diskgraph-pybind/src/lib.rs
• packages/kumo-core/src/kumo_core/pquery.py
• packages/kumo-core/src/kumo_core/training/table.py
• packages/kumo-core/tests/test_pquery.py
• packages/kumo-core/tests/test_train_table.py
• packages/kumo-engine/docs/mlops.md
• packages/kumo-engine/kumo/engine/__init__.py
• packages/kumo-engine/kumo/engine/_manifest.py
• packages/kumo-engine/kumo/engine/_mlops.py
• packages/kumo-engine/kumo/engine/_monitoring.py
• packages/kumo-engine/kumo/engine/_tracking.py
• packages/kumo-engine/kumo/engine/graph.py
• packages/kumo-engine/kumo/engine/integrations/mlflow.py
• packages/kumo-engine/kumo/engine/integrations/snowflake.py
• packages/kumo-engine/kumo/engine/pquery/prediction_table.py
• packages/kumo-engine/kumo/engine/pquery/predictive_query.py
• packages/kumo-engine/kumo/engine/trainer/_run_manifest.py
• packages/kumo-engine/kumo/engine/trainer/trainer.py
• packages/kumo-engine/test/test_job_templates.py
• packages/kumo-engine/test/test_mlflow_tracking.py
• packages/kumo-engine/test/test_mlops.py
• packages/kumo-engine/test/test_monitoring.py
• packages/kumo-engine/test/test_prediction_manifest.py
• packages/kumo-engine/test/test_predictive_query.py
• packages/kumo-engine/test/test_snapshot_manifest.py
• packages/kumo-engine/test/test_snowflake_tracking.py
• packages/kumo-engine/test/test_tracking_contract.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.

[coderabbitai]

Major · Data Integrity & Integration · Heavy lift

Pass the fitted model’s identity to Predict.

If two pipeline executions overlap, a later Fit can move champion before the earlier execution starts Predict. depends_on=[fit] orders steps within one execution, but Predict loads the shared alias through score_job.py; it does not load that execution’s fitted model. The earlier execution can therefore score with the wrong model. Pass an immutable model version or bundle path from Fit to Predict. (docs.aws.amazon.com)

[coderabbitai]

Major · Stability & Availability · Quick win

Preserve the running state when re-registering tasks.

If an operator runs register_jobs.py again after resuming the tasks, CREATE OR REPLACE TASK replaces both tasks in a suspended state. Scheduled training and scoring then stop until an operator resumes them again. Record each existing task’s state and restore resumed tasks after replacement, while leaving new tasks suspended. (docs.snowflake.com)

[coderabbitai]

Actionable comments posted: 2

---

Review info

Run configuration

• Source mode: Auto
• Repository: NVIDIA/kumo-engine
• Branch: vedaant/mlops-score-sketch
• File: .coderabbit.yaml
• Commit: 9390153
• Layers: Repository YAML
• Review profile: QUIET
• Plan: Enterprise
• Run ID: 9127c09b-8bca-43b2-b8d5-89363bf581e9

Commits

Reviewing files that changed from the base of the PR and between 41c548f and 9390153.

Files selected for processing: 1

• examples/mlops/databricks/README.md

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 10 remain after this review.

[coderabbitai]

Major · Security & Privacy · Quick win

Avoid schema-wide MODIFY for prediction output.

MODIFY on the schema is inherited by child tables. This gives the run-as principal write access to the three source tables, not only users_churn_predictions. Grant CREATE TABLE on the schema and MODIFY only on the prediction table after creation, or isolate predictions in a separate schema. (docs.databricks.com)

[coderabbitai]

Major · Data Integrity & Integration · Quick win

Document the privileges required for repeated model registration and alias updates.

kumo_train registers a new model version on each fit and updates champion, but this list grants only CREATE MODEL and EXECUTE. Unity Catalog requires model ownership or CREATE MODEL VERSION for new versions. Alias updates require model ownership plus USE SCHEMA and USE CATALOG. State these exact requirements, or require the run-as principal to own the model. (docs.databricks.com)

[coderabbitai]

Note

Quiet mode is enabled, so only the most important comments were posted inline. Other review comments are grouped below.

Other comments: 1

packages/kumo-engine/kumo/engine/trainer/_run_manifest.py: 1

58-68: Minor · Data Integrity & Integration · Quick win

The export sidecar records the run before the export result is known.

finish passes replace(manifest, outputs=records) to result._export as the sidecar. The sidecar written next to the exported Parquet parts therefore has these values:

• status=None.
• timings.finished_at=None.
• outputs that list every planned destination, including destinations that were never written.

After the export, finish writes the final manifest to the local _manifest_path. The artifact-root copy made by store also uses that final manifest.

The documented later-job workflow loads the run from the user's export with PredictionResult.load_from_connector. In that workflow, loaded.manifest.status is None. Assume embs/ fails after preds/ was written. Then the preds/_provenance.json file still lists both outputs as destinations. The prediction-runs section of docs/mlops.md says the manifest records status, timings, and the outputs as written. That description is not true for the copy that later jobs read.

Fix it in one of these ways:

• Document that sidecars in output_config destinations are the pre-export record. Only the local manifest and the artifact-root manifest are final.
• Rewrite each written destination's sidecar after the loop finishes. This requires a sink write for one file.

🧹 Nitpick comments: 1

packages/kumo-engine/kumo/engine/trainer/trainer.py: 1

335-339: Trivial · Performance & Scalability · Quick win

num_rows downloads every Parquet part of a loaded export.

For a result from load_from_connector, num_rows iterates self.files. That call runs _Export.files(), which copies all parts into local scratch. The class docstring says num_rows reads bounded producer metadata. load_from_connector also says it fetches parts only when files is first read. A later job that only checks the row count of a large export downloads the whole export.

Read the footers on the remote filesystem instead.

♻️ Proposed change

         if self._exported is not None:
             from pyarrow import parquet
-            return sum(
-                parquet.ParquetFile(path).metadata.num_rows
-                for path in self.files)
+            return self._exported.num_rows()

# In _Export:
def num_rows(self) -> int:
    from pyarrow import parquet
    from pyarrow.fs import FileSelector, FileType

    total = 0
    for info in self.filesystem.get_file_info(FileSelector(self.root)):
        if info.type == FileType.File and info.path.endswith('.parquet'):
            with self.filesystem.open_input_file(info.path) as handle:
                total += parquet.ParquetFile(handle).metadata.num_rows
    return total

---

Review info

Run configuration

• Source mode: Auto
• Repository: NVIDIA/kumo-engine
• Branch: vedaant/mlops-score-sketch
• File: .coderabbit.yaml
• Commit: 20b170a
• Layers: Repository YAML
• Review profile: QUIET
• Plan: Enterprise
• Run ID: fda43b38-5ba8-4b43-bba2-be1b36b518cb

Commits

Reviewing files that changed from the base of the PR and between 9390153 and 20b170a.

Files selected for processing: 15

• packages/diskgraph/diskgraph-pybind/python/diskgraph/tt_artifacts.py
• packages/diskgraph/diskgraph-pybind/python/tests/test_tt_artifacts_row_warnings.py
• packages/kumo-core/src/kumo_core/prediction.py
• packages/kumo-engine/docs/mlops.md
• packages/kumo-engine/kumo/engine/_manifest.py
• packages/kumo-engine/kumo/engine/_monitoring.py
• packages/kumo-engine/kumo/engine/_tracking.py
• packages/kumo-engine/kumo/engine/pquery/prediction_table.py
• packages/kumo-engine/kumo/engine/trainer/_run_manifest.py
• packages/kumo-engine/kumo/engine/trainer/trainer.py
• packages/kumo-engine/test/test_mlflow_model.py
• packages/kumo-engine/test/test_monitoring.py
• packages/kumo-engine/test/test_prediction_manifest.py
• packages/kumo-engine/test/test_predictive_query.py
• packages/kumo-engine/test/test_query_from_snapshot.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 10 remain after this review.

[coderabbitai]

Note

Quiet mode is enabled, so only the most important comments were posted inline. Other review comments are grouped below.

Other comments: 3

packages/kumo-engine/kumo/engine/pquery/prediction_table.py: 2

314-317: Minor · Functional Correctness · Quick win

Read the entity column for skipped-entity samples.

When DiskGraph reports a null_timestamp failure on TIMESTAMP, this code reads TIMESTAMP into skipped_sample. The sample then contains a null timestamp instead of the affected entity ID. Read ENTITY1 at the reported row location for every condition so the manifest can identify that row. DiskGraph reports both the failure column and sample locations. (raw.githubusercontent.com)

---

312-317: Minor · Functional Correctness · Quick win

Read the reported row as a row-group-relative index.

The sample format is <part> row group <n> row <m>, so <m> is the row index inside row group <n>. _sample_ids subtracts the preceding row-group row counts whenever row >= start. For a group-relative index, that subtraction reads the wrong row. In the example above, row 6 of group 1 (after a five-row group 0) reads row 1 instead of row 6. The recorded skipped_sample ID is then wrong.

-                start = sum(reader.metadata.row_group(index).num_rows
-                            for index in range(group))
                 column = reader.read_row_group(
                     group, columns=[entry['column']]).column(0)
-                ids.append(column[row - start if row >= start else row]
-                           .as_py())
+                ids.append(column[row].as_py())

packages/kumo-engine/docs/mlops.md: 1

133-133: Minor · Functional Correctness · Quick win

Document the all-skipped failure case.

If none of the supplied entities survives validation, custom prediction-table generation fails instead of returning a warning and an empty prediction result. Replace “never refused” with a statement that distinguishes partial skips from an all-skipped input. The prediction-table path rejects a zero-row result. (raw.githubusercontent.com)

---

Review info

Run configuration

• Source mode: Auto
• Repository: NVIDIA/kumo-engine
• Branch: vedaant/mlops-score-sketch
• File: .coderabbit.yaml
• Commit: cfdd76f
• Layers: Repository YAML
• Review profile: QUIET
• Plan: Enterprise
• Run ID: 0b03304f-1351-471f-92d0-9357ec6c57a6

Commits

Reviewing files that changed from the base of the PR and between 20b170a and cfdd76f.

Files selected for processing: 2

• packages/kumo-engine/docs/mlops.md
• packages/kumo-engine/kumo/engine/pquery/prediction_table.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 10 remain after this review.