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 Plus 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
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
…ROV-O document generation - Removed `prov.py` and `register_user.py` files. - Introduced `run.py` for training run PROV-O document generation. - Created `user.py` to handle user registration and PROV-O document building. - Updated `build_user_prov` and `register_new_user` functions to streamline user registration process. - Enhanced PROV document structure for better compatibility with MLflow and external tools.
…provenance module
…allback and run module
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
rationai/mlkit/provenance/user.py (1)
209-216: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueWrap temp-dir cleanup in
try/finally.Unlike
register_dataset, theshutil.rmtreehere is not guarded, so a failure injson.dumpormlflow.log_artifactleaves the_user_prov_*directory on disk.♻️ Proposed fix
prov_dir = f"_user_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) + try: + with open(prov_path, "w") as f: + json.dump(prov_doc, f, indent=2) + mlflow.log_artifact(prov_path, artifact_path="provenance") + finally: + shutil.rmtree(prov_dir, ignore_errors=True)🤖 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/user.py` around lines 209 - 216, Wrap the temporary provenance directory workflow in register_user (around prov_dir creation, JSON writing, and mlflow.log_artifact) with a try/finally block, and move shutil.rmtree(prov_dir, ignore_errors=True) into the finally clause so cleanup occurs when serialization or artifact logging fails.
🤖 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/run.py`:
- Around line 308-317: Update the verification metadata handling in the
`meta_entity` construction to coerce `verification["dataset_run_id"]` to a
string, including when its value is `None`, so `gen:dataset_run_id` is
serialized consistently as a string rather than `[null]`.
---
Nitpick comments:
In `@rationai/mlkit/provenance/user.py`:
- Around line 209-216: Wrap the temporary provenance directory workflow in
register_user (around prov_dir creation, JSON writing, and mlflow.log_artifact)
with a try/finally block, and move shutil.rmtree(prov_dir, ignore_errors=True)
into the finally clause so cleanup occurs when serialization or artifact logging
fails.
🪄 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: 249612b1-df44-492a-9f10-6497cbe67f20
📒 Files selected for processing (11)
.gitignorepyproject.tomlrationai/mlkit/__init__.pyrationai/mlkit/lightning/callbacks/dataset_verification.pyrationai/mlkit/lightning/callbacks/environment.pyrationai/mlkit/lightning/callbacks/provenance.pyrationai/mlkit/provenance/__init__.pyrationai/mlkit/provenance/common.pyrationai/mlkit/provenance/dataset.pyrationai/mlkit/provenance/run.pyrationai/mlkit/provenance/user.py
🚧 Files skipped from review as they are similar to previous changes (4)
- .gitignore
- rationai/mlkit/init.py
- rationai/mlkit/lightning/callbacks/dataset_verification.py
- rationai/mlkit/lightning/callbacks/environment.py
…nup in user registration
…ld_training_run_prov function
…proved organization and maintainability
…d provenance tracking
…cking in dataset registration
…hancing metadata capture
1 participant
Summary by CodeRabbit