[Vedaant-J]

Series: kumo-engine MLOps series: every snapshot, training table, model bundle and prediction run records where it came from; snapshots and models move between machines and registries; fits and prediction runs are tracked in MLflow or Snowflake; plain Python job scripts run on a scheduler through a small job helper; prediction runs are monitored for drift and realized performance.
Stack: Prediction runs and monitoring: #58 → #64 → #67 → #69 → #73
Merge order: one linear stack; merge bottom to top (#53 first)
This PR: Adds opt-in lineage columns on prediction outputs: KUMO_RUN_ID, KUMO_BUNDLE_ID and KUMO_SNAPSHOT_ID written into every output row by the scorer, so a row in a downstream table can be traced back to the prediction run, model bundle (#56) and graph snapshot (#50) that produced it.

Summary

This PR adds opt-in lineage columns on prediction outputs: KUMO_RUN_ID, KUMO_BUNDLE_ID and KUMO_SNAPSHOT_ID written into every output row by the scorer, so a row in a downstream table can be traced back to the prediction run, model bundle (#56) and graph snapshot (#50) that produced it.

Problem. Teams consuming Kumo predictions downstream need to answer questions like these:

• A row in the scores table looks wrong. Which prediction run wrote it?
• Which model bundle produced that score, and was it the one we promoted?
• Which graph snapshot (and so which source table versions) was it scored against?
• The table is appended to daily. Which rows came from which run?

Today none of these can be answered from the exported table.

Cause. An exported predictions Parquet carries only the prediction columns:

ENTITY: int64
TARGET_PRED: bool
False_PROB: float
True_PROB: float

The run, bundle and snapshot ids exist in the process that ran predict, but nothing writes them into the output.

Change. PredictionOutputConfig(lineage_columns=False) is a new flag. With it set to True and the config passed to Trainer.predict, the scorer's Parquet writer (kumo_core.train.io.StreamingFrameWriter, new constant_columns) appends the three ids as constant string columns to every row group it writes, for predictions and embeddings. There is no extra pass over the output: the columns are added in the same write, so result.files, result.frame and every export (predict(output_config=...) and later result.export(...)) carry them.

Exported predictions Parquet from a binary-classification fit on a small synthetic users/items/orders graph (PREDICT COUNT(orders.*, 0, 7, days) > 0 FOR EACH users.ID, 50 rows), with lineage_columns=True:

ENTITY: int64
TARGET_PRED: bool
False_PROB: float
True_PROB: float
KUMO_RUN_ID: string        # new
KUMO_BUNDLE_ID: string     # new
KUMO_SNAPSHOT_ID: string   # new

First rows (ids shortened):

ENTITY	TARGET_PRED	False_PROB	True_PROB	KUMO_RUN_ID	KUMO_BUNDLE_ID	KUMO_SNAPSHOT_ID
0	True	0.060519	0.939481	791f22bc…-predict-c26d64d5…	c584b62b…	52f642b6…@b747cbf1bf62f69d
8	True	0.061139	0.938861	791f22bc…-predict-c26d64d5…	c584b62b…	52f642b6…@b747cbf1bf62f69d
40	True	0.062343	0.937657	791f22bc…-predict-c26d64d5…	c584b62b…	52f642b6…@b747cbf1bf62f69d

The export directory itself is unchanged (part-00000-00000.parquet, _KUMO_EXPORT, _SUCCESS); only the Parquet columns differ.

What each column records:

• KUMO_RUN_ID is the prediction job id (<trainer id>-predict-<run uuid>). The export identity in _KUMO_EXPORT is derived from it.
• KUMO_BUNDLE_ID is the 64-hex bundle_id of the model that scored (from Record bundle_id and fit provenance in every model bundle #56). It matches Trainer.bundle_id and survives registry round trips, so it identifies the promoted model.
• KUMO_SNAPSHOT_ID is the snapshot id (<graph key>@<first 16 hex of the data version>), which matches snapshot-manifest.json (Record provenance for every artifact under one MLOps artifact root #50) and leads to the source table versions.

Details:

• The columns are constant per prediction run. An unknown id (for example no bundle_id on a bundle that predates Record bundle_id and fit provenance in every model bundle #56) is omitted rather than written as null.
• result.export(config) copies whatever columns the result already has; it does not add or remove them. Asking for lineage_columns=True on a result written without them emits a UserWarning and exports the result as it is.
• PredictionResult now records bundle_id and snapshot_id. Both are optional and excluded from equality and hashing.
• The lineage ids are part of the prediction's source identity only when lineage is on, so the default identity is unchanged.
• Not added to MetadataField, which is part of the hosted API.

Cost. Writing 10M rows of a 3-column predictions table through the scorer's writer with the three ids (run id 50 characters, bundle and snapshot ids 64):

	Writer time	Parquet size	Decoded (Arrow / warehouse row)
lineage off	0.91 s	89.7 MB	12.4 B/row
lineage on	2.59 s	90.0 MB (+0.3%, dictionary-encoded)	+190 B/row

Parquet files barely grow, but anything that decodes the rows (result.frame, CSV exports, warehouse tables) pays roughly 200 bytes per row for the three strings. This is the main reason the flag is off by default.

Why opt-in. Turning it on changes the destination's column set. On a warehouse table created without lineage columns:

Sink and write mode	Behaviour with lineage_columns=True
Snowflake OVERWRITE	Fails.
Snowflake APPEND	Succeeds, but silently drops the new columns.
BigQuery APPEND	Fails.
Databricks APPEND	Fails.
New table or directory	Lineage columns written.

So lineage should be enabled on a new table or directory. The README lists the per-sink behaviour above; the PredictionOutputConfig docstring says to enable it only on a new table or directory. Every other column is written exactly as without lineage, on every sink, because the export path is unchanged.

Back-compat

This PR only records and exposes information. With the flag at its default (False) the scorer writes exactly what it wrote before.

Situation	Behaviour
Existing predict/export call (flag not set)	Output files, result.frame and sink calls identical to before.
lineage_columns=True on a new destination	Three string columns appended after the existing ones; other columns keep their types; row count and row groups unchanged.
lineage_columns=True on an existing warehouse table	See the table above: most sinks fail, Snowflake APPEND drops the columns.
result.export(lineage_columns=True) on a result written without lineage	UserWarning; export proceeds without the columns.
PredictionResult built without ids (older code path)	Equal and hash-equal to one with ids.
Output read by an older engine or other tools	Extra columns are ordinary string columns; nothing in the engine reads them.

Known limitations

• Lineage is chosen at prediction time. A result scored without it cannot gain the columns on a later export; score again with the flag on.
• Each decoded row pays ~200 B for the three ids (see Cost).

Testing Done

• Engine (test_prediction_output_columns.py): one prediction with lineage on and one with it off, for predictions and embeddings. Column lists differ only by the three columns; every original column keeps its Arrow type; row counts are equal; values equal the run, bundle and snapshot ids; result.frame carries them; a later export without the flag copies them unchanged. A second test checks the warning for lineage on a result written without it, and that result equality and hashing ignore the new ids.
• kumo-core (test_train_output.py): the streaming writer appends the constant columns to every rotated part, keeps the source columns equal, and its recorded schema fingerprint and column count match the written files.
• On the current head (on a merge of the PRs below it): engine test_prediction_output_columns.py + test_prediction_completion.py + test_holdout_df.py: 52 passed; kumo-core test_train_output.py + test_prediction.py + test_prediction_artifact_stats.py: 121 passed (CPU).

Checklist

• No breaking changes (default off).
• README updated.
• Unit tests added.

🤖 Generated with Claude Code

[Vedaant-J]

Change map

File	Function / scope	What changed	Why
packages/kumo-core/src/kumo_core/prediction.py	_PredictionWorkerConfig, _score_prediction_shard, score_prediction	New constant_columns carried to each scoring worker	Spawned scorers write the same columns
packages/kumo-core/src/kumo_core/train/diskgraph_nested_trainer.py	DiskgraphNestedTrainer.predict, _predict_lp	New output_constant_columns, forwarded to every StreamingPredOutputWriter	Node and link prediction scorers both stamp
packages/kumo-core/src/kumo_core/train/io.py	StreamingFrameWriter.__init__, _file_schema, _stamped, _serialize, close	New constant_columns: string columns appended to each row group in the serializer before write_table; file schema and recorded stats schema include them	Stamp lineage in the one write the scorer already does, with no extra pass
packages/kumo-core/src/kumo_core/train/output.py	StreamingPredOutputWriter.__init__, write_embedding	Passes constant_columns to the predictions and embeddings writers	Both artifacts carry the same lineage
packages/kumo-core/tests/test_train_output.py	test_constant_columns_are_appended_to_every_part	New test: columns on every rotated part, source columns equal, stats fingerprint matches the files	Writer contract
packages/kumo-engine/README.md	Prediction output	Documents lineage_columns, per-sink effect on existing tables	Users must enable it on a new destination
packages/kumo-engine/kumo/engine/trainer/output_config.py	PredictionOutputConfig	New lineage_columns: bool = False and docstring	Opt-in switch
packages/kumo-engine/kumo/engine/trainer/trainer.py	PredictionResult, PredictionResult.export, _has_lineage_columns, Trainer.predict	predict builds the run/bundle/snapshot ids and passes them as constant_columns (and into the source identity only when on); result records bundle_id/snapshot_id; export warns when lineage is asked of a result without it	Ids come from the predicting process; export never re-encodes files
packages/kumo-engine/test/test_prediction_output_columns.py	test_lineage_columns_are_written_by_the_scorer, test_lineage_export_of_a_result_without_them_warns	Lineage on vs off for predictions and embeddings; later export copies columns; warning path; equality ignores ids	User-visible behaviour of the flag
test/train/test_kumo_pred_ttng.py	test_ttng_link_pred_predict_receives_runtime_and_top_k	expects output_constant_columns=None in the forwarded link-prediction call	the new keyword reaches _predict_lp

Requirement: exported predictions can be traced to their run, bundle and snapshot, opt-in. Smallest change: stamp constant columns in the scorer's existing Parquet writer (replaces the earlier export-time DiskGraph Parquet splicer, which is removed). Stopping condition: lineage on adds exactly three string columns everywhere; lineage off output is 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]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

• Source mode: Auto
• Repository: NVIDIA/kumo-engine
• Branch: vedaant/mlops-job-templates
• File: .coderabbit.yaml
• Commit: f21cf92
• Layers: Repository YAML
• Review profile: QUIET
• Plan: Enterprise
• Run ID: 5a9fe057-c3e5-416d-ab23-c3dd7a84f3a2

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

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

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

[Vedaant-J]

Closing: this PR's pieces landed elsewhere or were dropped in the final MLOps design.

• Lineage columns (KUMO_RUN_ID/KUMO_BUNDLE_ID/KUMO_SNAPSHOT_ID, PredictionOutputConfig.lineage_columns, the scorer's constant_columns plumbing) are dropped. Output rows stay exactly as on main.
• Lineage now lives in manifests: every prediction run writes _provenance.json (in the run directory, beside each Parquet export, and under <artifact root>/predictions/<id>/), which embeds the model and scoring-snapshot manifests and lists the run's outputs — Record a manifest and output score statistics for every prediction run #67.
• PredictionResult.bundle_id / .snapshot_id are back in Record a manifest and output score statistics for every prediction run #67, read from that manifest, with the run id as PredictionResult.id.
• The Databricks bare-table-name fix shipped separately in Write a bare Databricks table name into the connector schema #94.