[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 explicit model aliases, a bundle id pointer, and registration through kumo.MLOps(registry=...) for the Snowflake and Databricks model registries.

Summary

This PR adds explicit model aliases, a bundle id pointer, and registration through kumo.MLOps(registry=...) for the Snowflake and Databricks model registries. A version can be registered under an alias, promoted by moving the alias, and loaded by it; each version records which bundle it holds, and a fetch warns if the bundle it gets is a different one.

Problem. Teams that publish Kumo models to a warehouse registry need to answer questions like these:

• Which registered version is in production, and which bundle is it? Where is that bundle's lineage?
• Can we roll back to the previous version without a redeploy, by moving a stable name such as prod?
• Can a scoring job load "whatever is in production" instead of a hard-coded version number?

Today none of these can be answered from the registry, and loading by a stable name is not possible.

Cause. A Kumo publish records only a kumo_engine marker, which exists so that listing can tell Kumo versions apart from other tools' versions. The engine already passes metadata (spec_version, predictive_query, snapshot_id) and metrics to publish_model, but neither registry stored them. The call each registry received before this PR:

# Snowflake: Registry.log_model(...) arguments
metrics = {"kumo_engine": "1"}

# Databricks: POST /api/2.1/unity-catalog/models/versions body
tags = [{"key": "kumo_engine", "value": "1"}]

list_models and fetch_model returned no id. ModelRegistryLocation takes a model name and an optional exact version; there was no alias API. With MLOps on (#50), nothing registered a fit's model: only a plain export_config= published.

Change.

Registration and aliases:

import kumo.engine as kumo
from kumo.engine import ModelOutputConfig, ModelRegistryLocation, Trainer, set_model_alias

# registry: a SnowflakeConnector or DatabricksConnector; trainer, table, snapshot from your fit job.
cfg = kumo.MLOps("s3://<bucket>/kumo", registry=ModelOutputConfig(
    registry, location=ModelRegistryLocation("USERS_CHURN", alias="challenger")))
with kumo.mlops(cfg):
    result = trainer.fit(table, snapshot=snapshot)
published = result.registered_model                  # version assigned by the registry

# Promote; roll back by pointing "prod" at an earlier version the same way. No redeploy.
set_model_alias(registry, ModelRegistryLocation("USERS_CHURN", version=published.version), "prod")

# Scoring jobs load whichever version "prod" points at.
trainer = Trainer.load_from_connector(registry, location=ModelRegistryLocation("USERS_CHURN", alias="prod"))

• MLOps(registry=ModelOutputConfig(connector, location=...)) registers each fit's model after the plain export_config (if any), the root copy and tracking, and returns the version as the new TrainingResult.registered_model. Registration is independent of export=. If it fails, fit warns and returns the result with registered_model=None. A connector that cannot publish models, or an invalid alias, raises before training. TrainingResult.published_model stays the plain export_config result. (A bare ModelRegistryLocation, registered through MLflow, comes in Log, register and serve fitted models through MLflow pyfunc #65.)
• A plain fit(export_config=ModelOutputConfig(connector, location=...)) or Trainer.export(...) publishes as before, with or without a block, and also accepts an alias.
• ModelRegistryLocation.alias cannot be combined with version. On export, the version is published first and the alias is set on it afterwards. If setting it fails (including a sink without alias support), the call warns and still returns the published model.
• kumo.engine.set_model_alias goes through a new optional ModelAliasingSink protocol in kumoapi.plugin. A sink without it raises NotImplementedError for this new call and for fetching by alias.
• Snowflake moves an alias off the version that holds it (unset, then set). Snowflake cannot hold one alias on two versions, so the move is not atomic; if setting it on the new version fails, the alias is put back on the old version, a warning is logged and the error is raised. Snowflake allows one alias per version, so a version that already holds a different alias is refused rather than silently stripped. The system aliases DEFAULT, FIRST and LAST can be fetched by but not set.
• Databricks uses PUT /api/2.1/unity-catalog/models/{name}/aliases/{alias} to set an alias and GET on the same path to fetch by it.
• kumoapi.plugin.check_model_alias holds one alias syntax every registry accepts: a letter or underscore, then letters, digits or underscores (at most 255), and not DEFAULT, FIRST, LAST (Snowflake), LATEST or v<digits> (MLflow) in any case. Hyphens are refused because Snowflake would need a quoted identifier, which also excludes SageMaker's SageMakerLifecycle- aliases. Both registries call it, and fit (for export_config= and an MLOps registry) and Trainer.export call it before training or publishing, so an invalid export alias fails at once instead of after a fit has been published without it.

Bundle id pointer: each published version now records the marker and kumo_bundle_id, the bundle's content hash from #56. The bundle's own _provenance.json holds its lineage (snapshot, training table, source table versions), so the registry stores a pointer, not a copy. From a CPU fit on a small synthetic users/items/orders graph (PREDICT COUNT(orders.*, 0, 7, days) > 0 FOR EACH users.ID), hashes shortened:

# Snowflake: Registry.log_model(...) arguments
metrics = {
    "kumo_engine": "1",
    "kumo_bundle_id": "54c9eeb1…",              # new
}

# Databricks: POST /api/2.1/unity-catalog/models/versions body
tags = [
    {"key": "kumo_engine",    "value": "1"},
    {"key": "kumo_bundle_id", "value": "54c9eeb1…"},   # new
]

• Training metrics are not copied into either registry: they stay in the bundle's metrics.json and on the tracked fit run (Record fits in experiment tracking and publish models to registries #60, Record Trainer.fit in Snowflake Experiment Tracking through kumo.MLOps(tracking=) #68), where runs can be compared. Snowflake and Databricks versions record the same two keys.
• list_models and every fetch_model path (by version, default, or alias) return the recorded kumo_bundle_id in PublishedModel.metadata, read through the shared kumoapi.plugin.registry_lineage.
• Trainer.load_from_connector compares the recorded kumo_bundle_id with the fetched bundle's id on every fetch; a mismatch is a warning only.
• fit(export_config=), the MLOps registry step and Trainer.export (Add TrainingResult.with_scoring_defaults and Trainer.export #57) publish through the same _publish, so all record the pointer and set the alias. Exporting a derived bundle from TrainingResult.with_scoring_defaults publishes the child's own kumo_bundle_id, not the parent's.
• The ModelRegistryLocation(version=None) docstring now says what fetch loads: Snowflake's default version, or the highest ready Kumo version in Databricks.
• docs/mlops.md gains a "Registering models" section.

Also fixed: Databricks fetch of a real bundle. The Unity Catalog Files API lists directories with a trailing slash (.../1/model/), so the immediate-child check in the download rejected every real bundle and Trainer.load_from_connector(DatabricksConnector, location=...) raised Invalid model artifact path returned by Databricks. The slash is stripped before the check; the fake Files API in the tests now lists directories the same way. Found in a live Databricks run.

Not in this PR. Snowflake task= is not set, so Snowflake model monitors cannot attach to Kumo versions yet; it needs a live check first because a server-side rejection would fail every publish.

Not verified live: Snowflake log_model with the pointer record, alias set/unset, the rollback and the one-alias rule; Databricks alias PUT/GET. Registry tests use the packages' fakes. The Databricks trailing-slash bug was found in a live workspace run; the fix is tested against the fake.

Back-compat

This PR only records and exposes information on publish; aliases are a new opt-in field and a new call. No existing publish or fetch path gains a check that can fail.

Situation	Behaviour
Fit without MLOps, or MLOps without registry=	Nothing is registered; registered_model is None.
MLOps registry= and registration fails	Warning; the fit's result is returned with registered_model=None.
New publish	Marker plus kumo_bundle_id recorded. The assigned version and the returned PublishedModel are unchanged.
Version published before this PR	Listed and fetched as before, with no kumo_bundle_id in PublishedModel.metadata; the id is still inside its bundle. No mismatch warning (nothing recorded).
Fetch by version or with no version	Same version as before; metadata now carries the recorded kumo_bundle_id.
Version published by this PR, read by an older engine	The extra metrics and tag are ignored. Listing still keys on the marker.
Bundle without a bundle_id (written before #56)	Published with the marker only, as before.
Export with an invalid alias (new field)	ValueError before training or publishing.
Export with an alias, and setting the alias fails	A warning is logged; the version stays published and is returned.
Snowflake alias move fails half-way	The alias is restored on the version that held it, a warning is logged, and the error is raised (fit turns it into the warning above).
Recorded kumo_bundle_id differs from the fetched bundle	Warning; the fetched bundle loads.
Snowflake version metadata unreadable on fetch	Warning; fetch goes ahead with empty metadata.
Databricks fetch of a real bundle	Works (was: Invalid model artifact path).

Testing Done

• Fake-client tests: Snowflake metrics and Databricks tags carry the marker and kumo_bundle_id; list and fetch (by version and by alias) return it; Snowflake alias set, move, rollback of a failed move and the refusals Snowflake requires; Databricks alias needs a positive version; every Databricks fetch test lists directories with a trailing slash.
• kumo-api test_plugin_registry.py: check_model_alias syntax and reserved names, registry_lineage.
• Engine test_trainer_registry.py: an MLOps registry registers after the fit and sets the alias, and a failing registry warns and leaves registered_model unset; an invalid MLOps registry alias raises before training; the exact metadata passed to publish_model; export with an alias publishes then sets it, best-effort; the bundle-id mismatch warning; an alias fetch needs an aliasing sink; an invalid export alias is refused before publishing.
• Engine test_scoring_defaults.py: Trainer.export of a derived bundle publishes the child's kumo_bundle_id.
• On this head: test_trainer_registry.py, test_scoring_defaults.py, kumo-snowflake test_registry.py and kumo-api test_plugin_registry.py: 69 passed; kumo-databricks test_registry.py + test_uc_registry.py: 41 passed.
• The example values above come from a CPU run of this branch before the registry record was cut to the pointer; the records shown keep only the keys this head writes.

Checklist

• No breaking changes (existing assertions changed only where they encoded "marker only").
• Docs updated (docs/mlops.md).
• Unit tests added.

🤖 Generated with Claude Code

[Vedaant-J]

Change map

File	Function / scope	What changed	Why
packages/kumo-api/kumoapi/plugin.py	ModelRegistryLocation, ModelAliasingSink, ModelFetchingSink.fetch_model, ENGINE_MARKER, BUNDLE_ID_KEY, REGISTRY_KEY_PREFIX, check_model_alias, registry_lineage	alias field (export: set after publish; fetch: resolve), exclusive with version; new optional protocol; shared marker/key names, one alias syntax for every registry (letters, digits, underscores; DEFAULT/FIRST/LAST/LATEST/v reserved) and recorded-key reader; version=None docstring says what each registry loads	promotion and fetch by alias; one definition for both registries
packages/kumo-api/test/test_plugin_registry.py	new tests	alias syntax and reserved names refused, plain names accepted; registry_lineage keeps kumo_* keys as strings without the marker	coverage of the shared helpers
packages/kumo-databricks/src/kumo_databricks/registry.py	_alias_path, _version_number, publish_bundle, _download, fetch_bundle, list_models, fetch_model, set_model_alias	kumo_bundle_id tag next to the marker; recorded id returned from list and fetch; alias GET/PUT; trailing slash stripped from listed directories	bundle pointer, aliases; real UC bundles failed to download
packages/kumo-databricks/tests/test_registry.py, packages/kumo-databricks/tests/test_uc_registry.py	fake client, stand-in server, new tests	alias endpoints, kumo_bundle_id tag listed and fetched, directories listed with a trailing slash	coverage
packages/kumo-engine/docs/mlops.md	"Registering models"	new section: MLOps(registry=ModelOutputConfig(...)), registered_model, aliases, promotion, recorded pointers	user docs
packages/kumo-engine/kumo/engine/__init__.py, packages/kumo-engine/kumo/engine/trainer/__init__.py	exports	set_model_alias	public API
packages/kumo-engine/kumo/engine/trainer/trainer.py	TrainingResult.registered_model, fit, Trainer._register, _connector_registry, load_from_connector, set_model_alias, _require_publisher, _publish	MLOps registry=ModelOutputConfig(location=) registers after the fit (failure warns) into registered_model, checked before training; _require_publisher validates an alias; _publish adds kumo_bundle_id and sets the alias after publish (warning on failure); bundle-id mismatch warning on every fetch; alias fetch needs an aliasing sink	engine API
packages/kumo-engine/test/test_scoring_defaults.py	tests	exporting a with_scoring_defaults child publishes the child's kumo_bundle_id	export and fit share one publish step
packages/kumo-engine/test/test_trainer_registry.py	new tests	MLOps registry registers after the fit / warns on failure / checks the alias before training; exact metadata, export alias best-effort, invalid export alias refused before publishing, mismatch warning, non-aliasing sink	coverage
packages/kumo-snowflake/src/kumo_snowflake/registry.py	_version_record, _recorded_lineage, _alias_holders, _restore_alias, publish_model, list_models, fetch_model, set_model_alias	marker and kumo_bundle_id in version metrics (training metrics stay in the bundle and the tracked run, as on Databricks); recorded id returned from list and fetch; fetch by alias incl. system aliases; set/move alias with one-alias-per-version refusal; a failed move restores the alias on the old version and warns	traceability, aliases
packages/kumo-snowflake/tests/test_registry.py	publish/list/fetch/alias tests	new tests incl. rollback of a failed move; marker-only assertion now includes kumo_bundle_id	coverage

Requirement: registry versions must point at their bundle and support aliases (MLOps build). Existing: publish_model already received metadata/metrics, unused by both registries; lineage lives in the bundle's _provenance.json. Smallest change: store kumo_bundle_id through those calls on both registries; aliases via an optional protocol. Stopping condition: pointer recorded and returned, aliases settable at publish or later and fetchable, on both registries.

[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-scoring-defaults
• File: .coderabbit.yaml
• Commit: 206d88b
• Layers: Repository YAML
• Review profile: QUIET
• Plan: Enterprise
• Run ID: 9ad55644-b363-48a4-999d-0c1fdd062013

Commits

Reviewing files that changed from the base of the PR and between 71815cc and c0a5e72.

Files selected for processing: 13

• packages/kumo-api/kumoapi/plugin.py
• packages/kumo-api/test/test_plugin_registry.py
• packages/kumo-databricks/src/kumo_databricks/registry.py
• packages/kumo-databricks/tests/test_registry.py
• packages/kumo-databricks/tests/test_uc_registry.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/trainer.py
• packages/kumo-engine/test/test_scoring_defaults.py
• packages/kumo-engine/test/test_trainer_registry.py
• packages/kumo-snowflake/src/kumo_snowflake/registry.py
• packages/kumo-snowflake/tests/test_registry.py

Walkthrough

The change adds model aliases and bundle-lineage metadata to the shared registry API and the Databricks and Snowflake registry adapters. Trainer publishing can assign an alias after publication, and Trainer.fit can register a fitted model and expose it through registered_model. Alias-based loading checks sink support and warns when recorded and loaded bundle IDs differ.

Sequence Diagram(s)

sequenceDiagram
  participant TrainerFit as Trainer.fit
  participant RegistryPublisher as Registry publisher
  participant TrainingResult
  TrainerFit->>RegistryPublisher: Register fitted model
  RegistryPublisher-->>TrainerFit: Return PublishedModel
  TrainerFit->>TrainingResult: Set registered_model

Loading

Priority: Normal

Merge Risk: Low · up to 71815

The change appears mergeable with bounded follow-up: correct the export and rollback guidance so users know where to find a published version and when to check an alias after a failed move. Strengthen the alias-loading test to protect that behavior.

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 available. The supplied review context reports no returned review threads, and the repository contains only the…	Provide the full PR conversation comment starting with <!-- kumo-pr-change-map -->, then review each proposed change against the linked problem and repository code.

Passed checks: 2 passed

Check name	Status	Explanation
Pr Size		The pull request changes 13 files, based on the review-scoped diff. This does not exceed the 40-file limit.
Public Resource Access		No added or modified qualifying external resource reference was found. The added s3://bucket/kumo is an illustrative documentation example, and fake:// plus the relative Databricks API path are lo…

Full details: Public Resource Access

Explanation

No added or modified qualifying external resource reference was found. The added s3://bucket/kumo is an illustrative documentation example, and fake:// plus the relative Databricks API path are local test fixtures or implementation paths. No added public URL, repository location, or package source location requires anonymous access verification.

Full details: Change Map Design Review

Explanation

The required PR conversation comment beginning with <!-- kumo-pr-change-map --> is not available. The supplied review context reports no returned review threads, and the repository contains only the template and checker references, not the PR's Change map. The available summaries do not provide the proposed design, inspected behavior, stopping condition, or focused verification needed for this check.

Resolution

Provide the full PR conversation comment starting with <!-- kumo-pr-change-map -->, then review each proposed change against the linked problem and repository code.

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

packages/kumo-engine/docs/mlops.md: 2

87-87: Minor · Functional Correctness · Quick win

Correct the result described for Trainer.export.

Trainer.export(config) returns a PublishedModel. It does not set result.published_model. If a reader follows this sentence after calling Trainer.export, the stated field can remain None. Distinguish trainer.export(config).version from fit(...).published_model.version.

---

88-88: Minor · Data Integrity & Integration · Quick win

Describe Snowflake alias restoration as an attempt.

The statement that a failed alias move puts the alias back promises a successful rollback. The stated implementation only attempts to restore the previous assignment. If that restoration also fails, the alias can remain unset. Say that Snowflake attempts restoration and that callers should check the alias after a failed move.

🧹 Nitpick comments: 1

packages/kumo-engine/test/test_trainer_registry.py: 1

618-622: Trivial · Maintainability & Code Quality · Quick win

Make the fake resolve the requested alias.

_AliasingSink.fetch_model always copies CHURN_GNN@latest, and the alias-loading test never assigns prod. The test can therefore pass if Trainer.load_from_connector drops or changes the forwarded alias. Resolve location.alias through self.aliases, and assign prod before loading.

---

Review info

Run configuration

• Source mode: Auto
• Repository: NVIDIA/kumo-engine
• Branch: vedaant/mlops-scoring-defaults
• File: .coderabbit.yaml
• Commit: 248c49c
• Layers: Repository YAML
• Review profile: QUIET
• Plan: Enterprise
• Run ID: 0ccd3862-bc45-4b8c-99f6-a6b3e000a28d

Commits

Reviewing files that changed from the base of the PR and between 238a52e and 71815cc.

Files selected for processing: 7

• packages/kumo-api/kumoapi/plugin.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/trainer.py
• packages/kumo-engine/test/test_scoring_defaults.py
• packages/kumo-engine/test/test_trainer_registry.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.