[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 a self-contained manifest (_provenance.json) and output score statistics for every prediction run, embedding the model and scoring snapshot manifests; with MLOps on, the run is stored under <root>/predictions/<id>/.

Summary

This PR adds a self-contained manifest and output statistics for every prediction run. Every Trainer.predict call writes _provenance.json: the model and scoring snapshot it used (their manifests embedded), how it scored, where its outputs went, and what the scores looked like. With MLOps on, the run is also stored under the artifact root. In the series, it is the prediction-side link of the chain: snapshot (#50) → training table (#55) → model (#56, holdout statistics #64) → prediction run.

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?
• Did the score distribution move this week, compared with last week's run or with the model's holdout?
• Was this run scored with the default threshold, the model's default, or one passed by the caller?
• Where did run X write its outputs, and where is a durable copy of the run?

Today none of these can be answered from a prediction run's outputs.

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

<scratch>/trainers/<trainer>/predict_output/<id>/
  _selected_artifacts.json                  # pointer to the selected attempt
  _selections/  _attempts/                  # scorer attempts and parts

<out>/preds/
  part-00000-00000.parquet  _KUMO_EXPORT  _SUCCESS

None of these record the model, the snapshot, the source table versions, the scoring options and where they came from, the destinations, timings, or any statistic of the scores.

Change. Trainer.predict writes _provenance.json (a 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). With MLOps on (predict(..., mlops=cfg) or a kumo.mlops block), after the output_config export the run's outputs and manifest are stored under the artifact root:

<scratch>/trainers/<trainer>/predict_output/<id>/
  ...same files as above...
  _provenance.json                          # new

<out>/preds/                                # predict(output_config=...) or result.export(...)
  part-00000-00000.parquet  _KUMO_EXPORT  _SUCCESS
  _provenance.json                          # new, written before _SUCCESS

<root>/predictions/<id>/predictions/        # new, MLOps on only (export={"prediction"}, the default)
  part-00000-00000.parquet  _KUMO_EXPORT  _SUCCESS  _provenance.json

_provenance.json from a binary-classification run (PREDICT COUNT(orders.*, 0, 7, days) > 0 FOR EACH users.ID) on a small synthetic users/orders graph (500 users), fitted and scored inside a kumo.mlops block and exported to preds/ (ids shortened, percentiles elided, repeated environment blocks and the embedded generation plan elided, local paths replaced by <data>/<out>):

{
  "kind": "prediction", "schema_version": 1,
  "id": "7fe8ba6a…-predict-51689c3e…",
  "created_at": "2026-10-24T19:00:12+00:00",
  "environment": {"kumo_engine": "…", "build_commit": "8b211b2a", "diskgraph": "0.1.0", "torch": "2.12.0+cu130", "python": "3.14.7"},
  "parents": {
    "model": {
      "kind": "model", "id": "d6cd223f…", "fit_mode": "refit_trainval",
      "fit": {"started_at": "2026-10-24T18:59:53+00:00", "finished_at": "2026-10-24T19:00:11+00:00", "seconds": 17.631},
      "parents": {
        "training_table": {
          "kind": "training_table", "id": "tt-086b769a…",
          "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": "c1d54323…@3abe76f26b657db7", "…": "…"}}},
        "warm_start": null}},
    "snapshot": {
      "kind": "snapshot", "id": "c1d54323…@3abe76f26b657db7", "record_sources": true,
      "parents": {"source_tables": {
        "orders": {"connector": "LOCAL_FILE", "location": "<data>/orders.parquet", "version": "1792868392", "version_kind": "timestamp"},
        "users":  {"connector": "LOCAL_FILE", "location": "<data>/users.parquet",  "version": "1792868392", "version_kind": "timestamp"}}}}},
  "prediction_table": {"query_sha256": "69dab700…", "source": "pql", "anchors": ["2024-04-29T00:00:00+00:00"], "rows": 500},
  "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-24T19:00:11+00:00", "finished_at": "2026-10-24T19:00:12+00:00"},
  "stats": {"predictions": {"rows": 500, "predicted_positive_rate": 0.0,
    "columns": {"True_PROB": {"count": 500, "mean": 0.34659, "min": 0.33529, "max": 0.34877,
                              "percentiles": [0.33975, "…", 0.34702, "…", 0.34844]}}}}
}

What each field records:

• id is the prediction id, PredictionResult.id. Run directories and <root>/predictions/<id>/ are per id; run ids keep their format.
• parents embeds two manifests, each without its statistics and later-changing fields (_manifest.embedded):
  • model: the model manifest (Record bundle_id and fit provenance in every model bundle #56) 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 (Record provenance for every artifact under one MLOps artifact root #50), with each source table's connector, location and version. When it equals the training snapshot, both copies are kept.
  • A parent with no manifest (a bundle from before the 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 holds the query's SHA-256 (the text is in the embedded training table), whether the table came from PQL or a custom source, the anchor times in ISO UTC, and the row count.
• scoring holds the effective options and where each came from: argument (passed by the caller), bundle (the model's scoring defaults, Add TrainingResult.with_scoring_defaults and Trainer.export #57), default (filled in by the engine), query (from the PQL), or none.
• outputs lists the destinations of predict(output_config=...) as {connector, location}, in the same shape and 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 plain result.export(...) copies the files and the manifest as they are and records nothing new.
• stats holds the row count per artifact and, for predictions, the same score summary as the holdout baseline (Record prediction runs and monitor score drift #64): per score column count, mean, min, max and percentiles (q01..q99, approximate), from DiskGraph's compute_col_stats_for_parquet over the output Parquet (Rust, row groups in parallel, by-class statistics and histograms off):
  • binary: True_PROB, plus predicted_positive_rate at the threshold;
  • regression and forecast: TARGET_PRED;
  • multiclass, multilabel and link prediction: top_score, the highest score per entity and anchor (DiskGraph max_per, exact across row groups and files, rows in any order). Link prediction records no per-class counts: counting 2M recommended items cost 1.1 s and ~290 MB of the 1.3 s total, and nothing in the series reads it.
• status is succeeded, or failed when the output_config export raised; the manifest is written either way and predict re-raises. timings holds the start and finish of the call; environment follows the snapshot manifest (Record provenance for every artifact under one MLOps artifact root #50).
• In Python, PredictionResult.manifest is a PredictionManifest (public, kumo.PredictionManifest) with PredictionParents, PredictionTableRecord, RunScoring, PredictionOutput, RunTimings and ArtifactStats; for example run.manifest.stats['predictions'].columns. A newer file round-trips byte for byte. PredictionManifest.EMBED_OMITS is outputs, stats, status, timings.

Order inside predict: score → build the manifest → the output_config export (recording outputs) → write the manifest → with MLOps on, store <root>/predictions/<id>/ (whose manifest includes those outputs). Tracking the run comes last, in #73. The store is skipped if the export raised, and a failed store only warns. The example job test (#71's score_job.py inside a kumo.mlops block) now checks the predictions/<id>/ copy and that its manifest names the fitted model. docs/mlops.md has a new "Prediction runs" section.

Back-compat

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

Situation	Behaviour
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.
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. A predict(output_config=...) run's manifest still lists the destination.
PredictionResult from a run before this PR	manifest, bundle_id and snapshot_id are None; id is the run id it already had. export() works as before.
A parent without a manifest	Recorded as an id stub; bundle_id/snapshot_id still read.
Output statistics fail (unexpected column, unreadable part, DiskGraph error)	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; 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.
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.

Known limitations

• Without MLOps, a run exported only to a warehouse table keeps its manifest only in local scratch (PredictionResult.manifest). On an ephemeral machine, turn MLOps on to keep a copy under the artifact root.
• top_score treats each run of consecutive rows with equal entity and anchor as one prediction. The scorer writes each prediction's rows together.

Cost

Output statistics stay always on (no flag). They run once after scoring, in Rust (DiskGraph compute_col_stats_for_parquet, 8 threads). Measured inside a full Trainer.predict (A10G GPU scoring, 8 CPU cores, export to a local Parquet directory, two predicts per case):

Output	predict (end to end)	Output statistics	Share
Binary, 10M rows (10M users)	234–275 s	0.37–0.48 s	0.16–0.18%
Link prediction, 10M rows (1M users × top 10 of 50k items)	68–74 s	0.12–0.17 s	0.16–0.24%

The rest of the manifest adds at most 0.01 s. On their own, the statistics take 0.42–0.52 s / +27 MB for 10M binary rows and 1.48–1.72 s / +37 MB for 50M; link-prediction top scores take 0.12–0.18 s / +36 MB for 10M rows and 0.42–0.53 s / +45 MB for 50M (see #64). The artifact-root copy is one more copy of the output files, made only with MLOps on.

Testing Done

• test_prediction_manifest.py (10 passed, CPU):
  • fields of a run exported to two Parquet directories: kind, id, embedded model (without holdout) and scoring snapshot, bundle_id/snapshot_id, prediction table, scoring sources, {connector, location} outputs, timings, status;
  • each exported directory has a sidecar equal to the run's manifest;
  • percentiles within three percentile points of numpy's, exact min and max; predicted_positive_rate at the threshold;
  • top scores per entity and anchor with unsorted rows, across row groups and files (multiclass and link prediction give the same record);
  • a model without a manifest is an id stub parent;
  • with MLOps on, <root>/predictions/<id>/predictions/ holds the run's files and a manifest that lists the output_config destination; export={"model"} stores nothing;
  • a plain export() records nothing; no sidecar for a warehouse or CSV destination;
  • a bundle scoring default with no argument records bundle; an explicit argument records argument;
  • failing stats, destination records or the whole manifest each warn and leave exported bytes identical;
  • an export that raises on the second destination leaves a manifest with status: failed and the first destination; a newer manifest round-trips byte for byte.
• test_job_templates.py (12 passed): fit then score locally stores predictions/<id>/ with the model's id.
• Also run on the stack's head (Log prediction runs to MLflow and Snowflake through the tracking adapters #73): test_prediction_completion, test_query_from_snapshot, test_scoring_defaults.
• The example above was captured from a CPU run of this branch.
• flake8, isort and the prepush check are clean.

Checklist

• No breaking changes.
• Unit tests added.

🤖 Generated with Claude Code

[Vedaant-J]

Change map

File	Function / scope	What changed	Why
packages/kumo-engine/docs/mlops.md	docs (Prediction runs)	new section: the manifest's fields, sidecars, the artifact-root copy	user docs
packages/kumo-engine/kumo/engine/__init__.py	__all__	exports PredictionManifest	public manifest type, like the other kinds
packages/kumo-engine/kumo/engine/_manifest.py	PredictionManifest (+ EMBED_OMITS), PredictionParents, PredictionTableRecord, RunScoring, PredictionOutput, OutputDestination, RunTimings, ArtifactStats, manifest_bytes	prediction-run manifest dataclasses with embedded model and snapshot parents; one JSON encoding for manifest files and export sidecars	typed PredictionResult.manifest; no duplicated encoding
packages/kumo-engine/kumo/engine/pquery/prediction_table.py	PredictionTable.source, from_source_table	'pql' or 'custom'	prediction-table kind in the manifest
packages/kumo-engine/kumo/engine/trainer/_run_manifest.py	new: start, finish, store, write, output_records, destination_record, output_stats	builds the record (embedded model via embedded_model, scoring snapshot via embedded, stubs when missing; prediction table; scoring; stats from DiskGraph through _score_stats, incl. max_per top scores; no LP class counts); runs the output_config export recording completed destinations, writes the manifest, then stores <root>/predictions/<id>/ with MLOps on; every failure warns	prediction-run lineage, statistics and the root copy in one place
packages/kumo-engine/kumo/engine/trainer/trainer.py	PredictionResult.id, .manifest, .bundle_id, .snapshot_id, ._export, Trainer.predict(mlops=), Trainer._predict, _option_source, _sidecar_bytes, _holds_parquet	id and parent ids read from the manifest; Parquet directory exports carry _provenance.json; a plain export records nothing; mlops= resolved on the caller's thread; option sources recorded	record what each run scored and where it wrote
packages/kumo-engine/test/test_job_templates.py	test_fit_then_score_locally	asserts the predictions/<id>/ root copy and that it names the fitted model	the example score job stores its run with MLOps on
packages/kumo-engine/test/test_prediction_manifest.py	new	fields and embedded parents, sidecars, stub parent, MLOps root copy (and export off), plain export records nothing, no sidecar for warehouse/CSV, percentile accuracy, per-anchor top scores over unsorted rows, scoring sources, failure paths, status: failed after a raising export, byte-for-byte round trip	behaviour coverage

Requirement: every prediction run records its lineage, options, destinations and output score statistics, and with MLOps on is stored under the artifact root (final MLOps design, sections 1 and 3). Existing: the scorer's selection manifest, the sinks' export markers, the embedded-manifest helpers (#50/#56), and DiskGraph's batch-prediction statistics (#64 options). Smallest: one best-effort record per run, copied beside Parquet directory exports and to the root, never changing output bytes. Stopping condition: the record links run → model → training table → snapshot → source versions and carries the same score summary as the holdout baseline; every failure only warns.

[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-score-sketch
• File: .coderabbit.yaml
• Commit: d5b3c35
• Layers: Repository YAML
• Review profile: QUIET
• Plan: Enterprise
• Run ID: 7d53b27f-1e56-41c9-ab2a-228e9d2ecfda

Commits

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

Files selected for processing: 8

• 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/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_job_templates.py
• packages/kumo-engine/test/test_prediction_manifest.py

Walkthrough

The change adds prediction-run manifests with model and snapshot lineage, prediction-table metadata, scoring settings, output destinations, timing, and artifact statistics. Trainer.predict records run metadata and can store prediction artifacts through MLOps settings. PredictionResult exposes manifest data and attaches provenance sidecars to supported file exports. The package exports PredictionManifest. Tests and documentation cover manifest contents, export behavior, MLOps storage, and failure paths.

Sequence Diagram(s)

sequenceDiagram
  participant Trainer
  participant RunManifest
  participant PredictionResult
  participant ArtifactRoot
  Trainer->>RunManifest: start prediction manifest
  Trainer->>RunManifest: finish with output configuration and MLOps settings
  RunManifest->>PredictionResult: export configured outputs
  PredictionResult->>PredictionResult: attach provenance sidecars to supported file outputs
  RunManifest->>RunManifest: write manifest status and completion time
  RunManifest->>ArtifactRoot: store selected prediction artifacts when enabled

Loading

Priority: Normal

Merge Risk: Low · up to a8e18

Prediction exports remain usable on the established built-in paths, but provenance can be incomplete and some file-sink integrations may fail. Address these bounded issues or explicitly accept them before merging.

Pre-merge checks: 2 · 1

Failed checks: 1 inconclusive

Check name	Status	Explanation	Resolution
Change Map Design Review		I cannot assess the proposed Change map because the PR conversation comment beginning <!-- kumo-pr-change-map --> was not provided. The supplied context only summarizes the author’s change map and o…	Provide the full PR conversation comment that starts with <!-- kumo-pr-change-map -->. If the assessment depends on text omitted from the truncated PR description, provide its full text as well.

Passed checks: 2 passed

Check name	Status	Explanation
Pr Size		The pull request changes 8 files, based on the review-scoped base and head refs. The count does not exceed 40.
Public Resource Access		The pull request adds no external resource reference that requires public-access verification. The changed documentation and docstrings contain local artifact-root paths such as `/predictions/<i…

Full details: Public Resource Access

Explanation

The pull request adds no external resource reference that requires public-access verification. The changed documentation and docstrings contain local artifact-root paths such as <root>/predictions/<id>/predictions/; these describe output locations and are illustrative, not external resources. The added imports refer to repository modules. The diff contains no added or modified external URL, repository location, or package source location.

Full details: Change Map Design Review

Explanation

I cannot assess the proposed Change map because the PR conversation comment beginning <!-- kumo-pr-change-map --> was not provided. The supplied context only summarizes the author’s change map and omits its actual statements about the observed failure, existing behavior, smallest sufficient change, verification, and stopping condition. The available diff confirms the PR changed eight files, but it cannot establish whether the comment’s design claims are accurate. The change-map marker appears in repository templates and review guidance, not as the PR’s comment.

Resolution

Provide the full PR conversation comment that starts with <!-- kumo-pr-change-map -->. If the assessment depends on text omitted from the truncated PR description, provide its full text as well.

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: 3

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

61-63: Minor · Functional Correctness · Quick win

The sidecar next to exported Parquet files always has status: null and finished_at: null.

finish builds the sidecar from replace(manifest, outputs=records). At that point manifest.status is None and timings.finished_at is None. status has no omit_none(), so to_dict writes "status": null. The sidecar is written together with the Parquet parts, before the run ends. The final status and finish time reach only the scratch manifest and the artifact-root copy.

docs/mlops.md Line 94 and Lines 100-102 say that the manifest next to the outputs records status and timings. The test at Line 139 of test_prediction_manifest.py compares only outputs, so it does not detect the difference.

Do one of the following:

• Document that sidecars written during predict carry no final status or finished_at.
• Rewrite the sidecars after all exports succeed.

---

95-98: Minor · Data Integrity & Integration · Quick win

Remove partial artifact-root copies after a failed store.

store writes each selected artifact independently under predictions/<id>/<artifact>/. If a later write fails, earlier artifacts and their sidecars remain, while store only warns. The copy has no commit marker or partial status.

Compute base before the try block and delete that run directory when storage fails.

🧹 Suggested cleanup

+    base = f'predictions/{result.id}'
     try:
         selection = load_prediction_selection(result.selection,
                                               validate_objects=False)
         names = [item.artifact_type.value.lower()
                  for item in selection.artifacts]
-        base = f'predictions/{result.id}'
         result._export(
             PredictionOutputConfig(settings.root.connector, output_paths={
                 name: settings.root.location(f'{base}/{name}/')
                 for name in names
             }), manifest)
     except Exception as error:
+        try:
+            from kumo.engine.connector import file_store
+            filesystem, root = file_store(settings.root.connector,
+                                          settings.root.location(base))
+            filesystem.delete_dir_contents(root, missing_dir_ok=True)
+        except Exception:
+            pass
         warnings.warn(

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

279-284: Minor · Stability & Availability · Quick win

Add sidecars to the file-sink contract or make it optional at the call site.

Trainer._export passes sidecars= when _holds_parquet accepts a destination. FileArtifactSink.write_files does not declare that keyword, while is_file_destination accepts sinks through the separate FileStoreSink protocol. A strict-signature sink can satisfy both contracts, reach this branch, and raise TypeError before writing. Keep provenance best-effort by defining sidecar support in the protocol and implementations, or by skipping the sidecar when a sink does not support it.

---

Review info

Run configuration

• Source mode: Auto
• Repository: NVIDIA/kumo-engine
• Branch: vedaant/mlops-score-sketch
• File: .coderabbit.yaml
• Commit: cb72a69
• Layers: Repository YAML
• Review profile: QUIET
• Plan: Enterprise
• Run ID: 6916cc1d-61d7-46fb-a576-e6ea8792eae1

Commits

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

Files selected for processing: 8

• 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/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_job_templates.py
• packages/kumo-engine/test/test_prediction_manifest.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 0 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.