[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 MLflow model logging, registration and scoring for Kumo models.

Summary

This PR adds MLflow model logging, registration and scoring for Kumo models. With kumo.MLOps(tracking=MlflowTracking(), registry=ModelRegistryLocation(name, alias=...)), each fit logs its bundle as an MLflow pyfunc model on the fit run (#60), registers it, points the alias at the new version, and can be scored later from models:/name@alias. It is the MLflow counterpart of the connector-registry registration in #59.

Problem. Teams that run MLflow (Databricks Unity Catalog, SageMaker managed MLflow, a self-hosted server) want to:

• register each fitted Kumo model in the MLflow registry they already govern, and rank fits by metric;
• promote a version with @champion and roll back by moving the alias;
• score from a scheduled batch job by alias, without knowing where the bundle lives;
• see in Unity Catalog which tables a registered model was trained from. The fit run's inputs come from Record fits in experiment tracking and publish models to registries #60, and a model registered from that run inherits them.

Today none of these can be done. MLflow sees a fit's runs, metrics and source-table inputs (#60), but never the model.

Cause. With #60, a tracked fit records the run, its trials, its source tables as inputs and three artifacts. The bundle reaches the artifact root (#56) but never MLflow:

experiment/
  run "fit"                    # kumo.kind, kumo.run_id, kumo.bundle_id
    inputs/  source tables
    artifacts/  model_plan.json  _provenance.json  query.pql
    (no model, no registry version, no alias)

Change.

import mlflow
import pandas as pd
import kumo.engine as kumo
from kumo.engine.integrations.mlflow import MlflowTracking, load_model

cfg = kumo.MLOps("s3://<bucket>/kumo",
                 tracking=MlflowTracking(experiment_name="users-spend"),
                 registry=kumo.ModelRegistryLocation("users_spend", alias="champion"))  # catalog.schema.model on UC
with kumo.mlops(cfg):
    snapshot = graph.snapshot()            # stored under snapshots/<graph key>/<version>/
    result = trainer.fit(table, snapshot=snapshot)
result.registered_model     # PublishedModel(name='users_spend', version='1', uri='models:/users_spend/1')

# Scheduled scoring job, with kumo-engine installed.
model = mlflow.pyfunc.load_model("models:/users_spend@champion")
scores = model.predict(pd.DataFrame({"ENTITY": [3, 7, 42]}),
                       params={"snapshot": "s3://<bucket>/kumo/snapshots/<graph key>/<version>"})

# Or get the Trainer back, for explain or batch predict.
trainer = load_model("models:/users_spend@champion")

From a CPU fit of PREDICT ents.Y FOR EACH ents.ID on a synthetic 400-row table, local MLflow store (hashes shortened, local paths dropped):

run "fit" 9eed03ef…
  tags    kumo.kind=fit  kumo.run_id=8deb5944…-fit-5486380d…  kumo.bundle_id=b681cb48…    # from #60
  artifacts  model_plan.json  _provenance.json  query.pql                                 # from #60
LoggedModel m-1860a9b4…   source_run_id 9eed03ef…                                          # new
  metrics (model_id m-1860a9b4…)  val_mae=0.5095  test_mae=0.4316  test_refit_mae=0.6037 … seconds=12.54
registered model "users_spend"   aliases {champion: 1}                                     # new
  version 1   source models:/m-1860a9b4…

The logged model, downloaded from models:/users_spend@champion:

MLmodel                         # flavor, signature, metadata
_mlflow_model_code.py           # models-from-code loader: set_model(KumoModel())
artifacts/_provenance.json      # the bundle's ModelManifest
artifacts/published-model/      # the Kumo bundle, unchanged
  _provenance.json  metrics.json  prediction-schema.json  publication-manifest.json  training-spec.json
  manifests/  graph.json  input.json  task.json
  model/      checkpoint  config.json  encoders.json  resolved-model-plan.json  target-encoder.json
input_example.json  serving_input_example.json  requirements.txt  python_env.yaml  conda.yaml  registered_model_meta

Excerpt of MLmodel:

flavors:
  python_function:
    artifacts:
      bundle: {path: artifacts/published-model}
      provenance: {path: artifacts/_provenance.json}
    config: {columns: [ENTITY, TARGET_PRED], snapshot: null}
    loader_module: mlflow.pyfunc.loaders.code_model
metadata:
  kumo_bundle_id: cd4e28db…
signature:
  inputs:  [{name: ENTITY, type: long, required: true}, {name: TIMESTAMP, type: datetime, required: false}]
  outputs: [{name: ENTITY, type: long}, {name: TARGET_PRED, type: double}]
  params:  [snapshot: string, prediction_time: string, binary_classification_threshold: double,
            num_classes_to_return: long, population: boolean = false]

What each part records and does:

• Registration through MLOps(registry=ModelRegistryLocation(...)). It needs a tracker with the log_model hook (MlflowTracking). kumo.MLOps(...) raises in three cases: a bare location without such a tracker, a fixed version (MLflow assigns versions), or an alias that Add model registry aliases and a bundle id pointer on registered versions #59's check_model_alias refuses (which also refuses SageMaker's reserved SageMakerLifecycle- prefix). One alias is set with set_registered_model_alias, and ModelRegistryLocation.comment becomes the version's description. The version comes back as TrainingResult.registered_model (Add model registry aliases and a bundle id pointer on registered versions #59), and published_model stays the plain export_config result. MlflowTracking(log_model=True) logs the model without registering it.
• Pointers, not lineage. The model's metadata is kumo_bundle_id only, the same pointer Add model registry aliases and a bundle id pointer on registered versions #59 records on connector-registry versions. Lineage stays in the bundle's _provenance.json, which is both a fit-run artifact (Record fits in experiment tracking and publish models to registries #60) and a model artifact.
• Inputs of the fit run come from Record fits in experiment tracking and publish models to registries #60 unchanged, so a model registered from the run inherits them. log_model outside a fit logs the same inputs into the run it logs to, read with Record fits in experiment tracking and publish models to registries #60's fit_inputs(result.manifest).
• The model is a plain pyfunc logged from code with the bundle as an artifact: no custom flavor, no pickled object. requirements.txt pins kumo-engine[gcs,s3]==<training version>, the extras that loading a snapshot from S3 or GCS needs.
• signature is typed per bundle: ENTITY takes the type of the entity table's key, and TIMESTAMP is optional. Outputs are the task's prediction columns, from the new kumo_core.train.formatting.prediction_output_columns. For example, a binary model outputs ENTITY, TARGET_PRED, False_PROB, True_PROB.
• Params: snapshot is the full URI or local path of a stored or exported snapshot (s3://, gs://, file:// or local, through Record provenance for every artifact under one MLOps artifact root #50's connector_for_root). The others are prediction_time, binary_classification_threshold, num_classes_to_return and population.
• Metrics are the fit run's final metrics, logged with the LoggedModel's model_id, so mlflow.search_logged_models(..., order_by=[{"field_name": "metrics.val_mae"}]) ranks fits.
• Synchronous model logging. Record fits in experiment tracking and publish models to registries #60 runs tracker calls on a background thread and gives finish at most 30 s before the fit returns, so a scheduled job could exit mid-upload, leaving nothing registered and the run RUNNING. Model logging is therefore not part of finish.
  • After a fit that succeeded, Trainer.fit queues the tracker's optional log_model(result, snapshot, registry) on the tracking thread, so it never races the tracker's other calls. The fit waits for it while that thread is alive, up to 1 hour. Then finish(record, error) runs under the usual 30 s limit.
  • If logging or registering fails, the tracking thread stops, or the hour passes, fit warns (Could not log the model: ...) and finish gets the error, so the run ends FAILED. The fit still returns its result, with registered_model=None.
  • A model that finishes uploading after that is not registered, and a warning says so.
  • During the upload only, MLFLOW_HTTP_REQUEST_TIMEOUT is 600 s instead of the fit's 30 s, and only when the user has not set it.

Scoring and loading:

• Snapshots. predict loads each snapshot once per model instance and keeps every one it has loaded for the life of the process (no eviction). Each serving replica, and each Spark executor of spark_udf, downloads it once. With serving_snapshot=<snapshot URI> (on log_model or MlflowTracking), the model loads that snapshot in load_context, i.e. when a replica starts, and predict uses it when no snapshot param is given.

• Scratch. Each call scores under its own scratch root: a per-call temporary directory set through a context variable that get_scratch_dir reads (new internal scoped_scratch_dir(path)), removed when the call returns or raises.

• What a call scores. An empty frame returns no rows unless population=true, which scores the query's own population at prediction_time. Scoring goes through trainer.predict(snapshot, entities=..., prediction_time=...) (Score a snapshot with Trainer.predict(snapshot) and TrainingResult.predict #51).

• Where it runs. The model is for batch jobs where kumo-engine is installed, not for MLflow Model Serving. MLflow's serving-input check, which cannot pass without a snapshot, is silenced during logging.

• Outside a fit. kumo.engine.integrations.mlflow.log_model(trainer, snapshot=...) logs a fitted or loaded model in the given, active or a new run. load_model(uri) -> Trainer restores it and warns if the model's kumo_bundle_id differs from the bundle it loaded.

• docs/mlops.md gains an "MLflow models" section and a "Which registry" table:

  Platform	Registry
  Databricks	MLflow on Unity Catalog, or the DatabricksConnector registry
  SageMaker	SageMaker managed MLflow
  Snowflake	the Snowflake Model Registry
  Plain VM	any MLflow registry

  It also states the MLflow minimum: a 3.13 or later client, because DiskGraph pins pyarrow 24.

Back-compat

Model logging is opt-in. It only adds records to MLflow; it never changes what a fit trains, writes or returns.

Situation	Behaviour
fit() without MLOps, or MLOps without registry= / log_model=True	Unchanged; no model logged.
MLOps(tracking=MlflowTracking(), registry=ModelRegistryLocation(...))	Model logged, registered and aliased after a fit that succeeded, before fit returns; fit waits up to 1 hour. A failed fit logs no model.
Logging or registering fails, the tracking thread stops, or the hour passes	A warning. Metrics and tags stay on the run, the run ends FAILED, and the fit returns its result with registered_model=None.
Bare ModelRegistryLocation without an MLflow tracker, a fixed version, or an invalid alias	ValueError when kumo.MLOps is built, before any work.
Custom Tracking adapters	log_model(result, snapshot, registry) is optional; adapters without it are unaffected.
Bundle without _provenance.json logged with log_model	Logs; no inputs and no provenance artifact.
MLFLOW_HTTP_REQUEST_TIMEOUT set by the user	Kept as set, also during the model upload.
A logged model loaded where kumo-engine is a different version	requirements.txt pins the training version; load_model does not check it.

Known limitations

• A model registered through a connector registry (Add model registry aliases and a bundle id pointer on registered versions #59) is linked to the MLflow run only through the shared kumo_bundle_id / kumo.bundle_id; MLflow has no native link to a registry it did not create.
• A model upload slower than 1 hour is abandoned: the run ends FAILED and the model, if it finishes uploading later, is not registered.
• SageMaker managed MLflow and scoring from live S3/GCS snapshots were not run on this head. An earlier head of this PR, which configured registration on the tracker instead of kumo.MLOps(registry=...), registered and scored live on SageMaker and Databricks.

Testing Done

CPU, mlflow 3.17.1:

• test_mlflow_model.py:
  • Real fit. A fit inside MLOps(tracking=MlflowTracking(), registry=ModelRegistryLocation('m', alias='champion')):
    • one LoggedModel on the fit run, carrying the fit's metrics;
    • the @champion version's source is that model, and result.registered_model names it;
    • the run's kumo.* tags are exactly kind, run id and bundle id;
    • the model metadata is {kumo_bundle_id}, and its _provenance.json equals result.manifest;
    • signature types; pyfunc.predict returns ENTITY, TARGET_PRED per entity;
    • load_model returns the same bundle id and warns on a mismatched kumo_bundle_id.
  • Scoring with log_model. A model logged with log_model:
    • pins kumo-engine[gcs,s3]==<version> and loads its serving_snapshot when it is loaded;
    • scores the query population and given entities at a prediction time exactly like Trainer.predict, also with no snapshot param;
    • returns no rows for an empty input;
    • scores each call under its own scratch root and removes only that.
  • Waits and timeouts.
    • Model logging completes before finish, and a failing log_model warns and reaches finish as the error.
    • An upload that outlives the wait bound, or a dead tracking thread, warns and reaches finish as the error.
    • The upload sees a 600 s HTTP timeout and the rest of the run 30 s.
    • A model whose run ended during the upload is not registered.
  • Configuration checks. A bare location without an MLflow tracker, a fixed version and a reserved alias are refused by kumo.MLOps.
• On this head, test_mlflow_model.py + test_mlops.py gave 14 passed and test_mlflow_tracking.py + test_tracking_contract.py passed. kumo-core test_prediction_output_columns.py passed. flake8 and isort are clean.
• The records above were captured from a CPU run of this head with a local MLflow store.

Checklist

• No breaking changes (opt-in).
• Unit tests added.
• Docs updated (docs/mlops.md).

🤖 Generated with Claude Code

[Vedaant-J]

Change map

File	Function / scope	What changed	Why
packages/kumo-core/src/kumo_core/train/formatting.py	prediction_output_columns, format_prediction	new: prediction columns per task type, sharing the long-format task set format_prediction now uses	typed MLflow signature outputs from one definition
packages/kumo-core/tests/test_prediction_output_columns.py	new tests	columns for every task type	coverage
packages/kumo-engine/docs/mlops.md	"MLflow models", "Which registry", Tracking protocol paragraph	registration through MLOps(registry=ModelRegistryLocation), waits, serving snapshot, MLflow minimum; registry choice per platform; log_model(result, snapshot, registry) hook	user docs
packages/kumo-engine/kumo/engine/_mlops.py	MLOps.__post_init__	a bare ModelRegistryLocation needs a tracker with log_model, no fixed version, a valid alias	fail before any work instead of after the fit
packages/kumo-engine/kumo/engine/_scratch.py	scoped_scratch_dir, get_scratch_dir	new internal context manager: a per-call scratch root set through a context variable that get_scratch_dir reads	each pyfunc predict removes only its own scratch
packages/kumo-engine/kumo/engine/_tracking.py	Tracking, TrackingSession.log_model, MODEL_SECONDS	optional log_model(result, snapshot, registry) tracker member, run on the tracking thread while the fit waits as long as that thread is alive, up to 1 hour; returns the registered version and its error (a dead thread or the timeout too), warned	model upload/registration completes before fit returns, outside the 30 s finish limit
packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py	new: KumoModel, predict, snapshot_location, signature, _entity_type	pyfunc wrapper; serving snapshot loaded in load_context; snapshot URI → connector; scoring through trainer.predict (#51) under the call's own scratch root; per-bundle signature and example	MLflow-native batch scoring
packages/kumo-engine/kumo/engine/integrations/_mlflow_model_code.py	new module	set_model(KumoModel())	models from code instead of a pickle
packages/kumo-engine/kumo/engine/integrations/mlflow.py	MlflowTracking.__init__, log_model (member), start, finish, _log_metrics, log_model, load_model, _log_model, _log_inputs, _register, _in_run	log_model=/serving_snapshot=; logs the model on the fit run and registers it at the given location (one alias, comment as description); final metrics carry the model_id; model metadata kumo_bundle_id, _provenance.json artifact, pip requirement kumo-engine[gcs,s3]==<version>; standalone log_model logs fit_inputs(manifest) as run inputs; 600 s HTTP timeout during the upload only; a model finishing after its run ended is not registered	MLflow registry
packages/kumo-engine/kumo/engine/trainer/trainer.py	fit, _mlflow_registry, _publish, load_from_connector, warn_if_other_bundle, registry_pointers	fit waits for model logging before finish and sets registered_model; pointer metadata and bundle-id mismatch warning shared with MLflow log_model/load_model	one pointer builder and one mismatch check
packages/kumo-engine/test/test_mlflow_model.py	new	MLOps fit → register → alias → pyfunc predict; exact kumo.* tags; metadata kumo_bundle_id; manifest artifact; pinned requirement; serving snapshot; population/timed rows equal Trainer.predict; per-call scratch cleanup; model logging waits past the finish limit and a failure reaches finish; bounded wait; 600 s upload timeout, late model not registered; invalid MLOps registries refused	coverage

Requirement: MLflow pyfunc + registry with signature, LoggedModel and alias loading, registered through MLOps(registry=ModelRegistryLocation). Existing: #60 run records, inputs and background tracking thread, #59 pointer metadata, alias checks and registered_model. Smallest: one pyfunc wrapper, one optional tracker member for the model, one synchronous session call, one construction check. Stopping condition: register, alias, load and score round trip; tags and model metadata read back from MLflow.

[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-registry-aliases
• File: .coderabbit.yaml
• Commit: c0a5e72
• Layers: Repository YAML
• Review profile: QUIET
• Plan: Enterprise
• Run ID: ae7aa48b-ec97-4e00-94c8-da346c1e1e75

Commits

Reviewing files that changed from the base of the PR and between 9b8a83a and a8aad9a.

Files selected for processing: 11

• packages/kumo-core/src/kumo_core/train/formatting.py
• packages/kumo-core/tests/test_prediction_output_columns.py
• packages/kumo-engine/docs/mlops.md
• packages/kumo-engine/kumo/engine/_mlops.py
• packages/kumo-engine/kumo/engine/_scratch.py
• packages/kumo-engine/kumo/engine/_tracking.py
• packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py
• packages/kumo-engine/kumo/engine/integrations/_mlflow_model_code.py
• packages/kumo-engine/kumo/engine/integrations/mlflow.py
• packages/kumo-engine/kumo/engine/trainer/trainer.py
• packages/kumo-engine/test/test_mlflow_model.py

Walkthrough

The changes add task-specific prediction output columns and MLflow model logging, registration, loading, and serving. Trainer.fit passes model-logging results and errors through the tracking session. MLflow model logging supports registry configuration, model metadata, and upload timeout handling. The pyfunc integration loads trainer bundles and graph snapshots, constructs prediction signatures, and serves predictions. Tests cover output columns, model lifecycle and serving behavior, registry validation, and bounded tracking waits.

Sequence Diagram(s)

sequenceDiagram
  participant TrainerFit
  participant TrackingSession
  participant MlflowTracking
  participant MLflowModelRegistry
  TrainerFit->>TrackingSession: Queue model logging with result, snapshot, and registry
  TrackingSession->>MlflowTracking: Call log_model
  MlflowTracking->>MLflowModelRegistry: Register model when registry is supplied
  MlflowTracking-->>TrackingSession: Return registered model or error
  TrackingSession-->>TrainerFit: Return model and error

Loading

Priority: Normal

Merge Risk: Moderate · up to 9b8a8

Resolve the registration and tracking lifecycle inconsistencies and restrict snapshot selection before merging. The prediction signature should also match the values returned to consumers.

Pre-merge checks: 2 · 1

Failed checks: 1 inconclusive

Check name	Status	Explanation	Resolution
Change Map Design Review		The required PR conversation comment beginning <!-- kumo-pr-change-map --> is not available. The repository contains only the template and review instructions, and the supplied review context report…	Provide the complete authored change-map conversation comment, including its observed failure or release requirement, existing behavior, smallest sufficient change, focused verification, stopping condition, and rationale for each changed fi…

Passed checks: 2 passed

Check name	Status	Explanation
Pr Size		The pull request changes 11 files. This does not exceed the limit of 40 files.
Public Resource Access		No added or modified concrete external URL, repository location, or package source location requires public-access verification. The diff contains only illustrative s3://bucket/... and models:/...…

Full details: Public Resource Access

Explanation

No added or modified concrete external URL, repository location, or package source location requires public-access verification. The diff contains only illustrative s3://bucket/... and models:/... examples, local test paths, and package/import names.

Full details: Change Map Design Review

Explanation

The required PR conversation comment beginning <!-- kumo-pr-change-map --> is not available. The repository contains only the template and review instructions, and the supplied review context reports no posted threads. The exact proposed design, rationale, focused verification, stopping condition, and file/function rows therefore cannot be assessed. The authoritative diff is 949 lines, so the missing function/scope-level map matters. The code also requires explicit design review: KumoModel.predict lets params['snapshot'] override the configured default, and _snapshots grows without eviction while loading each URI.

Resolution

Provide the complete authored change-map conversation comment, including its observed failure or release requirement, existing behavior, smallest sufficient change, focused verification, stopping condition, and rationale for each changed file/function scope. Then reassess the snapshot trust boundary and cache behavior against that design.

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

Note

Quiet mode is enabled, so only the most important comments were posted inline. Other review comments are grouped below.

Other comments: 3

packages/kumo-engine/kumo/engine/integrations/mlflow.py: 1

385-396: Minor · Functional Correctness · Quick win

Warn instead of failing when the alias or comment step fails after registration.

mlflow.register_model creates the version first. If update_model_version or set_registered_model_alias then raises, _register raises too. TrackingSession.log_model reports "Could not log the model" and returns None. fit returns registered_model=None and the run ends FAILED. The version still exists in the registry, so the result does not match the registry state.

The connector path in trainer._publish handles this case differently. It warns about the missing alias and still returns the published model. Use the same behavior here.

🛠️ Proposed fix

     version = str(mlflow.register_model(info.model_uri, name).version)
     client = mlflow.MlflowClient()
-    if location.comment:
-        client.update_model_version(name, version,
-                                    description=location.comment)
-    if location.alias:
-        client.set_registered_model_alias(name, location.alias, version)
+    try:
+        if location.comment:
+            client.update_model_version(name, version,
+                                        description=location.comment)
+        if location.alias:
+            client.set_registered_model_alias(name, location.alias, version)
+    except Exception as error:
+        warnings.warn(
+            f'Registered {name} version {version} without its comment or '
+            f'alias: {error}', UserWarning, stacklevel=2)
     return PublishedModel(name=name, version=version,
                           uri=f'models:/{name}/{version}')

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

584-586: Minor · Stability & Availability · Quick win

Finish the tracking session if the model-logging wait is interrupted.

session.log_model can block for up to an hour. The try/except BaseException earlier in fit covers only self._fit. If a user presses Ctrl-C, or another BaseException arrives, during this wait, session.finish never runs. The results are:

• The MLflow fit run stays RUNNING.
• The process-global MLFLOW_HTTP_REQUEST_* variables set by _limited_http are not removed.
• The tracking thread keeps running.

Wrap the call so the session always finishes.

🛠️ Proposed fix

-        registered, error = session.log_model(result, snapshot,
-                                              _mlflow_registry(settings))
-        session.finish(result, error)
+        try:
+            registered, error = session.log_model(
+                result, snapshot, _mlflow_registry(settings))
+        except BaseException as interrupt:
+            session.finish(result, interrupt)
+            raise
+        session.finish(result, error)

packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py: 1

19-19: Minor · Functional Correctness · Quick win

Align the declared PREDICTED signature type with classification output.

_DTYPES declares PREDICTED as boolean at packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py:19. The classification path in packages/kumo-core/src/kumo_core/train/formatting.py builds PREDICTED with df.apply(lambda x: 1 if x["CLASS"] == x[target_col_name] else 0, axis=1), then renames the column to PREDICTED. That produces integer 0/1 values, not booleans.

MLflow does not enforce output schemas in pyfunc.predict, so this mismatch is not caught at scoring time. Consumers that derive result types from signature.outputs can then receive values that do not match the declared type. One example is mlflow.pyfunc.spark_udf when no result_type is given.

Declare the classification PREDICTED output with an integer type, such as long, rather than boolean. Change only the classification entry. Before editing _DTYPES, confirm that the entry is not shared with other task types. The multilabel path sets PREDICTED to None (formatting.py:477), so one type may not fit every task.

---

Review info

Run configuration

• Source mode: Auto
• Repository: NVIDIA/kumo-engine
• Branch: vedaant/mlops-registry-aliases
• File: .coderabbit.yaml
• Commit: 71815cc
• Layers: Repository YAML
• Review profile: QUIET
• Plan: Enterprise
• Run ID: 8127603b-c8f6-4624-84a5-45c9bd99d389

Commits

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

Files selected for processing: 11

• packages/kumo-core/src/kumo_core/train/formatting.py
• packages/kumo-core/tests/test_prediction_output_columns.py
• packages/kumo-engine/docs/mlops.md
• packages/kumo-engine/kumo/engine/_mlops.py
• packages/kumo-engine/kumo/engine/_scratch.py
• packages/kumo-engine/kumo/engine/_tracking.py
• packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py
• packages/kumo-engine/kumo/engine/integrations/_mlflow_model_code.py
• packages/kumo-engine/kumo/engine/integrations/mlflow.py
• packages/kumo-engine/kumo/engine/trainer/trainer.py
• packages/kumo-engine/test/test_mlflow_model.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 4 remain after this review.

[coderabbitai]

Major · Security & Privacy · Heavy lift

Restrict client-supplied snapshot URIs on a served model.

predict takes params['snapshot'] from the caller. _snapshot passes it unchecked to snapshot_location and connector_for_root. It then calls GraphSnapshot.load_from_connector. This causes three problems for any serving endpoint:

• A client can make the replica read from any s3://, gs:// or local path, using the server's credentials. It then loads the downloaded index into DiskGraph.
• self._snapshots caches every loaded URI, and nothing evicts entries. Each distinct URI adds disk and memory use for the life of the process. A client can exhaust both.
• The serving_snapshot default does not constrain this, because a snapshot param always overrides it.

Add a way to pin serving to the configured snapshot. For example, take an allow-list of URI prefixes from model_config, or ignore the snapshot param when serving_snapshot is set. Optionally bound the cache.

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