Repository navigation
Conversation
Contributor
Author
Change map
|
Vedaant-J
temporarily deployed
to
gha-kumo-workflow-tests
October 20, 2026 22:13 — with 
GitHub Actions
Inactive
GitHub Actions
Inactive
|
Note Reviews pausedIt 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 Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. ℹ️ Recent review info
Walkthrough
Merge Risk:
 |
| 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. |
-
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 mlflowandnvidia-ml-pyinpackages/kumo-engine/pyproject.toml:67-70anduv.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
mlflowandnvidia-ml-pyinpackages/kumo-engine/pyproject.toml:67-70anduv.lock; the 86 newfiles.pythonhosted.orgartifact 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
-
Fix failing CI checks
-
- Commit to this branch
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
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 winDo not mark the session as started when
context()raises.
startsetsself._started = Truebefore it builds theFitContext. The build happens inside the_calllambda, and it can raise. Two examples are aKeyErrorontrain_table_spec['tt_id']and aTypeErrorfromto_json_dict. When the build raises,_calllogs the error andTracking.startis never called. The session still forwards later events, andTrainer.fitstill callsfinish. A third-partyTrackingthen getson_eventandfinishwithoutstart, which breaks the lifecycle that theTrackingprotocol documents.MlflowTrackinghandles this case only because_reset()clears its state.Build the context first. Set
_startedonly after the build succeeds. Keep_started = Truebeforeself._tracking.start(...), so thatfinishstill runs cleanup whenstartitself 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 winHandle
loggervalues that are lists or booleans.
create_pl_trainerwraps every suppliedloggervalue in a list with the TensorBoard logger. A logger list becomes nested, andlogger=Falsebecomes a logger collection containingFalse. These values do not produce a valid Lightning logger collection.logger=Falsealso 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
⛔ Files ignored due to path filters: 1
uv.lockis excluded by!**/*.lock,!**/*.lock
📒 Files selected for processing: 15
packages/kumo-core/src/kumo_core/automl/auto_trainer.pypackages/kumo-core/src/kumo_core/train/trainer.pypackages/kumo-core/src/kumo_core/training/orchestration.pypackages/kumo-core/src/kumo_core/training/runtime.pypackages/kumo-core/tests/test_training.pypackages/kumo-engine/kumo/engine/__init__.pypackages/kumo-engine/kumo/engine/_tracking.pypackages/kumo-engine/kumo/engine/integrations/__init__.pypackages/kumo-engine/kumo/engine/integrations/mlflow.pypackages/kumo-engine/kumo/engine/trainer/_runtime.pypackages/kumo-engine/kumo/engine/trainer/trainer.pypackages/kumo-engine/pyproject.tomlpackages/kumo-engine/test/test_mlflow_tracking.pypackages/kumo-engine/test/test_optional_imports.pysetup.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.
3 tasks done
Vedaant-J
force-pushed
the
vedaant/mlops-mlflow-tracking
branch
from
October 20, 2026 23:29
74477db to
c4c4d43
Compare
Vedaant-J
temporarily deployed
to
gha-kumo-workflow-tests
October 20, 2026 23:29 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
Vedaant-J
temporarily deployed
to
gha-kumo-workflow-tests
October 21, 2026 00:14 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to GitHub limitations.
🟡 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 winHandle concurrent experiment creation without disabling tracking.
When
create_experimentraisesRESOURCE_ALREADY_EXISTS, perform a freshget_experiment_by_namelookup 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 winCover global-rank gating in the tracking test.
The test covers only
LOCAL_RANK=1. It does not detect a regression that removesRANKgating while retaining local-rank gating. Add theRANK=1, LOCAL_RANK=0case.🐛 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
📒 Files selected for processing: 10
packages/kumo-core/src/kumo_core/automl/auto_trainer.pypackages/kumo-core/src/kumo_core/train/trainer.pypackages/kumo-core/tests/test_training.pypackages/kumo-engine/README.mdpackages/kumo-engine/kumo/engine/__init__.pypackages/kumo-engine/kumo/engine/_tracking.pypackages/kumo-engine/kumo/engine/integrations/mlflow.pypackages/kumo-engine/kumo/engine/trainer/trainer.pypackages/kumo-engine/test/test_mlflow_tracking.pypackages/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.
This was referenced Oct 21, 2026
Vedaant-J
temporarily deployed
to
gha-kumo-workflow-tests
October 21, 2026 01:56 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
There was a problem hiding this comment.
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
📒 Files selected for processing: 5
packages/kumo-engine/README.mdpackages/kumo-engine/kumo/engine/_tracking.pypackages/kumo-engine/kumo/engine/integrations/mlflow.pypackages/kumo-engine/test/test_mlflow_tracking.pypackages/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.
2 tasks done
Vedaant-J
force-pushed
the
vedaant/mlops-mlflow-tracking
branch
from
October 21, 2026 08:47
90ddb38 to
50dda4a
Compare
Vedaant-J
temporarily deployed
to
gha-kumo-workflow-tests
October 21, 2026 08:47 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
2 tasks done
Vedaant-J
force-pushed
the
vedaant/mlops-mlflow-tracking
branch
from
October 21, 2026 09:58
50dda4a to
11f23bd
Compare
Vedaant-J
temporarily deployed
to
gha-kumo-workflow-tests
October 21, 2026 09:58 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
Vedaant-J
added a commit
that referenced
this pull request
Oct 21, 2026
Vedaant-J
force-pushed
the
vedaant/mlops-mlflow-tracking
branch
from
October 21, 2026 15:57
11f23bd to
ec7ea6e
Compare
Vedaant-J
temporarily deployed
to
gha-kumo-workflow-tests
October 21, 2026 15:57 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
Vedaant-J
added a commit
that referenced
this pull request
Oct 21, 2026
Vedaant-J
force-pushed
the
vedaant/mlops-mlflow-tracking
branch
from
October 21, 2026 19:08
ec7ea6e to
be04f00
Compare
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…eaches Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… sinks do Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Fits are fit-<model>-<start>, trials trial-<n>-<fit>, a user run_name names the fit fit-<name>, and Snowflake numbers a taken name _2, _3. Trackers report the experiment and run name on TrackingRef. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Adds Note, Table.brief and Scalar.role; MLflow logs a run's tables with one loggedArtifacts update and digests datasets by name, source and schema. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Trackers and registries get the shipped model's holdout (before_refit/ for the pre-refit one); fits drop seconds and the alias tag. Snowflake keeps each trial's best tune metric under TRIAL_<n>/ plus a run's duration_seconds, MLflow times runs from the record, and a Snowflake registry version defaults to the next V<n> and records the holdout. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…r the model register() reuses the model the fit run logged, else logs it with no run in the configured or model-named experiment; load_model resolves an alias to its version first. Callers look up the held version once. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
SageMaker scopes also stop MLflow printing run links, and the registry version error names the tracker rather than MLflow. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… type A single-trial fit logs its epochs on the fit run alone, and searches no longer log a search/best_* series. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Fit params skip unset plan values, and the tune metric uses the logged metric name. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…em notes Data warnings become one Note per problem with a count tag; a manifest's artifacts are a locations table, its datasets are named by kind and date, and the snapshot and training table get schema and size details. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…as given A model whose assets sit on a volume, stage or bucket now loads outside a kumo.mlops block through load_model's target, load_from_connector's connector or the pyfunc's model_config['assets_from']; when none can read them, the error names the argument to pass. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…set loading Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This branch was successfully deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stack: #53 → #112 → #50 → #60 → #64; #114 on #112; #115 on #64. Merge bottom to top.
Summary
This PR adds experiment tracking and model registration for
Trainer.fit. Before it, a fit reported progress only as progress-bar text on stdout and could not record itself anywhere a team already looks, and a chosen threshold had nowhere to live in the bundle.kumo.MLOps(tracking=..., registry=...)records each fit as one run in MLflow (Databricks, SageMaker managed MLflow, any server) or Snowflake Experiment Tracking. The run gets params, per-split metrics, the search trials, and data summaries (split, health and column tables, notes, and the source tables as datasets) read from the stats files the builds already write. The model is then registered under an alias in that tracker's registry.kumo.engine.integrations.mlflow.log_model(trainer, snapshot=..., experiment_name=...)logs one outside a fit, creating the experiment if it is missing.fit(callbacks=)andgraph.snapshot(callbacks=)receive progress events.fit(scoring_defaults=)stores a threshold or top-k in the bundle.Usage
Snowflake uses
SnowflakeTracking(connector, experiment_name=...). A custom tracker implementskumo.engine.tracking.Tracking.Back-compat
kumo.MLOpswithouttrackingandregistryon_error="warn"(default) warns once;"raise"raises after the result is stored.predictwithout a thresholdTARGET_PRED; before, it raised.scoring-defaults.json, nodependencies)scoring-defaults.jsonread by an older enginekumo_bundle_id.mlfloworsnowflakeextraimport kumo.engineworks. Building a tracker raises anImportErrornaming the extra.Testing
test_fit_events.py,test_tracking_contract.py,test_mlflow_tracking.py,test_mlflow_target.py,test_mlflow_model.py,test_snowflake_tracking.py,test_scoring_defaults.py,test_data_summaries.py,test_summary_tables.py,test_model_agnostic_mlops.py;test_trainer_registry.pygrows with registration, alias andregistry=cases.Trainer.predict. Snowflake tests use a recording fake ofsnowflake.ml.experimentwhose signatures are checked against the installed library.mlflow-floorjob that runs the MLflow suites on mlflow-skinny 3.5.0.Also fixes
pidfd_open;encoder_selectioninstall hint; OptunaExperimentalWarningfilter scoped toanalyze().Change map: see the
kumo-pr-change-mapcomment.🤖 Generated with Claude Code