[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: #53 → #112 → #50 → #60 → #64; #114 on #112; #115 on #64
Merge order: one linear stack; merge bottom to top (#53 first); #114 sits on #112 and merges any time after it; #115 (draft) sits on #64 and merges after it
This PR: Adds experiment tracking and model registration for Trainer.fit: kumo.MLOps(tracking=..., registry=...) records every fit in MLflow or Snowflake Experiment Tracking, registers its model in a Snowflake, Unity Catalog or MLflow registry under an alias, and logs it as an MLflow pyfunc model that batch jobs and SageMaker endpoints can score.

Update (10-29, harness pass): no warning for overridden text columns

• Problem. A fit whose plan set users.BIO to TextEmbedding(...) still logged Text column 'users.BIO' will not be encoded by default: encoder inference runs for every column before the plan's overrides apply.
• Change. Fit drops encoder-inference warnings for the columns in encoder_overrides and keeps them for the rest (no_encoder_inference_warnings(overridden); no kumo-core change).
• Verified. test_embedding_models.py 4 passed: the fixture's fit overrides the warned column and logs nothing; the fit without overrides still warns. With the trainer change reverted, the test fails.

Update (10-29): usability fix pass

From the first-hour usability tests on Databricks, SageMaker and Snowflake:

• MLflow floor. Loading by alias passes download_artifacts(registry_uri=), added in MLflow 3.5; Databricks Runtime 17.3 ships mlflow-skinny 3.0.1, so it failed with a TypeError in notebooks. The mlflow extra now requires mlflow-skinny>=3.5 (%pip upgrades the runtime's), and a new CI job mlflow-floor runs the MLflow suites against mlflow 3.5.0 (four tests that read logs are left out: MLflow 3.5's SQLite migrations reconfigure logging). Locally: 101 passed at 3.5.0.
• Own MLflow calls as the tracker's identity. with tracking.client() as client: yields an MlflowClient for the tracker's server with its credentials and tracking URI set for the block (mlflow.search_runs inside it works too). With SageMakerMLflow(identity=...) the docs' raw MlflowClient() got 403.
• Promotion recipes. MLflow: checks client.get_registered_model(name).aliases for the first promotion. Snowflake: VERSION champion UNSET ALIAS (skipped on the first promotion), VERSION <v> UNSET ALIAS (drops candidate), then SET ALIAS = champion; the old two-line recipe failed both times.
• Exports. kumo.TextEmbedding, kumo.WriteMode (Stype stays in kumoapi.typing; the docs pass stype="text"); cfg.snapshot_uri(snapshot) gives a stored snapshot's location (the pyfunc's snapshot param).
• Text columns. A text column left unencoded now says when no word-vector model is installed. Docs: GloVeVectors("glove.6B.zip", dim=100), with_column("bio", stype="text").
• Docs. S3 table example; starting a stopped SageMaker tracking server; SnowflakeConnector(role=); torch CUDA note; Databricks install line; TMPDIR next to scratch_dir.

Update (10-28): one registry style, platform run ids, scoring defaults at fit, live-run fixes

• One registry style: kumo.MLOps(registry=ModelRegistryLocation(name, alias=...)) only, and the tracker's registry registers it: MlflowTracking an MLflow pyfunc (Unity Catalog, SageMaker MLflow, any MLflow server), SnowflakeTracking(connector) the Snowflake Model Registry version through its connector. registry=ModelOutputConfig(...) is gone from MLOps; plain fit(export_config=ModelOutputConfig(...)) is unchanged.
• Load by alias from the settings: Trainer.load_from_connector(cfg, location=ModelRegistryLocation(name, alias="champion")), one way on every platform. Re-register: result.export(cfg) / trainer.export(cfg) registers a model through the tracker; a registry that already holds the bundle under that name returns its version (MLflow, Unity Catalog, SageMaker, Snowflake).
• No Kumo run ids: the kumo.run_id and kumo.parent_run tags, trial ids <fit run id>-trial-… and Event.run_id are gone. The model manifest records the tracker's own run as tracking: {platform, run_id, url}, read as TrainingResult.tracking; runs are found in the platform UI or by kumo.bundle_id.
• fit(scoring_defaults=ScoringDefaults(...)) writes the defaults into the bundle before it is stored and registered; with_scoring_defaults stays for retuning.
• Signatures: MlflowTracking(target=) also takes a tracking URI (no tracking_uri=) and needs experiment_name only to record; serving_snapshot moves off MlflowTracking (it stays on log_model, which now also takes a GraphSnapshot or snapshot id); load_model/log_model take target=; SnowflakeTracking(connector_or_session, experiment_name=..., schema_name=...).
• Fixes from the live validation: every store failure (snapshot, training table, model, asset) follows on_error like tracking and registration, and a model whose bundle or assets were not stored is never registered; assets are written to object stores without bucket-level calls (prefix-scoped S3 rights suffice); a second MLflow registration into a Unity Catalog model no longer fails on ModelVersionSearch.tags; an ambient Databricks target sets the process-wide tracking URI per call, so a ~/.databrickscfg profile works with no environment variables; a service-principal client sets auth_type="oauth-m2m", so the token lent to MLflow on another thread cannot make its auth ambiguous.
• docs/mlops.md is restructured: quickstart per platform, id table, train/register/promote, score, credentials, then Reference.

Captured outputs further down (event JSON with run_id, run trees with kumo.run_id/kumo.parent_run, registry=ModelOutputConfig(...), SnowflakeTracking(connector=...), tracking_uri=) predate this update; the bullets above supersede them.

Summary

This PR adds experiment tracking and model registration for Trainer.fit: kumo.MLOps(tracking=..., registry=...) records every fit in MLflow or Snowflake Experiment Tracking, registers its model in a Snowflake, Unity Catalog or MLflow registry under an alias, and logs it as an MLflow pyfunc model that batch jobs and SageMaker endpoints can score. Underneath are structured fit progress events (fit(callbacks=...)), and scoring defaults (a threshold or class count) stored in the model bundle.

It builds on #50, which adds the _provenance.json manifests, the artifact root and the kumo.MLOps switch with only artifact_root and store. #64, the PR above, adds prediction-run manifests, records prediction runs in the same trackers, monitors drift and realized performance, and ships job templates.

Problem. Teams running Kumo fits as scheduled Databricks, SageMaker or Snowflake jobs need to answer questions like these:

• Is last night's training job still making progress? Which search trial is best so far, and with which hyperparameters?
• Which of this week's fits had the best validation AUROC? Did any fit fail, and with what error? We want to see this in the tracker we already use (MLflow or Snowflake), not in logs.
• Which registered version is in production, and which bundle (and so which snapshot and source table versions) does it hold?
• Can we promote a new version or roll back by moving a stable name such as champion, without redeploying the scoring job?
• The threshold we picked after looking at the holdout lives in a job script. Can it travel with the model?
• Can a SageMaker endpoint or a Spark batch job score models:/users_churn@champion without the training code?

Today none of these can be answered.

Cause. At #50, fit() reports progress only as Lightning progress-bar text, overwritten in place on stdout. The engine's training bindings hard-wire NoOpTrainingInstrumentation, and fit() takes no callbacks. The MLOps switch knows only where to store artifacts:

kumo.MLOps(artifact_root="s3://bucket/kumo", export=True)   # no tracking, no registry

Trial history stays in the trainer's scratch directory, which no tracker reads:

<scratch>/trainers/<trainer>/runs/<fit run id>/models/<fit run id>/automl/
  auto_train_progress.json  best_trainer.json  encoders.json
  version_0/  version_1/  version_0_refit_trainval/     # tensorboard/, train_progress.json, checkpoint/ per trial

A publish through fit(export_config=ModelOutputConfig(connector, location=...)) records only a marker, so listing can tell Kumo versions from other tools' versions. ModelRegistryLocation has no alias, and a fetch is by exact version or the registry's default:

# Snowflake: Registry.log_model(..., metrics=...)
{"kumo_engine": "1"}
# Databricks: POST /api/2.1/unity-catalog/models/versions, tags
[{"key": "kumo_engine", "value": "1"}]

A bundle is sealed by its publication manifest (the bundle id is a hash of it), so a threshold chosen after the fit has no place in it. predict needs binary_classification_threshold on every binary call. Publishing happens only inside fit(). Nothing logs a model to MLflow.

Change.

kumo/engine/
  _events.py                  # new: Event, EventEmitter (fit progress events)
  tracking.py                 # new, public: Event, Tracking, RunRecord (for writing an adapter)
  _tracking.py                # new: FitContext, RunRecord, run_record, trial_record, fit_inputs, Tracking, TrackingSession
  _mlops.py                   # + MLOps.tracking, MLOps.registry
  _manifest.py                # + TrackingRef; ModelManifest.tracking
  _scratch.py                 # + scoped_scratch_dir (per-call scratch for the pyfunc model)
  connector/snowflake.py      # + SnowflakeConnector.session()
  integrations/
    __init__.py               # new
    mlflow.py                 # new: MlflowTracking, log_model, load_model
    _mlflow_model.py          # new: KumoModel (the pyfunc), signature, snapshot loading
    _mlflow_model_code.py     # new: models-from-code entry point
    snowflake.py              # new: SnowflakeTracking
  trainer/
    _scoring.py               # new: ScoringDefaults, scoring-defaults.json
    trainer.py                # fit(callbacks=), Trainer.export, TrainingResult.registered_model / .scoring_defaults / .with_scoring_defaults
docs/mlops.md                 # + Scoring defaults, Registering models, Promote a Kumo model, Progress events, Experiment tracking, MLflow models
kumoapi/plugin.py             # + ModelRegistryLocation.alias, ModelAliasingSink, check_model_alias, registry_lineage, BUNDLE_ID_KEY
kumo_snowflake/registry.py, kumo_databricks/registry.py   # + kumo_bundle_id, set_model_alias, fetch by alias
kumo_core: instrumentation.event hook, trial/epoch events, prediction_output_columns, bundle_task_type, remove_bundle_file, pidfd fallback

All examples below come from real CPU runs of this branch on a synthetic graph: users (200 rows) and orders (3000 rows, user_id foreign key, ts time column) read through a FileConnector, query PREDICT COUNT(orders.*, 0, 7, days) > 0 FOR EACH users.user_id, two search trials of two epochs. Hashes are shortened and floats rounded; local paths are <data> / <scratch>.

1. Progress events

Trainer.fit(..., callbacks=[...]) calls each callback with a kumo.engine.tracking.Event (frozen dataclass: name, run_id, time, trial, epoch, metrics, params, plus to_dict()). Order: fit_started; per trial attempt trial_started, one epoch_ended per epoch, trial_ended; then fit_ended. JSON lines for platform log capture are one callback:

trainer.fit(table, snapshot=snapshot, callbacks=[lambda e: print(json.dumps(e.to_dict()))])

Output (lines trimmed: trial 1 only, metrics cut to val_auroc/train_loss, the searched params cut to three):

{"name": "fit_started", "run_id": "1611ba2c…-fit-816ff38d…", "time": 1793026044.44, "trial": null, "epoch": null, "metrics": {}, "params": {"snapshot_id": "608b12db…@6084e62f9be989de", "query_sha256": "cdf3deff…"}}
{"name": "trial_started", "run_id": "1611ba2c…-fit-816ff38d…", "time": 1793026044.61, "trial": 1, "epoch": null, "metrics": {}, "params": {"optimization.base_lr": 0.001, "optimization.batch_size": 512, "model_architecture.activation": "leaky_relu", "...": "...", "attempt": 1}}
{"name": "epoch_ended", "run_id": "1611ba2c…-fit-816ff38d…", "time": 1793026049.58, "trial": 1, "epoch": 1, "metrics": {"train_loss": 0.7117, "val_auroc": 0.502, "...": "..."}, "params": {}}
{"name": "epoch_ended", "run_id": "1611ba2c…-fit-816ff38d…", "time": 1793026052.48, "trial": 1, "epoch": 2, "metrics": {"train_loss": 0.6832, "val_auroc": 0.5015, "...": "..."}, "params": {}}
{"name": "trial_ended", "run_id": "1611ba2c…-fit-816ff38d…", "time": 1793026053.72, "trial": 1, "epoch": null, "metrics": {"val_auroc": 0.502, "seconds": 9.1058}, "params": {"status": "succeeded", "attempt": 1}}
{"name": "fit_ended", "run_id": "1611ba2c…-fit-816ff38d…", "time": 1793026069.26, "trial": null, "epoch": null, "metrics": {"val_auroc": 0.517, "test_auroc": 0.502, "test_refit_auroc": 0.5303, "...": "...", "seconds": 24.82}, "params": {"status": "succeeded"}}

• run_id is the fit's run id, the name of its scratch directory; every event of one fit shares it. trial and epoch count from 1.
• fit_started.params: the snapshot id and a SHA-256 of the query (never the query text).
• trial_started.params: only the hyperparameters the search varied, keyed by model-plan path (non-scalars as strings), plus attempt. An out-of-memory retry is a new attempt of the same trial.
• epoch_ended.metrics: the epoch's train_* and val_* metrics. trial_ended: the tune metric and seconds. fit_ended: the final TrainingResult.metrics plus seconds.
• A failed attempt or fit ends with status: failed and error ("<Type>: <message>", at most 200 characters) before the exception propagates.
• Metrics are finite floats only. A callback that raises is logged and not called again; the fit continues.
• kumo-core gains an additive TrainingInstrumentation.event(name, **fields) hook (a no-op by default) and local_training_bindings(instrumentation=). AutoTrainer emits the trial events, AutoTrainingProgressLogger emits epoch_ended, and the searcher records sampled_params.

2. Experiment tracking

from kumo.engine.integrations.mlflow import MlflowTracking          # kumo-engine[mlflow]
from kumo.engine.integrations.snowflake import SnowflakeTracking    # kumo-engine[snowflake]

cfg = kumo.MLOps(artifact_root="s3://<bucket>/kumo",
                 tracking=MlflowTracking(experiment_name="users-churn", run_name="fit",
                                         tags={"team": "growth"}))
# or tracking=SnowflakeTracking(connector.session(), experiment_name="users_churn",
#                               database="ML", schema="EXPERIMENTS", run_name="fit_1")
with kumo.mlops(cfg):                    # or fit(..., mlops=cfg)
    result = trainer.fit(table, snapshot=snapshot)

Tracking is on only through MLOps settings; there is no per-call tracking= argument. SnowflakeConnector.session() (new) returns the Snowpark session the connector reads through; SnowflakeTracking(connector=...) opens it from the connector's identity.

Credentials. Trackers authenticate with the same Identity objects as connectors, so a fit on a laptop or VM can record to Databricks, SageMaker or Snowflake:

from kumo.engine.integrations.mlflow import MlflowTracking, SageMakerMLflow

MlflowTracking(target=DatabricksConnector(..., identity=DatabricksIdentity(host=..., client_id=..., client_secret=...)),
               experiment_name="/Users/<you>/users_churn")          # Databricks MLflow + Unity Catalog
MlflowTracking(target=SageMakerMLflow("<tracking server ARN>", identity=S3Identity(role_arn="<role ARN>")),
               experiment_name="users-churn")                         # kumo-engine[mlflow,sagemaker]
SnowflakeTracking(connector=SnowflakeConnector(..., identity=SnowflakeIdentity(...)), experiment_name="USERS_CHURN")
MlflowTracking()                                                      # MLflow's own tracking URI, read once here

• Every MLflow call goes through one MlflowClient(tracking_uri, registry_uri) per tracker; the engine no longer calls mlflow.start_run, mlflow.pyfunc.log_model or mlflow.register_model, which only reach the process-wide tracking URI. Model logging is pyfunc.save_model plus create_logged_model / log_model_artifacts / finalize_logged_model / log_outputs; registration is create_registered_model plus create_model_version; system metrics come from MLflow's CPU/disk/network/GPU monitors on a kumo thread that logs through the client.
• MLflow takes no per-client credentials, so with explicit ones each call sets, under one lock, exactly what MLflow reads, and restores it: Databricks gets DATABRICKS_HOST, a short-lived DATABRICKS_TOKEN minted by the identity's own client and the process-wide tracking URI databricks (MLflow's Unity Catalog store downloads the logged model through the global URI; without it, registration from a laptop failed with "Logged model ... not found") (other DATABRICKS_* auth variables unset; the client secret never enters the environment); SageMaker gets the process-wide tracking URI (the sagemaker-mlflow plugin signs and routes with it), the identity's boto3 session through sagemaker_mlflow.session.use_session (per thread), and SAGEMAKER_MLFLOW_ASSUME_ROLE_ARN unset; SageMaker model artifacts go through an S3ArtifactRepository built with the identity's credentials. Ambient identities (DatabricksIdentity(), S3Identity()) set no credentials; a SageMaker target still sets the process-wide tracking URI for each call (the plugin signs and routes with it), so no MLFLOW_TRACKING_URI is needed, and other threads calling MLflow during a fit see the SageMaker URI. The docs say not to call MLflow or Databricks on other threads during a fit with explicit credentials.
• Checked when kumo.MLOps is built (Tracking.check(registry), new optional member): MLflow looks up the experiment (or lists one), lists one registered model when it registers, and on Unity Catalog reads the model's schema (GET /api/2.1/unity-catalog/schemas/{catalog}.{schema}, no warehouse) so a missing catalog or schema, or one the principal cannot use, fails before training; Snowflake runs SHOW SCHEMAS LIKE ... IN DATABASE .... An unreachable server or refused credentials raise RuntimeError: MLflow at <uri> is not usable with these credentials: <error>.

Layers and boundaries.

progress Events + FitContext (query, resolved plan, source tables) + TrainingResult
     │
run_record / trial_record  ->  RunRecord(name, params, metrics, steps, tags, artifacts, inputs)
     │      the only place parameter, metric and tag names are defined
     │
TrackingSession            rank 0 only; one background thread, queue of 1000; 30 s end-of-fit drain
     │
Tracking protocol          start(record) · on_trial(record, event) · finish(record | None, error)
     │                     optional: prepare(), tags, log_model(result, snapshot, registry)
adapters                   MlflowTracking · SnowflakeTracking · your own

• Public: kumo.engine.tracking exports Event, Tracking and RunRecord, for writing an adapter. The adapters live in kumo.engine.integrations.mlflow and .snowflake. Nothing is added to the top-level kumo.engine namespace except ScoringDefaults.
• Private: kumo.engine._tracking (records, session, fit_inputs) and kumo.engine._events. Adapters translate ready-built records into platform calls and never compute names or values. A platform that lacks a feature records less; nothing is emulated.
• Off the training thread. Every tracker call is queued on one background thread, so a slow or unreachable tracking server never stalls training or the ranks waiting on rank 0. Only prepare() runs on the fit's thread (MLflow uses it to nest the fit run under the caller's active run). The fit waits at most 30 s at the end; anything left finishes in the background. If the queue overflows, the rest of the fit is untracked, with one warning. While a fit is tracked, MLflow HTTP calls use 2 retries and a 30 s timeout unless MLFLOW_HTTP_REQUEST_MAX_RETRIES / MLFLOW_HTTP_REQUEST_TIMEOUT are set.
• Failures never fail the fit. A tracker whose hook raises is logged once, gets one finish(None, None) that ends its run as failed, and is not called again for that fit.

Runs carry pointers, not lineage. Lineage (snapshot, training table, source table versions, query) lives once, in the model's _provenance.json (#50). A run repeats none of it as tags or params:

Tag	On	Value
kumo.kind	every run	fit or trial
kumo.run_id	every run	the engine's fit run id; trials <fit run id>-trial-<n>-attempt-<m>
kumo.bundle_id	a fit run that succeeded	the model's bundle id, i.e. models/<bundle id>/ under the artifact root
kumo.parent_run, kumo.status, kumo.error	trial runs; kumo.error also on a failed fit	outcome
kumo.dropped_params, kumo.dropped_metrics	when a platform limit dropped some	count

Params are options: the resolved model plan flattened to dotted keys on the fit run; the searched hyperparameters and attempt on a trial. User tags= are added under their own names; one named kumo.* or like a param is dropped with a warning. Artifacts: query.pql, model_plan.json and, at finish, the model's _provenance.json.

The model's _provenance.json points back at the run. A tracked fit's ModelManifest gains (from the run above):

"tracking": {"run_id": "1611ba2c…-fit-816ff38d…"}

tracking.run_id is the kumo.run_id tag of the fit run, not the tracker's own id. It is written only when the fit was tracked, and dropped where the model is embedded in a child manifest (warm start, derived model). After a fit, the tracker object's own run_id attribute (cfg.tracking.run_id) is the MLflow run id, or the Snowflake run name, for job summaries.

MLflow. One run per fit, nested under an already active run (for example a Databricks notebook's), and one child run per trial attempt linked by mlflow.parentRunId. Read back with MlflowClient.search_runs / get_metric_history from a local store:

experiment "users-churn"
└── run "fit"                      FINISHED   3577867c…
    ├── tags       team=growth  kumo.kind=fit  kumo.run_id=1611ba2c…-fit-816ff38d…  kumo.bundle_id=52b1ba56…
    ├── params     (70) optimization.base_lr=[0.01, 0.001]  optimization.batch_size=[512, 1024]  …
    ├── metrics    val_auroc=0.517  test_auroc=0.502  test_refit_auroc=0.5303  train_loss=0.7243  … seconds=24.82
    ├── artifacts  _provenance.json  model_plan.json  query.pql
    ├── inputs     orders  local  {"uri": "<data>/orders.parquet"}
    │              users   local  {"uri": "<data>/users.parquet"}
    ├── run "trial-1-attempt-1"    FINISHED   258202ce…
    │   ├── tags     mlflow.parentRunId=3577867c…  kumo.kind=trial  kumo.status=succeeded
    │   │            kumo.run_id=1611ba2c…-fit-816ff38d…-trial-1-attempt-1  kumo.parent_run=1611ba2c…-fit-816ff38d…
    │   ├── params   (9) attempt=1  optimization.base_lr=0.001  optimization.batch_size=512  model_architecture.activation=leaky_relu  …
    │   └── metrics  best_val_auroc=0.502  seconds=9.11
    │                val_auroc   (step = epoch)  1: 0.502   2: 0.5015
    │                train_loss  (step = epoch)  1: 0.7117  2: 0.6832      (+ the other train_/val_ metrics)
    └── run "trial-2-attempt-1"    FINISHED   1758a530…
        ├── params   (9) attempt=1  optimization.base_lr=0.001  optimization.batch_size=1024  …
        └── metrics  best_val_auroc=0.5207  seconds=7.35   val_auroc  1: 0.5207  2: 0.517

• Source tables as run inputs. fit_inputs(manifest) reads them from the training table's embedded snapshot manifest (plus a custom training table's own source) without opening a file; a stub parent adds none. RunRecord.inputs maps each alias to {connector, location, version, version_kind}. MLflow logs DATABRICKS tables as delta_table sources (unquoted name, plus the Delta version when the snapshot recorded an exact one), so Unity Catalog links each table to a model registered from the run; LOCAL_FILE, S3, GCS by URI (local, s3, gs). Snowflake, BigQuery and Custom SQL have no MLflow source type and are not logged; they stay in _provenance.json.
• MLflow rejects @ in metric names, so the adapter logs link-prediction metrics such as test_map@10 as test_map_at_10 (integrations.mlflow.metric_key, also used for the logged model's metrics).
• Params MLflow rejects (bad key, value over 6000 characters) are counted in kumo.dropped_params; the full plan is in model_plan.json. Optional system metrics (CPU, memory, disk, network, GPU with nvidia-ml-py) go on the fit run; log_system_metrics=False turns them off.
• A user's mlflow.pytorch.autolog() is suspended once per fit, so it adds no runs; a user's disable_autologging() is left alone.
• Where runs go: MLFLOW_TRACKING_URI (automatic on Databricks; the tracking server ARN plus the sagemaker-mlflow plugin for SageMaker managed MLflow; MLflow's local store otherwise).

Snowflake. Snowflake Experiment Tracking (snowflake-ml-python) has one active run per process, no nested runs and no tags, and start_run(name) resumes a run whose name matches as a SQL LIKE pattern. So SnowflakeTracking records one run per fit with each trial attempt inside it under a trial-<n>-attempt-<m>/ prefix, logs tags as params of the same name, and never resumes a run: a taken run_name gets _<8 hex of the fit run id> with a warning. Two real fits with the same run_name="fit_1", recorded into the test fake that stands in for snowflake.ml.experiment (no Snowflake account was used):

ML.EXPERIMENTS
└── experiment USERS_CHURN
    ├── run FIT_1                                   FINISHED
    │   ├── params (98)
    │   │   optimization.base_lr=[0.01, 0.001]  optimization.batch_size=[512, 1024]  … (resolved plan)
    │   │   team=growth  kumo.kind=fit  kumo.run_id=3779956f…-fit-3e3ec2b5…  kumo.bundle_id=b33bbcca…
    │   │   trial-1-attempt-1/optimization.base_lr=0.001  trial-1-attempt-1/attempt=1  …
    │   │   trial-1-attempt-1/kumo.kind=trial  trial-1-attempt-1/kumo.status=succeeded
    │   │   trial-1-attempt-1/kumo.run_id=…-trial-1-attempt-1  trial-1-attempt-1/kumo.parent_run=3779956f…-fit-3e3ec2b5…
    │   │   trial-2-attempt-1/…
    │   ├── metrics (37 names)
    │   │   trial-1-attempt-1/val_auroc  step 1: 0.512  step 2: 0.4976    trial-1-attempt-1/best_val_auroc 0.512
    │   │   trial-2-attempt-1/val_auroc  step 1: 0.5067 step 2: 0.5196    trial-2-attempt-1/best_val_auroc 0.5196
    │   │   val_auroc 0.5196  … seconds 114.07
    │   └── artifacts  query.pql  model_plan.json  _provenance.json
    └── run FIT_1_3AC7C899                          FINISHED   (second fit; warning "Snowflake run fit_1 already exists …; recording this fit as fit_1_3ac7c899")

Calls for the first fit: set_experiment("users_churn", database_name="ML", schema_name="EXPERIMENTS") → show_runs_in_experiment → start_run("fit_1") → log_params (71) → log_artifact ×2 → per trial log_params (12), log_metrics(step=1), log_metrics(step=2), log_metrics (summary), log_params (kumo.status) → log_metrics (17, final) → log_params (kumo.bundle_id) → log_artifact (_provenance.json) → end_run(status="FINISHED").

• database / schema default to the session's current ones (PUBLIC when only database is given), passed explicitly on every fit.
• Limits (1000 params, 200 metric names per run): per-epoch metrics stop taking new names at 140 and trial summaries at 170, so the final metrics always fit; params stop at 980, keeping room for finish-time tags. Anything dropped, and params containing $$ (they would end Snowflake's $$ SQL literal), are counted in kumo.dropped_params / kumo.dropped_metrics.
• If the process's ExperimentTracking singleton already has an active run or belongs to another session, the fit is not tracked there, with a warning, and nothing on it is touched. The adapter reads the private _instance, _run, _session and _sql_client.show_runs_in_experiment, so kumo-snowflake[sdk] now pins snowflake-ml-python>=2.0,<3 (2.0 for end_run(status=)).
• native_model_lineage=True holds the run open as a context for the whole fit, so a model version registered meanwhile is linked to it through Snowflake lineage. Off by default.
• No ALTER SESSION; the session's parameters are not changed. Snowflake has no dataset inputs or system metrics.

Record field	MLflow	Snowflake
fit params, user tags, kumo.* tags	params; tags	params; params with the same names
trial record	child run trial-1-attempt-1	trial-1-attempt-1/<name> params and metrics in the fit run
trial steps	child-run metrics, step = epoch	log_metrics(step=epoch)
artifacts	log_text	log_artifact of temp files
inputs	native run inputs	not recorded (no API)

3. Scoring defaults

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

with_scoring_defaults validates the defaults against the bundle's task with the rules predict applies (a threshold only for binary and multilabel classification, in [0, 1]; a count only for multiclass, multilabel and link prediction), normalises them (1, 1.0 and np.float32(1) give one file and one id), copies the bundle with hardlinks so the parent is never written, adds the new file to the publication manifest, and writes the copy's _provenance.json:

<derived bundle>/
  model/  manifests/  metrics.json  prediction-schema.json  training-spec.json   # hardlinks to the parent
  scoring-defaults.json                # new
  publication-manifest.json            # rewritten: lists scoring-defaults.json (role scoring_defaults, sha256 5bd069d9…, 78 bytes)
  _provenance.json                     # new record: own id, parents.derived_from

// scoring-defaults.json
{
  "binary_classification_threshold": 0.5,
  "num_classes_to_return": null
}
// _provenance.json (shortened)
{"kind": "model", "id": "b213bcb0…", "fit": null, "fit_mode": null,
 "parents": {"derived_from": {"kind": "model", "id": "52b1ba56…", "fit_mode": "refit_trainval",
                              "fit": {"started_at": "2026-10-26T14:47:24+00:00", "seconds": 24.809, "...": "..."},
                              "parents": {"training_table": {"kind": "training_table", "id": "tt-3fa3be7f…", "...": "..."}, "...": "..."}},
             "training_table": null, "warm_start": null},
 "created_at": "2026-10-26T14:47:52+00:00", "schema_version": 1, "environment": {"...": "..."}}

• parents.derived_from embeds the parent's manifest; the parent's own model parents are cut to {kind, id, reconstructed: true}, so re-tuning does not grow the chain. A parent without a manifest reads as that stub.
• ScoringDefaults() removes the file. Clearing b213bcb0… gave back the parent's id 52b1ba56….
• predict precedence: an explicit argument, then (for the count) the query's own TOP k, then the bundle default, then today's behaviour (binary raises without a threshold; multilabel uses 0.5).
• Trainer.export(ModelOutputConfig) (new) publishes a fitted or loaded bundle through the same step as fit(export_config=). It is a plain command: it writes only to the given config, never to the artifact root.
• kumo-core gains remove_bundle_file and bundle_task_type; the warm-start checks now share bundle_task_type.

4. Registering models

Two registration paths, chosen by the type of registry:

# Connector registry: Snowflake Model Registry, or Unity Catalog through a DatabricksConnector.
cfg = kumo.MLOps("s3://<bucket>/kumo", registry=kumo.ModelOutputConfig(
    registry_connector, location=kumo.ModelRegistryLocation("USERS_CHURN", alias="challenger")))

# MLflow registry (Unity Catalog with MLFLOW_REGISTRY_URI=databricks-uc, SageMaker managed MLflow, any MLflow server).
cfg = kumo.MLOps("s3://<bucket>/kumo", tracking=MlflowTracking(experiment_name="users-churn"),
                 registry=kumo.ModelRegistryLocation("users_churn", alias="champion"))

with kumo.mlops(cfg):
    result = trainer.fit(table, snapshot=snapshot)
result.registered_model   # PublishedModel(name='users_churn', version='1', uri='models:/users_churn/1', ...)

• TrainingResult.registered_model (new) is the version the MLOps registry created; published_model stays the plain export_config result. Registration is independent of export=.
• Order inside fit: plain export_config publish, root copy, tracking finish, then the connector registry. A failed registration warns and returns the result with registered_model=None. A connector that cannot publish models, or an invalid alias, raises before training.
• kumo.MLOps(...) raises at construction for a bare ModelRegistryLocation without a tracker that has log_model (MlflowTracking), for a fixed version (MLflow assigns versions), or for an invalid alias.
• Aliases. ModelRegistryLocation.alias (new; exclusive with version). On publish, the version is created first and the alias set afterwards; a failure to set it warns and the version stays published. kumoapi.plugin.check_model_alias defines one syntax every registry accepts: a letter or underscore, then letters, digits or underscores (at most 255), and not DEFAULT, FIRST, LAST, LATEST or v<digits> in any case (which also excludes SageMaker's SageMakerLifecycle- aliases). Trainer.load_from_connector(connector, location=ModelRegistryLocation(name, alias="champion")) loads whichever version an alias points at.
• Connector registries set aliases through a new optional kumoapi.plugin.ModelAliasingSink protocol. Snowflake moves an alias by unsetting it on the holder, then setting it; if setting fails it puts it back and raises. Snowflake allows one alias per version, so a version that already holds a different alias is refused. Unity Catalog uses PUT /api/2.1/unity-catalog/models/{name}/aliases/{alias} and GET on the same path. File, S3 and GCS connectors have no aliases.
• There is no engine call to move an alias later. Promotion uses the registry's own API (MlflowClient().set_registered_model_alias, or ALTER MODEL … VERSION … SET ALIAS on Snowflake); docs/mlops.md has a "Promote a Kumo model" recipe that compares the new model's test metrics with the current champion's.
• TrainingResult.export(ModelOutputConfig) (new) publishes that result's bundle exactly as Trainer.export does for the trainer's latest result, so every artifact has <artifact>.export(config) and a load_from_connector; fit(export_config=...) stays as the shorthand.
• No duplicate Unity Catalog versions through a DatabricksConnector. Publishing a bundle whose kumo_bundle_id a READY Kumo version of the model already carries returns that version (an alias then moves to it) and logs <name> version <n> already holds bundle <id>; reused it; no version is created or uploaded.
• No duplicate MLflow versions. Registering a bundle whose kumo_bundle_id a version of the same name already carries returns that version (the alias and comment still apply) and logs <name> version <n> already holds bundle <id>; reused it; no new version is created.
• Pointer, not lineage. Each version now also records kumo_bundle_id. For the bundle above:

# Snowflake: Registry.log_model(..., metrics=...)
{"kumo_engine": "1", "kumo_bundle_id": "52b1ba56…"}
# Databricks: POST /api/2.1/unity-catalog/models/versions, tags
[{"key": "kumo_engine", "value": "1"}, {"key": "kumo_bundle_id", "value": "52b1ba56…"}]

Training metrics are not copied into either registry; they stay in metrics.json and on the tracked run. list_models and every fetch_model path return the recorded kumo_bundle_id in PublishedModel.metadata (through kumoapi.plugin.registry_lineage), and load_from_connector warns when the fetched bundle's id differs.

• Task type. A Snowflake version also records its Snowflake ML task (task=Task.TABULAR_BINARY_CLASSIFICATION, TABULAR_MULTI_CLASSIFICATION or TABULAR_REGRESSION, from the bundle's prediction schema), which Snowflake Model Monitor requires. Other tasks, or a schema that cannot be read, register without one; nothing raises.
• Also fixed: the Unity Catalog Files API lists directories with a trailing slash, which made the Databricks download reject every real bundle (Invalid model artifact path returned by Databricks.). The slash is stripped before the immediate-child check.

5. MLflow models and serving

With registry=ModelRegistryLocation(...), or MlflowTracking(log_model=True) to log without registering, the fit logs its bundle as a plain pyfunc model (logged from code, no pickled object, no custom flavor) on the fit run, registers it and sets the alias. From the run above:

LoggedModel m-a8132620…   source_run_id 3577867c…
  metrics (model_id m-a8132620…)  val_auroc=0.517  test_auroc=0.502  test_refit_auroc=0.5303 … seconds=24.82
registered model "users_churn"   aliases {champion: 1}
  version 1   source models:/m-a8132620…

The model as downloaded from models:/users_churn@champion:

MLmodel                         # flavor, signature, metadata
_mlflow_model_code.py           # 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
requirements.txt                # mlflow==3.17.1, kumo-engine[gcs,s3]==<training version>
python_env.yaml  conda.yaml  input_example.json  serving_input_example.json  registered_model_meta

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

Scoring it:

model = mlflow.pyfunc.load_model("models:/users_churn@champion")
model.predict(pd.DataFrame({"ENTITY": [3, 7, 42]}),
              params={"snapshot": "<exported snapshot>", "prediction_time": "2024-04-01",
                      "binary_classification_threshold": 0.5})
#    ENTITY  TARGET_PRED  False_PROB  True_PROB
# 0       3         True      0.4440     0.5560
# 1       7         True      0.4392     0.5608
# 2      42         True      0.4230     0.5770

trainer = kumo.engine.integrations.mlflow.load_model("models:/users_churn@champion")   # the Trainer back

• Metadata is kumo_bundle_id only, the same pointer the connector registries record; load_model warns if it differs from the loaded bundle.
• Signature is typed per bundle: ENTITY takes the entity key's type, TIMESTAMP is optional, and outputs are the task's prediction columns from the new kumo_core.train.formatting.prediction_output_columns.
• Params: snapshot is a stored or exported snapshot's URI or path (s3://, gs://, file://, local). An empty frame returns no rows unless population=true, which scores the query's population at prediction_time. A temporal query needs prediction_time or a TIMESTAMP column.
• Snapshots load once per model instance and stay loaded for the life of the process (no eviction). With serving_snapshot=<URI> (on log_model or MlflowTracking), each replica loads it when it starts and predict uses it when no snapshot param is given.
• Each predict call works in its own scratch directory (new internal scoped_scratch_dir, a context variable get_scratch_dir reads) and removes it before returning.
• Requirements: kumo-engine[s3,gcs] of the training version, plus [text] when the bundle encodes a text column (requirements.txt lists extras sorted as gcs,s3).
• Synchronous model logging. The background drain gives finish only 30 s, so a scheduled job could exit mid-upload. After a fit that succeeded, fit queues the tracker's log_model(result, snapshot, registry) on the tracking thread and waits for it while that thread is alive, up to an hour, with a 600 s MLflow HTTP timeout during the upload (unless MLFLOW_HTTP_REQUEST_TIMEOUT is set). On failure, a dead tracking thread, a tracker already dropped earlier in the fit, or the hour passing, fit warns (Could not log the model: ...), the run ends FAILED, and the fit returns its result with registered_model=None; an upload that finishes later is not registered.
• Outside a fit: log_model(trainer, snapshot=...) logs a fitted or loaded model in the given, active or a new run, with the bundle's source tables as inputs.
• Serving needs an image with kumo-engine installed: mlflow models serve -m <model dir> --env-manager local (on SageMaker endpoints, -m /opt/ml/model -p 8080). MLflow's stock serving containers pip-install the requirements at start and cannot install kumo-engine. SageMaker batch transform takes JSON (dataframe_split plus params), not CSV.
• GPU real-time endpoints on SageMaker need an inference image whose NVIDIA driver supports CUDA 12.6; the default GPU hosting image ships driver 470, so set a newer InferenceAmiVersion on the production variant (for example al2023-ami-sagemaker-inference-gpu-4-1). Not yet tested live.
• kumo-core fix for SageMaker hosting: SageMaker hosting containers return ENOSYS for pidfd_open, so every prediction failed with Could not capture stable worker handles: OSError(38, 'Function not implemented'). _WorkerHandle.open now falls back to the duplicated process sentinel fd on ENOSYS or EPERM (other errors still raise), and send_signal uses os.kill whenever there is no pidfd; the live sentinel guards against PID reuse.

Platform	Track	Register
Databricks	MlflowTracking (automatic URI)	Unity Catalog through MLflow, or DatabricksConnector
SageMaker	MlflowTracking + managed MLflow	SageMaker managed MLflow
Snowflake	SnowflakeTracking	SnowflakeConnector (Snowflake Model Registry)
Plain VM	any MLflow server	any MLflow registry

Packaging. New extra kumo-engine[mlflow] installs mlflow-skinny>=3.5,<4 and nvidia-ml-py. The client alone is enough (3.5 for download_artifacts(registry_uri=)), and it does not pull the server's dependencies, so a runtime's pyarrow and cryptography stay as they are (full MLflow upgraded cryptography and broke the preinstalled pyOpenSSL on a Databricks runtime). The pyarrow range DiskGraph accepts is widened in its companion, #124 (pyarrow>=16,<25). New extra kumo-engine[sagemaker] = kumo-engine[mlflow,s3] plus sagemaker-mlflow>=0.5.0, the first release with sagemaker_mlflow.session.use_session (0.4.0 has no session.py); it is a separate extra because the plugin's request-header provider runs on every MLflow request and warns when the process tracking URI is not a SageMaker ARN. mlflow and snowflake.ml are imported lazily, so import kumo.engine loads neither. The root test extra in setup.py adds full mlflow>=3.13,<4, a superset of the extra that the tests need for MLflow's local SQLite store, because engine tests run with --no-deps in the root environment. uv.lock adds mlflow, mlflow-skinny and their dependencies and the snowflake-ml-python specifier; no existing locked version changes.

Scheduled jobs and finding things

• kumo.MLOps(on_error="warn"|"raise"), default "warn". A tracking or registration failure is logged (traceback only at DEBUG) and reported once: in warn mode as one UserWarning (Could not log the model: ..., Could not register the model: ..., or Trainer.fit (its model is in trainer.result) finished, but MLOps tracking or registration failed: ...), so notebooks and job logs both show it; with "raise", as one RuntimeError chained to the cause, after the fit has finished and its result is stored, with no warning before it. The cause's full text is kept (it was cut at 200 characters, which hid AWS's "not authorized to perform ..." part).
• kumo.error on failed fits. A fit whose model logging or registration failed ends its tracked run FAILED with kumo.error (MLflow tag; Snowflake param). A connector registration failure now also ends the run FAILED.
• Bundle id on MLflow-registered versions. registry=ModelRegistryLocation(...) with MlflowTracking registers each version with a kumo_bundle_id tag and returns it in result.registered_model.metadata, as the connector registries do, so the Unity Catalog model page shows which bundle a version is.
• Run names. MLflow fit runs are named fit-<UTC start> (for example fit-20261027T061838Z) unless run_name is given. docs/mlops.md has a "Which id is which" table: kumo run id (kumo.run_id, manifest tracking.run_id), bundle id (result.id, kumo.bundle_id), tracker run id (tracking.run_id on the tracker), registry version (registered_model.version), with a search_runs example from a bundle id to its fit run.
• pyfunc row order. Scoring reorders rows, so the pyfunc model now returns one row per input row in the input's order, matching on ENTITY and, for temporal queries, the anchor time (repeated rows pair up in order; entities not in the snapshot are left out). Trainer._predict takes the internal extra_fields it needs for the anchor.
• pyfunc inputs. A temporal query scored through the pyfunc model without params={"prediction_time": ...} or a TIMESTAMP column fails with temporal query: pass params={'prediction_time': ...} or a TIMESTAMP column; a binary model without a threshold or scoring defaults with binary classification: pass params={'binary_classification_threshold': ...}. The docs example shows both.
• Databricks runtime. The docs recommend Databricks Runtime 17.3 LTS ML or later: its torch already satisfies kumo-core, so installing kumo-engine keeps the runtime's torch, torchvision and MLflow, and (since main's Accept pyarrow 16 through 24 in DiskGraph; explain custom S3 CA bundles below pyarrow 21 #124 accepts pyarrow 16 through 24) its pyarrow 19. On 16.4 LTS ML pip replaces torch 2.6 and breaks the preinstalled torchvision. On Databricks the docs recommend the MLflow model over the connector registry unless only kumo-engine jobs load it.

Registry names and SageMaker

• Names fail fast. A registry model name is checked when kumo.MLOps is built, against the target registry: SageMaker managed MLflow (a tracking-server or MLflow-app ARN) takes letters, digits and hyphens (^[a-zA-Z0-9](-*[a-zA-Z0-9]){0,56}$, e.g. users-churn); Unity Catalog takes up to catalog.schema.model, each part at most 255 characters without spaces, /, . or control characters; Snowflake takes an identifier (letter or _, then letters, digits, _, $, or a quoted name); other MLflow registries add no rule. Before, a bad name failed only after training. The rules conflict (SageMaker forbids the . Unity Catalog needs and the _ Snowflake uses), so the docs say which form each registry takes.
• SageMaker UI. MLflow's printed run links do not open on SageMaker; after a fit tracked on a SageMaker ARN, kumo-engine warns (UserWarning, so notebooks and job logs show it) with the aws sagemaker create-presigned-mlflow-tracking-server-url command (or Studio).
• Docs. A SageMaker managed MLflow section: creating a tracking server, the caller's sagemaker-mlflow:* permissions, pip install "kumo-engine[mlflow,sagemaker]" and MlflowTracking(target=SageMakerMLflow(arn), ...), the name rule, opening the UI; S3 examples for the artifact root, source tables (data["orders.parquet"]), prediction outputs and the monitoring store; and "Running from a laptop" with one example per platform and the explicit-credentials caveat. examples/mlops/smoke.py is a manual live check per platform (not run in CI).

6. Internal seams: any bundle, run and predictor

No new public API. The built-in Trainer is now one client of internal, model-agnostic seams, so a later user-supplied training loop can reuse tracking, root storage, registries and serving:

• Run records. _tracking.run_record(RunSpec) builds every tracked record from a neutral spec: kind, run id, flat params, metrics, per-step histories, user tags, inputs, parents (pointer tags such as kumo.bundle_id), status and error. TrackingSession takes spec callables, and trainer/_run_spec.py adapts a fit into one. Nothing in _tracking.py needs a TrainingResult.

• Store, log and register. trainer._store_and_register(bundle_dir, manifest, settings, ...) stores any bundle under <root>/models/<id>/, logs it in the session's run, and registers it. The tracker's log_model and MLflow's _log_model take (bundle_dir, ModelManifest, snapshot). A bundle without a training spec is logged without a signature.

• Dependencies. ModelManifest.dependencies lists what predicting needs besides the bundle and a snapshot:

  • the predictor entry point;
  • the kumo-engine package and its extras;
  • content-addressed assets, by asset_kind, name, version and sha256.

  Bundles without the field read as none, and get the built-in predictor. _deps.Dependencies.require(asset) records an asset during fit. With MLOps on, it also stores the asset once under <root>/assets/<asset_kind>/<sha256>/ with its own _provenance.json. resolve(asset_kind, name) returns a sha256-checked local directory. It looks in explicit locations first, then the recorded location, then the enclosing block's root. A missing asset raises MissingDependencyError, naming the kind, name and hash. Text encoders are the first assets: a fit with MLOps on stores its word vectors, and Trainer.load, so also the pyfunc and score jobs, resolves them before falling back to the embedding-model store. A bundle never embeds the vectors. FileAsset covers plain files.

• Predictors. The MLflow pyfunc loads load_predictor(bundle_dir): the bundle's predictor dependency, with .load(bundle_dir) returning an object with .predict(rows, snapshot, params). TrainerPredictor wraps today's scoring. A predictor that cannot be imported raises a ValueError that names it.

Back-compat

Everything here is opt-in and only records: tracking, registration and model logging never change what a fit trains or returns, and their failures are warnings.

Situation	Behaviour
fit() with MLOps off, or MLOps with neither tracking nor registry	Unchanged. No tracking block in _provenance.json, registered_model is None.
fit() without callbacks	Unchanged; nothing is attached to kumo-core and nothing is printed.
Tracked fit	Bundle files and metric keys as untracked, so the bundle id is too; only _provenance.json (not part of the id) gains tracking.run_id.
Tracker raises, server errors or hangs	Logged once, tracker dropped for that fit, its run ended as failed; the fit is unaffected. The end-of-fit drain waits at most 30 s.
Training fails	fit_ended with status: failed, runs end FAILED with kumo.error, the same exception propagates; no model is logged or registered.
Registration or model logging fails	Warning; the fit returns its result with registered_model=None.
MlflowTracking() with no target	Records where MLflow's tracking and registry URIs point when it is built (before: when each fit started).
kumo.MLOps(..., tracking=...)	Now contacts the tracker once when built (check), and raises there if the server or registry target is unusable; MLOps off, or without a tracker, is unchanged.
A user's active MLflow run	The fit run nests under it only when it is on the same tracking server.
SnowflakeTracking(session, ...)	Unchanged; connector= is the new alternative.
kumo.MLOps(root, ...) options	Everything after artifact_root (tracking, registry, store, on_error) is keyword-only, so adding fields never shifts a positional argument.
Existing predict, 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 as before; its id is unchanged.
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.
scoring-defaults.json unlisted, unreadable, or invalid for the task	Ignored (with a warning when listed); the bundle behaves as one without defaults.
Publishing the same bundle again through a DatabricksConnector	Returns the existing READY version holding its kumo_bundle_id instead of creating another; a new bundle id, or no id, publishes as before.
Registry version published before this PR	Listed and fetched as before, with no kumo_bundle_id in metadata and no mismatch warning.
Version published by this PR, read by an older engine	The extra metric or tag is ignored; listing still keys on the marker.
Fetch by version or with no version	Same version as before; metadata now carries kumo_bundle_id.
Snowflake alias move fails half-way	The alias is restored on the version that held it, a warning is logged, the error is raised (fit and export turn it into a warning).
Fetch by alias from a sink without ModelAliasingSink	NotImplementedError (new call only).
fit() with an export_config whose connector cannot publish models	Still raises NotImplementedError, now before fit_started is emitted.
Environment with snowflake-ml-python < 2.0 or >= 3	kumo-snowflake[sdk] now asks for >=2.0,<3; the registry publish works on 2.x.
kumo-engine without the mlflow or snowflake extra	import kumo.engine works; MlflowTracking() / SnowflakeTracking(...) raise an ImportError naming the extra.
Bundle without dependencies in _provenance.json (all earlier bundles)	Loads and serves as before, with the built-in predictor; its models come from the embedding-model store.
Fit with MLOps on and a text column	Its word vectors are also stored once under <root>/assets/text_encoder/<sha256>/; the bundle and its id are unchanged.
Trainer.load of a bundle whose stored assets cannot be found, and that the store does not hold either	MissingDependencyError at load (only for bundles that record a stored copy); otherwise unchanged.
Model storage under the root	Runs after fit_ended instead of before it, so fit_ended's seconds no longer includes the upload.
Linux host where pidfd_open works	Unchanged: workers are still tracked and signalled through pidfds.

Known limitations

• Snowflake tracking was not run against a live account. It was checked against a recording fake. The 8 tests that compare the fake's signatures with an installed snowflake-ml-python skip when the library is absent, as it was here and is in CI; privileges, Snowflake's undocumented per-run byte limit and name lookup are unexercised. native_model_lineage=True is off by default for the same reason.
• The Snowflake and Databricks alias set, move and rollback, and the kumo_bundle_id record, were tested against the packages' fakes only.
• Concurrent fits in one process are not supported with MLflow autolog on, and not with Snowflake tracking at all (one ExperimentTracking per process).
• If the tracking queue overflows (1000 pending calls), the rest of the fit is untracked. A model upload slower than an hour is abandoned.
• The pyfunc model has a fixed cost per call (it stages the rows, builds a prediction table and starts a loader worker): about 0.6 s on a CPU endpoint. It suits batch-like calls, not low-latency online scoring. A temporal query needs an explicit prediction_time.
• Snapshots loaded by a serving replica are never evicted, so its disk must hold every snapshot it is asked to score against.
• A model registered through a connector registry is linked to its MLflow run only through the shared kumo_bundle_id / kumo.bundle_id.
• Snowflake model monitors cannot attach to Kumo versions yet: log_model does not set task=.
• With explicit credentials, MLflow calls swap process-wide state (environment variables; the SageMaker tracking URI) for their duration, so another thread calling MLflow or Databricks at that moment would see it. The SageMaker swap uses MLflow's private mlflow.tracking._tracking_service.utils._tracking_uri; test_mlflow_target.py pins it and the other internals relied on (mlflow.tracking.fluent._get_experiment_id, the system-metric monitor classes, the plugin's per-thread session, the Databricks SDK's variable names).
• Unity Catalog's version search returns no tags, so the duplicate check reads each version of the name once. Connector registries (Databricks, Snowflake) do not deduplicate in this PR.
• The live smoke scripts were not run for this push; the credential paths are tested with fakes.

Testing Done

Run at this head, one file at a time, CPU, mlflow 3.17.1, without snowflake-ml-python installed (sqlite MLflow store in the tests; Snowflake through the recording fake):

Test file	Passed
test/connector/test_custom_source.py	20
test_bundle_provenance.py	7
test_custom_sql.py	11
test_filestore_sink.py	4
test_fit_events.py	4
test_holdout_df.py	11
test_mlflow_model.py	18
test_mlflow_target.py	15
test_mlflow_tracking.py	15
test_mlops.py	4
test_optional_imports.py	48
test_prediction_completion.py	14
test_prediction_workers.py	55
test_predictive_query.py	67
test_query_from_snapshot.py	4
test_scoring_defaults.py	7
test_semantic_carry_forward.py	20
test_snapshot_export.py	36
test_snapshot_manifest.py	21
test_snowflake_connector.py	8
test_snowflake_tracking.py	13 (8 skipped)
test_tracking_contract.py	8
test_trainer_registry.py	38
test_training_explanations.py	16
test_training_table_export.py	18
test_warm_start.py	6

• Rerun after the credentials and retest changes (on the tree before the rebase onto main, which adds Push down feature projections through DiskGraph and Kumo loaders #46 and Accept pyarrow 16 through 24 in DiskGraph; explain custom S3 CA bundles below pyarrow 21 #124): test_mlflow_target.py, test_mlflow_tracking.py, test_mlflow_model.py, test_snowflake_tracking.py, test_tracking_contract.py, test_trainer_registry.py, with sagemaker-mlflow 0.5.0 installed. test_mlflow_target.py covers the credential scopes with fake Databricks clients and real boto3 sessions (no network), the construction checks against a closed port and a faked Unity Catalog 404, registry reuse on a SQLite store, and pins the MLflow, plugin and Databricks SDK internals the scopes rely on.
• kumo-core changed tests: 235 passed, 1 skipped; test_worker_handles.py: 3 passed. kumo-api: 12 passed. kumo-snowflake, kumo-databricks and kumo-bigquery changed tests pass.
• What the new tests cover:
  • test_fit_events.py: event order and numbering on a two-trial CPU fit, JSON round trip, a raising callback, a trial that fails after training reports one failed trial_ended, a failed fit's fit_ended, non-finite metrics dropped.
  • test_tracking_contract.py: a recorded fake fit (three source tables; two trials, one with an out-of-memory attempt) is replayed through TrackingSession into both adapters; each store is read back and compared with run_record / trial_record (params, metrics, step histories, kumo.* and user tags, parent links, artifacts). Also a raising tracker, a failed fit, rank > 0, and a hung tracker (stalls the fit for under 1.5 s).
  • test_mlflow_tracking.py: link-prediction metric names (test_map@10) logged as test_map_at_10; a real two-trial fit nested under a user's run with autolog on; a failed fit; log_batch errors; a user's disabled autologging; a hung MLflow HTTP server; source tables as local, s3 and delta_table inputs (Snowflake skipped), with the Delta input parsed back the way Unity Catalog registration does; fit_inputs on embedded, custom, warm-start and stub parents.
  • test_snowflake_tracking.py: the fake's signatures match the installed library (skipped where it is not installed); native lineage; errors; an instance in use; location defaults; a taken or LIKE-matching run name; limits; the user's session is not altered; a hung call.
  • test_mlflow_model.py: a real fit with MLOps(tracking=MlflowTracking(), registry=ModelRegistryLocation('m', alias='champion')) gives one logged model with the fit's metrics, the @champion version, metadata {kumo_bundle_id}, the provenance artifact equal to result.manifest, typed signature, served predictions and load_model's mismatch warning; log_model scoring equals Trainer.predict for the population, given entities and a prediction time, under a per-call scratch root; the wait bound, the 600 s upload timeout, a late upload not registered; the [text] extra; a tracker dropped earlier in the fit warns that the model was not logged; invalid MLflow registries refused.
  • test_scoring_defaults.py, test_trainer_registry.py, kumo-snowflake and kumo-databricks registry tests, kumo-api test_plugin_registry.py: as described in sections 3 and 4.
  • kumo-core test_worker_handles.py: with pidfd_open refused by ENOSYS or EPERM, a worker is still observed alive, killed and seen to exit; other errors still raise.
• Live (earlier run, same final API): on a SageMaker real-time endpoint (CPU instance), the pyfunc served by mlflow models serve --env-manager local in an image with kumo-engine preinstalled scored both with a snapshot param and with serving_snapshot, at about 0.6 s per call for 1 to 100 entities; that run found the pidfd_open failure fixed above. MLflow's stock SageMaker container failed because it cannot install kumo-engine. Batch transform completed with a JSON request of 2000 entities; CSV input is refused.
• Seam tests (test_model_agnostic_mlops.py, 6 tests, CPU): a toy run, bundle, predictor and FileAsset go through each seam. One test trains a tiny model with non-default word vectors (a GloVe-format directory built in the test) and MLOps on. A clean subprocess then serves it from the bundle, an exported snapshot and the stored assets alone, with no store, an empty HOME, another cwd and scratch, and proxies pointing nowhere, and matches the training process's scores. After the assets are moved, it fails with MissingDependencyError. Rerun with the change: test_tracking_contract.py, test_snowflake_tracking.py, test_mlflow_tracking.py, test_mlflow_model.py, test_trainer_registry.py, test_bundle_provenance.py, test_embedding_models.py, test_mlops.py (CPU, all pass), and the kumo-snowflake and kumo-databricks registry tests. Tests that call the changed internal signatures were adapted mechanically; their assertions are unchanged. test_mlflow_model.py::test_logged_model_scores_the_query_population_and_timed_rows is flaky on GPU, before and after this change.
• The examples above were captured from CPU runs of this head on synthetic data.

Checklist

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

🤖 Generated with Claude Code

[Vedaant-J]

Change map

File	What changed	Why
.github/workflows/testing.yml	new job mlflow-floor: the MLflow suites against mlflow 3.5.0 (log-reading tests left out)	the extra's floor is tested, not assumed
examples/mlops/smoke.py	new: manual live smoke check per platform (fit, track, register, load back with an explicit identity); install lines per platform; the ambient SageMaker path needs no MLFLOW_TRACKING_URI; 10-28: Snowflake registers too; loads back with Trainer.load_from_connector(cfg, location=...); prints result.tracking	not run in CI; One load route on every platform
packages/kumo-api/kumoapi/plugin.py	ModelRegistryLocation.alias (exclusive with version); new ModelAliasingSink protocol (set_model_alias); check_model_alias (one alias syntax and the reserved names of every registry); registry_lineage; constants ENGINE_MARKER, BUNDLE_ID_KEY, REGISTRY_KEY_PREFIX; fetch docstrings say what version=None selects	registry aliases and the kumo_bundle_id pointer, defined once for every registry plugin
packages/kumo-api/test/test_plugin_registry.py	tests for check_model_alias (reserved names, accepted names) and registry_lineage	covers the shared alias rule and pointer reading
packages/kumo-core/src/kumo_core/artifacts.py	remove_bundle_file(bundle_dir, role) beside add_bundle_file; bundle_task_type(artifacts) reads the task type from prediction_schema	removing scoring defaults from a bundle; one task-type reader for scoring defaults and warm start
packages/kumo-core/src/kumo_core/automl/auto_trainer.py	emits trial_started (searched params + attempt) and trial_ended (tune metric, seconds, status, error) per attempt through the instrumentation's event hook; passes the trial number to the progress logger; _emit never raises	fit progress events per trial attempt
packages/kumo-core/src/kumo_core/automl/optuna_searcher.py	copies the fallback sampler's sampled_params	Optuna trials report their searched hyperparameters too
packages/kumo-core/src/kumo_core/automl/runtime.py	resolve_warm_start uses bundle_task_type	shared task-type reader
packages/kumo-core/src/kumo_core/automl/searcher.py	Searcher.sampled_params: the fields the search varied and their sampled values (non-scalars as strings)	trial_started params name only searched hyperparameters
packages/kumo-core/src/kumo_core/callbacks/train_progress_logger.py	AutoTrainingProgressLogger(on_event=); emits epoch_ended with the epoch's train/val metrics for the current trial	per-epoch progress events
packages/kumo-core/src/kumo_core/diskgraph/lightning_data.py	_WorkerHandle.open falls back to the sentinel fd when pidfd_open raises ENOSYS/EPERM (other errors still raise); send_signal uses os.kill whenever there is no pidfd	SageMaker hosting refuses pidfd_open, which made every pyfunc prediction fail
packages/kumo-core/src/kumo_core/encoder_selection.py	a text column left unencoded also says when no word-vector model is installed	the message blamed the data when the model was missing
packages/kumo-core/src/kumo_core/instrumentation.py	TrainingInstrumentation.event(name, **fields) (no-op default, also on NoOpTrainingInstrumentation); format_error ("<Type>: <message>", 200 chars)	additive hook that carries events out of kumo-core; one error format for events and tags
packages/kumo-core/src/kumo_core/train/formatting.py	prediction_output_columns(task_type); the long-format task set is a module constant shared with format_prediction	typed MLflow signature outputs that match what predict returns
packages/kumo-core/src/kumo_core/training/runtime.py	local_training_bindings(instrumentation=) instead of a hard-wired NoOpTrainingInstrumentation	lets the engine receive events
packages/kumo-core/tests/test_artifacts.py	test for bundle_task_type	covers the new reader
packages/kumo-core/tests/test_encoder_selection.py	the unencoded-text message with and without an installed model	covers the hint
packages/kumo-core/tests/test_prediction_output_columns.py	new: output columns per task type	covers the signature source
packages/kumo-core/tests/test_searched_params.py	new: only varied fields are reported, as JSON-serialisable values	covers sampled_params
packages/kumo-core/tests/test_worker_handles.py	new: with pidfd_open refused by ENOSYS/EPERM a worker is observed, killed and seen to exit; EBADF still raises	covers the pidfd fallback
packages/kumo-databricks/src/kumo_databricks/registry.py	records the kumo_bundle_id tag; list_models and fetch_model return it in metadata; fetch by alias (GET .../aliases/{alias}); set_model_alias (PUT .../aliases/{alias}); trailing slash stripped from Files API directory entries; shared ENGINE_MARKER; check_model_location: coordinates and Unity Catalog name rules; publish_bundle returns the READY Kumo version that already holds the bundle id instead of creating another, with an info log (_registered_version)	Unity Catalog aliases and pointer; real bundles could not be fetched
packages/kumo-databricks/tests/test_registry.py	bundle id tagged, listed and fetched by version and alias; alias needs a positive version; fake Files API lists directories with a trailing slash; publishing a registered bundle again reuses its version with no new POST, and another bundle id creates one	covers the new registry paths and the download fix
packages/kumo-databricks/tests/test_uc_registry.py	the publish assertion accepts the marker among other tags	publishes now carry kumo_bundle_id too
packages/kumo-engine/docs/mlops.md	new sections: Scoring defaults, Registering models, Promote a Kumo model, Progress events, Experiment tracking (MLflow and Snowflake, tag table, Tracking protocol), MLflow models (registration, scoring, serving, SageMaker), Which registry; tracking/registry in "Turning MLOps on"; GPU SageMaker endpoints need an inference AMI with a CUDA 12.6 driver; on_error, "Which id is which", registry tags per style, the pyfunc temporal example, the mlflow-skinny extra and Databricks Runtime 17.3 LTS ML; registry name rules per registry; SageMaker managed MLflow setup and S3 examples; "Running from a laptop" (one example per platform, explicit-credentials caveat, smoke script); MlflowTracking(target=...) and the construction checks; kumo-engine[sagemaker] and SageMakerMLflow in the SageMaker section; warn-once / "Could not register the model" wording; duplicate bundles reuse their version; pyfunc row order; data["orders.parquet"]; no pyarrow claims; result.export(...) on a TrainingResult and fit(export_config=) as shorthand; Databricks connector reuse of a version holding the same bundle; a Databricks recommendation under Which registry; pyarrow 16 through 24 kept on Databricks 17.3; SageMaker installs kumo-engine[mlflow,sagemaker]; ambient SageMaker needs no MLFLOW_TRACKING_URI, with the other-threads caveat; Running from a laptop: the Databricks scope also sets the tracking URI; new section Artifacts, bundles and dependencies (parents vs dependencies, assets under <root>/assets/, resolution order, registering your own word vectors); 10-28: restructured: quickstart per platform, id table, train/register/promote, trackers, credentials, Reference; 10-29: S3 table example; promotion recipes for the first promotion and a version holding candidate, run through tracking.client(); start a stopped SageMaker server; GloVe release zip with dim=, kumo.TextEmbedding, text stype; role=; torch CUDA note; mlflow-skinny 3.5; Databricks install line; cfg.snapshot_uri; TMPDIR	user docs for everything in this PR; Main path about a third of the old guide; usability tests on all three platforms
packages/kumo-engine/kumo/engine/__init__.py	exports ScoringDefaults; 10-29: exports TextEmbedding, WriteMode	public type for scoring defaults; no kumoapi imports in user code
packages/kumo-engine/kumo/engine/_deps.py	new: Dependencies (require(asset) records an asset and stores it once under <root>/assets/<asset_kind>/<sha256>/ with its _provenance.json; resolve(asset_kind, name) returns a sha256-checked local directory from given locations, the recorded location or the root, cached under ~/.cache/kumo/assets/); Asset protocol, FileAsset, MissingDependencyError, asset_locations; 10-28: asset store failures are collected (Dependencies.errors) for on_error; object stores get no create_dir	what a model needs at predict time besides its bundle and a snapshot, for any model; Live run: assets failed under prefix-scoped S3 rights and only warned
packages/kumo-engine/kumo/engine/_embeddings.py	TextEncoderAsset (word vectors as an asset); require_text_encoders records a fitted model's text encoders; with_text_encoders adds stored vectors to a loaded trainer's store	serving a model with non-default word vectors without registering them again
packages/kumo-engine/kumo/engine/_events.py	new: Event, EventEmitter (fans events out to callbacks, drops a callback that raises, keeps finite metrics only), query_sha256; 10-28: Event.run_id removed	fit progress events; No Kumo-invented run ids
packages/kumo-engine/kumo/engine/_manifest.py	TrackingRef; ModelManifest.tracking (omitted when None, dropped when embedded); ModelParents.derived_from (omitted when None); Dependency and ModelManifest.dependencies (the predictor entry point, the kumo-engine package and its extras, content-addressed assets; omitted when None, so older manifests read as none); AssetManifest (kind=asset, id = sha256); 10-28: TrackingRef(platform, run_id, url)	model → run link in _provenance.json; The tracker's own run, not a Kumo id
packages/kumo-engine/kumo/engine/_mlops.py	MLOps.tracking, MLOps.registry (keyword-only); a ModelRegistryLocation registry is checked at construction (needs a tracker with log_model, no fixed version, valid alias); on_error, "warn" or "raise" (keyword-only, default "warn") and the shared raise helper; a registry model name is checked when kumo.MLOps is built (MLflow tracker or connector registry); MLOps calls the tracker's check(registry) when built; raise_on_error warns once in warn mode (warn_failure) with the untruncated cause (full_error); 10-28: registry takes a ModelRegistryLocation the tracker registers; snapshot_uri; store failures in raise_on_error; 10-29: MLOps.snapshot_uri(snapshot)	tracking and registration through the MLOps switch, failing before any work; One registry style; snapshots by object; one on_error path; the pyfunc's snapshot param without rebuilding the path
packages/kumo-engine/kumo/engine/_predictor.py	new: Predictor protocol, load_predictor(bundle_dir) (the bundle's predictor dependency, default the built-in one; an unimportable one raises a ValueError naming it), TrainerPredictor wrapping today's Trainer scoring	the pyfunc and Snowflake models load any bundle's predictor
packages/kumo-engine/kumo/engine/_scratch.py	scoped_scratch_dir(path) context variable read by get_scratch_dir	each pyfunc predict call works in, and removes, its own scratch directory
packages/kumo-engine/kumo/engine/_tracking.py	new: FitContext, RunRecord, run_record, trial_record (the only place run names, params, metrics and tags are defined), fit_inputs, the Tracking protocol, TrackingSession (rank 0, one background thread, 1000-call queue, 30 s drain, synchronous log_model wait up to 1 h); log_model warns when the tracker was dropped earlier in the fit; every tracking warning names the error; the session keeps the first failure for on_error="raise"; a fit ended by an error gets kumo.error; RegistrationError tells a registration failure from a logging one; tracking tracebacks only at DEBUG; failures reported once (warn=False under on_error="raise"); records are built from a neutral RunSpec (kind, run id, flat params, metrics, steps, user tags, inputs, parents pointer tags, status, error) by run_record(spec), and TrackingSession takes spec callables and log_model(bundle_dir, manifest, snapshot, registry), so nothing here needs a TrainingResult; FitContext moved to trainer/_run_spec.py; 10-28: no kumo.run_id/kumo.parent_run; TrackingSession.pointer(); non-callable log_model is absent	tracker-neutral records; a slow tracker never stalls training; Platform run ids only
packages/kumo-engine/kumo/engine/connector/snowflake.py	SnowflakeConnector.session() returns the reader's Snowpark session	pass the connector's identity to SnowflakeTracking
packages/kumo-engine/kumo/engine/integrations/__init__.py	new package	home of the tracker adapters
packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py	new: KumoModel pyfunc (loads the bundle, caches snapshots, optional serving_snapshot), predict through Trainer.predict, snapshot_location, typed signature; a temporal query without prediction_time/TIMESTAMP, or a binary model without a threshold, fails naming the missing param; predict returns rows in input order (_in_input_order, matching on ENTITY and the anchor time); KumoModel scores through the predictor the bundle's manifest names (load_predictor)	Kumo models as MLflow pyfunc models
packages/kumo-engine/kumo/engine/integrations/_mlflow_model_code.py	new: models-from-code entry (set_model(KumoModel()))	no pickled object in the logged model
packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py	new: SageMakerMLflow; Target (tracking and registry URIs, scope() that sets explicit credentials under one lock and restores them, artifact_repository with the identity's AWS credentials); resolve; Databricks token and SageMaker URI/session scopes; an ambient SageMaker target sets only the process-wide tracking URI around each call; the Databricks explicit-credentials scope also sets the process-wide tracking URI to databricks (Unity Catalog registration downloads the logged model through it); 10-28: resolve(target) takes a URI string; Target.platform/run_url; ambient Databricks sets the tracking URI per call	MLflow has no per-client credentials; Fold tracking_uri=; run links; live run: profile auth needed MLFLOW_TRACKING_URI
packages/kumo-engine/kumo/engine/integrations/mlflow.py	new: MlflowTracking (fit run nested under the active run, child run per trial attempt, source tables as native inputs, capped HTTP retries/timeout, autolog suspended, log_model hook that logs and registers with an alias), log_model, load_model; metric_key maps @ to _at_ in metric names, which MLflow rejects; fit runs default to fit-<UTC start>; registered versions get a kumo_bundle_id tag and PublishedModel.metadata; failures name their cause; check_model_name applies the SageMaker or Unity Catalog name rule from the registry URI; on SageMaker a fit logs how to open the UI (presigned URL or Studio); target= / tracking_uri=; one MlflowClient per tracker, no fluent calls (_log_model = pyfunc.save_model + logged-model client calls; _register = client registration that reuses a version already holding the bundle); check(); own system-metrics sampler; check_registry_schema (Unity Catalog schema GET); UI hint as a UserWarning; log_model/load_model(tracking=); no construction-time MLFLOW_TRACKING_URI requirement for ambient SageMaker; the tracker's log_model and _log_model take (bundle_dir, manifest, snapshot) instead of a TrainingResult; a bundle without a training spec is logged without a signature; the pip requirement reads the bundle's listed encoders; 10-28: no tracking_uri=/serving_snapshot on MlflowTracking; load_model/register; log_model(target=, experiment_name=); UC tags read when search returns a method; 10-29: MlflowTracking.client(): an MlflowClient with the tracker's credentials and tracking URI set for the block	MLflow adapter, model logging and registration; Consistent signatures; one load and re-register route; live run: second UC registration failed; users' own MLflow calls got 403 with an explicit identity
packages/kumo-engine/kumo/engine/integrations/snowflake.py	new: SnowflakeTracking (one run per fit, trial attempts under a prefix, tags as params, tiered limits with drop counts, never resumes a run, leaves an in-use singleton alone, optional native model lineage); connector= opens the session from the connector identity; check() runs SHOW SCHEMAS LIKE when kumo.MLOps is built; 10-28: SnowflakeTracking(source, ..., schema_name=), connector, load_model, random run-name suffix	Snowflake Experiment Tracking adapter; Registry follows the tracker; no Kumo ids
packages/kumo-engine/kumo/engine/tracking.py	new public module exporting Event, Tracking, RunRecord	stable import path for custom adapters
packages/kumo-engine/kumo/engine/trainer/__init__.py	exports ScoringDefaults	public type
packages/kumo-engine/kumo/engine/trainer/_provenance.py	write_provenance(tracked_run_id=) writes tracking.run_id; record_derivation writes a derived bundle's manifest with parents.derived_from; an embedded model's derived_from is cut to a stub like warm_start; writes the model's dependencies; a derived bundle keeps its parent's; 10-28: manifest tracking is the tracker's TrackingRef	run link and derived-bundle lineage; Platform run ids
packages/kumo-engine/kumo/engine/trainer/_run_spec.py	new: FitContext, fit_spec, fitted_spec: a fit as a RunSpec; 10-28: FitContext.run_id removed	the Trainer is one client of the neutral tracking spec; No Kumo run ids
packages/kumo-engine/kumo/engine/trainer/_runtime.py	forwards instrumentation to kumo-core's bindings	events reach the engine
packages/kumo-engine/kumo/engine/trainer/_scoring.py	new: ScoringDefaults, validation and normalisation against the task, scoring-defaults.json write/read (unreadable or invalid files ignored with a warning)	scoring defaults carried in the bundle
packages/kumo-engine/kumo/engine/trainer/_spec.py	SCORING_DEFAULTS_ROLE, excluded from the scoring artifact paths; SCORING_DEFAULTS_ROLE excluded beside #50's holdout roles	the new bundle role is not a scoring input
packages/kumo-engine/kumo/engine/trainer/output_config.py	check_model_location runs a connector registry's name and alias checks; kumo.MLOps and fit call it	Registry names fail when kumo.MLOps is built, not after training
packages/kumo-engine/kumo/engine/trainer/trainer.py	fit(callbacks=) with events and a TrackingSession; MLflow model logging then connector registration after a fit; TrainingResult.registered_model, .scoring_defaults, .with_scoring_defaults; Trainer.export; one _publish that records kumo_bundle_id and sets an alias best-effort; fetch by alias with a bundle-id mismatch warning; predict falls back to bundle defaults; connector registration and alias failures name their cause, end the fit run FAILED with kumo.error, and raise with on_error="raise"; registration warnings via warn_failure, skipped under on_error="raise"; _predict(extra_fields=) for the pyfunc anchor time; TrainingResult.export(ModelOutputConfig); a loaded bundle keeps its holdout alongside its scoring defaults (#50); _store_and_register(bundle_dir, manifest, settings, ...) stores, logs and registers any bundle and runs after fit_ended; _publish takes the bundle dir, id and metrics and reads a training spec only if the bundle has one; fit records the predictor, package and text-encoder dependencies (storing word vectors under the root with MLOps on); Trainer.load resolves stored text-encoder assets (recorded location, then the enclosing block's root) before falling back to the store; 10-28: fit(scoring_defaults=); TrainingResult.tracking; export(cfg); registry through the tracker; unstored models never registered	wires events, tracking, registration and scoring defaults into the trainer; Defaults travel with the registered model; one re-register route; live run
packages/kumo-engine/pyproject.toml	new extra mlflow (mlflow>=3.13,<4, nvidia-ml-py); the mlflow extra is mlflow-skinny>=3.0.1,<4; new extra sagemaker (kumo-engine[mlflow,s3], sagemaker-mlflow>=0.5.0); 10-29: mlflow extra: mlflow-skinny>=3.5	MLflow adapter dependency; 3.13 is the first MLflow that allows DiskGraph's pyarrow 24; loading by alias uses download_artifacts(registry_uri=) (3.5); Databricks 17.3 ships 3.0.1
packages/kumo-engine/test/_snowflake_experiment_fake.py	new: recording stand-in for snowflake.ml.experiment.ExperimentTracking with its observable rules	Snowflake tests without an account
packages/kumo-engine/test/test_bundle_provenance.py	the stub fit stores through _store_and_register	storing moved out of _finish_fit
packages/kumo-engine/test/test_embedding_models.py	imports TextEmbedding from kumo.engine; the new hint in the expected message	covers the exports and the hint
packages/kumo-engine/test/test_fit_events.py	new: event order and contents on a real fit, a raising callback, failed fit and trial, non-finite metrics; 10-28: events carry no run id	covers progress events; No Kumo run ids
packages/kumo-engine/test/test_mlflow_model.py	new: fit logs, registers and serves the model; scoring equals Trainer.predict; wait bound and upload timeout; a tracker dropped earlier in the fit warns that the model was not logged; [text] requirement; invalid MLflow registries refused; on_error warn and raise with a registration failure, a tracking-start failure, the version tag and metadata, the pyfunc error texts; pyfunc input order with shuffled and repeated rows; "Could not register the model" warned once with the full cause, or raised once with no warning; a tracking failure logged without traceback and warned once; calls the spec-based TrackingSession and bundle-based log_model; the pyfunc's trainer is _predictor.trainer; 10-28: load by alias through the settings; export(cfg) returns the held version	covers MLflow models; No duplicate versions
packages/kumo-engine/test/test_mlflow_target.py	new: URI resolution; ambient identities set nothing; Databricks token scope; SageMaker URI, session and artifact credentials; calls take turns; construction checks (SageMaker ambient URI, unreachable server, Unity Catalog schema); registry reuse of a bundle's version; pins of the MLflow, plugin and SDK internals used; the smoke script compiles; ambient SageMaker calls see the ARN with a non-ARN global URI, which is restored after; the Databricks scope sets and restores the global tracking URI; the ambient SageMaker test skips without sagemaker-mlflow; 10-28: URI targets, run URLs, ambient Databricks scope, service-principal auth type, UC tags as a method; 10-29: client() reaches the tracker's server as its identity, and restores the global URI	covers the targets with fakes, no network; Reproduce the live-run failures; covers client()
packages/kumo-engine/test/test_mlflow_tracking.py	new: link-prediction metric names, real fit run tree, failed fit, MLflow errors, autolog, hung server, source tables as native inputs, fit_inputs; fit run names and a user run name winning; name rules through check_model_name; the SageMaker UI hint as a warning; builds records with run_record(fit_spec(CONTEXT)); 10-28: manifest records the MLflow run and result.tracking	covers the MLflow adapter; Platform run ids
packages/kumo-engine/test/test_mlops.py	on_error is validated at construction; 10-28: the tracker registers; ModelOutputConfig registry refused; 10-29: snapshot_uri for a path and an S3 root	Opt-in failure mode; One registry style; covers snapshot_uri
packages/kumo-engine/test/test_model_agnostic_mlops.py	new: a toy run tracked from its spec; a toy bundle stored, logged and registered; a toy predictor through load_predictor and the pyfunc, and an unknown one's error; a FileAsset stored once and resolved by hash, with the missing-dependency error; a manifest without dependencies; a model with non-default word vectors served in a clean subprocess from bundle + snapshot + stored assets only, and its error once the assets are gone; 10-28: prefix-scoped S3 asset store; unstored models not registered; connector-tracker registration	proves each seam takes any model; Reproduce the live-run failures
packages/kumo-engine/test/test_optional_imports.py	mlflow is optional; a bare install's MlflowTracking() / SnowflakeTracking() name their extras	bare install keeps importing
packages/kumo-engine/test/test_prediction_completion.py	stub training results carry scoring_defaults	predict now reads them
packages/kumo-engine/test/test_prediction_workers.py	stub training results carry scoring_defaults	predict now reads them
packages/kumo-engine/test/test_scoring_defaults.py	new: derived bundle thresholds unless explicit, equal ids for equal defaults, stub parent, export and connector load, child bundle id published, validation, ignored files; 10-28: a fit stores its defaults in the bundle	covers scoring defaults; fit(scoring_defaults=)
packages/kumo-engine/test/test_semantic_carry_forward.py	stub training results carry scoring_defaults	predict now reads them
packages/kumo-engine/test/test_snowflake_connector.py	session() returns the reader's pooled session	covers the new accessor
packages/kumo-engine/test/test_snowflake_tracking.py	new: fake matches the installed library, native lineage, errors, in-use instance, location, run names, limits, session untouched, hung call; connector= and the schema check; replays through the spec-based TrackingSession; 10-28: positional connector, schema_name, load through the settings	covers the Snowflake adapter; Consistent signatures
packages/kumo-engine/test/test_tracking_contract.py	new: a recorded fit replayed into every adapter matches the core records; raising, failed, non-zero-rank and hung trackers; replays through the spec-based TrackingSession (fit_spec, fitted_spec); 10-28: runs keyed by name, no Kumo ids	one contract for all adapters; Platform run ids
packages/kumo-engine/test/test_trainer_registry.py	pointer-only publish metadata, export with alias, alias fetch and mismatch warning, aliasing sink required, invalid alias refused, MLOps registry registers after the fit and checks its alias before it; TrainingResult.export publishes and sets the alias like Trainer.export; 10-28: connector registries through a tracker	covers registration through the engine; One registry style
packages/kumo-snowflake/pyproject.toml	snowflake-ml-python>=2.0,<3 in the sdk extra	end_run(status=) needs 2.0; the adapter reads 2.x internals
packages/kumo-snowflake/src/kumo_snowflake/registry.py	records kumo_bundle_id in version metrics; list and fetch return it in metadata; fetch by alias; set_model_alias moves an alias, refuses a version holding another, restores on failure; the carrier model's docstring and refusal name Trainer.load_from_connector (kumo-engine) alongside kumoai_studio.import_model() (hosted service); log_model gets task= (binary, multiclass or regression) from the bundle's prediction schema, unset otherwise; check_model_location: Snowflake identifier rule	Snowflake aliases and pointer
packages/kumo-snowflake/tests/test_registry.py	pointer record, listing metadata, alias set/move/refusals, fetch by alias, failed move restored; the task passed for binary classification and left unset for link prediction	covers the Snowflake registry paths
setup.py	mlflow>=3.13,<4 in the root test extra; the test requirement notes it is a superset of the mlflow-skinny extra	engine wheel tests run with --no-deps in the root environment
uv.lock	adds mlflow 3.17.1 and its dependencies, the mlflow extra, and the snowflake-ml-python>=2.0,<3 specifier; no existing locked version changes; the mlflow extra resolves to mlflow-skinny; the sagemaker extra and sagemaker-mlflow 0.5.0; 10-29: mlflow-skinny specifier >=3.5,<4	lockfile for the new extras; follows pyproject
packages/kumo-databricks/src/kumo_databricks/session.py	10-28: service-principal clients set auth_type="oauth-m2m"	Live run: the token lent to MLflow made a volume store fail with "oauth and pat"
packages/kumo-engine/kumo/engine/_snapshot_io.py	10-28: store raises; the caller applies on_error	Store failures were warn-only
packages/kumo-engine/kumo/engine/graph.py	10-28: snapshot store failure goes through on_error	One on_error path
packages/kumo-engine/kumo/engine/pquery/training_table.py	10-28: training-table store failure goes through on_error	One on_error path
packages/kumo-engine/test/test_snapshot_export.py	10-28: a failed store follows on_error	One on_error path
packages/kumo-engine/kumo/engine/_embeddings.py	10-29: no_encoder_inference_warnings(overridden) drops inference warnings for the given columns only	fit warned about columns the plan overrides
packages/kumo-engine/kumo/engine/trainer/trainer.py	10-29: fit builds the task dataset under no_encoder_inference_warnings(encoder_overrides)	same
packages/kumo-engine/test/test_embedding_models.py	10-29: the fixture's fit overrides the warned column; that fit logs nothing	covers the override case

Requirement: a scheduled fit must be visible in the team's experiment tracker, its model registered under a movable alias with a pointer back to its bundle, its post-fit threshold carried by the model, and the model scorable from MLflow without the training code. Existing behaviour: at #50, fit() emits only progress-bar text, kumo.MLOps has only artifact_root and export, registry publishes record only a marker with no alias, bundles cannot hold scoring arguments, and nothing logs a model to MLflow. Smallest change: an additive kumo-core event hook; tracker-neutral run records with thin MLflow and Snowflake adapters on one background thread; tracking= and registry= on kumo.MLOps; an alias field, an optional aliasing protocol and a kumo_bundle_id pointer in the registry plugins; a derived bundle file for scoring defaults; a code-logged pyfunc model; and the pidfd fallback that SageMaker hosting needs. Verification: the tests listed in the PR body, run one file at a time at this head, plus real CPU fits on synthetic data tracked in a local MLflow store and the Snowflake fake, registered and scored from models:/...@champion. Stopping condition: every fit inside an MLOps block with a tracker yields one run with the documented names and a kumo.bundle_id, its model is registered and aliased (or a warning says why not), and every call without MLOps, callbacks or defaults behaves as at #50.

[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: 50dda4a
• Latest commit, not reviewed: 339b306

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

No actionable comments were generated in the recent review.

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/kumo-engine/.coderabbit.yaml
• Review profile: QUIET
• Plan: Enterprise
• Run ID: 69b8fedb-1ddc-416f-b70d-45bab0cef3e7

📥 Commits

Reviewing files that changed from the base of the PR and between 90ddb38 and 50dda4a.

⛔ Files ignored due to path filters: 1

• uv.lock is excluded by !**/*.lock, !**/*.lock

📒 Files selected for processing: 8

• packages/kumo-core/src/kumo_core/automl/auto_trainer.py
• packages/kumo-engine/README.md
• packages/kumo-engine/kumo/engine/_tracking.py
• packages/kumo-engine/kumo/engine/integrations/mlflow.py
• packages/kumo-engine/kumo/engine/tracking.py
• packages/kumo-engine/kumo/engine/trainer/trainer.py
• packages/kumo-engine/test/test_mlflow_tracking.py
• packages/kumo-engine/test/test_tracking_contract.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.

Walkthrough

The change adds opt-in fit and trial tracking through a public tracking API and an MLflow integration. Trainer.fit starts and finishes tracking and forwards fit events. Trial logger callbacks pass through the training runtime to AutoTrainer, which supplies trial and attempt numbers. MLflow records fit and trial runs, metrics, parameters, and artifacts. The change also adds optional dependency handling, tests, and documentation.

Sequence Diagram(s)

sequenceDiagram
  participant Trainer.fit
  participant TrackingSession
  participant MlflowTracking
  participant AutoTrainer
  participant DiskgraphNestedTrainer
  Trainer.fit->>TrackingSession: start fit tracking
  TrackingSession->>MlflowTracking: create parent run
  AutoTrainer->>TrackingSession: request logger for trial and attempt
  TrackingSession->>MlflowTracking: create nested trial logger
  TrackingSession-->>AutoTrainer: return Lightning logger
  AutoTrainer->>DiskgraphNestedTrainer: pass Lightning logger
  Trainer.fit->>TrackingSession: forward events and finish
  TrackingSession->>MlflowTracking: record events and finish runs

Loading

Priority: Normal

Merge Risk: Low · previous review 50dda

Tracking is opt-in and its failures do not fail training. Some edge cases can still leave MLflow child runs marked RUNNING or disable tracking for a fit. These are bounded issues and worth fixing, so the change is mergeable with owner awareness.

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. Repository search found only the template and review tooling, not the authored Change map. The supplied P…	Provide the complete authored PR conversation comment beginning <!-- kumo-pr-change-map -->, including its Change map and supporting prose, or make the PR conversation available for inspection.

Passed checks: 2 passed

Check name	Status	Explanation
Pr Size		The pull request changes 18 files. This does not exceed the failure threshold of 40 files.
Public Resource Access		All added package sources are publicly accessible. The PR adds the PyPI sources for mlflow and nvidia-ml-py in packages/kumo-engine/pyproject.toml:67-70 and uv.lock; the 86 new `files.pythonho…

Full details: Public Resource Access

Explanation

All added package sources are publicly accessible. The PR adds the PyPI sources for mlflow and nvidia-ml-py in packages/kumo-engine/pyproject.toml:67-70 and uv.lock; the 86 new files.pythonhosted.org artifact URLs tested successfully with HTTP 200. The relevant PyPI indexes also returned HTTP 200. No added repository URL was found, and the README references are illustrative documentation.

Full details: Change Map Design Review

Explanation

The required PR conversation comment beginning <!-- kumo-pr-change-map --> is not available. Repository search found only the template and review tooling, not the authored Change map. The supplied PR description is also explicitly truncated. The code diff is available, but it cannot establish whether the proposed map identified necessity, abstraction-layer placement, existing behavior, focused verification, and a stopping condition. This check therefore cannot assess the proposed design or its explanations.

Resolution

Provide the complete authored PR conversation comment beginning <!-- kumo-pr-change-map -->, including its Change map and supporting prose, or make the PR conversation available for inspection.

✨ 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/kumo/engine/_tracking.py: 1

59-63: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

Do not mark the session as started when context() raises.

start sets self._started = True before it builds the FitContext. The build happens inside the _call lambda, and it can raise. Two examples are a KeyError on train_table_spec['tt_id'] and a TypeError from to_json_dict. When the build raises, _call logs the error and Tracking.start is never called. The session still forwards later events, and Trainer.fit still calls finish. A third-party Tracking then gets on_event and finish without start, which breaks the lifecycle that the Tracking protocol documents. MlflowTracking handles this case only because _reset() clears its state.

Build the context first. Set _started only after the build succeeds. Keep _started = True before self._tracking.start(...), so that finish still runs cleanup when start itself raises.

Proposed fix

     def start(self, context: Callable[[], FitContext]) -> None:
         if self._tracking is None:
             return
-        self._started = True
-        self._call(lambda: self._tracking.start(context()))
+        fit_context = self._call(context)
+        if fit_context is None:
+            return
+        self._started = True
+        self._call(lambda: self._tracking.start(fit_context))

packages/kumo-core/src/kumo_core/train/trainer.py: 1

1443-1445: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

Handle logger values that are lists or booleans.

create_pl_trainer wraps every supplied logger value in a list with the TensorBoard logger. A logger list becomes nested, and logger=False becomes a logger collection containing False. These values do not produce a valid Lightning logger collection. logger=False also no longer disables logging. The fix must preserve that boolean behavior.

Suggested fix

-from collections.abc import Callable
+from collections.abc import Callable, Iterable
...
         extra_logger = self.kwargs.pop("logger", None)
-        if extra_logger is not None:
+        if isinstance(extra_logger, bool):
+            pl_logger = extra_logger
+        elif isinstance(extra_logger, Iterable):
+            pl_logger = [pl_logger, *extra_logger]
+        elif extra_logger is not None:
             pl_logger = [pl_logger, extra_logger]

---

ℹ️ Review info

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/kumo-engine/.coderabbit.yaml
• Review profile: QUIET
• Plan: Enterprise
• Run ID: 7e4f692d-f9c2-4dc1-9c3a-15926fab9fd4

📥 Commits

Reviewing files that changed from the base of the PR and between d45e5ed and 74477db.

⛔ Files ignored due to path filters: 1

• uv.lock is excluded by !**/*.lock, !**/*.lock

📒 Files selected for processing: 15

• packages/kumo-core/src/kumo_core/automl/auto_trainer.py
• packages/kumo-core/src/kumo_core/train/trainer.py
• packages/kumo-core/src/kumo_core/training/orchestration.py
• packages/kumo-core/src/kumo_core/training/runtime.py
• packages/kumo-core/tests/test_training.py
• packages/kumo-engine/kumo/engine/__init__.py
• packages/kumo-engine/kumo/engine/_tracking.py
• packages/kumo-engine/kumo/engine/integrations/__init__.py
• packages/kumo-engine/kumo/engine/integrations/mlflow.py
• packages/kumo-engine/kumo/engine/trainer/_runtime.py
• packages/kumo-engine/kumo/engine/trainer/trainer.py
• packages/kumo-engine/pyproject.toml
• packages/kumo-engine/test/test_mlflow_tracking.py
• packages/kumo-engine/test/test_optional_imports.py
• setup.py

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

[coderabbitai]

Caution

Some comments are outside the diff and can’t be posted inline due to GitHub limitations.

⚠️ Outside diff range comments (1)

🟡 Minor · Handle concurrent experiment creation without disabling tracking. · mlflow.py:153-158

packages/kumo-engine/kumo/engine/integrations/mlflow.py:153-158
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win

Handle concurrent experiment creation without disabling tracking.

When create_experiment raises RESOURCE_ALREADY_EXISTS, perform a fresh get_experiment_by_name lookup and use the returned experiment. Catch only this duplicate-creation error. Re-raise unrelated MLflow failures.

🛡️ Suggested fix

         if self.experiment_name is not None:
             experiment = self._client.get_experiment_by_name(
                 self.experiment_name)
-            experiment_id = (experiment.experiment_id
-                             if experiment else self._client.create_experiment(
-                                 self.experiment_name))
+            if experiment is None:
+                try:
+                    experiment_id = self._client.create_experiment(
+                        self.experiment_name)
+                except Exception as error:
+                    exceptions = import_provider(
+                        'mlflow.exceptions', 'mlflow')
+                    databricks_pb2 = import_provider(
+                        'mlflow.protos.databricks_pb2', 'mlflow')
+                    if not isinstance(error, exceptions.MlflowException) or (
+                            error.error_code !=
+                            databricks_pb2.RESOURCE_ALREADY_EXISTS):
+                        raise
+                    experiment = self._client.get_experiment_by_name(
+                        self.experiment_name)
+                    if experiment is None:
+                        raise
+                    experiment_id = experiment.experiment_id
+            else:
+                experiment_id = experiment.experiment_id

🧹 Nitpick comments: 1

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

378-388: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Cover global-rank gating in the tracking test.

The test covers only LOCAL_RANK=1. It does not detect a regression that removes RANK gating while retaining local-rank gating. Add the RANK=1, LOCAL_RANK=0 case.

🐛 Suggested fix

 def test_only_rank_zero_tracks(monkeypatch):
     monkeypatch.setenv('LOCAL_RANK', '1')
     recorder = _Recorder()
     session = TrackingSession(recorder)
 
     session.start(lambda: CONTEXT)
     session.on_event(EVENTS[0])
     session.finish(RESULT, None)
 
     assert not session.enabled
     assert recorder.calls == []
 
+def test_only_global_rank_zero_tracks(monkeypatch):
+    monkeypatch.setenv('RANK', '1')
+    monkeypatch.setenv('LOCAL_RANK', '0')
+    recorder = _Recorder()
+    session = TrackingSession(recorder)
+
+    session.start(lambda: CONTEXT)
+    session.on_event(EVENTS[0])
+    session.finish(RESULT, None)
+
+    assert not session.enabled
+    assert recorder.calls == []

---

ℹ️ Review info

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/kumo-engine/.coderabbit.yaml
• Review profile: QUIET
• Plan: Enterprise
• Run ID: 413901b2-1304-4d92-8208-5f67666de42c

📥 Commits

Reviewing files that changed from the base of the PR and between c4c4d43 and 5e0624d.

📒 Files selected for processing: 10

• packages/kumo-core/src/kumo_core/automl/auto_trainer.py
• packages/kumo-core/src/kumo_core/train/trainer.py
• packages/kumo-core/tests/test_training.py
• packages/kumo-engine/README.md
• packages/kumo-engine/kumo/engine/__init__.py
• packages/kumo-engine/kumo/engine/_tracking.py
• packages/kumo-engine/kumo/engine/integrations/mlflow.py
• packages/kumo-engine/kumo/engine/trainer/trainer.py
• packages/kumo-engine/test/test_mlflow_tracking.py
• packages/kumo-engine/test/test_tracking_contract.py

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

[coderabbitai]

Actionable comments posted: 1

---

ℹ️ Review info

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/kumo-engine/.coderabbit.yaml
• Review profile: QUIET
• Plan: Enterprise
• Run ID: 4675a596-3919-4607-940f-c95403683072

📥 Commits

Reviewing files that changed from the base of the PR and between 5e0624d and 90ddb38.

📒 Files selected for processing: 5

• packages/kumo-engine/README.md
• packages/kumo-engine/kumo/engine/_tracking.py
• packages/kumo-engine/kumo/engine/integrations/mlflow.py
• packages/kumo-engine/test/test_mlflow_tracking.py
• packages/kumo-engine/test/test_tracking_contract.py

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

[coderabbitai]

🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win

Close child runs after an epoch-log failure.

If _log_steps raises after _start_trial creates a child run, _safe disables tracking. finish then skips child-run termination because it checks _enabled at Lines 143–145. A transient logging failure can leave the child run marked RUNNING after the fit ends. Terminate known child runs during cleanup even when tracking is disabled. MLflow requires explicit termination for runs created through MlflowClient.create_run. (mlflow.org)