Feature/prov#8
Conversation
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
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:
📝 WalkthroughWalkthroughThe toolkit now includes MLflow dataset and user registries, dataset verification and environment/provenance Lightning callbacks, expanded lazy public exports, PROV document generation, revised MLflow run resumption, and documentation covering setup, APIs, integrations, and experiment artifacts. ChangesUnified ML Provenance Toolkit
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant LightningTrainer
participant EnvironmentCallback
participant DatasetVerificationCallback
participant ProvenanceCallback
participant MLflow
LightningTrainer->>EnvironmentCallback: invoke on_fit_start
LightningTrainer->>DatasetVerificationCallback: verify manifest and create split
EnvironmentCallback->>MLflow: log environment and hardware metadata
DatasetVerificationCallback->>MLflow: log verification and split metadata
LightningTrainer->>ProvenanceCallback: invoke on_fit_end
ProvenanceCallback->>MLflow: log run_summary.json and prov.json
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (3 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces rationai.mlkit, a unified ML Provenance & Metrics Toolkit featuring automatic PROV-O-aware experiment tracking, dataset and user registration, and integration with PyTorch Lightning and MLflow. Key feedback from the review highlights several robustness issues in the newly added provenance.py and dataset_to_mlflow.py files. Specifically, the reviewer noted that the MLflow run is never ended, which could leave runs active indefinitely. Additionally, there are potential crashes if git is missing or if /proc/self/cgroup access is restricted, a potential KeyError when searching for user runs, dead code in the PROV document builder, and a potential file leak of dataset_provenance.csv if exceptions occur during dataset registration.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 9
🧹 Nitpick comments (2)
demo.py (1)
61-66: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low valueUnclosed file handle when counting manifest rows.
open(manifest)is never closed (Ruff SIM115). Low risk in a short-lived demo script, but easy to fix.♻️ Proposed fix
- n_rows = sum(1 for _ in open(manifest)) - 1 if manifest.exists() else "?" + if manifest.exists(): + with open(manifest) as f: + n_rows = sum(1 for _ in f) - 1 + else: + n_rows = "?"🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@demo.py` around lines 61 - 66, Update the manifest row-counting expression in the directory iteration to open manifest.csv with a context manager, ensuring the file handle is closed after counting rows while preserving the existing n_rows behavior for present and missing manifests.rationai/mlkit/__init__.py (1)
4-7: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winEager provenance import defeats the lazy-loading design.
rationai.mlkit.provenanceimportstorch,mlflow, andpandasat module scope, so this eager import pulls all of them in wheneverrationai.mlkitis imported — undermining the__getattr__mechanism you added specifically to defer heavy submodules (lightning,metrics,data.*). Consider resolvingprovenance_autolog/register_datasetlazily via__getattr__as well.♻️ Move provenance re-exports into
__getattr__from rationai.mlkit.stream import StreamCapture, StreamLogger, StreamModifier -from rationai.mlkit.provenance import ( - autolog as provenance_autolog, - register_dataset, -)if name in ("MetaTiledSlides", "OpenSlideTilesDataset"): import importlib _mod = importlib.import_module("rationai.mlkit.data.datasets") return getattr(_mod, name) + if name in ("provenance_autolog", "register_dataset"): + import importlib + _mod = importlib.import_module("rationai.mlkit.provenance") + return getattr(_mod, "autolog" if name == "provenance_autolog" else name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rationai/mlkit/__init__.py` around lines 4 - 7, Remove the module-scope provenance import from rationai.mlkit and extend its existing __getattr__ lazy-resolution logic to load provenance only when provenance_autolog or register_dataset is accessed. Preserve the current re-export names and behavior while preventing torch, mlflow, and pandas from loading during a plain rationai.mlkit import.</codeવું
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@dataset_to_mlflow.py`:
- Around line 6-48: The standalone register_dataset_as_provenance path diverges
from the Dataset_Registry tagging and SHA-256 manifest contract used by
rationai.mlkit.provenance.register_dataset. Remove this orphaned registration
implementation, or replace its MLflow registration logic with a call to the
canonical register_dataset function so manifest hashing and verification tags
remain compatible with autolog and demo.py.
In `@demo.py`:
- Around line 353-401: The artifact summary lookup in the run-reporting loop
must search nested directories, not only root entries. Update the artifact
handling around client.list_artifacts and the prov_path lookup to reuse the
directory download and recursive glob approach from tests/test_all.py, ensuring
nested provenance/run_summary.json files are found and processed by the existing
provenance and dataset verification output.
In `@dummy_dataset_create.py`:
- Around line 90-125: Update create_dummy_datasets so clean=False preserves
existing datasets: choose the dataset index from the first non-colliding
dummy_dataset_N directory rather than always starting at ds_idx + 1, and append
new rows to any existing manifest.csv instead of overwriting it. Keep clean=True
behavior unchanged, including starting from dummy_dataset_1 and creating a fresh
manifest.
In `@rationai/mlkit/lightning/autolog.py`:
- Around line 60-61: Update the func-is-None branch of autolog to forward
log_hyperparams through the returned partial alongside log_config and
log_stream, preserving the caller’s explicitly provided value in the
parameterized decorator form.
In `@rationai/mlkit/provenance.py`:
- Around line 1252-1263: Update the fail-fast branch in the provenance
verification flow to clean up all tracked temporary directories, including
artifact_dir and split_dir, before raising the RuntimeError. Reuse the existing
_temp_dirs cleanup mechanism and preserve the mlflow.end_run(status='FAILED')
call and verification error details.
- Around line 1275-1381: Update the cleanup flow in the surrounding finally
block to end the active MLflow run after all provenance artifact handling and
temporary-directory cleanup completes. Ensure the run started by
mlflow.start_run() is closed on both success and failure paths, while preserving
the existing cleanup and fail-fast behavior.
In `@README.md`:
- Around line 247-261: Update the README StreamCapture example to match the
actual API: define a small StreamLogger subclass, pass its instance as the first
argument and the stdout/stderr streams tuple to StreamCapture, and demonstrate
output through logger.log_stream(...). Remove the unsupported stream keyword and
get_text()/get_clean_text() calls.
- Around line 76-101: Update the README dataset structure and Options table to
consistently document test_data/ as the default and output directory for
--data-dir, replacing the current data/ references while leaving the other
option descriptions unchanged.
In `@user_to_mlflow.py`:
- Around line 28-37: Replace the real personal data in the __main__ invocation
of register_new_user with clearly fictional placeholder values, and accept
actual registration values through CLI arguments or environment variables
instead of source-code literals. Keep register_new_user’s existing behavior
unchanged while ensuring the example no longer commits names, email addresses,
or other identifying data.
---
Nitpick comments:
In `@demo.py`:
- Around line 61-66: Update the manifest row-counting expression in the
directory iteration to open manifest.csv with a context manager, ensuring the
file handle is closed after counting rows while preserving the existing n_rows
behavior for present and missing manifests.
In `@rationai/mlkit/__init__.py`:
- Around line 4-7: Remove the module-scope provenance import from rationai.mlkit
and extend its existing __getattr__ lazy-resolution logic to load provenance
only when provenance_autolog or register_dataset is accessed. Preserve the
current re-export names and behavior while preventing torch, mlflow, and pandas
from loading during a plain rationai.mlkit import.</codeવું
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 7b24e717-abaa-4fcc-9ee3-b2eab7e6b6d4
📒 Files selected for processing (18)
.gitignoreREADME.mddataset_to_mlflow.pydemo.pydummy_dataset_create.pyrationai/__init__.pyrationai/mlkit/__init__.pyrationai/mlkit/autolog.pyrationai/mlkit/data/__init__.pyrationai/mlkit/lightning/__init__.pyrationai/mlkit/lightning/autolog.pyrationai/mlkit/lightning/loggers/mlflow.pyrationai/mlkit/lightning/with_cli_args.pyrationai/mlkit/provenance.pyrationai/mlkit/stream/__init__.pyrationai/mlkit/with_cli_args.pytests/test_all.pyuser_to_mlflow.py
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
rationai/mlkit/autolog.py (1)
60-61: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
log_hyperparamsis dropped whenautologis used with keyword args.When called as
@autolog(log_hyperparams=False)(func isNone), the returnedpartialomitslog_hyperparams, so it silently reverts to the defaultTrueon the second call. The user's setting is lost.🐛 Forward all keyword options
if func is None: - return partial(autolog, log_config=log_config, log_stream=log_stream) + return partial( + autolog, + log_config=log_config, + log_stream=log_stream, + log_hyperparams=log_hyperparams, + )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rationai/mlkit/autolog.py` around lines 60 - 61, Update the func-is-None branch of autolog to preserve and forward the log_hyperparams option when returning the decorator partial, so `@autolog`(log_hyperparams=False) retains that value on the subsequent call.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@rationai/mlkit/provenance/provenance.py`:
- Around line 82-96: Handle absent git executables in both helpers by broadening
the exception handling in _get_git_info at
rationai/mlkit/provenance/provenance.py lines 82-96 to catch FileNotFoundError
and preserve the existing unknown-value fallback, and in _lookup_user_run at
lines 106-113 to catch FileNotFoundError alongside the existing git failure
exception when reading git config user.name.
- Around line 226-233: The samples hash currently uses host-dependent absolute
paths. In rationai/mlkit/provenance/provenance.py lines 226-233 within
register_dataset, keep full for file hashing or existence checks but store rel
in each sample’s hashed path; apply the identical change in lines 342-349 within
_verify_dataset so both registration and verification hash relative paths
consistently.
---
Outside diff comments:
In `@rationai/mlkit/autolog.py`:
- Around line 60-61: Update the func-is-None branch of autolog to preserve and
forward the log_hyperparams option when returning the decorator partial, so
`@autolog`(log_hyperparams=False) retains that value on the subsequent call.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: da1f6a41-50f4-4b51-bf52-9ece780f8723
📒 Files selected for processing (6)
rationai/mlkit/__init__.pyrationai/mlkit/autolog.pyrationai/mlkit/lightning/__init__.pyrationai/mlkit/provenance/dataset_to_mlflow.pyrationai/mlkit/provenance/provenance.pyrationai/mlkit/provenance/user_to_mlflow.py
💤 Files with no reviewable changes (1)
- rationai/mlkit/lightning/init.py
🚧 Files skipped from review as they are similar to previous changes (1)
- rationai/mlkit/init.py
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
rationai/mlkit/__init__.py (1)
31-35: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winFix lazy routing for the Lightning exports.
rationai.mlkit.lightningonly definesTrainer, so resolvingMLFlowLogger,MultiloaderLifecycle,autolog, orwith_cli_argsthrough this module raisesAttributeError. Route each symbol to its defining module, or re-export all five symbols fromrationai/mlkit/lightning/__init__.py; otherwise these documented top-level imports are broken.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rationai/mlkit/__init__.py` around lines 31 - 35, Update the __getattr__ lazy-routing logic for Trainer, MLFlowLogger, MultiloaderLifecycle, autolog, and with_cli_args so each name resolves from the module that actually defines it, or ensure rationai.mlkit.lightning re-exports all five symbols. Preserve the documented top-level imports without AttributeError.
🧹 Nitpick comments (4)
rationai/mlkit/provenance/__init__.py (1)
1-39: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAddress the reported Ruff findings.
If
ruff checkis a merge gate, replace ambiguous en dashes, use the configured absolute-import style, and sort__all__to avoid lint failures.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rationai/mlkit/provenance/__init__.py` around lines 1 - 39, Update the module docstring and imports in the provenance package to satisfy Ruff: replace ambiguous en dashes with the project’s accepted punctuation, use the configured absolute-import style for register_dataset and register_user, and alphabetize the entries in __all__ while preserving the same public exports.Source: Linters/SAST tools
rationai/mlkit/lightning/callbacks/provenance.py (2)
44-45: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
storage/metaPROV prefixes hardcodehttp://localhost:8083.These URIs are baked into every generated PROV bundle, so documents produced on any machine reference a local-only endpoint and are not portable across environments. Consider sourcing the base URL from config/env with
localhost:8083only as a fallback default.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rationai/mlkit/lightning/callbacks/provenance.py` around lines 44 - 45, Update the PROV prefix configuration for "storage" and "meta" to derive the base URL from the existing application configuration or environment, while retaining http://localhost:8083 as the fallback default. Ensure generated provenance bundles use the configured endpoint consistently for both document and metadata URIs.
680-696: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
log_streamis dead configuration —_StreamCaptureis never instantiated.
self.log_streamis stored (Line 692) and_StreamCapture(Lines 286-330) is fully defined, but neitheron_fit_startnoron_fit_endever wraps execution in it, so stdout/stderr are never captured toconsole.log. Either wire the capture into the fit lifecycle or drop the unused param and class to avoid advertising a feature that does nothing.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rationai/mlkit/lightning/callbacks/provenance.py` around lines 680 - 696, Remove the unused log_stream configuration from the callback constructor and its self.log_stream assignment, and delete the unreferenced _StreamCapture class since the fit lifecycle does not instantiate it. Update any callers or related interfaces that pass log_stream so the API no longer advertises unsupported stream capture.rationai/mlkit/lightning/callbacks/dataset_verification.py (1)
38-59: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReuse the public
verify_datasethelper instead of re-wiring internals.Lines 38-59 duplicate the exact detect-manifest → derive-data_root → lookup-run →
_verify_datasetsequence thatverify_dataset()inregister_dataset.py(Lines 59-89) already encapsulates. Callingverify_dataset(self._manifest_path)removes the reliance on three underscore-prefixed internals and keeps the two code paths from drifting.♻️ Proposed simplification
- # Import here so the callback doesn't require provenance as a hard dep - from rationai.mlkit.provenance.register_dataset import ( - _detect_manifest, - _lookup_dataset_run, - _verify_dataset, - ) - - manifest_path = self._manifest_path - data_root = None - if manifest_path is None: - manifest_path, data_root = _detect_manifest() - - if manifest_path is None: - print(" [DatasetVerificationCallback] No manifest.csv found — skipping") - return - - if data_root is None: - import os - data_root = os.path.dirname(os.path.abspath(manifest_path)) - - dataset_run_id = _lookup_dataset_run() - verification = _verify_dataset(manifest_path, data_root, dataset_run_id) + # Import here so the callback doesn't require provenance as a hard dep + from rationai.mlkit.provenance.register_dataset import verify_dataset + + verification = verify_dataset(self._manifest_path)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rationai/mlkit/lightning/callbacks/dataset_verification.py` around lines 38 - 59, Update the dataset verification callback method containing the shown manifest handling to import and call the public verify_dataset helper with self._manifest_path, replacing the local _detect_manifest, _lookup_dataset_run, and _verify_dataset sequence. Preserve the callback’s existing handling of the returned verification result and avoid relying on underscore-prefixed internals.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@rationai/mlkit/lightning/callbacks/provenance.py`:
- Around line 730-736: Update the manifest resolution flow around manifest_path
and data_root so a user-supplied manifest_path with no data_root derives the
dataset root from that manifest, matching the existing behavior in
verify_dataset. Preserve _detect_manifest() for cases where manifest_path is
absent, and ensure the resulting values allow the existing manifest verification
and split flow guarded by manifest_path and data_root to run.
- Around line 87-88: Update the git metadata fallback in
rationai/mlkit/lightning/callbacks/provenance.py:73-88 to catch
FileNotFoundError alongside subprocess.CalledProcessError so commit, remote, and
branch remain "unknown"; likewise update the user lookup fallback at
rationai/mlkit/lightning/callbacks/provenance.py:92-103 to catch both exceptions
and continue to the existing environment/USER fallback.
In `@rationai/mlkit/provenance/register_dataset.py`:
- Around line 22-39: Update _lookup_dataset_run so manifest_path participates in
selecting the Dataset_Registry run by deriving the expected manifest identifier
and filtering or matching the corresponding MLflow run tag before choosing the
latest match. Preserve the existing None behavior when the experiment or
matching runs are absent; alternatively, remove manifest_path and update its
callers and documentation to explicitly define latest-run selection.
---
Outside diff comments:
In `@rationai/mlkit/__init__.py`:
- Around line 31-35: Update the __getattr__ lazy-routing logic for Trainer,
MLFlowLogger, MultiloaderLifecycle, autolog, and with_cli_args so each name
resolves from the module that actually defines it, or ensure
rationai.mlkit.lightning re-exports all five symbols. Preserve the documented
top-level imports without AttributeError.
---
Nitpick comments:
In `@rationai/mlkit/lightning/callbacks/dataset_verification.py`:
- Around line 38-59: Update the dataset verification callback method containing
the shown manifest handling to import and call the public verify_dataset helper
with self._manifest_path, replacing the local _detect_manifest,
_lookup_dataset_run, and _verify_dataset sequence. Preserve the callback’s
existing handling of the returned verification result and avoid relying on
underscore-prefixed internals.
In `@rationai/mlkit/lightning/callbacks/provenance.py`:
- Around line 44-45: Update the PROV prefix configuration for "storage" and
"meta" to derive the base URL from the existing application configuration or
environment, while retaining http://localhost:8083 as the fallback default.
Ensure generated provenance bundles use the configured endpoint consistently for
both document and metadata URIs.
- Around line 680-696: Remove the unused log_stream configuration from the
callback constructor and its self.log_stream assignment, and delete the
unreferenced _StreamCapture class since the fit lifecycle does not instantiate
it. Update any callers or related interfaces that pass log_stream so the API no
longer advertises unsupported stream capture.
In `@rationai/mlkit/provenance/__init__.py`:
- Around line 1-39: Update the module docstring and imports in the provenance
package to satisfy Ruff: replace ambiguous en dashes with the project’s accepted
punctuation, use the configured absolute-import style for register_dataset and
register_user, and alphabetize the entries in __all__ while preserving the same
public exports.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 6e1a1cdf-17f4-4960-9d26-0ce26dc9bfa3
📒 Files selected for processing (8)
.gitignorerationai/mlkit/__init__.pyrationai/mlkit/lightning/callbacks/__init__.pyrationai/mlkit/lightning/callbacks/dataset_verification.pyrationai/mlkit/lightning/callbacks/provenance.pyrationai/mlkit/provenance/__init__.pyrationai/mlkit/provenance/register_dataset.pyrationai/mlkit/provenance/register_user.py
🚧 Files skipped from review as they are similar to previous changes (1)
- .gitignore
- Added EnvironmentCallback to capture hardware, docker, git, user, and environment snapshot at training start. - Refactored ProvenanceCallback to depend on EnvironmentCallback for environment data, enhancing modularity. - Removed legacy CSV-based dataset registration from register_dataset.py for cleaner codebase. - Introduced load_manifest function to streamline manifest loading across multiple modules. - Updated documentation and logging for better clarity and error handling.
There was a problem hiding this comment.
Actionable comments posted: 6
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
rationai/mlkit/lightning/callbacks/dataset_verification.py (1)
100-120: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
fail_fastabort is gated behindmlflow.active_run(), silently disabling the safety check when no run is active.The
RuntimeErrorguarding against a verified-mismatch dataset only fires insideif mlflow.active_run():. If no MLflow run happens to be active at this point, verification still runs but a failed result is neither logged nor raised —fail_fast=True(the default) is silently bypassed. Data-integrity enforcement shouldn't depend on whether logging infrastructure happens to be active.🐛 Proposed fix
if mlflow.active_run(): mlflow.log_params({ "dataset_verified": verification["verified"], "dataset_file_sizes_match": verification["file_sizes_match"] is True, "dataset_files_missing": verification["files_missing"], "dataset_files_total": verification["files_total"], }) if verification["verified"]: mlflow.set_tag("dataset_verification", "VERIFIED") else: mlflow.set_tag("dataset_verification", "MISMATCH") mlflow.set_tag( "dataset_verification_details", "; ".join(verification.get("details", [])), ) - if self.fail_fast and not verification["verified"]: - raise RuntimeError( - "Dataset verification failed — aborting training.\n" - + "\n".join(f" {d}" for d in verification.get("details", [])) - ) + if self.fail_fast and not verification["verified"]: + raise RuntimeError( + "Dataset verification failed — aborting training.\n" + + "\n".join(f" {d}" for d in verification.get("details", [])) + )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rationai/mlkit/lightning/callbacks/dataset_verification.py` around lines 100 - 120, Move the fail-fast check from inside the mlflow.active_run() block to the surrounding verification flow so it executes regardless of MLflow run state. Preserve the existing logging and tagging under mlflow.active_run(), while ensuring self.fail_fast with an unverified verification result always raises the existing RuntimeError.rationai/mlkit/lightning/callbacks/provenance.py (1)
557-729: 📐 Maintainability & Code Quality | 🟠 Major | 🏗️ Heavy lift
_fallback_on_fit_startduplicates ~150 lines already implemented inEnvironmentCallback/DatasetVerificationCallbackinstead of delegating to them.Git-info retrieval, tag construction, hardware/docker detection, and environment snapshotting are all re-implemented here nearly verbatim from
EnvironmentCallback.on_fit_start, and dataset verification/split logic is re-implemented fromDatasetVerificationCallback.on_fit_start. This isn't just a style concern: it has already caused a real drift bug — thedata_rootderivation fix landed inDatasetVerificationCallbackbut not here (see the Line 603-609 comment), precisely because the logic lives in two places that must be kept in sync by hand.Consider having the "no siblings" fallback instantiate
EnvironmentCallback()/DatasetVerificationCallback()internally and delegate to theiron_fit_start, then read the results back the same way_gather_from_siblingsdoes — eliminating the duplicate implementation entirely.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rationai/mlkit/lightning/callbacks/provenance.py` around lines 557 - 729, The _fallback_on_fit_start method duplicates EnvironmentCallback and DatasetVerificationCallback behavior, allowing fixes to drift between implementations. Replace its inline Git, environment, dataset verification, split, and logging logic with internally created EnvironmentCallback and DatasetVerificationCallback instances, delegate through their on_fit_start methods, and collect their outputs using the same result handling as _gather_from_siblings; preserve the no-sibling fallback behavior and existing callback configuration.
♻️ Duplicate comments (1)
rationai/mlkit/lightning/callbacks/provenance.py (1)
603-609: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
data_rootstill not derived from a user-suppliedmanifest_path— unresolved from prior review.When
self.manifest_pathis set butself.data_rootisNone,_detect_manifest()is skipped (manifest_path is not None), sodata_rootstaysNone. Theif manifest_path and data_root:guard then fails and the code falls into the misleading "No manifest.csv found" warning, silently skipping verification, split, and split artifacts — exactly the bug flagged on a previous commit of this file.DatasetVerificationCallback.on_fit_start(in the sibling file) already carries the fix; this fallback path still doesn't.🐛 Proposed fix
if manifest_path is None: manifest_path, data_root = _detect_manifest() + elif data_root is None: + data_root = os.path.dirname(os.path.abspath(manifest_path))🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rationai/mlkit/lightning/callbacks/provenance.py` around lines 603 - 609, Update the manifest resolution logic in the callback method around manifest_path and data_root so a user-supplied manifest_path with no self.data_root derives data_root from that manifest path before the availability guard. Reuse the same resolution behavior as DatasetVerificationCallback.on_fit_start, then preserve the existing verification, split, and artifact flow when both values are resolved.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@rationai/mlkit/lightning/callbacks/environment.py`:
- Around line 44-53: Update `_lookup_user_run` to handle a missing `git`
executable by suppressing both `subprocess.CalledProcessError` and
`FileNotFoundError` around the `subprocess.check_output` call. Add the
`contextlib` import and use `contextlib.suppress` instead of the current
try/except/pass, preserving the existing `USER` fallback when lookup fails.
In `@rationai/mlkit/lightning/callbacks/provenance.py`:
- Around line 591-592: Wrap the `_lookup_user_run()` call in the fallback path
with `try/except Exception`, matching `EnvironmentCallback.on_fit_start`. On any
lookup failure, preserve the `strict` behavior by treating the user run ID and
tags as unavailable and allowing dataset verification and split logging to
continue.
- Around line 733-754: Update on_fit_start to handle has_env and has_verify
independently instead of routing all partial sibling configurations through
_gather_from_siblings. Ensure the missing environment or dataset-verification
work, including split-data population, is performed via the appropriate fallback
logic when only one sibling callback is present; preserve sibling-derived data
for callbacks that exist and avoid silently omitting it.
- Around line 61-78: Update _get_prov_prefixes to validate that the parsed
PROV_BASE_URI JSON is a mapping/object before merging it with
_DEFAULT_PROV_PREFIXES. Treat valid non-object JSON values like malformed JSON
by logging the existing warning and returning the defaults, while preserving the
override and valid-object merge behavior.
In `@rationai/mlkit/lightning/loggers/mlflow.py`:
- Around line 60-76: Update the uninitialized branch of the experiment property
to call the valid MlflowClient termination API, replacing set_terminated_status
with set_terminated using the existing run ID, status, and timestamp. Then
restore the fluent active run with mlflow.start_run(run_id=self.run_id) or an
equivalent supported mechanism before setting _active_run_id and _initialized,
so callbacks observe the run as active.
In `@rationai/mlkit/provenance/__init__.py`:
- Line 5: Replace the Unicode en dash in the register_dataset docstring with an
ASCII hyphen, preserving the existing wording and meaning so Ruff RUF002 passes.
---
Outside diff comments:
In `@rationai/mlkit/lightning/callbacks/dataset_verification.py`:
- Around line 100-120: Move the fail-fast check from inside the
mlflow.active_run() block to the surrounding verification flow so it executes
regardless of MLflow run state. Preserve the existing logging and tagging under
mlflow.active_run(), while ensuring self.fail_fast with an unverified
verification result always raises the existing RuntimeError.
In `@rationai/mlkit/lightning/callbacks/provenance.py`:
- Around line 557-729: The _fallback_on_fit_start method duplicates
EnvironmentCallback and DatasetVerificationCallback behavior, allowing fixes to
drift between implementations. Replace its inline Git, environment, dataset
verification, split, and logging logic with internally created
EnvironmentCallback and DatasetVerificationCallback instances, delegate through
their on_fit_start methods, and collect their outputs using the same result
handling as _gather_from_siblings; preserve the no-sibling fallback behavior and
existing callback configuration.
---
Duplicate comments:
In `@rationai/mlkit/lightning/callbacks/provenance.py`:
- Around line 603-609: Update the manifest resolution logic in the callback
method around manifest_path and data_root so a user-supplied manifest_path with
no self.data_root derives data_root from that manifest path before the
availability guard. Reuse the same resolution behavior as
DatasetVerificationCallback.on_fit_start, then preserve the existing
verification, split, and artifact flow when both values are resolved.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: c28cde80-0493-410d-940f-1c9d5d92a6f8
📒 Files selected for processing (7)
rationai/mlkit/lightning/callbacks/__init__.pyrationai/mlkit/lightning/callbacks/dataset_verification.pyrationai/mlkit/lightning/callbacks/environment.pyrationai/mlkit/lightning/callbacks/provenance.pyrationai/mlkit/lightning/loggers/mlflow.pyrationai/mlkit/provenance/__init__.pyrationai/mlkit/provenance/register_dataset.py
🚧 Files skipped from review as they are similar to previous changes (1)
- rationai/mlkit/lightning/callbacks/init.py
|
/gemini summary |
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
|
/gemini review |
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
|
@coderabbitai review |
✅ Action performedReview finished.
|
There was a problem hiding this comment.
Actionable comments posted: 8
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
README.md (2)
371-371: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winFix the
prov_mlflowdocumentation link.The referenced repository currently returns HTTP 404, so readers cannot follow the compatibility reference. Update it to the canonical repository or remove the link until it is available. ()
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@README.md` at line 371, Update the prov_mlflow Markdown link in the README to point to the canonical, accessible repository; if no valid canonical repository is available, remove the hyperlink while retaining the tool reference text.Source: MCP tools
20-24: 🔒 Security & Privacy | 🟠 Major | ⚡ Quick winBind MLflow to localhost by default.
0.0.0.0exposes the UI on every interface; for the local setup here, use127.0.0.1or document the auth/firewall boundary before exposing provenance data.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@README.md` around lines 20 - 24, Update the README’s MLflow startup command to bind the server to 127.0.0.1 instead of 0.0.0.0, preserving the documented local URL and avoiding exposure on all network interfaces.rationai/mlkit/provenance/register_dataset.py (1)
242-301: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winRun/temp-dir cleanup isn't exception-safe.
register_dataset()manages the run manually (mlflow.start_run()...mlflow.end_run()) with notry/finally. The new PROV-O and legacy artifact steps (Lines 260-297) add several more failure-prone calls (build_dataset_prov, file writes, twomlflow.log_artifactcalls) inside this unprotected window. If any of them raises, the run is leftRUNNINGforever in MLflow and the temp dirs (prov_dir,legacy_prov_dir) are never cleaned up.🛡️ Suggested fix
run = mlflow.start_run(run_name=f"Dataset_{dataset_name}_{version}") run_id = run.info.run_id + prov_dir = None + legacy_prov_dir = None + try: mlflow.log_params({...}) mlflow.set_tags({...}) ... prov_dir = f"_dataset_prov_{uuid.uuid4().hex[:8]}" os.makedirs(prov_dir, exist_ok=True) prov_path = os.path.join(prov_dir, "prov.json") with open(prov_path, "w") as f: json.dump(prov_doc, f, indent=2) mlflow.log_artifact(prov_path, artifact_path="provenance") - shutil.rmtree(prov_dir, ignore_errors=True) legacy_prov_dir = f"_dataset_legacy_{uuid.uuid4().hex[:8]}" ... mlflow.log_artifact(legacy_path, artifact_path="provenance") - shutil.rmtree(legacy_prov_dir, ignore_errors=True) - - mlflow.end_run() + finally: + if prov_dir: + shutil.rmtree(prov_dir, ignore_errors=True) + if legacy_prov_dir: + shutil.rmtree(legacy_prov_dir, ignore_errors=True) + status = "FINISHED" if <no exception> else "FAILED" # or use except/else + mlflow.end_run(status=status)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rationai/mlkit/provenance/register_dataset.py` around lines 242 - 301, Make register_dataset exception-safe by wrapping the MLflow run and both provenance artifact workflows in try/finally, ensuring every created temporary directory is removed with cleanup that tolerates missing paths and mlflow.end_run() is attempted whenever start_run() succeeds. Preserve the existing provenance generation, artifact logging, return value, and success output while preventing failures in build_dataset_prov, file writes, or log_artifact from leaving resources active.
🧹 Nitpick comments (3)
README.md (1)
307-307: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winSpecify the language for the project-tree fence.
Markdownlint reports MD040 here. Change the opening fence to
```text.Proposed fix
-``` +```text🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@README.md` at line 307, Update the project-tree code fence in the README to specify the text language by changing its opening delimiter to ```text, resolving the MD040 lint violation while leaving the fenced content unchanged.Source: Linters/SAST tools
rationai/mlkit/provenance/register_dataset.py (1)
275-281: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRepeated "temp-dir → write JSON → log_artifact → rmtree" boilerplate across three sites. Extract a shared helper (e.g. in
prov.py) that takes the doc/dict, filename, and MLflow artifact path, and internally handles the temp dir + cleanup (ideally with atry/finallyso cleanup and run-status happen even on failure).
rationai/mlkit/provenance/register_dataset.py#L275-L281: replace the prov.json temp-dir/write/log/rmtree block with a call to the shared helper.rationai/mlkit/provenance/register_dataset.py#L284-L297: replace the legacy dataset_provenance.json block with the same helper call.rationai/mlkit/provenance/register_user.py#L64-L71: replace the user prov.json block with the same helper call.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rationai/mlkit/provenance/register_dataset.py` around lines 275 - 281, Extract the repeated temporary-directory, JSON-writing, MLflow artifact logging, and cleanup flow into a shared helper in prov.py that accepts the document, filename, and artifact path and guarantees cleanup via try/finally. Replace the blocks at rationai/mlkit/provenance/register_dataset.py lines 275-281 and 284-297, and rationai/mlkit/provenance/register_user.py lines 64-71, with calls to this helper while preserving each existing filename, document, and MLflow artifact path.rationai/mlkit/provenance/prov.py (1)
65-66: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the unused
_typed_valuetype parameters. This helper always serializes plain strings, and all current call sites use the default string case. If typed PROV literals are needed here, encodetype_prefix:type_localin the value shape; otherwise drop the misleading parameters.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rationai/mlkit/provenance/prov.py` around lines 65 - 66, Update _typed_value to remove the unused type_prefix and type_local parameters, retaining its existing plain-string serialization behavior. Verify its current call sites and adjust any invocations that rely on the removed arguments; do not add typed-literal handling unless the value representation explicitly requires it.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In @.gitignore:
- Line 177: Update the `example*` entry in `.gitignore` to scope matching to the
intended repository-root directory, such as `/example*/`, rather than ignoring
matching files or directories throughout the tree.
In `@rationai/mlkit/provenance/__init__.py`:
- Around line 4-6: Replace the Unicode en dashes in the module docstring entries
for prov, register_dataset, and register_user with Ruff-compliant ASCII
punctuation, preserving the documented descriptions and meanings.
In `@rationai/mlkit/provenance/prov.py`:
- Line 69: Update _qualified_name, build_user_prov, build_dataset_prov, and
every flagged local annotation in prov.py to use explicit generic type
parameters for all dict and list types, including nested containers. Preserve
the existing data shapes while replacing bare dict/list annotations and return
types with accurately typed equivalents so the file passes strict MyPy checking.
- Around line 73-78: Update _iso_timestamp to use datetime.UTC instead of
timezone.utc for both timestamp conversion and current-time creation, and remove
the now-unused timezone import while preserving the existing UTC ISO formatting.
In `@rationai/mlkit/provenance/register_user.py`:
- Around line 42-49: Remove all real names and email addresses from the __main__
example and replace them with non-identifying placeholders or externally
supplied values. Update register_new_user so MLflow tags do not persist
plaintext name or email fields, using anonymized values or an explicit opt-in
mechanism. Remove or sanitize the trailing print output so it cannot expose the
real name or username.
- Around line 30-39: Update the experiment setup in the user registration flow
to handle concurrent creation of "User_Registry": when no experiment is found,
call create_experiment and catch the already-exists error, retaining the
resulting or existing experiment ID without performing a second
get_experiment_by_name lookup. Pass that reused ID to mlflow.start_run.
In `@README.md`:
- Around line 65-72: Replace the real person names, organization, and email
addresses in the register_new_user example with clearly synthetic placeholder
values, including example.com email addresses, while preserving the example’s
structure and parameter names.
- Around line 103-120: Update the README training example by importing hydra,
DictConfig, and MLFlowLogger from their appropriate packages, and make the
Trainer invocation valid Python by removing the ellipsis after the callbacks
keyword and supplying any needed arguments or placing explanatory text outside
the call. Apply the same corrections to the corresponding example section
referenced by the comment.
---
Outside diff comments:
In `@rationai/mlkit/provenance/register_dataset.py`:
- Around line 242-301: Make register_dataset exception-safe by wrapping the
MLflow run and both provenance artifact workflows in try/finally, ensuring every
created temporary directory is removed with cleanup that tolerates missing paths
and mlflow.end_run() is attempted whenever start_run() succeeds. Preserve the
existing provenance generation, artifact logging, return value, and success
output while preventing failures in build_dataset_prov, file writes, or
log_artifact from leaving resources active.
In `@README.md`:
- Line 371: Update the prov_mlflow Markdown link in the README to point to the
canonical, accessible repository; if no valid canonical repository is available,
remove the hyperlink while retaining the tool reference text.
- Around line 20-24: Update the README’s MLflow startup command to bind the
server to 127.0.0.1 instead of 0.0.0.0, preserving the documented local URL and
avoiding exposure on all network interfaces.
---
Nitpick comments:
In `@rationai/mlkit/provenance/prov.py`:
- Around line 65-66: Update _typed_value to remove the unused type_prefix and
type_local parameters, retaining its existing plain-string serialization
behavior. Verify its current call sites and adjust any invocations that rely on
the removed arguments; do not add typed-literal handling unless the value
representation explicitly requires it.
In `@rationai/mlkit/provenance/register_dataset.py`:
- Around line 275-281: Extract the repeated temporary-directory, JSON-writing,
MLflow artifact logging, and cleanup flow into a shared helper in prov.py that
accepts the document, filename, and artifact path and guarantees cleanup via
try/finally. Replace the blocks at rationai/mlkit/provenance/register_dataset.py
lines 275-281 and 284-297, and rationai/mlkit/provenance/register_user.py lines
64-71, with calls to this helper while preserving each existing filename,
document, and MLflow artifact path.
In `@README.md`:
- Line 307: Update the project-tree code fence in the README to specify the text
language by changing its opening delimiter to ```text, resolving the MD040 lint
violation while leaving the fenced content unchanged.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 9002e967-ea54-46dc-80b2-9ecb67161665
📒 Files selected for processing (9)
.gitignoreREADME.mdrationai/mlkit/__init__.pyrationai/mlkit/lightning/callbacks/provenance.pyrationai/mlkit/lightning/loggers/mlflow.pyrationai/mlkit/provenance/__init__.pyrationai/mlkit/provenance/prov.pyrationai/mlkit/provenance/register_dataset.pyrationai/mlkit/provenance/register_user.py
🚧 Files skipped from review as they are similar to previous changes (2)
- rationai/mlkit/init.py
- rationai/mlkit/lightning/callbacks/provenance.py
… ProvenanceCallback
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
rationai/mlkit/lightning/callbacks/provenance.py (1)
873-877: 🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy liftPartial sibling presence still silently drops the missing callback's work.
has_env or has_verifyroutes to_gather_from_siblingswhenever either sibling exists. If onlyEnvironmentCallbackis added (noDatasetVerificationCallback), verification/split are silently omitted with no warning. Trackhas_env/has_verifyindependently, or at minimum log when only one sibling is present.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rationai/mlkit/lightning/callbacks/provenance.py` around lines 873 - 877, The callback dispatch around _gather_from_siblings must handle partial sibling presence without dropping the missing callback’s work. Track has_env and has_verify independently, invoke sibling gathering only for the callbacks that exist, and run the corresponding fallback work for any missing callback; preserve the existing all-siblings and no-siblings behavior.
♻️ Duplicate comments (1)
rationai/mlkit/lightning/callbacks/environment.py (1)
49-57: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winMissing
gitbinary still raises an uncaughtFileNotFoundError.
contextlib.suppressonly coverssubprocess.CalledProcessError; whengitisn't onPATH,check_outputraisesFileNotFoundError, which escapes_lookup_user_run. Suppress both so theUSER/"unknown"fallback applies.🐛 Proposed fix
- with contextlib.suppress(subprocess.CalledProcessError): + with contextlib.suppress(subprocess.CalledProcessError, FileNotFoundError): username = ( subprocess.check_output( ["git", "config", "user.name"], stderr=subprocess.DEVNULL, ) .decode() .strip() )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rationai/mlkit/lightning/callbacks/environment.py` around lines 49 - 57, Update the exception suppression around the git lookup in _lookup_user_run to also catch FileNotFoundError, while preserving the existing CalledProcessError handling so the USER/"unknown" fallback applies when git is unavailable or fails.
🧹 Nitpick comments (1)
rationai/mlkit/lightning/callbacks/provenance.py (1)
361-368: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRedundant nested
ifinside the prefix-stripping loop.The inner
if clean.startswith(...)duplicates thewhilecondition and is dead. The whole block can be simplified.♻️ Proposed simplification
for key, val in params.items(): if key.startswith(("opt_", "sch_")): clean = key while clean.startswith(("opt_", "sch_")): - if clean.startswith(("opt_", "sch_")): - clean = clean[4:] + clean = clean[4:] if f"gen:{clean}" not in run_activity: run_activity[f"gen:{clean}"] = _typed_value(val)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rationai/mlkit/lightning/callbacks/provenance.py` around lines 361 - 368, In the parameter normalization loop, remove the redundant inner startswith check inside the while loop. Keep the while condition responsible for repeatedly stripping “opt_” or “sch_” prefixes, then preserve the existing run_activity update and _typed_value handling.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@rationai/mlkit/provenance/prov.py`:
- Around line 42-49: Update get_prov_prefixes to validate that the parsed
PROV_BASE_URI JSON is an object/mapping before merging it with
_DEFAULT_PROV_PREFIXES, matching the guard used by
provenance.py::_get_prov_prefixes. Treat valid non-object JSON such as arrays as
invalid and return the defaults without raising.
---
Outside diff comments:
In `@rationai/mlkit/lightning/callbacks/provenance.py`:
- Around line 873-877: The callback dispatch around _gather_from_siblings must
handle partial sibling presence without dropping the missing callback’s work.
Track has_env and has_verify independently, invoke sibling gathering only for
the callbacks that exist, and run the corresponding fallback work for any
missing callback; preserve the existing all-siblings and no-siblings behavior.
---
Duplicate comments:
In `@rationai/mlkit/lightning/callbacks/environment.py`:
- Around line 49-57: Update the exception suppression around the git lookup in
_lookup_user_run to also catch FileNotFoundError, while preserving the existing
CalledProcessError handling so the USER/"unknown" fallback applies when git is
unavailable or fails.
---
Nitpick comments:
In `@rationai/mlkit/lightning/callbacks/provenance.py`:
- Around line 361-368: In the parameter normalization loop, remove the redundant
inner startswith check inside the while loop. Keep the while condition
responsible for repeatedly stripping “opt_” or “sch_” prefixes, then preserve
the existing run_activity update and _typed_value handling.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 353d2977-84bd-472e-955d-22004af79e19
📒 Files selected for processing (11)
.gitignoreREADME.mdrationai/mlkit/__init__.pyrationai/mlkit/lightning/callbacks/dataset_verification.pyrationai/mlkit/lightning/callbacks/environment.pyrationai/mlkit/lightning/callbacks/provenance.pyrationai/mlkit/lightning/loggers/mlflow.pyrationai/mlkit/provenance/__init__.pyrationai/mlkit/provenance/prov.pyrationai/mlkit/provenance/register_dataset.pyrationai/mlkit/provenance/register_user.py
🚧 Files skipped from review as they are similar to previous changes (6)
- .gitignore
- rationai/mlkit/lightning/loggers/mlflow.py
- rationai/mlkit/provenance/init.py
- rationai/mlkit/init.py
- rationai/mlkit/lightning/callbacks/dataset_verification.py
- README.md
| env_json = os.environ.get("PROV_BASE_URI", "") | ||
| if env_json: | ||
| try: | ||
| merged = {**_DEFAULT_PROV_PREFIXES, **json.loads(env_json)} | ||
| return merged | ||
| except json.JSONDecodeError: | ||
| pass | ||
| return _DEFAULT_PROV_PREFIXES |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
get_prov_prefixes crashes on valid-but-non-object PROV_BASE_URI JSON.
Only json.JSONDecodeError is caught. A value like PROV_BASE_URI='[1,2,3]' parses fine but raises TypeError when spread via **json.loads(...), escaping this handler. The same issue was already fixed in provenance.py::_get_prov_prefixes; apply the identical guard here.
🐛 Proposed fix
if env_json:
try:
- merged = {**_DEFAULT_PROV_PREFIXES, **json.loads(env_json)}
+ parsed = json.loads(env_json)
+ if not isinstance(parsed, dict):
+ raise TypeError("expected a JSON object")
+ merged = {**_DEFAULT_PROV_PREFIXES, **parsed}
return merged
- except json.JSONDecodeError:
+ except (json.JSONDecodeError, TypeError):
pass📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| env_json = os.environ.get("PROV_BASE_URI", "") | |
| if env_json: | |
| try: | |
| merged = {**_DEFAULT_PROV_PREFIXES, **json.loads(env_json)} | |
| return merged | |
| except json.JSONDecodeError: | |
| pass | |
| return _DEFAULT_PROV_PREFIXES | |
| env_json = os.environ.get("PROV_BASE_URI", "") | |
| if env_json: | |
| try: | |
| parsed = json.loads(env_json) | |
| if not isinstance(parsed, dict): | |
| raise TypeError("expected a JSON object") | |
| merged = {**_DEFAULT_PROV_PREFIXES, **parsed} | |
| return merged | |
| except (json.JSONDecodeError, TypeError): | |
| pass | |
| return _DEFAULT_PROV_PREFIXES |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@rationai/mlkit/provenance/prov.py` around lines 42 - 49, Update
get_prov_prefixes to validate that the parsed PROV_BASE_URI JSON is an
object/mapping before merging it with _DEFAULT_PROV_PREFIXES, matching the guard
used by provenance.py::_get_prov_prefixes. Treat valid non-object JSON such as
arrays as invalid and return the defaults without raising.
1 participant
Summary by CodeRabbit
provenance/run_summary.json,provenance/prov.json).