[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: Prediction runs and monitoring: #64 → #67 → #69 → #73
Merge order: one linear stack; merge bottom to top (#53 first)
This PR: Adds realized performance and score drift for prediction runs.

Summary

This PR adds realized performance and score drift for prediction runs. Trainer.evaluate scores a past run against labels that have since arrived and writes the evaluated rows that warehouse model monitors read. Trainer.score_drift compares a run's scores with the model's holdout and with a previous run. PredictionResult.load_from_connector lets a later job do both from an exported run. With MLOps tracking on, each call is logged as a monitoring run. It builds on the prediction-run manifest (#67), the holdout score statistics (#64) and the snapshot manifest (#50).

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

• We scored three weeks ago. How good were those scores, now that the outcomes are in?
• Did this week's score distribution drift from the training baseline, or from last week's run, and by how much?
• Did the inputs change between the training snapshot and the scoring snapshot?
• 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?

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

Cause. The engine can generate labels only for training. Scoring a past run against outcomes means regenerating the query's labels at that run's anchor times on a newer snapshot, and nothing exposes that. Label generation also refuses tables that could not train a model (a single class, fewer classes than TOP K), even though those are valid outcomes to read. The run's score statistics (#67), the model's holdout statistics (#64) and the snapshot manifests (#50) were recorded, but nothing compared them, and a PredictionResult existed only inside the process that ran predict.

Before:

trainer.fit(...)                           -> model (+ holdout score statistics, #64)
trainer.predict(snapshot, ...)             -> PredictionResult (+ _provenance.json, #67)
# realized metrics, drift, monitor tables, monitoring from another job: write it yourself

Change. After:

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="predictions/2024-03-30/") -> PredictionResult       # new
result.evaluate(...), result.score_drift(...)                          # TrainingResult passthroughs, like result.predict

Trainer.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, hash-partitions them on ENTITY (and ANCHOR_TIMESTAMP with several anchors), 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. The evaluated rows are written next to the run as a Parquet directory (EvaluationResult.actuals): the label columns plus the matched prediction's columns. Metrics use the fit-time test_* definitions:

• binary: auroc, ap over True_PROB, exact (torchmetrics definitions, ties counted half; scores compared at float32);
• regression: mae, rmse over TARGET_PRED;
• multiclass: acc of the top-scored CLASS per entity and anchor;
• link prediction: map@k, 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;
• multilabel and forecast: no metrics, with a reason.

Trainer.score_drift compares the run's score summary with the model's holdout (#64) and with a previous run (#67), 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ε (ε: sketch rank error)
wasserstein	∫₀¹ |Q_a − Q_b| dp, exact per segment (a segment adds 0.01·(|d_k|+|d_{k+1}|)/2 when the knot gaps d have one sign, else 0.01·(d_k²+d_{k+1}²)/(2(|d_k|+|d_{k+1}|)))	≤ 0.01·(both score ranges) + sketch error

wasserstein is in score units: comparable over time for probabilities, scale-dependent for regression. The report also has per-table row and per-edge count changes between the training and scoring snapshots (reverse edges skipped).

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 (FileNotFoundError otherwise), then reads only its _provenance.json. The Parquet parts are downloaded in parallel the first time files is read, into a private scratch directory removed when the result is garbage-collected, so score_drift from another job fetches nothing else. PredictionResult.files stays as on main.

Example from a binary query (PREDICT COUNT(orders.*, 0, 7, days) > 0 FOR EACH users.ID) on a synthetic users/orders graph (60 users). The model was trained and run on a snapshot whose orders stop at 2024-03-31, exported to predictions/2024-03-30/, loaded back with load_from_connector, then evaluated on a later snapshot with 410 more orders (ids shortened; dataclasses.asdict of the frozen results).

EvaluationResult:

{
  "prediction_id": "1bb6e3b1…-predict-9c94ade2…", "bundle_id": "02f5a396…",
  "snapshot_id": "66570cbc…@c85adb16754ce481",
  "task_type": "binary_classification",
  "anchors": ["2024-03-30T00:00:00+00:00"],
  "metrics": {"auroc": 0.4270, "ap": 0.7665},
  "rows": 60, "labelled": 60, "labelled_fraction": 1.0,
  "matured": true, "bundle_matches": true, "reasons": {},
  "actuals": "<scratch>/evaluations/1bb6e3b1…/actuals/66570cbc…@c85adb16754ce481"
}

actuals holds 60 rows such as {"ENTITY": 0, "ANCHOR_TIMESTAMP": "2024-03-30", "TARGET": true, "True_PROB": 0.66656, "TARGET_PRED": true}.

DriftReport (with previous set to the same run, so the previous_* statistics are 0):

{
  "prediction_id": "1bb6e3b1…-predict-9c94ade2…", "bundle_id": "02f5a396…",
  "score": "True_PROB",
  "psi": 0.5288, "ks": 0.2407, "wasserstein": 0.00087,
  "holdout_truncated": false, "bundle_matches": true,
  "previous_prediction_id": "1bb6e3b1…-predict-9c94ade2…",
  "previous_psi": 0.0, "previous_ks": 0.0, "previous_wasserstein": 0.0,
  "input_changes": {"tables.orders.rows": [1090, 1500], "edges.orders_USER_ID_to_users_ID": [1090, 1500]},
  "reasons": {}
}

(The tiny model was trained on two steps; the numbers are real but say nothing about the model.)

What the fields mean. Nothing is refused; the flags say how far to trust a number:

• prediction_id and bundle_id are the run's own ids (from its manifest); snapshot_id is the snapshot labels were read from; previous_prediction_id is the comparison run.
• matured is False while the label window ends after the earliest latest-event time among the query's target tables (a conservative proxy for label completeness). A static query reads labels as of the snapshot, so anchors is empty, matured is None, and the evaluated rows carry the run's prediction time as ANCHOR_TIMESTAMP.
• bundle_matches compares the run's model with the trainer's. Evaluating a run with a different model is allowed and flagged.
• labelled_fraction is labelled predictions / run predictions. For link prediction it is the share of entities with any future link.
• holdout_truncated repeats the holdout record's truncated flag (Record prediction runs and monitor score drift #64).
• actuals: ENTITY, ANCHOR_TIMESTAMP, TARGET, and the prediction (binary True_PROB and the predicted label TARGET_PRED; regression TARGET_PRED; multiclass the top CLASS and its SCORE; link prediction none). None when no labels were joined.
• reasons: when a step cannot be computed (no anchors, no holdout statistics, no manifest, unsupported task), the field is None or empty and reasons says why, instead of an exception.

Tracking. With kumo.MLOps(tracking=...) (a block or mlops=), each evaluate and score_drift call opens a new run in the tracker's experiment, named monitor-<prediction id> and tagged kumo.kind=monitoring, kumo.prediction_run=<prediction id>, kumo.bundle_id; the metrics, the drift statistics (incl. previous_*) and the input count changes (input.<key>.before/after) are its metrics. A new run is used because Snowflake cannot reliably reopen a finished run. MLflow also tags the prediction run (#73) with kumo.monitored=<monitoring run id>. The hook is Tracking.log_monitoring(record), optional and best-effort (logged on a tracking thread for up to 30 s, never raises); the MLflow and Snowflake adapters implement it, and both now name a run after its record when it has a name (fit records have none, so fit runs are unchanged).

Warehouse model monitors. docs/mlops.md ("Monitoring prediction runs") gives two recipes over the evaluated rows, exported with evaluation.export(...):

• 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, look up the table and schema UUIDs (w.tables.get(...).table_id, w.schemas.get(...).schema_id), then one 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="ANCHOR_TIMESTAMP", granularities=[1 day])))) and w.data_quality.create_refresh(...). The inference log has no probability column.
• Snowflake Model Monitor: one CREATE MODEL MONITOR ... MODEL=<m> VERSION=<v> FUNCTION=predict SOURCE=<evaluated rows> BASELINE=<holdout evaluated rows> ... TIMESTAMP_COLUMN=ANCHOR_TIMESTAMP PREDICTION_SCORE_COLUMNS=(True_PROB) ACTUAL_CLASS_COLUMNS=(TARGET), with the per-task column choices. The model must be a Snowflake Model Registry version (kumo.MLOps(registry=ModelOutputConfig(sf, location=...))).

Both live checks are pending (to run on the test accounts): creating the Databricks monitor on a table written by evaluation.export, and whether Snowflake accepts FUNCTION=predict on a Kumo registry version, whose predict has a placeholder signature (if not, a real predict signature must be logged at registration). CI verifies the evaluated-rows columns per task.

Follow-up (not in this PR): input-feature drift. This PR compares scores and table and edge counts; per-feature distributions of the scoring snapshot against the training snapshot need column statistics on both snapshots and are a separate change.

Supporting changes:

• 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, by arrow::compute::interleave over the partition's kept batches), and constant_timestamp adds the static query's prediction time. It reuses the projected row-group reader of the batch-prediction statistics (bp_col_stats::row_group_batches, Record prediction runs and monitor score drift #64) and arrow-rs building blocks (RowConverter for join keys and class identity, take_record_batch to partition, Arrow IPC to spill).
• kumo-core generate_train_table(..., validate_targets=True): False keeps rows that could not train a model, so outcomes can be read. The default is unchanged.
• kumo-core pquery: one private _target_maximum(query, index_path, min|max) behind latest_temporal_timestamp; evaluate uses its minimum for matured.
• PredictiveQuery: one _generate helper for training-table and label generation, and a private _label_parts context manager that keeps label parts on disk for the join.
• kumo.engine exports EvaluationResult and DriftReport. _tracking.log_run(tracking, hook, record) logs one standalone run through an optional tracker hook (Log prediction runs to MLflow and Snowflake through the tracking adapters #73 reuses it for prediction runs).

Back-compat

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

Situation	Behaviour
Existing fit, predict, export, training-table generation	Unchanged. validate_targets defaults to True, which is today's behaviour.
Fit runs in MLflow or Snowflake	Unchanged: fit records have no name, so run_name applies as before. A taken Snowflake name now logs "recording this run as" instead of "this fit as".
New calls: evaluate / score_drift	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.
mlops= / a kumo.mlops block with tracking	A new monitoring run per call; tracking errors are logged, never raised. Without tracking, nothing is logged.
New call: PredictionResult.load_from_connector	Needs a connector whose sink holds files at the path (TypeError otherwise) and a completed export (_SUCCESS; FileNotFoundError otherwise). A completed directory without a manifest loads, with no ids or anchors. Downloaded parts are removed with the result.
EvaluationResult.export without evaluated rows	Raises ValueError with the reasons (new call).
Run without a manifest (before #67)	score_drift reports None statistics with reasons; evaluate of a temporal query has no anchors, so metrics is empty with a reason.
Model without holdout statistics (before #64, refit_full, link prediction)	Holdout statistics are None with a reason.
Snapshots not passed	input_changes is None with a reason.
Unsupported task (multilabel, forecast)	metrics is empty, with a reason.

Performance

evaluate's join and metrics, before (pandas: concat of label parts, run.frame, merge, per-row Python ranking) and after (DiskGraph). Synthetic predictions and label parts on local NVMe, one anchor, 90% of predictions labelled, LP with k = 10 and 1–5 targets per entity; 8 threads; wall time and peak RSS of the process. Label generation (DiskGraph, unchanged) is not included. Measured before the join emitted prediction columns, which adds one column copy per matched row.

Output	Rows	Before	After
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 (8 × about 1M rows of predictions; labels stream), plus 8 bytes per labelled row for the exact binary ranking. score_drift's statistics are O(101) on manifest metadata.

Known limitations

• matured is a proxy: it uses the earliest latest-event time across target tables.
• Results are returned, not written (except actuals and the tracked monitoring run).
• 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); neither warehouse monitor supports link prediction or ranking.
• The two warehouse recipes are pending live verification (above).

Testing Done

• Rust (cargo test -p diskgraph-pybind --lib bp_, 34 passed incl. 3 bp_evaluate; clippy -D warnings and rustfmt clean): binary AUROC equals a brute-force pair count in one and five spilled partitions, and each evaluated row carries its own entity's score; the AP step definition; regression with a constant timestamp; multiclass emits the top-scored class and score per entity and anchor; link-prediction metrics.
• test_monitoring.py (10 passed; one real CPU fit on a two-snapshot users/orders fixture, other cases on stub runs):
  • a matured binary run: actuals labels against a pandas recomputation, the emitted True_PROB/TARGET_PRED equal the run's, AUROC against scikit-learn; the older snapshot is flagged immature;
  • several anchors from the manifest are joined; a run without anchors gets a reason;
  • regression metrics; a static query's rows carry the prediction time and the score;
  • EvaluationResult.export writes the evaluated rows;
  • realized metrics equal kumo-core BinaryAUROC / BinaryAveragePrecision and PyG LinkPredMAP / LinkPredPrecision;
  • PSI, KS and Wasserstein-1 from percentile curves: KS within 0.02 of the exact two-sample statistic and W1 within 2% of the exact value on 100k-sample normals; point masses; 0 for identical summaries; None without curves;
  • drift against the holdout, a previous run and the snapshot counts; drift without a holdout;
  • a tracked score_drift and evaluate each log a monitoring run with the right tags and metrics;
  • a run exported and loaded back: score_drift fetches no parts, evaluate and score_drift match the original, the downloaded parts are removed with the result, an export without _SUCCESS is refused.
• test_mlflow_tracking.py: a monitoring run is a new finished run with the record's metrics, and the prediction run is tagged kumo.monitored; test_snowflake_tracking.py, test_tracking_contract.py: fit runs unchanged.
• kumo-core test_pquery.py + test_train_table.py: 287 passed; test_predictive_query.py passed.
• The examples above were captured from a CPU run of this branch.
• flake8, isort and the prepush check are clean.

Checklist

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

🤖 Generated with Claude Code

[Vedaant-J]

Change map

File	Function / scope	What changed	Why
packages/diskgraph/diskgraph-pybind/python/diskgraph/__init__.py	exports	evaluate_prediction_parquet	public binding
packages/diskgraph/diskgraph-pybind/src/bp_col_stats.rs	row_group_batches	made pub(crate)	the evaluate join reuses its projected row-group reader
packages/diskgraph/diskgraph-pybind/src/bp_evaluate.rs	new: EvaluateSpec (+ prediction_columns, constant_timestamp), evaluate_prediction_parquet, join, ranking_metrics, metrics, tests	streaming hash-partitioned join of prediction Parquet with label parts (keys via arrow RowConverter, partitions spilled to Arrow IPC, at most 8 joins in flight); fit-time metric definitions (exact AUROC/AP, MAE/RMSE, top-1 acc, PyG MAP@k/precision@k); evaluated rows written with the label columns, the matched (top-scored) prediction's columns via interleave, and an optional constant timestamp	evaluate scales with output size; the rows warehouse monitors read
packages/diskgraph/diskgraph-pybind/src/lib.rs	evaluate_prediction_parquet, module registration	Python binding taking a JSON spec	thin Python API
packages/kumo-core/src/kumo_core/pquery.py	latest_temporal_timestamp, _target_maximum	one private min/max helper over the target tables' source maxima	the matured flag without a second copy
packages/kumo-core/src/kumo_core/training/table.py	generate_train_table	validate_targets=True keyword; False keeps single-class / short-TOP-K label rows	read outcomes that could not train a model
packages/kumo-core/tests/test_pquery.py, packages/kumo-core/tests/test_train_table.py	tests	target maxima; validate_targets=False keeps staging	coverage
packages/kumo-engine/docs/mlops.md	docs (Monitoring prediction runs)	new section: realized performance, score drift (PSI/KS/W1 table), tracking, a later job, Databricks inference-log and Snowflake Model Monitor recipes	user docs
packages/kumo-engine/kumo/engine/__init__.py	exports	EvaluationResult, DriftReport	public API
packages/kumo-engine/kumo/engine/_monitoring.py	new: EvaluationResult (+ export), DriftReport, evaluate, score_drift, track, _psi, _ks, _wasserstein, _curves, _count_changes, _actuals_dir	labels joined and scored by DiskGraph with prediction columns emitted; anchors from the manifest; matured/bundle_matches/labelled_fraction flags; PSI, KS and Wasserstein-1 from percentile curves vs the model's holdout and a previous run; snapshot count changes; one reasons map instead of raises; each call logged as a new monitoring run when tracking is on	realized performance and drift for prediction runs
packages/kumo-engine/kumo/engine/_tracking.py	log_run	logs one standalone run through an optional tracker hook on a tracking thread, up to 30 s, never raising	monitoring runs (and #73's prediction runs) share one path
packages/kumo-engine/kumo/engine/integrations/mlflow.py	MlflowTracking.start, MlflowTracking.log_monitoring	names a run after its record when set; a monitoring run is one finished run, and the prediction run it monitors is tagged kumo.monitored	the monitoring hook for MLflow
packages/kumo-engine/kumo/engine/integrations/snowflake.py	SnowflakeTracking.start, log_monitoring, _new_run_name	names a run after its record when set (taken names get a suffix); a monitoring run is one finished run; warnings say which kind of run is skipped	the monitoring hook for Snowflake
packages/kumo-engine/kumo/engine/pquery/predictive_query.py	PredictiveQuery._generate, _label_parts, generate_training_table	one generation helper for training tables and labels; label parts kept on disk for the join	label regeneration for evaluation without a second workflow block
packages/kumo-engine/kumo/engine/trainer/trainer.py	Trainer.evaluate(mlops=), Trainer.score_drift(mlops=), TrainingResult.evaluate, TrainingResult.score_drift, PredictionResult.load_from_connector, PredictionResult.id, _manifest_path, num_rows, files, _Export	entry points delegating to _monitoring with settings resolved on the caller's thread; an exported run loaded through file_store once its _SUCCESS exists, _provenance.json first, parts downloaded in parallel on first use into a directory removed with the result	public entry points; monitoring from a later job
packages/kumo-engine/test/test_mlflow_tracking.py	test_a_monitoring_run_is_new_and_tags_its_prediction_run	a monitoring run is a new finished run with the record's metrics; the prediction run gets kumo.monitored	adapter contract
packages/kumo-engine/test/test_monitoring.py	new	evaluate (matured, immature, anchors, regression, static with prediction time), emitted prediction columns, EvaluationResult.export, realized metrics vs kumo-core/PyG, PSI/KS/W1 from curves (vs exact values, point masses), drift (holdout, previous, counts, missing records), tracked monitoring runs, a run loaded lazily from its export	behaviour coverage
packages/kumo-engine/test/test_predictive_query.py	tests	Generation fakes accept validate_targets	Training-table and label generation share one _generate helper that passes it
packages/kumo-engine/test/test_snowflake_tracking.py	test_an_existing_run_is_never_resumed	expects the run-neutral rename warning	message now covers prediction and monitoring runs

Requirement: realized performance and drift for prediction runs that scale with output size, logged through the tracker and feeding warehouse monitors (final MLOps design, section 9). Existing: the run manifest and its score statistics (#67), the holdout statistics (#64), snapshot counts, DiskGraph training-table generation for labels, the tracking adapters (#60/#68). Smallest: two Trainer methods, one loader, the join and metrics in one DiskGraph function, statistics on stored curves, one optional tracker hook. Stopping condition: metrics, drift and input changes computable in-process or from an exported run at 100M rows within bounded memory, with every missing input explained; recipes documented, live checks pending.

[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]

Warning

Review limit reached

Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available.

Next included review available in 24 minutes.

Check out review usage here.

View limit details

Limit details: You’ve used all 12 included reviews currently available.

Learn how review limits work.

Review configuration:

Run configuration

• Source mode: Auto
• Repository: NVIDIA/kumo-engine
• Branch: vedaant/mlops-prediction-manifest
• File: .coderabbit.yaml
• Commit: ef2cb21
• Layers: Repository YAML
• Review profile: QUIET
• Plan: Enterprise
• Run ID: 6704c2da-8473-4cf4-96c4-5e0022b8940e

Commits

Reviewing files that changed from the base of the PR and between 0c39510 and 4aba059.

Files selected for processing: 20

• 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/_monitoring.py
• packages/kumo-engine/kumo/engine/_tracking.py
• packages/kumo-engine/kumo/engine/integrations/mlflow.py
• packages/kumo-engine/kumo/engine/integrations/snowflake.py
• packages/kumo-engine/kumo/engine/pquery/predictive_query.py
• packages/kumo-engine/kumo/engine/trainer/trainer.py
• packages/kumo-engine/test/test_mlflow_tracking.py
• packages/kumo-engine/test/test_monitoring.py
• packages/kumo-engine/test/test_predictive_query.py
• packages/kumo-engine/test/test_snowflake_tracking.py

Walkthrough

The change adds Parquet-based evaluation for binary, regression, multiclass, and link predictions. It adds trainer APIs for evaluating prediction runs and measuring score drift, plus support for loading exported prediction runs. Label generation can use specified timestamps without target validation. MLflow and Snowflake tracking can record monitoring runs.

Sequence Diagram(s)

sequenceDiagram
  participant Trainer
  participant Monitoring as monitoring.evaluate
  participant Labels as PredictiveQuery._label_parts
  participant Evaluator as evaluate_prediction_parquet
  Trainer->>Monitoring: Evaluate prediction run against snapshot
  Monitoring->>Labels: Generate labels at prediction timestamps
  Labels-->>Monitoring: Return staged Parquet parts
  Monitoring->>Evaluator: Join predictions and labels
  Evaluator-->>Monitoring: Return counts and task metrics
  Monitoring-->>Trainer: Return EvaluationResult

Loading

Priority: Normal

Merge Risk: Low · up to 0c395

Monitoring runs overwrite the tracker's run_id, and the MLflow path can raise a KeyError when a prediction run has no id. Failed evaluations can leave large temporary spill files on disk, and the docs misdescribe the link-prediction labelled fraction. None of these are severe, but they are worth fixing before or shortly after merge.

Pre-merge checks: 2 · 1

Failed checks: 1 inconclusive

Check name	Status	Explanation	Resolution
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 map, and the full PR description is explicitly …	Provide the complete marked change-map comment and the omitted PR description, or make the linked conversation data available. Then review each proposed change against the stated problem, repository behavior, verification, and stopping cond…

Passed checks: 2 passed

Check name	Status	Explanation
Pr Size		The pull request changes 20 files. This does not exceed the 40-file limit.
Public Resource Access		The diff adds no URL, repository location, or package source location. It adds documentation references to Databricks and Snowflake APIs, plus databricks-sdk >= 0.68; these are names and API example…

Full details: Public Resource Access

Explanation

The diff adds no URL, repository location, or package source location. It adds documentation references to Databricks and Snowflake APIs, plus databricks-sdk >= 0.68; these are names and API examples, not restricted resource locations. The referenced SDK is publicly available on PyPI (anonymous request returned HTTP 200).

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 map, and the full PR description is explicitly truncated. The repository diff is available, but it cannot establish the map's stated problem, rationale for each change, focused verification, or stopping condition.

Resolution

Provide the complete marked change-map comment and the omitted PR description, or make the linked conversation data available. Then review each proposed change against the stated problem, repository behavior, verification, and stopping condition.

Finishing Touches: 💡 1

🛠️ Fix failing CI checks: 💡

• Commit to this branch
• Create a new PR

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

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

[coderabbitai]

Note

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

Other comments: 4

packages/kumo-engine/kumo/engine/integrations/snowflake.py: 1

119-122: Minor · Functional Correctness · Quick win

Restore the fit run's run_id after logging a monitoring run.

SnowflakeTracking has the same run_id problem as MlflowTracking. start sets self.run_id to the monitoring run's name, and finish does not restore it. The class docstring says run_id "is the latest fit run's name". After a tracked evaluate or score_drift call, run_id names the monitoring run.

🐛 Proposed fix

     def log_monitoring(self, record: RunRecord) -> None:
+        fit_run_id = self.run_id
         self.start(record)
         self.finish(record, None)
+        self.run_id = fit_run_id

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

600-600: Minor · Stability & Availability · Quick win

Remove spill_dir on every exit path, not only on success.

evaluate_prediction_parquet removes spill_dir only after both scatter calls and all joins succeed. Several steps can fail with ? before Line 600:

• scatter (for example, a column is missing or a cast fails).
• unspill.
• join.
• ArrowWriter creation.

When any of them fails, the Arrow IPC partitions stay on disk. These partitions hold the projected prediction and label columns for the whole run. The Python caller puts spill_dir under the run's directory as .spill-<uuid>. Each failed evaluate call therefore leaves a new full-size copy behind. At the 100M-row target, this can fill the disk.

Use a drop guard so the directory is removed on both success and error.

♻️ Proposed fix

     std::fs::create_dir_all(&spec.spill_dir)?;
+    struct RemoveOnDrop<'a>(&'a Path);
+    impl Drop for RemoveOnDrop<'_> {
+        fn drop(&mut self) {
+            let _ = std::fs::remove_dir_all(self.0);
+        }
+    }
+    let _spill_guard = RemoveOnDrop(&spec.spill_dir);
 ...
-    let _ = std::fs::remove_dir_all(&spec.spill_dir);
     let (rows, labelled) = (acc.rows, acc.labelled);

packages/kumo-engine/kumo/engine/integrations/mlflow.py: 1

156-167: Minor · Functional Correctness · Quick win

Keep tracking.run_id on the fit run, and stop requiring the kumo.prediction_run tag.

log_monitoring reuses start, and start sets self.run_id to the new monitoring run. finish → _reset does not restore it. The class docstring and mlops.md state that run_id is "the MLflow id of the latest fit run". After trainer.evaluate(...) or trainer.score_drift(...) runs with tracking on, tracking.run_id names the monitoring run. Code that reads tracking.run_id after a fit to find the fit run then gets the wrong run.

record.tags['kumo.prediction_run'] can also be missing. _monitoring.track leaves the tag out when run.id is None, for example on a loaded export without _provenance.json. The lookup then raises KeyError after the monitoring run was already logged. log_run reports the run as "not tracked", which is not true.

🐛 Proposed fix

     def log_monitoring(self, record: RunRecord) -> None:
+        fit_run_id = self.run_id
         self.start(record)
         experiment_id, run_id = self._experiment_id, self.run_id
         self.finish(record, None)
+        self.run_id = fit_run_id
+        prediction_run = record.tags.get('kumo.prediction_run')
+        if prediction_run is None:
+            return
         client = self._mlflow.MlflowClient()
         for run in client.search_runs(
                 [experiment_id], filter_string=(
                     "tags.`kumo.kind` = 'prediction' AND tags.`kumo.run_id` "
-                    f"= '{record.tags['kumo.prediction_run']}'"),
+                    f"= '{prediction_run}'"),
                 max_results=1):

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

275-275: Minor · Functional Correctness · Quick win

Fix the link-prediction labelled_fraction description.

The docs say that for link prediction, labelled_fraction is "the share of entities with any future link". The code counts more than that.

In bp_evaluate.rs join, a label row with an empty list still gives Some(vec![]). The row then sets matured[row] = true, so it counts toward labelled. Only the metric accumulation skips empty label lists.

Two tests show this behavior:

• In test_realized_metrics_match_fit_time_definitions, some entities get zero true items, yet the test asserts joined['labelled'] == entities.
• The Rust test comment says entity (1, 1) "has no labels and is skipped". That entity is skipped for the metric, but it is still counted as labelled.

For link tasks, labelled_fraction is therefore the share of predicted entities and anchors that have a label row. Update the docs to say this. If the documented meaning is the intended one, change join instead so that empty lists do not set matched.

📝 Proposed doc fix

-- `rows`, `labelled` and `labelled_fraction` (for link prediction, the share of entities with any future link);
+- `rows`, `labelled` and `labelled_fraction` (the share of predicted entities, and anchors, that have a label row; for link prediction this includes entities with no future link, which the metrics skip);

---

Review info

Run configuration

• Source mode: Auto
• Repository: NVIDIA/kumo-engine
• Branch: vedaant/mlops-prediction-manifest
• File: .coderabbit.yaml
• Commit: a8e1884
• Layers: Repository YAML
• Review profile: QUIET
• Plan: Enterprise
• Run ID: 711748bf-64b6-4811-a6c2-c82e8d67dc79

Commits

Reviewing files that changed from the base of the PR and between a8e1884 and 0c39510.

Files selected for processing: 20

• 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/_monitoring.py
• packages/kumo-engine/kumo/engine/_tracking.py
• packages/kumo-engine/kumo/engine/integrations/mlflow.py
• packages/kumo-engine/kumo/engine/integrations/snowflake.py
• packages/kumo-engine/kumo/engine/pquery/predictive_query.py
• packages/kumo-engine/kumo/engine/trainer/trainer.py
• packages/kumo-engine/test/test_mlflow_tracking.py
• packages/kumo-engine/test/test_monitoring.py
• packages/kumo-engine/test/test_predictive_query.py
• packages/kumo-engine/test/test_snowflake_tracking.py

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

[Vedaant-J]

Folded into #64 (prediction runs & monitoring), which now contains this PR's final code: the consolidated stack shows only the end state. Review continues there.