[Vedaant-J]

Stack: #53 → #112 → #50 → #60 → #64; #114 on #112; #115 on #64. Merge bottom to top.

Summary

Teams that score on a schedule cannot tell which model and data produced a table of scores, whether the scores moved, or how good they were once outcomes arrived. This PR records and monitors prediction runs:

• Every Trainer.predict writes a PredictionManifest (_provenance.json): model, snapshot and source-table versions, scoring options, destinations, skipped entities, and exact counts of its scores on fixed bins. With MLOps on, the run is also stored under <root>/predictions/<id>/ and tracked.
• Every model manifest keeps exact counts of its holdout scores as the baseline. Trainer.score_drift gives PSI, KS and Wasserstein against the holdout and a previous run, plus input-column drift; Trainer.evaluate gives realized metrics once outcomes are in a newer snapshot, with a DiskGraph join that keeps memory bounded. For link prediction, PSI against the holdout is a plain number with no shift verdict, since the holdout covers only entities with a future target; every other task, multilabel included, gets one.
• Every tracker logs prediction and monitoring runs through shared hooks (SingleRunHooks), with a Markdown run card and build time and size as prep.* tags. On MLflow, a monitoring call folds into its prediction run, and the fit's test/* holdout metrics are also logged on its logged model without a dataset link, so the Models tab and Compare versions show one test/<metric> column per version; the training table remains the run's input dataset. Snowflake runs carry the card and a kumo.model param, and tracking stops with a clear error at the experiment's 500-run limit.
• examples/mlops ships fit and score job templates for Databricks Jobs, SageMaker Pipelines and Snowflake Tasks; docs/mlops.md has recipes that feed exported evaluations to Databricks and Snowflake warehouse monitors.

Usage

cfg = kumo.MLOps("s3://<bucket>/kumo", tracking=MlflowTracking("churn"))
with kumo.mlops(cfg):
    run = trainer.predict(snapshot, output_config=kumo.PredictionOutputConfig(conn, path="scores"))
    drift = trainer.score_drift(run, previous=earlier_run)        # psi, ks, wasserstein, features, reasons
    evaluation = trainer.evaluate(run, labels_snapshot=newer)     # metrics, matured, labelled_fraction
evaluation.export(kumo.PredictionOutputConfig(conn, path="scores_evaluated", write_mode=kumo.WriteMode.APPEND))
later = kumo.PredictionResult.load_from_connector(cfg, path=f"predictions/{run.id}")

Back-compat

Situation	Behaviour
Table destination	Replaced, as before. WriteMode.APPEND adds PREDICTION_ID and BUNDLE_ID.
File destination that already holds data	Refused before scoring by default or with APPEND; OVERWRITE empties it first.
Bundle from before this PR	Loads and scores as before; drift vs holdout is None with a reason.
Manifest, statistics or store write fails	Warning, or the error under on_error="raise" after outputs are written; predictions unchanged.
Custom Tracking with fit hooks only	Still valid; prediction and monitoring hooks are optional.
diskgraph.compute_col_stats_for_parquet callers	New options are keyword-only and off; default output is byte-identical.

Testing

• kumo-core: test_score_counts.py, test_evaluation_counts.py, test_prediction_counts.py, test_score_drift.py check every count and drift statistic against a brute-force recount.
• kumo-engine: test_monitoring.py (25 tests: maturity, anchors, exports, metric definitions, exact drift), test_prediction_manifest.py (12), test_tracking_contract.py, test_mlflow_tracking.py, test_snowflake_tracking.py, test_card.py, test_prediction_table_output.py.
• test_job_templates.py runs fit_job.py then score_job.py end to end on local Parquet.
• DiskGraph: test_col_stats_options.py pins the default statistics output by SHA-256; Rust tests cover the evaluation join and value bins.

Change map: see the kumo-pr-change-map comment.

🤖 Generated with Claude Code

[Vedaant-J]

Change map

File	Function / scope	What changed	Why
packages/diskgraph/diskgraph-ingest/src/col_stats.rs	ValueBins, FileColStats	Exact value counts on fixed bins per numerical column	Input drift without a second data pass
packages/diskgraph/diskgraph-pybind/src/bp_col_stats.rs	ColStatsOptions, compute_parquet_col_stats_with_options	Percentiles, per-class stats, max_per, fixed-edge counts as options	Exact score counts at Parquet scale; default output unchanged
packages/diskgraph/diskgraph-pybind/src/bp_evaluate.rs	new: EvaluateSpec, evaluate_prediction_parquet	Partitioned join of predictions with labels; fit-time metrics	evaluate with memory bounded by partition size
packages/diskgraph/diskgraph-pybind/src/lib.rs	compute_col_stats_for_parquet, evaluate_prediction_parquet	Keyword-only stats options; new evaluation binding	Python access to both
packages/diskgraph/diskgraph-pybind/python/diskgraph/__init__.py	exports	Exports evaluate_prediction_parquet	Public binding
packages/diskgraph/diskgraph-pybind/python/diskgraph/tt_artifacts.py	prediction-table builder	No skipped-row warning when building a prediction table	Caller reports skipped entities once
packages/diskgraph/diskgraph-train-table/src/split_stats.rs	TargetHistogram, KumoStageStats	50-bin target histogram per split for node regression	Fixed bins for regression score counts
packages/diskgraph/diskgraph-train-table/src/pass_b.rs	pass B stage stats	Passes the target histogram into each split's stats	Histogram filled in the existing pass
packages/diskgraph/diskgraph-train-table/src/ttgen_stats.rs	target_range	Target min and max from the Welford stats	Histogram bin range
packages/diskgraph/diskgraph-train-table/src/ttgen_executor.rs	custom-source skip warning	Not logged when plan.is_prediction	It said "training" while scoring
packages/kumo-core/src/kumo_core/score_counts.py	new: Counts, bin_bounds, merged	Exact score counts on fixed bins, mergeable across batches	One count format for holdout and predictions
packages/kumo-core/src/kumo_core/evaluation_counts.py	new: evaluation_counts, evaluation_summaries	Exact test-split counts and holdout/ summaries per task	Holdout baseline and fit summaries without sampling
packages/kumo-core/src/kumo_core/prediction_counts.py	new: count_predictions, prediction_summaries	Per-run score, class and item counts and their summaries	Prediction run summaries and drift input
packages/kumo-core/src/kumo_core/score_drift.py	new: score_drift, psi, ks, wasserstein, psi_shift, holdout_shift	Exact drift between two runs' counts on the same bins	Score drift vs holdout and previous run; no holdout verdict for link prediction
packages/kumo-core/src/kumo_core/summaries.py	Scalar.role, Series.chart, Histogram.extent, Table.brief, Note	New summary fields and a Note type	Run cards and charts through the sinks seam
packages/kumo-core/src/kumo_core/prediction.py	GeneratedPredictionTable.row_failures, count_edges	Skipped-row report; scorers count scores into one file	Skipped entities and exact score counts
packages/kumo-core/src/kumo_core/pquery.py	_target_maximum	Min or max of the target tables' latest events	Label maturity for evaluate
packages/kumo-core/src/kumo_core/object_store_sink.py	ObjectStoreSinkBackend.preflight_write_files, write_files	Occupied directory refused before scoring unless OVERWRITE; covers FilesystemSink	Fail before scoring, with the full URI
packages/kumo-core/src/kumo_core/train/trainer.py	Trainer.test, _test_counts	Test split returns exact counts	Holdout baseline recorded at fit
packages/kumo-core/src/kumo_core/train/diskgraph_nested_trainer.py	test, predict	Test counts incl. link top scores; prediction batches counted	Same counts on the DiskGraph path
packages/kumo-core/src/kumo_core/train/output.py	EvalOutput.counts	New field for the test split's counts	Carries counts to the engine
packages/kumo-core/src/kumo_core/train/base_link_module.py	holdout_top_scores	Each test entity's top score	Link prediction holdout counts
packages/kumo-core/src/kumo_core/train/link_gnn_module.py	test step	Returns top scores	Link prediction holdout counts
packages/kumo-core/src/kumo_core/train/hybrid_link_gnn_module.py	test step	Returns top scores	Link prediction holdout counts
packages/kumo-core/src/kumo_core/train/id_link_gnn_module.py	test step	Returns per-entity max score	Link prediction holdout counts
packages/kumo-core/src/kumo_core/train/shallow_rhs_link_gnn_module.py	test step	Returns top scores	Link prediction holdout counts
packages/kumo-core/src/kumo_core/training/orchestration.py	TrainingOptions, result	Passes threshold and target edges; keeps the published model's counts	Baseline of the shipped model
packages/kumo-core/src/kumo_core/training/table.py	generate_train_table(validate_targets=)	False skips classification-target checks	Read outcome labels for evaluate
packages/kumo-engine/kumo/engine/__init__.py	exports	PredictionManifest, EvaluationResult, DriftReport	Public API
packages/kumo-engine/kumo/engine/_manifest.py	HoldoutStats, BuildCost, PredictionManifest and records	Holdout baseline, build cost, prediction run records	Typed record of every prediction run
packages/kumo-engine/kumo/engine/_monitoring.py	new: evaluate, score_drift, holdout_evaluation, EvaluationResult, DriftReport, track	Realized metrics, score and input drift, reasons for gaps	Monitoring of prediction runs
packages/kumo-engine/kumo/engine/_score_stats.py	new: output_counts	Counts a model's own prediction Parquet through DiskGraph	Score counts for models outside kumo's loop
packages/kumo-engine/kumo/engine/_card.py	new: heading, body, holdout, model_line	Markdown run card for fit and prediction runs	Readable run summary in each tracker
packages/kumo-engine/kumo/engine/_summary_tables.py	route(charts=), Routed.charts, confusion_table	Series and histograms become tables or chart metrics; confusion matrices become tables	Prediction and monitoring runs need these views
packages/kumo-engine/kumo/engine/_prep_cost.py	new: build_cost, prep_summaries	Build time and disk size as prep.* tags	Shown by every tracker on runs that use them
packages/kumo-engine/kumo/engine/_data_summaries.py	target_edges, _train_target	Target bin edges from the training split statistics	Holdout bins match the training target
packages/kumo-engine/kumo/engine/_tracking.py	SingleRunHooks, prediction_record, track_prediction, model_metrics, system_metrics, log_run	Shared prediction and monitoring hooks and records	No per-tracker copies
packages/kumo-engine/kumo/engine/integrations/mlflow.py	MlflowTracking: log_monitoring, _write_card, _update_card, _describe_version, _log_model_input, sample_system_metrics	Nested monitoring runs, run and version cards, test/* metrics on the logged model without a dataset, system metrics	MLflow side of prediction runs; Models tab compares versions
packages/kumo-engine/kumo/engine/integrations/snowflake.py	SnowflakeTracking: SingleRunHooks, _log_card, _check_room, _model_param	Run card, scoring model param, run-limit warning and refusal	Snowflake side of prediction runs
packages/kumo-engine/kumo/engine/graph.py	report_other_graph(strict=), _changed_columns, snapshot build cost	Names changed columns; raises under on_error='raise'	A model scores only its training graph in strict mode
packages/kumo-engine/kumo/engine/connector/base.py	Connector._monitored_files	Hook typing exported files for the platform monitor	Default leaves files unchanged
packages/kumo-engine/kumo/engine/connector/snowflake.py	SnowflakeConnector._monitored_files, _for_model_monitor	TARGET as 0/1 and UTC times without a zone for table destinations	Snowflake Model Monitor accepts the files
packages/kumo-engine/kumo/engine/connector/file.py	_FileSinkBackend.preflight_write_files	OVERWRITE accepts an occupied directory	OVERWRITE replaces a file destination
packages/kumo-engine/kumo/engine/pquery/prediction_table.py	PredictionTable.source, skipped_rows, skipped_entities	Source and skipped entities with one warning	Recorded in the prediction manifest
packages/kumo-engine/kumo/engine/pquery/predictive_query.py	_generate, _label_parts	Shared generation; label parts at given anchors	Labels for evaluate without a second workflow
packages/kumo-engine/kumo/engine/pquery/training_table.py	_record_manifest	Records build time and disk size	Data preparation cost on runs
packages/kumo-engine/kumo/engine/trainer/_provenance.py	holdout_counts, count_edges, write_provenance	Holdout counts in the model manifest; bins for scoring	Drift baseline travels with the model
packages/kumo-engine/kumo/engine/trainer/_run_manifest.py	new: start, finish, store, output_records	Builds, stores and tracks the prediction manifest	Owns the prediction run lifecycle
packages/kumo-engine/kumo/engine/trainer/_run_spec.py	fitted_spec, holdout_metrics, model_metrics	Holdout metrics named test/<metric>	Fit runs show the shipped model's quality
packages/kumo-engine/kumo/engine/trainer/output_config.py	PredictionOutputConfig.path, write_mode	path= for file or table; tables replace unless APPEND	One argument for every destination; ids on append
packages/kumo-engine/kumo/engine/trainer/trainer.py	predict, evaluate, score_drift, holdout_evaluation, PredictionResult	Prediction runs recorded; monitoring entry points	Public API for monitoring
packages/kumo-engine/pyproject.toml	databricks extra	Adds databricks-sdk>=0.68	Data Quality monitor recipe imports it
uv.lock	databricks extra	Lock follows pyproject	Lock in sync
tests/public/Dockerfile.dockerignore	image context	Admits kumo_core/artifact_export.py	Table-append tests import it
packages/diskgraph/diskgraph-pybind/python/tests/test_col_stats_options.py	new	Stats options from Python; default output pinned by hash	Python contract; default unchanged
packages/diskgraph/diskgraph-pybind/python/tests/test_tt_artifacts_row_warnings.py	builders	Training builder warns; prediction builder does not	Warning ownership
packages/diskgraph/diskgraph-train-table/tests/custom_tt_integration.rs	custom TT stats	No histogram for a custom table	Covers the histogram
packages/diskgraph/diskgraph-train-table/tests/ttgen_e2e.rs	ttgen stats	50-bin histogram counts every train row	Covers the histogram
packages/kumo-core/tests/test_score_counts.py	new	Batch and worker counts equal brute force	Counts contract
packages/kumo-core/tests/test_evaluation_counts.py	new	Holdout summaries per task match brute force	Exact holdout summaries
packages/kumo-core/tests/test_prediction_counts.py	new	Prediction counts per task match brute force	Exact prediction summaries
packages/kumo-core/tests/test_score_drift.py	new	Drift equals brute force; bins must match	Drift contract
packages/kumo-core/tests/test_prediction.py	scorer counts	Every scorer's counts merge into one file	Multi-worker counts
packages/kumo-core/tests/test_pquery.py	_target_maximum	Order independent	Covers the helper
packages/kumo-core/tests/test_train_table.py	validate_targets	False keeps degenerate labels	Covers the keyword
packages/kumo-core/tests/test_training.py	fake trainer	test takes keyword options	Follows the new signature
packages/kumo-databricks/tests/test_write.py	append	Appended files stamped with id columns	Table appends carry run ids
packages/kumo-snowflake/tests/test_sink.py	append	Appended files stamped with id columns	Table appends carry run ids
packages/kumo-engine/test/test_monitoring.py	evaluate, score_drift, export	Maturity, anchors, exports, metric definitions, exact drift	Monitoring contract
packages/kumo-engine/test/test_prediction_manifest.py	manifest	Fields, store, counts, failure paths, tracking	Prediction manifest contract
packages/kumo-engine/test/test_score_stats.py	new	Output counts match Counts on edges and top scores	Score count contract
packages/kumo-engine/test/test_fit_summaries.py	new	Holdout summaries count every row per task	Fit summary contract
packages/kumo-engine/test/test_card.py	new	Card headings, body, notes and locations	Card contract
packages/kumo-engine/test/test_summary_tables.py	route, histogram tables	Histogram tables close end bins at the extent	Table routing contract
packages/kumo-engine/test/test_data_summaries.py	target_edges	Even target bins without a histogram	Data summary contract
packages/kumo-engine/test/test_model_agnostic_predictions.py	new	A toy model's run is recorded and monitored	Seams take any model
packages/kumo-engine/test/test_tracking_contract.py	prediction_record, sinks	Prediction records, names, every summary type has a sink	Tracker contract
packages/kumo-engine/test/test_mlflow_tracking.py	MLflow hooks	Nesting, system metrics, tables, charts, cards, renamed fit	MLflow behaviour
packages/kumo-engine/test/test_mlflow_model.py	alias, logged model, version card	Alias resolves once; test/* on the logged model; version description; partial scoring warns	MLflow model behaviour
packages/kumo-engine/test/test_snowflake_tracking.py	Snowflake hooks	Card, duration, prediction names, run limit, monitoring runs	Snowflake behaviour
packages/kumo-engine/test/test_job_templates.py	new	Fit then score locally; alias moves; appends; task and step args	Templates keep working
packages/kumo-engine/test/test_prediction_table_output.py	new	Table replace vs append with ids; path=	Output contract
packages/kumo-engine/test/test_prediction_output_preflight.py	preflight	Occupied destination fails before scoring; OVERWRITE passes	Checks before scoring
packages/kumo-engine/test/test_prediction_output_columns.py	counts	Forecast, item and top-score counts	Per-task counts
packages/kumo-engine/test/test_filestore_sink.py	file store	Occupied directory refused by full location	Error names the destination
packages/kumo-engine/test/test_query_from_snapshot.py	skipped entities	Counts, sample, reason, warning, metrics	Skipped entities
packages/kumo-engine/test/test_graph_mismatch.py	strict mode	Raises on schema change; names changed columns	Covers strict mode
packages/kumo-engine/test/test_holdout_df.py	holdout record	Matches the holdout frame per fit mode	Holdout record contract
packages/kumo-engine/test/test_artifacts.py	prediction artifacts	A prediction run's snapshot input and outputs	Artifact listing
packages/kumo-engine/test/test_bundle_provenance.py	stub fit	Passes holdout_truncated	Follows the fit record
packages/kumo-engine/test/test_scoring_defaults.py	derived model	Keeps the parent's holdout	Derived-model baseline
packages/kumo-engine/test/test_trainer_registry.py	registered model	Loaded model keeps its registered version	Registry contract
packages/kumo-engine/test/test_training_table_export.py	build cost	Records disk bytes under MLOps only	Covers build cost
packages/kumo-engine/test/test_snapshot_manifest.py	build cost	Snapshot under MLOps records its build	Covers build cost
packages/kumo-engine/test/test_predictive_query.py	fakes	Accept keyword args; stub row_failures	Follows new signatures
packages/kumo-engine/test/test_mlops.py	kinds	Includes prediction	Export kinds
packages/kumo-engine/docs/mlops.md	Score, Monitor, Reference	Prediction runs, monitoring, warehouse monitors, drift definitions	User docs for this PR
docs/development/testing.md	DiskGraph build	Rebuild DiskGraph from the same tree	Stale builds miss new APIs
examples/mlops/README.md	new	Job scripts, platforms, flags	Scheduler entry point
examples/mlops/jobs.py	new	Shared flags, graph, connectors, MLOps settings	One module both jobs share
examples/mlops/fit_job.py	new	Snapshot, table, fit and register in an MLOps block	Scheduled training template
examples/mlops/score_job.py	new	Load by alias or path, score, store the run	Scheduled scoring template
examples/mlops/databricks/README.md	new	Deploy and run the Asset Bundle	Databricks setup
examples/mlops/databricks/databricks.yml	new	Train and daily score jobs with variables	Databricks Jobs template
examples/mlops/sagemaker/README.md	new	Build and run the pipeline with managed MLflow	SageMaker setup
examples/mlops/sagemaker/pipeline.py	new	Fit and score pipeline steps (SDK v3)	SageMaker Pipelines template
examples/mlops/sagemaker/Dockerfile	new	GPU or CPU image from a wheel directory	The pipeline's image
examples/mlops/snowflake/README.md	new	Register jobs and Tasks; experimental	Snowflake setup
examples/mlops/snowflake/register_jobs.py	new	ML Job definitions and scheduled Tasks	Snowflake schedule template
packages/kumo-engine/kumo/engine/_text.py	number	Number formatting for cards	Used by card and trackers
packages/kumo-engine/test/test_prediction_completion.py	stub training result	Adds manifest=None to the stub	Matches the new training result field
packages/kumo-engine/test/test_prediction_workers.py	stub training result	Adds manifest=None to the stub	Matches the new training result field
packages/kumo-engine/test/test_semantic_carry_forward.py	stub results	Adds manifest=None and published_counts=None to stubs	Matches the new result fields
packages/kumo-engine/test/test_training_explanations.py	stub result	Adds published_counts=None to the stub	Matches the new result field
packages/kumo-engine/kumo/engine/_predictor.py	predict_frame	Scores under the active MLOps settings	Pyfunc scoring is recorded like Trainer.predict

[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: 1b2d873

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

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