[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 prediction-run tracking: with kumo.MLOps(tracking=...), each prediction run is recorded in MLflow or Snowflake under the same kumo.* tags as the fit run, built from its manifest.

Summary

This PR adds prediction-run tracking. With kumo.MLOps(tracking=...), every Trainer.predict is recorded as a run in MLflow or Snowflake Experiment Tracking, in the same experiment as the training runs and through the same Tracking adapters as fits. It closes the prediction side of the MLOps series: #67 writes a manifest for every prediction run, #60 and #68 record fits, and this PR sends the prediction manifest to those trackers.

Problem. Teams that score on a schedule need to answer questions like these:

• Which model produced today's scores? Can I see that in MLflow or Snowsight next to the training runs?
• Did the score distribution or predicted positive rate move between daily runs?
• Which source tables did today's run read?

Since #67 the answers exist, but only in each run's _provenance.json. Today none of them can be answered from the experiment tracker where the training runs already live.

Cause. The Tracking protocol (#60, #68) records fits: start(record), on_trial(record, event) and finish(record, error), fed by a TrackingSession on a background thread. It lists log_prediction_run(record) as an optional hook, but nothing calls it and nothing maps a prediction manifest into a RunRecord.

before                                         after
with kumo.mlops(MLOps(..., tracking=t)):       with kumo.mlops(MLOps(..., tracking=t)):
    trainer.fit(...)        # tracked              trainer.fit(...)        # unchanged
    trainer.predict(...)    # not tracked          trainer.predict(...)    # new: one run per prediction

Change.

1. Core mapping (_tracking.prediction_record(manifest, error=None) -> RunRecord, no platform imports), the only place prediction-run names are defined:

   Part	Contents
   name	the prediction id
   params	the manifest's scoring options (binary_classification_threshold, threshold_source, num_classes_to_return, count_source, num_workers) plus anchors; None options are left out
   metrics	the manifest stats, flattened with the same helper as fit params: prediction_table.rows; <artifact>.rows; per score column (True_PROB, TARGET_PRED or top_score) count, mean, min, max and seven percentile levels q01, q05, q25, q50, q75, q95, q99 (the manifest keeps all 99); predictions.predicted_positive_rate (binary)
   inputs	the scoring snapshot's source tables, read from the embedded parent by the same fit_inputs helper fits use (now also for prediction manifests); MLflow logs them as run inputs
   tags	pointers only: 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

2. When. Inside predict, after the output_config export, the manifest write and the artifact-root copy (Record a manifest and output score statistics for every prediction run #67), _run_manifest.finish calls track_prediction(manifest, settings.tracking, error). It goes through Evaluate prediction runs and report score drift #69's log_run: the tracker's optional log_prediction_run hook runs on a tracking thread, only on global rank 0, for up to 30 s, so a hung tracker cannot hold predict, and its errors are logged, never raised. If the export raises, the run is still recorded as failed (kumo.error = <Type>: <message>, the same formatting as fit runs) and predict raises the export's error. Settings are resolved on the caller's thread (Record provenance for every artifact under one MLOps artifact root #50).

3. Adapters. Each implements the hook as its own start(record) then finish(record, record_failure(record)), so params, metrics, tags, inputs and artifacts go through the same code as a fit run.

   • MlflowTracking: one run named by the prediction id in the configured experiment, nested under the caller's active run if there is one (the same rule as fits). Params MLflow rejects are counted in kumo.dropped_params.
   • SnowflakeTracking: one run named by the (quoted) prediction 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 adapter's run_id attribute is now the most recent tracked run, fit or prediction.

4. Docs. docs/mlops.md (Experiment tracking) describes prediction runs and adds prediction / monitoring to the tag table.

from kumo.engine.integrations.mlflow import MlflowTracking

cfg = kumo.MLOps(artifact_root="s3://<bucket>/kumo",
                 tracking=MlflowTracking(experiment_name="users-churn"))
with kumo.mlops(cfg):
    scores = trainer.predict(snapshot, output_config=output,
                             binary_classification_threshold=0.5)

Real example. Synthetic graph: 500 users, 4,000 orders, query PREDICT COUNT(orders.*, 0, 7, days) > 0 FOR EACH users.ID. A short fit and a predict inside one kumo.mlops block with MlflowTracking, read back with MlflowClient.search_runs from a local SQLite store (ids shortened, metrics rounded):

{
  "run_name": "972e73fe…-predict-550b108f…",
  "status": "FINISHED",
  "tags": {"kumo.kind": "prediction", "kumo.run_id": "972e73fe…-predict-550b108f…", "kumo.bundle_id": "fa2d033d…"},
  "params": {"anchors": "[\"2024-04-29T00:00:00+00:00\"]", "binary_classification_threshold": "0.5",
             "threshold_source": "argument", "count_source": "none", "num_workers": "1"},
  "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.3485,
              "predictions.True_PROB.min": 0.3176, "predictions.True_PROB.max": 0.3663,
              "predictions.True_PROB.q01": 0.324, "predictions.True_PROB.q05": 0.329, "predictions.True_PROB.q25": 0.3418,
              "predictions.True_PROB.q50": 0.3504, "predictions.True_PROB.q75": 0.356,
              "predictions.True_PROB.q95": 0.3615, "predictions.True_PROB.q99": 0.3643},
  "inputs": ["orders", "users"],
  "artifacts": ["_provenance.json"]
}

The fit run in the same experiment carries the same kumo.bundle_id (#65), so filtering on it returns the training run and every prediction run of that model; kumo.kind = prediction separates scoring runs from training runs. Every daily run has the same metric names, so predictions.True_PROB.q50 or predictions.predicted_positive_rate can be charted across runs. For PSI, KS and Wasserstein against the holdout or another run, Trainer.score_drift (#69) reads the full curves from the manifests and logs a monitoring run that MLflow links back with kumo.monitored.

Back-compat

This PR only records and exposes information. Tracking is opt-in, and a tracking failure never changes or fails a prediction.

Situation	Behaviour
predict() with MLOps off, or on without tracking	Unchanged: no tracker code runs.
predict() with MLOps tracking on	One run per call if the tracker has log_prediction_run, after the export and the artifact-root copy. A test compares the frames of tracked and untracked runs.
The output_config export raises	The run is recorded and ended failed with kumo.error; the export's error is raised as before.
Existing user Tracking implementations (fit hooks only)	Still satisfy the protocol; they get no call for predictions.
Fit runs	Unchanged.
No manifest (it could not be built)	Not recorded, with a warning. The result is returned as before.
Tracker error or hang (MLflow store down, Snowflake error, …)	One warning; the prediction is unaffected; predict waits at most 30 s.
Snowflake instance in use, or run name taken	Not recorded there / recorded under a suffixed name, with a warning.
Non-zero rank	Nothing is logged.

Known limitations

• MLflow and Snowflake time a run from when it is logged. The prediction's own start and finish times are in the manifest artifact.
• No live Snowflake run was made; the Snowflake path is covered by the recording fake from Record Trainer.fit in Snowflake Experiment Tracking through kumo.MLOps(tracking=) #68.

Testing Done

• test_tracking_contract.py (13 passed):
  • the record of a recorded binary prediction manifest: params, metrics with seven percentile levels, pointer tags, inputs from the embedded scoring snapshot, the manifest artifact; a partial manifest (stub parents, num_classes_to_return, a top_score column without percentiles, embeddings rows, a failure tag);
  • the recorded manifest fed to both adapters (MLflow on a local store, the Snowflake fake), with what each stored read back and compared with prediction_record;
  • no manifest, or a tracker without the hook, records nothing.
• test_prediction_manifest.py: with MLOps(tracking=recorder) the recorder gets prediction_record(result.manifest) and the frame equals an untracked run's; a raising tracker only warns; an export that raises in a tracked predict re-raises, records a failed run with kumo.error=OSError: disk full and the manifest's id, and stores nothing in the root.
• test_mlflow_tracking.py: a prediction run nests under an active user run, which stays active, and otherwise lands top-level in the named experiment under the prediction id.
• test_snowflake_tracking.py: with the instance in use, the prediction run is not recorded and the instance is untouched; logging the same record twice creates a suffixed run instead of resuming.
• On this head (extension rebuilt from it): the suites above plus test_monitoring, test_score_stats, test_bundle_provenance, test_holdout_df, test_scoring_defaults, test_prediction_completion, test_job_templates, test_query_from_snapshot, test_mlflow_model: 125 passed, 8 skipped (checks against the installed snowflake-ml, which is not installed); test_predictive_query + test_optional_imports: 118 passed.
• 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 (Experiment tracking)	prediction runs paragraph; prediction/monitoring kinds and kumo.prediction_run in the tag table; hook signatures	user docs
packages/kumo-engine/kumo/engine/_tracking.py	prediction_record, record_failure, track_prediction, _logged_percentiles, fit_inputs	maps a prediction manifest to a RunRecord: scoring options + anchors as params; stats flattened (seven percentile levels) as metrics; the scoring snapshot's source tables as inputs (fit_inputs now reads prediction manifests too); pointer tags kumo.kind/run_id/bundle_id plus kumo.error when the export failed; _provenance.json as the artifact. track_prediction logs it through #69's log_run and never raises	prediction-run names defined once in core; tracking never stalls predict
packages/kumo-engine/kumo/engine/integrations/mlflow.py	MlflowTracking.log_prediction_run, docstrings	start(record) + finish(record, record_failure(record)), FAILED when the record has kumo.error	MLflow adapter reusing the fit path
packages/kumo-engine/kumo/engine/integrations/snowflake.py	SnowflakeTracking.log_prediction_run, docstrings	start(record) + finish(record, record_failure(record)) with the fit runs' never-resume rule	Snowflake adapter reusing the fit path
packages/kumo-engine/kumo/engine/trainer/_run_manifest.py	finish	after the export, the manifest write and the root copy, tracks the run with settings.tracking, failed when the export raised	the ordering of section 3: export → store → track
packages/kumo-engine/test/test_mlflow_tracking.py	test_prediction_run_nests_under_the_active_run	nesting and experiment rule	MLflow specifics
packages/kumo-engine/test/test_prediction_manifest.py	_Recorder, test_failed_export_still_writes_and_tracks_the_run, test_tracking_records_the_run_and_never_changes_it	tracked vs untracked frame equality; the recorder gets the record; a raising tracker warns; a failing export records a failed run, re-raises and stores nothing	API behaviour on a real fit
packages/kumo-engine/test/test_snowflake_tracking.py	test_an_instance_in_use_is_left_alone, test_prediction_runs_are_never_resumed	in-use instance untouched; a re-logged run gets a suffixed name	Snowflake specifics
packages/kumo-engine/test/test_tracking_contract.py	prediction manifest fixture, Snowflake observer run names, 4 tests	record of a full and a partial manifest; both adapters vs prediction_record; no manifest / no hook → no calls	contract test

Requirement: each prediction run must appear in the experiment tracker that holds the training runs, pointing to the model that produced it (final MLOps design, section 3). Existing: the #67 manifest has the facts, the #60/#68 adapters log fits, #69's log_run logs a standalone run. Smallest change: one core mapping from manifest to RunRecord, called at the end of predict with MLOps tracking on, and adapter hooks that reuse their fit start/finish. Stopping condition: one run per tracked predict with pointer tags, options, stats, inputs and the manifest; predict without tracking 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]

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-monitoring
• File: .coderabbit.yaml
• Commit: 4aba059
• Layers: Repository YAML
• Review profile: QUIET
• Plan: Enterprise
• Run ID: 73b7644a-c672-42fe-a712-471faaf45bec

Commits

Reviewing files that changed from the base of the PR and between 5b51082 and 1d3acd6.

Files selected for processing: 9

• packages/kumo-engine/docs/mlops.md
• 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/trainer/_run_manifest.py
• packages/kumo-engine/test/test_mlflow_tracking.py
• packages/kumo-engine/test/test_prediction_manifest.py
• packages/kumo-engine/test/test_snowflake_tracking.py
• packages/kumo-engine/test/test_tracking_contract.py

Walkthrough

Prediction runs are built from manifests and sent to tracking integrations. Prediction completion records export failures and tracks runs when settings are present. Prediction artifacts are stored only after successful export. MLflow and Snowflake integrations now handle prediction records and their failure status.

Sequence Diagram(s)

sequenceDiagram
  participant RunManifestFinish
  participant TrackPrediction
  participant MlflowTracking
  participant MLflow
  RunManifestFinish->>TrackPrediction: pass manifest and export error
  TrackPrediction->>MlflowTracking: submit prediction record
  MlflowTracking->>MLflow: start run and finish with record failure status

Loading

Priority: Normal

Merge Risk: Moderate · up to 5b510

Prediction tracking is opt-in, and predictions themselves are not affected by tracking errors. Two bounded issues should be fixed before merge. A slow tracking call can overlap a later one and mis-record or leave runs open. A tracking failure partway through starting a run can leave an MLflow or Snowflake run open and block later tracking.

Pre-merge checks: 2 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Change Map Design Review		The proposed design uses the correct core layer: _tracking.py builds a platform-neutral RunRecord, _run_manifest.finish submits it after export, and adapters provide the optional hook. The desig…	Serialize standalone tracking operations per tracker, or give each operation an independent adapter state. Keep the 30-second caller timeout, but hold the serialization until the worker actually exits. Wrap adapter startup so any exception …

Passed checks: 2 passed

Check name	Status	Explanation
Pr Size		The pull request changes 9 files. The limit is 40 files, so the check passes.
Public Resource Access		The pull-request diff adds no external URL, repository location, or package source location. The MLflow and Snowflake pip install references remain unchanged text within modified docstrings, and the…

Full details: Public Resource Access

Explanation

The pull-request diff adds no external URL, repository location, or package source location. The MLflow and Snowflake pip install references remain unchanged text within modified docstrings, and the new test identifiers and internal imports are not external resource references.

Full details: Change Map Design Review

Explanation

The proposed design uses the correct core layer: _tracking.py builds a platform-neutral RunRecord, _run_manifest.finish submits it after export, and adapters provide the optional hook. The design does not preserve the stated stopping condition of one correct tracked run per prediction. log_run starts a new worker for every call and returns after 30 seconds, while MlflowTracking and SnowflakeTracking mutate shared run state in start and finish. A later prediction can therefore reset or finish the first prediction's run after the first worker continues in the background. The adapters also do not clean up a run when start creates it and _log_params, _log_artifacts, or _log_inputs raises: log_prediction_run calls finish only after start returns. This leaves an MLflow run and HTTP guard open, or leaves Snowflake's singleton run active and blocks later tracking. The changed tests cover normal runs and a hung call, but not these two failure paths.

Resolution

Serialize standalone tracking operations per tracker, or give each operation an independent adapter state. Keep the 30-second caller timeout, but hold the serialization until the worker actually exits. Wrap adapter startup so any exception after run creation calls cleanup (finish(None, error) or an equivalent adapter-specific cleanup) before re-raising for the worker logger. Add focused tests for a timed-out prediction followed by another prediction, and for failures in MLflow and Snowflake parameter, artifact, and input logging; verify that each run is finished correctly and that a later prediction can track.

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]

Actionable comments posted: 2

---

Review info

Run configuration

• Source mode: Auto
• Repository: NVIDIA/kumo-engine
• Branch: vedaant/mlops-monitoring
• File: .coderabbit.yaml
• Commit: 0c39510
• Layers: Repository YAML
• Review profile: QUIET
• Plan: Enterprise
• Run ID: 14c0939a-0a91-42de-abb5-0e04c2cd0201

Commits

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

Files selected for processing: 9

• packages/kumo-engine/docs/mlops.md
• 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/trainer/_run_manifest.py
• packages/kumo-engine/test/test_mlflow_tracking.py
• packages/kumo-engine/test/test_prediction_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; 1 remain after this review.

[coderabbitai]

Major · Stability & Availability · Heavy lift

Prevent overlapping runs on one tracker.

If a prediction’s tracking call exceeds the 30-second wait, its worker continues. A later predict call can start another worker on the same tracker. Both adapters reset and update run state on that instance, so the workers can assign records to the wrong run or leave a run unfinished. Serialize calls per tracker or give each prediction an independent run session while preserving the timeout.

[coderabbitai]

Major · Stability & Availability · Heavy lift

Finish a run when start fails after creating it.

If MLflow creates the run but _log_params, _log_artifacts, or _log_inputs raises, this hook never calls finish. The tracking worker suppresses the exception, leaving the run open and the HTTP-settings guard unclosed. SnowflakeTracking.log_prediction_run has the same failure path; its active run can also block later tracking attempts. Clean up a partially started run in both hooks, then let the tracking worker report the failure.

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