[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: Experiment tracking: #52 → #60 → #68
Merge order: one linear stack; merge bottom to top (#53 first)
This PR: Adds native Snowflake Experiment Tracking for Trainer.fit: kumo.engine.integrations.snowflake.SnowflakeTracking, the Snowflake counterpart of #60's MlflowTracking.

Summary

This PR adds native Snowflake Experiment Tracking for Trainer.fit: kumo.engine.integrations.snowflake.SnowflakeTracking, the Snowflake counterpart of #60's MlflowTracking. It is a thin adapter over #60's run records (run_record / trial_record) and implements the same Tracking protocol, so fits show up in Snowsight under AI & ML » Experiments without MLflow.

Problem. Teams that work only in Snowflake need to answer questions like these:

• Where do we see this week's training runs without standing up MLflow?
• Which fit had the best validation MAE, and with which search space?
• How do two hyperparameter-search trials compare, per epoch?
• Did last night's fit finish or fail, and with what error?
• Which experiment run produced the model version we published to the Snowflake Model Registry? (The run carries the model's kumo.bundle_id, and this PR also links them with Snowflake's native lineage when asked.)

Today none of these can be answered in Snowflake. kumo.MLOps(tracking=...) (#60) only has an MLflow adapter, and Snowflake ML Experiments is not MLflow.

Cause. #60 records a fit as a run tree: one MLflow run per fit and one child run per trial attempt. Snowflake Experiment Tracking (GA since 2026-02-19, snowflake-ml-python) has a different shape. These facts were checked in the docs and in the installed snowflake-ml-python 2.4.0 source (snowflake/ml/experiment/experiment_tracking.py):

• ExperimentTracking(session, *, database_name, schema_name, ...) is a process singleton. A second construction ignores the database and schema it is given, so the location must be passed again to set_experiment.
• There is one active run per process. start_run raises RuntimeError("A run is already active ..."). There are no nested runs, no parent-run field and no tags.
• start_run(name) looks the name up as a SQL LIKE pattern and resumes a matching run that is still running, or raises if it has ended.
• log_params(dict) converts values to strings. log_metrics(dict, step=0) takes one step per call. Both are sent as JSON inside $$…$$ SQL. log_artifact(local_path) takes files only.
• end_run(status='FINISHED'|'FAILED') exists only from 2.0.0. In 1.19 to 1.29 a run cannot be marked FAILED.
• Limits per the docs: 500 experiments per schema, 500 runs per experiment, 1000 unique params and 200 unique metrics per run.
• Run → model link: while a Run is open as a context manager, snowflake-ml patches the registry's ModelManager.log_model so the model version records the experiment and run as a lineage source. ExperimentTracking.list_model_versions(run) and Snowsight lineage read it.
• No system metrics. Built-in training callbacks exist only for LightGBM, Keras, XGBoost and CatBoost, not for Lightning.

Change.

kumo/engine/integrations/
  mlflow.py          # #60
  snowflake.py       # new: SnowflakeTracking
docs/mlops.md        # + Snowflake part of "Experiment tracking"

import kumo.engine as kumo
from kumoapi.model_plan import ModelPlan
from kumo.engine import Trainer
from kumo.engine.integrations.snowflake import SnowflakeTracking  # kumo-engine[snowflake]

trainer = Trainer(ModelPlan())
# `session` is a Snowpark Session; `table` and `snapshot` come from
# PredictiveQuery.generate_training_table(snapshot) and Graph.snapshot()
cfg = kumo.MLOps(artifact_root="s3://bucket/kumo", tracking=SnowflakeTracking(
    session,                       # Snowpark session
    experiment_name="users_spend", # created if missing
    database="ML", schema="EXPERIMENTS",   # default: the session's current ones
    run_name="fit_1",              # a Snowflake identifier; default: Snowflake picks one
    tags={"team": "growth"},
    native_model_lineage=False))   # default; True links a model published during the fit via Snowflake lineage
with kumo.mlops(cfg):              # or fit(..., mlops=cfg)
    result = trainer.fit(table, snapshot=snapshot)

As for MLflow, tracking is on only through MLOps settings; there is no per-call tracking= argument.

One run per fit, with the trials inside it. Snowflake cannot hold a child run open while the fit run is active, so each trial attempt goes inside the fit run. Its record name prefixes everything it records: trial-1-attempt-1/<name>, the same name trial_record gives the MLflow child run. The attempt is part of the prefix, so an out-of-memory retry with a different batch size does not overwrite the first attempt. All names come from the core records; the adapter only adds prefixes.

A new run every fit. An existing run is never resumed. Before start_run, the adapter lists the experiment's runs; if run_name is taken (or matches an existing run as a LIKE pattern, since _ is a wildcard), the fit is recorded as <run_name>_<first 8 hex of the kumo fit run id> with a warning, so a scheduled job that reuses run_name="fit_1" every night gets one run per night instead of an error or a resumed run.

Example

What Snowflake receives from two real CPU fits with two search trials of two epochs each (PREDICT users.spend_next_30d FOR EACH users.user_id on a one-table synthetic 400-row users graph), both with the SnowflakeTracking(...) call above. The fits ran for real; no Snowflake account was used: the Snowflake side is the test fake (test/_snowflake_experiment_fake.py) installed in place of snowflake.ml.experiment, which records each call and payload and applies the library's run, naming and identifier rules (see Testing Done). The tree below is that fake's recorded state. Ids are shortened and values rounded:

ML.EXPERIMENTS
└── experiment USERS_SPEND
    ├── run FIT_1                               FINISHED
    │   ├── params (97)
    │   │   optimization.base_lr=[0.01, 0.001]  optimization.batch_size=[512, 1024]  … (resolved plan, 68 keys)
    │   │   kumo.kind=fit  kumo.run_id=bb59cc1e…-fit-9aab43a2…  team=growth
    │   │   trial-1-attempt-1/optimization.base_lr=0.01  trial-1-attempt-1/optimization.batch_size=1024
    │   │   trial-1-attempt-1/model_architecture.activation=relu  trial-1-attempt-1/attempt=1  … (8 searched + attempt)
    │   │   trial-1-attempt-1/kumo.kind=trial  trial-1-attempt-1/kumo.status=succeeded
    │   │   trial-1-attempt-1/kumo.run_id=bb59cc1e…-fit-9aab43a2…-trial-1-attempt-1
    │   │   trial-1-attempt-1/kumo.parent_run=bb59cc1e…-fit-9aab43a2…
    │   │   trial-2-attempt-1/…  (same keys; base_lr=0.001, activation=elu)
    │   ├── metrics (37 names)
    │   │   trial-1-attempt-1/val_mae      step 1: 1.8510   step 2: 0.9875
    │   │   trial-1-attempt-1/train_loss   step 1: 2.4466   step 2: 1.2356   (+ val/train loss, mse, rmse)
    │   │   trial-1-attempt-1/best_val_mae 0.9875   trial-1-attempt-1/seconds 5.7
    │   │   trial-2-attempt-1/val_mae      step 1: 1.7011   step 2: 1.0750   …
    │   │   trial-2-attempt-1/best_val_mae 1.0750   trial-2-attempt-1/seconds 5.1
    │   │   val_mae 0.9875  test_mae 0.9424  test_refit_mae 0.9021  train_loss 1.2356  … seconds 17.0
    │   └── artifacts  query.pql (51 B)  model_plan.json (3806 B)  (+ _provenance.json at finish on the current head)
    └── run FIT_1_F209E8BA                      FINISHED   (second fit, same run_name; same keys)

Recorded call sequence of the first fit: ExperimentTracking(session) → set_experiment("users_spend", database_name="ML", schema_name="EXPERIMENTS") → show_runs_in_experiment("USERS_SPEND") → start_run("fit_1") (entered as a context only with native_model_lineage=True) → log_params (71) → log_artifact ×2 → for each trial: log_params (12), log_metrics(step=1) (8), log_metrics(step=2) (8), log_metrics (2, summary), log_params (1, kumo.status) → log_metrics (17, final) → end_run(status="FINISHED"). The second fit reuses the process's instance (no construction), finds FIT_1 in show_runs_in_experiment, warns Snowflake run fit_1 already exists in experiment users_spend; recording this fit as fit_1_f209e8ba, and starts fit_1_f209e8ba.

What each part records:

• Fit params are the fit record's params: the resolved model plan flattened to dotted keys. The record's tags are also params under the same names, because Snowflake has no tags: kumo.kind, kumo.run_id (the engine's fit run id) and, at finish, kumo.bundle_id (the model's id). Lineage (snapshot, training table, source versions) is not copied onto the run: it lives in the model's _provenance.json and the manifests, which record the run back as tracking.run_id (Record fits in experiment tracking and publish models to registries #60), and the run only points to the fit and the model. User tags are params under their given names too (core drops a user tag named kumo.* or like a param, with a warning). The full plan is in the model_plan.json artifact and the query in query.pql; the model's _provenance.json is logged too, at finish.
• Trial params are each attempt's searched hyperparameters, attempt, and its kumo.kind, kumo.run_id and kumo.parent_run tags, all under the attempt's prefix. When the attempt ends, its kumo.status (and kumo.error for a failed attempt) follow under the same prefix, as MLflow logs them as child-run tags.
• Per-epoch metrics are each epoch_ended event's metrics (the values the trial record keeps in steps), logged with step = epoch as the events arrive, so Snowsight's metric-history chart plots them per trial. Trial summaries (best_<tune metric>, seconds) use step 0. The final metrics are the fit record's metrics, unprefixed: metrics.json plus seconds.
• Status: the run ends FINISHED, or FAILED when the fit raises; a failed fit also gets the kumo.error param (error type and message). The fit's exception is re-raised unchanged.
• Published model: with native_model_lineage=True the run is held open as a context for the whole fit, and a model version published meanwhile is linked to the run through Snowflake lineage, with no extra call. This is opt-in because it puts the run into the model manifest that the existing publish path writes, which has not been verified live (see Known limitations), and while the context is open snowflake-ml applies it to every model the process registers. At finish the adapter also logs every tag of the finish-time record that was not logged at the start (kumo.bundle_id, and kumo.error on a failed fit) as params with the same names, and any artifact it has not logged yet (_provenance.json).

Limits:

• Per-epoch metrics stop taking new names at 140, trial summaries at 170. The final metrics can use all 200, so they always fit.
• Params stop at 980; the remaining slots are kept for the finish-time tags and drop counts logged at the end.
• Whatever did not fit is counted in the kumo.dropped_metrics and kumo.dropped_params params (names from core; MLflow records kumo.dropped_params the same way). Neither is logged when nothing was dropped.
• Params whose key or value contains $$ are skipped and counted, because they would end the SQL literal.
• Non-finite metric values are not logged.

Edge cases:

• The process already has an ExperimentTracking (snowflake-ml keeps one per process): it is reused, never re-constructed, because re-constructing ignores the new session and resets the user's capture_source_info / warnings_as_errors. If it holds an active run (for example the user's own start_run in a notebook) or belongs to another Snowpark session, the fit is not tracked in Snowflake, a warning says so, and nothing on it is touched: no run is ended or detached, no setting changes. The library has no public "is a run active" or "list runs" call, so the adapter reads the private _instance, _run, _session and _sql_client.show_runs_in_experiment. It fails closed if _run is missing, and (when run_name is set) a missing _sql_client turns tracking off with one warning. A probe test pins these on the installed class.
• Location: database / schema left as None resolve to the session's current database and schema (PUBLIC when only database is given or the session has none). They are passed explicitly every time, so a later fit never inherits the location of an earlier one.
• The user's session is not changed. The adapter issues no ALTER SESSION and sets no session parameter. Snowflake calls run on Record fits in experiment tracking and publish models to registries #60's background tracking thread, so training is not blocked: per-epoch cost on the training thread measured with the recording fake went from 0.08 ms to 0.012 ms, and with a 0.5 s delay per Snowflake call from 500 ms to 0.016 ms. A hung Snowflake call stalls a replayed fit for under 1.5 s (the fit waits at most 30 s at the end for tracking), and the run still ends FINISHED once the call returns. snowflake-ml-python offers no per-statement timeout for its calls, so a stuck statement keeps only the background thread busy.
• Any Snowflake error logs one warning and stops tracking for the rest of the fit. The run is still ended with the fit's status. The fit result is unchanged.
• Rank: only global rank 0 tracks. The check is in Record fits in experiment tracking and publish models to registries #60's session, so this adapter never sees other ranks.
• Not recorded: no system metrics, because Snowflake has no such API and this adapter does not emulate one.

Packaging. kumo-engine[snowflake] already installs kumo-snowflake[sdk], which already depends on snowflake-ml-python for the model registry. This PR sets it to >=2.0,<3 in kumo-snowflake[sdk]: 2.0 for end_run(status=), and below 3 because the adapter reads ExperimentTracking internals (present unchanged from 2.0.0 to 2.6.0, the latest release; checked in both wheels). No new extra. uv.lock changes only that specifier; the locked version stays 2.1.0. snowflake.ml is imported lazily through import_provider, so import kumo.engine does not load it.

How the two backends relate

Both adapters translate the same RunRecords from run_record / trial_record. The parameter, metric and tag names and values are identical; only where they live differs:

Record field	MLflow (#60)	Snowflake (this PR)
fit params	params of the fit run	params of the fit run
fit tags (kumo.kind, kumo.run_id, kumo.bundle_id, kumo.error)	tags	params with the same names
user tags	tags	params with the same names
trial record	child run named trial-1-attempt-1	trial-1-attempt-1/<name> params and metrics in the fit run
trial tags (kumo.kind, kumo.status, kumo.error, …)	child-run tags	trial-1-attempt-1/kumo.* params
trial steps (step = epoch)	child-run metrics with step = epoch, same names	log_metrics(step=epoch)
artifacts (query.pql, model_plan.json, _provenance.json)	log_text	log_artifact of temp files
inputs (source tables)	native MLflow run inputs	not recorded (no API)
published model	MLflow model link with #65	native lineage opt-in (native_model_lineage=True)

#60's contract test (test_tracking_contract.py) now replays the same recorded fake fit into both adapters, reads each one back and compares it with the records. What each adapter records beyond params, metrics and tags (nested runs: MLflow only; per-epoch steps and artifacts: both) is per-adapter test data in that test.

Back-compat

This PR only adds an opt-in adapter; nothing changes for a fit that does not pass it.

Situation	Behaviour
fit() with MLOps off, or with MlflowTracking	Unchanged.
After a fit with MLOps(tracking=SnowflakeTracking(...)) returns	tracking.run_id is that fit run's name (Snowflake identifies runs by name; it includes any _<8 hex> suffix added for a taken run_name), set when the run starts and kept after fit returns, for job summaries.
MLOps(tracking=SnowflakeTracking(session, experiment_name=...))	One new Snowflake run as above. Metric keys and the bundle id are the same as an untracked fit; the model's _provenance.json gains tracking.run_id (#60).
Same run_name as an existing run in the experiment	The existing run is left as it is; this fit is recorded as <run_name>_<8 hex> with a warning.
fit(export_config=...) with SnowflakeTracking(...) on (default)	The registry publish runs exactly as without tracking (no lineage source added).
Same with native_model_lineage=True	The published version also carries the run as a Snowflake lineage source (opt-in; unverified live).
A Snowflake run is already active, or ExperimentTracking was created with another session	Not tracked, with one warning; the instance and its run are left as they are.
The user's Snowpark session	Unchanged: no ALTER SESSION, no parameter set.
Snowflake errors (permissions, network, limits)	One warning, tracking off for the rest of the fit, run ended with the fit's status, fit succeeds.
Failure inside training	Run ended FAILED with kumo.error; the same exception is raised.
kumo-engine without the snowflake extra	import kumo.engine works; SnowflakeTracking(...) raises an ImportError naming kumo-engine[snowflake].
snowflake-ml-python < 2.0 or >= 3 already installed	kumo-snowflake[sdk] now asks for >=2.0,<3; the existing registry publish works on 2.x (registry tests pass on 2.4.0).

Known limitations

• No live Snowflake run. Everything Snowflake-side was checked against a recording fake whose method signatures match the installed snowflake-ml-python 2.4.0; behaviour against a real account (privileges, the undocumented byte limit below, name lookup) has not been exercised.
• Process-wide side effects of the library. With native_model_lineage=True, while the fit's run is open, snowflake-ml links every model version the process registers to that run. A second Snowpark session in the same process is not tracked: snowflake-ml's single ExperimentTracking stays bound to the first one.
• Native lineage is opt-in and unverified live. With native_model_lineage=True, Registry.log_model writes the experiment run into the model manifest as a lineage source. This was not tested against a registry whose role has no privileges on the experiment, or that lives in another database. If Snowflake rejected such a manifest, a fit that both tracks and publishes would fail at publish. That is why it is off by default: the default never touches the publish path.
• Undocumented byte limit. Run metadata also has a byte-size limit that Snowflake does not document. When it is hit, the library warns once instead of raising, and later params or metrics for that run are lost.
• Run-name check is not atomic. The unused-name check and start_run are two calls; a run created with the same name in between, by another process, would be resumed or refused by the library (a refusal turns tracking off with a warning).
• Snowflake records no dataset inputs (the experiment API has none).
• Run names must be Snowflake identifiers (fit_1; unquoted names are upper-cased, so fit-1 is refused and tracking turns off with a warning).
• Concurrent fits in one process are not supported.

Testing Done

CPU only, no Snowflake account. snowflake-ml-python 2.4.0 is installed in the dev venv; the CI environment does not install it, so there the signature and probe checks skip and the other tests use the fake.

• Fake: test/_snowflake_experiment_fake.py reproduces the library's observable rules: a singleton whose re-construction resets settings and keeps the first session, one active run, set_experiment to another experiment or location dropping the run, auto-start on log_*, identifier upper-casing and refusal, start_run(name) resuming or refusing a LIKE-matching run, FINISHED/FAILED only, and a lineage stack while a Run context is open. Its method signatures are checked to equal the installed ExperimentTracking's (7 checks). A probe of the real class with a mock session pins the _instance, _run, _session and _sql_client.show_runs_in_experiment members the adapter reads.
• Contract: test_tracking_contract.py runs the recorded two-trial fake fit (with one out-of-memory attempt) through both MLflow and Snowflake: runs, names, params, metric names and values, (step, value) histories, kumo.* and user tags (including trial kumo.status / kumo.error and the failed fit's kumo.error) and artifacts, including a _provenance.json and the kumo.bundle_id tag that first appear in the finish record, all match the core records, read back from Snowflake's prefixed params under their core names. A tracker that raises is ended once and not called again.
• Native lineage: replaying the recorded fit with native_model_lineage=True puts the run on the lineage stack for the whole fit (where a publish would happen) and takes it off at the end; set_experiment gets the given database and schema.
• Run names: a taken name, and a name that matches a still-running run only through the _ wildcard, each give a new <run_name>_<8 hex> run with a warning, and tracking.run_id is that new name; the existing run gets no params.
• Errors: log_metrics failing gives one warning, and the run is still ended FINISHED. Instance in use: with an active user run or another session, no call reaches the instance, the user's run stays RUNNING and its experiment and warnings_as_errors are untouched. Location: a fit without database after one with it uses the session's current database and schema. Failed fit: run ended FAILED with the kumo.error param. Limits: tier order (final, then summaries, then per-epoch), drop counts, $$ params and NaN metrics skipped, non-string tag values converted, finish-time tags logged at the end. Session: no statement is issued on the user's session; a hung call stalls a replayed fit for under 1.5 s. Bare install: SnowflakeTracking(...) names kumo-engine[snowflake].
• Re-run on the current head (rebased onto Record fits in experiment tracking and publish models to registries #60's MLOps-driven tracking): test_snowflake_tracking.py, test_tracking_contract.py, test_optional_imports.py: 73 passed. uv lock --check passes. kumo-snowflake registry tests were not re-run; the kumo-snowflake change is only the snowflake-ml-python>=2.0,<3 specifier.
• The example above was captured from two real CPU fits on an earlier head (before kumo.MLOps and the kumo.bundle_id tag); Snowflake receives the same calls and names, plus the finish-time kumo.bundle_id param and _provenance.json artifact.

Checklist

• No breaking changes (opt-in).
• Unit tests added.

🤖 Generated with Claude Code

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

[Vedaant-J]

Change map

File	Function / scope	What changed	Why
packages/kumo-engine/docs/mlops.md	"Experiment tracking"	Snowflake example through kumo.MLOps(tracking=SnowflakeTracking(...)); trials inside the fit run, tags as params, limits, instance-in-use and run-name rules, no dataset inputs, session untouched	user docs
packages/kumo-engine/kumo/engine/integrations/snowflake.py	SnowflakeTracking, SnowflakeTracking.run_id, _suffixed, _taken (new module)	Tracking adapter; public run_id = the latest fit run's name. start reuses or creates snowflake-ml's ExperimentTracking (skipped with one warning if it has an active run or another session), calls set_experiment with the location resolved from the session, picks a run name no existing run matches (<run_name>_<8 hex> with a warning when taken; never resumes), starts the fit run (held open as a context across the fit only with native_model_lineage=True), and logs fit params and the record's tags (kumo.* and user tags, same names) as params plus query.pql / model_plan.json. on_trial logs each trial attempt's params and start tags, per-epoch metrics (step = epoch), summaries and end tags (kumo.status, kumo.error) under trial-<n>-attempt-<a>/. finish logs final metrics, every finish-time tag not logged at start (kumo.error), any artifact not yet logged (provenance.json) and the core kumo.dropped_* counts, then end_run(FINISHED/FAILED). Metric tiers 140/170/200, param cap 980+20, $$ and non-finite filters. Issues nothing on the user's Snowpark session; records no dataset inputs	Native Snowflake Experiment Tracking for fit. Names come only from run_record / trial_record. Snowflake has one active run per process, no nesting and no tags, resumes runs by LIKE name, and caps 200 metrics / 1000 params per run. Native model lineage is opt-in because its effect on the publish path is unverified live
packages/kumo-engine/test/_snowflake_experiment_fake.py	FakeExperimentTracking, FakeRun, FakeSession, install (new)	Recording test double (the session records any statement) of snowflake.ml.experiment.ExperimentTracking with the library's singleton, run, LIKE resume, identifier, status, location and lineage rules and _sql_client.show_runs_in_experiment; install swaps it in through sys.modules	Tests run without a Snowflake account or snowflake-ml-python (not installed in CI)
packages/kumo-engine/test/test_optional_imports.py	_BARE_INSTALL	SnowflakeTracking(...) without snowflake-ml raises an ImportError naming kumo-engine[snowflake]	Optional-import guard for the new integration
packages/kumo-engine/test/test_snowflake_tracking.py	new tests	Fake signatures equal the installed class; probe of the real singleton's _instance / _run / _session / _sql_client; two-trial CPU fit with native lineage (call order, location, per-epoch steps, final metrics, run on the lineage stack at publish); errors warn once; instance in use is left alone; location from the session without lineage context; taken and wildcard-matching run names (and run_id = the suffixed name); failed fit with kumo.error; limits; no statement on the user's session; a hung call does not stall the fit	Covers the adapter's behaviour beyond the shared contract
packages/kumo-engine/test/test_tracking_contract.py	snowflake_adapter fixture, adapter params	Snowflake joins the contract replay; observe splits trial-…/ prefixes and reads kumo.* and user-tag params back as tags under their core names, step-0 points included; per-epoch steps and artifacts expected, no nested runs	Proves both adapters record the same RunRecord facts
packages/kumo-snowflake/pyproject.toml	sdk extra	snowflake-ml-python → snowflake-ml-python>=2.0,<3	end_run(status=) exists from 2.0; capped below 3 because the adapter reads ExperimentTracking internals (unchanged 2.0.0–2.6.0); kumo-engine[snowflake] already installs this extra
uv.lock	kumo-snowflake requires-dist	Specifier >=2.0,<3 recorded; locked version unchanged (2.1.0)	Lock consistency

Requirement: Snowflake-only teams see fits in Snowflake Experiment Tracking with the same names and shapes as MLflow, through #60's Tracking protocol; opt-in and best-effort. Existing behaviour: only an MLflow adapter. Smallest change: one adapter module translating the core records, plus a recording fake so CI covers it without an account. Stopping condition: the shared contract test passes for Snowflake; a real fit records one new run per fit with trials, per-epoch metrics, status and errors; nothing outside the adapter changes behaviour.

[coderabbitai]

Warning

Review limit reached

Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available.

Next included review available in 24 minutes.

Check out review usage here.

View limit details

Limit details: You’ve used all 12 included reviews currently available.

Learn how review limits work.

Review configuration:

Run configuration

• Source mode: Auto
• Repository: NVIDIA/kumo-engine
• Branch: vedaant/mlops-mlflow-tracking
• File: .coderabbit.yaml
• Commit: d97b285
• Layers: Repository YAML
• Review profile: QUIET
• Plan: Enterprise
• Run ID: 2758ae9f-4992-413f-aa54-9afd37fa26c3

Commits

Reviewing files that changed from the base of the PR and between 61b7ba6 and ef5ed4b.

Files ignored due to path filters: 1

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

Files selected for processing: 7

• packages/kumo-engine/docs/mlops.md
• packages/kumo-engine/kumo/engine/integrations/snowflake.py
• packages/kumo-engine/test/_snowflake_experiment_fake.py
• packages/kumo-engine/test/test_optional_imports.py
• packages/kumo-engine/test/test_snowflake_tracking.py
• packages/kumo-engine/test/test_tracking_contract.py
• packages/kumo-snowflake/pyproject.toml

Walkthrough

Adds SnowflakeTracking to record each Trainer.fit as one Snowflake run. The adapter logs fit data and trial events, supports optional native model lineage, and handles run-name collisions and Snowflake logging limits. It skips tracking when another active process run or a different session prevents tracking. The change also adds a fake Snowflake API, adapter and contract tests, optional-import coverage, SDK version constraints, and tracking documentation.

Sequence Diagram(s)

sequenceDiagram
  participant Trainer
  participant SnowflakeTracking
  participant SnowflakeExperimentTracking
  Trainer->>SnowflakeTracking: Start fit tracking
  SnowflakeTracking->>SnowflakeExperimentTracking: Select experiment and start run
  Trainer->>SnowflakeTracking: Send trial events
  SnowflakeTracking->>SnowflakeExperimentTracking: Log trial and fit data
  Trainer->>SnowflakeTracking: Finish fit
  SnowflakeTracking->>SnowflakeExperimentTracking: Log final data and end run

Loading

Priority: Normal

Merge Risk: Low · up to 61b7b

Snowflake-tracked fits may stop recording after a run-closing error, and tests can fail with a supported SDK version. Address these before relying on the integration; the tag documentation also needs clarification.

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 inspection found only the template marker, and the supplied PR description is truncated. The d…	Provide the complete author-authored change-map conversation comment, or make the full PR conversation available. Then compare each mapped scope with the changed code and verify necessity, abstraction layer, behavior preservation, and focus…

Passed checks: 2 passed

Check name	Status	Explanation
Pr Size		The pull request changes 8 files. This does not exceed the 40-file limit.
Public Resource Access		No added external URL or repository location is inaccessible. The changed dependency snowflake-ml-python>=2.0,<3 is publicly available from PyPI; anonymous requests to PyPI returned HTTP 200, and ve…

Full details: Public Resource Access

Explanation

No added external URL or repository location is inaccessible. The changed dependency snowflake-ml-python>=2.0,<3 is publicly available from PyPI; anonymous requests to PyPI returned HTTP 200, and versions 2.0.0 through 2.6.0 are available. The s3://bucket/kumo reference is an illustrative example and is excluded by the check.

Full details: Change Map Design Review

Explanation

The required PR conversation comment beginning <!-- kumo-pr-change-map --> is not available. Repository inspection found only the template marker, and the supplied PR description is truncated. The diff and code are available, but the check requires the map’s stated requirement, existing behavior, smallest sufficient change, focused verification, and stopping condition to assess the proposed design.

Resolution

Provide the complete author-authored change-map conversation comment, or make the full PR conversation available. Then compare each mapped scope with the changed code and verify necessity, abstraction layer, behavior preservation, and focused validation.

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.

Caution

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

⚠️ Outside diff range comments (1)

Minor · Release the Snowflake run before resetting adapter state. · snowflake.py:154-162

packages/kumo-engine/kumo/engine/integrations/snowflake.py:154-162
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win

Release the Snowflake run before resetting adapter state.

When end_run raises, finish resets only SnowflakeTracking state. The process-wide singleton can still hold the active run. TrackingSession contains the exception, so the fit can succeed. A later Trainer.fit then sees the active singleton and skips Snowflake tracking.

Ensure cleanup releases the active SDK run before _reset(), and add a regression case that performs a second fit after an end_run failure.

Other comments: 2

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

31-32: Minor · Functional Correctness · Quick win

Limit signature checks to the supported API surface.

The SDK constraint permits Snowflake ML Python 2.0, but FakeExperimentTracking.__init__ includes warnings_as_errors, which Snowflake added in 2.4.0. This exact comparison fails on supported older versions even though the adapter does not pass that option. Compare only the parameters the adapter uses, or make the check version-aware. (docs.snowflake.com)

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

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

Qualify which tags Snowflake records as params.

Line 139 says all tags become params of the same name. That holds for the fit run's record tags, including kumo.* run tags such as kumo.kind and kumo.run_id. It does not hold for user-supplied tags= entries. Those entries that start with kumo. or match a param name are dropped before the Snowflake adapter receives them, with a warning. Document that exclusion here, because the Snowflake section is where users look for the tag mapping.

📝 Proposed documentation fix

-- Snowflake has no child runs or tags, so trial attempts are logged in the fit run as `trial-1-attempt-1/<name>` params and metrics, and tags are params of the same name.
+- Snowflake has no child runs or tags, so trial attempts are logged in the fit run as `trial-1-attempt-1/<name>` params and metrics. The fit run's tags are logged as params of the same name. A `tags=` entry that starts with `kumo.` or matches a param name is not recorded.

---

Review info

Run configuration

• Source mode: Auto
• Repository: NVIDIA/kumo-engine
• Branch: vedaant/mlops-mlflow-tracking
• File: .coderabbit.yaml
• Commit: e625375
• Layers: Repository YAML
• Review profile: QUIET
• Plan: Enterprise
• Run ID: 60582b8c-c009-482a-974e-c3227dcb7e12

Commits

Reviewing files that changed from the base of the PR and between 8dc7f3f and 61b7ba6.

Files ignored due to path filters: 1

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

Files selected for processing: 3

• packages/kumo-engine/docs/mlops.md
• packages/kumo-engine/test/test_snowflake_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; 2 remain after this review.

[Vedaant-J]

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