[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: Models, registry and job templates: #57 → #59 → #65 → #71
Merge order: one linear stack; merge bottom to top (#53 first)
This PR: Adds scoring defaults carried in the model bundle: a threshold or class count chosen after fit is stored in a derived bundle with its own id, and predict uses it when a call does not pass one.

Summary

This PR adds scoring defaults carried in the model bundle: a threshold or class count chosen after fit is stored in a derived bundle with its own id, and predict uses it when a call does not pass one. It also adds Trainer.export, so a fitted or loaded model can be published outside fit().

Problem. Teams scoring on a schedule run into these:

• The threshold chosen after looking at the holdout lives in a job script, so promoting a new model means changing that script.
• Two jobs scoring with the same model use different thresholds, and nothing in the model says which one is intended.
• A model that was loaded, inspected and tuned cannot be published to a registry; only fit() can publish.

Today none of these can be solved without a code change in every job that calls predict.

Cause. The threshold depends on holdout predictions, so it exists only after fit() has written and sealed the bundle. The bundle has no place for it:

<bundle>/
  model/  checkpoint  config.json  encoders.json  resolved-model-plan.json  target-encoder.json
  manifests/  graph.json  input.json  task.json
  metrics.json  prediction-schema.json  training-spec.json  _provenance.json
  publication-manifest.json            # lists and hashes every file above; the bundle id is derived from it

Editing a file inside it would change the bundle without anything recording the edit. predict therefore needs binary_classification_threshold on every binary call, and publishing happens only inside fit().

Change.

result = trainer.fit(table, snapshot=snapshot)
tuned = result.with_scoring_defaults(
    kumo.ScoringDefaults(binary_classification_threshold=0.5))
tuned.scoring_defaults   # ScoringDefaults(binary_classification_threshold=0.5, num_classes_to_return=None)
tuned.predict(snapshot)  # thresholds at 0.5

• kumo.engine.ScoringDefaults(binary_classification_threshold, num_classes_to_return). Field names match the predict arguments.
• TrainingResult.scoring_defaults: the bundle's defaults (all None when it has none). It is set for a fit's result and for Trainer.load/load_from_connector (trainer.result).
• TrainingResult.with_scoring_defaults(defaults) -> TrainingResult:
  • validates the defaults against the bundle's task type with the rules predict applies: a threshold only for binary and multilabel classification and only in [0, 1]; a count only where predict accepts one (multiclass, multilabel, link prediction);
  • normalises values first (threshold to float, count to int, numpy scalars accepted), so 1, 1.0 and np.float32(1) give the same file and the same id;
  • copies the bundle, hardlinking unchanged files, so the parent bundle is never written;
  • writes scoring-defaults.json, lists it in the publication manifest and writes the copy's _provenance.json;
  • returns the result of loading the copy. New defaults replace any the bundle had; ScoringDefaults() removes the file, so the id returns to the parent's.
• Trainer.export(ModelOutputConfig) -> PublishedModel publishes a fitted or loaded model through the same step fit(export_config=) uses. It is a plain command: inside or outside a kumo.mlops block it writes only to the given config, never to the artifact root.
• kumo-core gains remove_bundle_file(bundle_dir, role) next to add_bundle_file. It also gains bundle_task_type(artifact_paths), which the scoring defaults, the trainer's warm-start check and kumo-core's warm-start resolver now share instead of each reading prediction-schema.json.
• docs/mlops.md gains a "Scoring defaults" section.

Derived bundle from a binary fit on a small synthetic graph (PREDICT COUNT(events.*, 0, 7, days) > 0 FOR EACH ents.ID, 50 entities):

<derived bundle>/
  model/  manifests/  metrics.json  prediction-schema.json  training-spec.json   # hardlinks to the parent
  _provenance.json                     # new record: own id, parents.derived_from
  scoring-defaults.json                # new
  publication-manifest.json            # rewritten: lists scoring-defaults.json

scoring-defaults.json:

{
  "binary_classification_threshold": 0.5,
  "num_classes_to_return": null
}

Its publication-manifest entry, so it is hash-checked at load and part of the bundle id:

{"relative_path": "scoring-defaults.json", "role": "scoring_defaults", "sha256": "5bd069d9…", "size_bytes": 78}

The derived bundle's _provenance.json (ModelManifest, #56) records no fit of its own. Its only parent is the original model's manifest, embedded in derived_from, which carries the training table and snapshot (shortened):

{
  "kind": "model",
  "id": "5c4f281c…",
  "fit": null,
  "fit_mode": null,
  "parents": {
    "derived_from": {
      "kind": "model",
      "id": "c728c4af…",
      "fit": {"started_at": "2026-10-24T18:19:46+00:00", "finished_at": "2026-10-24T18:20:04+00:00", "seconds": 17.913},
      "fit_mode": "refit_trainval",
      "parents": {
        "training_table": {"kind": "training_table", "id": "tt-3d20699a…", "query": "PREDICT COUNT(events.*, 0, 7, days) > 0 FOR EACH ents.ID", "parents": {"snapshot": {"kind": "snapshot", "id": "bca8c681…@2f054e40…", "...": "..."}}, "...": "..."},
        "warm_start": null
      },
      "...": "..."
    },
    "training_table": null,
    "warm_start": null
  },
  "created_at": "2026-10-24T18:20:04+00:00",
  "environment": {"build_commit": "14d25143", "diskgraph": "0.1.0", "...": "..."},
  "schema_version": 1
}

• id: the derived bundle's id.
• parents.derived_from: the parent's manifest, embedded the way Record bundle_id and fit provenance in every model bundle #56 embeds a warm-start model. Its own model parents (warm_start, derived_from) are cut to {kind, id, reconstructed: true}, so a chain of re-tunings does not grow each manifest.
• When the parent has no manifest (a bundle from main), derived_from is the stub {kind: "model", id: <parent id>, reconstructed: true}.

In that run the parent's id was c728c4af… and the derived one's 5c4f281c…. Clearing the defaults on the derived result gave back c728c4af….

How predict picks each value:

Source, highest priority first	Threshold	Count (num_classes_to_return)
Explicit predict argument	used	used
The query's own TOP k	n/a	used
Bundle default (scoring-defaults.json)	used	used
Nothing set	binary raises, multilabel uses 0.5, as today	today's default

Back-compat

This PR only adds new calls and a new optional bundle file, and records rather than enforces: no existing call changes behaviour.

Situation	Behaviour
Existing predict, fit, load, load_from_connector calls	Unchanged. An explicit argument always wins over a bundle default.
Bundle without scoring-defaults.json (all existing bundles)	Loads and predicts exactly as before; binary still raises without a threshold. The id is unchanged.
with_scoring_defaults on a bundle from main (no _provenance.json)	Works; derived_from is a stub with the parent's id.
with_scoring_defaults on a bundle whose _provenance.json is unreadable or from a newer engine	Works; derived_from is a stub.
Derived bundle read by an older engine	scoring-defaults.json is a listed, hash-checked extra role, which older loaders verify and ignore. Predict then needs explicit arguments, as today.
scoring-defaults.json present but not listed in the manifest	Ignored.
Listed file unreadable, or invalid for the bundle's task	Ignored with a warning; predict behaves as for a bundle without defaults.
Invalid defaults passed to with_scoring_defaults	ValueError before anything is copied.
export to a connector that cannot publish models	NotImplementedError, the same check fit(export_config=) makes.

Testing Done

• test_scoring_defaults.py (6 tests, CPU fit):
  • default vs explicit threshold: identical outputs, explicit wins; parent bundle bytes and id unchanged; the derived manifest's id and its embedded derived_from; re-tuning cuts the grandparent to a stub; clearing the defaults restores the parent's id;
  • 1 and np.float32(1) give one id;
  • a parent without a manifest reads as a stub;
  • export, then load_from_connector, keeps the defaults and the id;
  • invalid defaults for the task; a multilabel threshold-plus-count default;
  • stray or invalid defaults files are ignored, and the bundle then behaves like an old one.
• kumo-core test_artifacts.py. Together: 17 passed on this head.
• The example above was captured from a CPU run of this head.

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/kumo-core/src/kumo_core/artifacts.py	remove_bundle_file, _write_bundle_manifest, bundle_task_type	remove a listed file and its manifest entry; manifest writer shared with add_bundle_file; task type read from prediction-schema.json in one place	clearing defaults restores the parent's id; one writer and one task-type reader
packages/kumo-core/src/kumo_core/automl/runtime.py	resolve_warm_start	reads the parent bundle's task type through bundle_task_type	one reader shared with scoring defaults and the trainer's warm-start check
packages/kumo-core/tests/test_artifacts.py	tests	bundle_task_type reads the task type a bundle records	coverage of the shared reader
packages/kumo-engine/docs/mlops.md	"Scoring defaults"	new section: with_scoring_defaults, scoring_defaults, precedence, Trainer.export	user docs
packages/kumo-engine/kumo/engine/__init__.py, packages/kumo-engine/kumo/engine/trainer/__init__.py	exports	ScoringDefaults	public API
packages/kumo-engine/kumo/engine/trainer/_provenance.py	record_derivation	new: the derived bundle's ModelManifest with its own id and the parent embedded in parents.derived_from (embedded_model; a stub via parent_model when the parent has none)	lineage of a derived model
packages/kumo-engine/kumo/engine/trainer/_scoring.py	ScoringDefaults, normalized_scoring_defaults, write_scoring_defaults, read_scoring_defaults	new: validate and normalise against the task type (threshold for binary/multilabel, count where predict takes one); write/remove the file; read only when listed, warn and ignore when unreadable or invalid	storage and checks
packages/kumo-engine/kumo/engine/trainer/_spec.py	SCORING_DEFAULTS_ROLE, read_bundle_artifact_paths	role constant; excluded from artifact paths	keeps it out of staging and warm start
packages/kumo-engine/kumo/engine/trainer/trainer.py	TrainingResult.scoring_defaults, TrainingResult.with_scoring_defaults, Trainer.export, _fitted, _require_publisher, _publish, _link_or_copy, fit, _finish_fit, _predict, _load_bundle	new API; publish check and publish step shared by fit and export; predict falls back to the defaults after explicit args and TOP k; load reads the defaults	scoring defaults and export
packages/kumo-engine/test/test_prediction_completion.py	prepared_prediction fixture	stub TrainingResult carries an empty ScoringDefaults	predict now reads scoring_defaults from the training result
packages/kumo-engine/test/test_prediction_workers.py	test_valid_class_count_reaches_scorer	stub TrainingResult carries an empty ScoringDefaults	predict now reads scoring_defaults from the training result
packages/kumo-engine/test/test_scoring_defaults.py	new	behaviour, normalisation and back-compat	coverage
packages/kumo-engine/test/test_semantic_carry_forward.py	trainer stubs in prediction tests	stub TrainingResult carries an empty ScoringDefaults	predict now reads scoring_defaults from the training result

Requirement: scoring defaults live in the bundle and can be overridden per predict call (MLOps build). Existing: the sealed bundle with its publication manifest and add_bundle_file from #56. Smallest: one optional listed file in a derived copy, read at load, used only when an argument is omitted. Stopping condition: defaults carried, overridable, exported and reloaded; bundles without them 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-snowflake-tracking
• File: .coderabbit.yaml
• Commit: ef5ed4b
• Layers: Repository YAML
• Review profile: QUIET
• Plan: Enterprise
• Run ID: caca471f-8909-4764-94b8-e7ccbb42bfbc

Commits

Reviewing files that changed from the base of the PR and between 248c49c and 206d88b.

Files selected for processing: 14

• packages/kumo-core/src/kumo_core/artifacts.py
• packages/kumo-core/src/kumo_core/automl/runtime.py
• packages/kumo-core/tests/test_artifacts.py
• packages/kumo-engine/docs/mlops.md
• packages/kumo-engine/kumo/engine/__init__.py
• packages/kumo-engine/kumo/engine/trainer/__init__.py
• packages/kumo-engine/kumo/engine/trainer/_provenance.py
• packages/kumo-engine/kumo/engine/trainer/_scoring.py
• packages/kumo-engine/kumo/engine/trainer/_spec.py
• packages/kumo-engine/kumo/engine/trainer/trainer.py
• packages/kumo-engine/test/test_prediction_completion.py
• packages/kumo-engine/test/test_prediction_workers.py
• packages/kumo-engine/test/test_scoring_defaults.py
• packages/kumo-engine/test/test_semantic_carry_forward.py

Walkthrough

The change adds task-aware scoring defaults to model bundles. Training results can create derived bundles with provenance, and Trainer.export publishes fitted or loaded bundles. Prediction uses bundle defaults when relevant arguments are omitted. Bundle loading restores those defaults. Warm-start task validation now reads task types through bundle artifacts.

Sequence Diagram(s)

sequenceDiagram
  participant Caller
  participant Trainer._predict
  participant TrainingResult
  Caller->>Trainer._predict: Predict with scoring arguments omitted
  Trainer._predict->>TrainingResult: Read scoring_defaults
  TrainingResult-->>Trainer._predict: Return configured threshold and class count
  Trainer._predict-->>Caller: Return predictions

Loading

Priority: Normal

Merge Risk: Low · up to 248c4

Link predictions using both a bundle default and TOP k can fail unexpectedly. Fix this narrow prediction path before merging, or explicitly accept the limitation.

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 present in the supplied review context, and the repository contains only the template and review instructions. …	Provide the complete PR conversation comment beginning with <!-- kumo-pr-change-map --> (including its Change map and rationale), and the omitted PR description or linked problem context if the map depends on it. Re-run this design review…

Passed checks: 2 passed

Check name	Status	Explanation
Pr Size		The pull request changes 14 files. The limit is 40 files, so the check passes.
Public Resource Access		The pull-request diff adds no external URLs, repository locations, or package source locations. The documentation adds only local Python imports and API examples. No accessibility check is required.

Full details: Change Map Design Review

Explanation

The required PR conversation comment beginning with <!-- kumo-pr-change-map --> is not present in the supplied review context, and the repository contains only the template and review instructions. The available description is explicitly truncated. The actual diff is available, but this check requires the proposed Change map to assess necessity, abstraction layer, existing behavior, focused verification, and stopping condition for each change. The supplied objectives are not a substitute for that comment.

Resolution

Provide the complete PR conversation comment beginning with <!-- kumo-pr-change-map --> (including its Change map and rationale), and the omitted PR description or linked problem context if the map depends on it. Re-run this design review with that evidence.

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

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

962-963: Minor · Functional Correctness · Quick win

Apply the bundle's link count when a query includes TOP k.

For link predictions, the bundle's stored count should fill an omitted num_classes_to_return, even when the query has TOP k. The current condition skips the default whenever TOP k is present, so prediction raises. Only the TOP k precedence rule is scoped to class counts, not link counts. Keep TOP k in the resolver call so its existing handling is unchanged.

🐛 Suggested fix

-        if num_classes_to_return is None and parsed.query.top_k is None:
+        if num_classes_to_return is None and (
+                parsed.query.top_k is None or task_type.is_link_pred):
             num_classes_to_return = training_result.scoring_defaults.num_classes_to_return

---

Review info

Run configuration

• Source mode: Auto
• Repository: NVIDIA/kumo-engine
• Branch: vedaant/mlops-snowflake-tracking
• File: .coderabbit.yaml
• Commit: 61b7ba6
• Layers: Repository YAML
• Review profile: QUIET
• Plan: Enterprise
• Run ID: ceacb288-4e90-4d12-af34-7c504d7ba756

Commits

Reviewing files that changed from the base of the PR and between 248091b and 248c49c.

Files selected for processing: 12

• packages/kumo-core/src/kumo_core/artifacts.py
• packages/kumo-core/tests/test_artifacts.py
• packages/kumo-engine/docs/mlops.md
• packages/kumo-engine/kumo/engine/__init__.py
• packages/kumo-engine/kumo/engine/trainer/__init__.py
• packages/kumo-engine/kumo/engine/trainer/_provenance.py
• packages/kumo-engine/kumo/engine/trainer/_scoring.py
• packages/kumo-engine/kumo/engine/trainer/trainer.py
• packages/kumo-engine/test/test_prediction_completion.py
• packages/kumo-engine/test/test_prediction_workers.py
• packages/kumo-engine/test/test_scoring_defaults.py
• packages/kumo-engine/test/test_semantic_carry_forward.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.

[Vedaant-J]

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