Feature/implement k fold train split#3
Conversation
|
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:
📝 WalkthroughWalkthroughAdds k‑fold stratified split: new dataset MLflow artifact reference ( ChangesK-fold split feature
Sequence Diagram(s)sequenceDiagram
participant Config as Hydra Config
participant MLflow as MLflow Artifacts
participant Data as Parquet Loader
participant Split as K-fold Processor
participant Storage as Artifact Upload
Config->>MLflow: resolve `filter_tiles_run_id` & artifact path
MLflow->>Data: download train tiles parquet
Data-->>Split: provide rows with `roi_coverage_*` and `slide_id`
Split->>Split: derive `label`, `tissue_prop`, `slide_id`
Split->>Split: collapse rare classes and build stratification labels
Split->>Split: assign folds via StratifiedKFold
Split->>MLflow: log per-fold metrics and label distribution tables
Split->>Storage: write augmented parquet and upload to MLflow artifact path
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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 a stratified k-fold splitting utility for tissue classification datasets, including new configuration files, dependency updates, and a Kubernetes job submission script. The core logic in split/kfold_split.py handles label derivation, rare label management, and MLFlow logging. Review feedback identifies potential crashes due to missing ROI columns or placeholder values in the submission script and suggests a more efficient approach for adding multiple columns to the dataset to avoid unnecessary data copying.
Previously, rare classes (fewer than n_folds tiles) were collapsed into 'background' in-place, and the mutated labels flowed all the way into the stored 'label' column — losing the original semantic class for those samples. Split the array in two: original 'labels' stays untouched and is what gets persisted and reported, while 'stratification_labels' (a copy with rare classes merged) is used solely as the StratifiedKFold target. Log both distributions to MLflow for visibility.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
split/kfold_split.py (1)
125-127:⚠️ Potential issue | 🟠 Major | ⚡ Quick win
configs/split.yamlprimary config is still missing — runtime will fail withMissingConfigException.
@hydra.main(config_name="split")requires a file atconfigs/split.yamlto exist. The+split=kfold_splitoverride injected by@with_cli_argsis a config-group append that Hydra processes after locating and parsing the primary config; it cannot substitute for the missing primary file.Create a minimal
configs/split.yaml:# `@package` _global_ defaults: - _self_Then select the group via CLI (
+split=kfold_split) or set a default:defaults: - split: kfold_split - _self_Verify startup with
python -m split.kfold_split --cfg job.🤖 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 `@split/kfold_split.py` around lines 125 - 127, The Hydra runtime fails because `@hydra.main`(config_name="split") requires a primary configs/split.yaml which is missing; add a minimal configs/split.yaml (with an empty defaults or a defaults list selecting split: kfold_split) so Hydra can load the primary config before group overrides applied by with_cli_args(defaults=["+split=kfold_split"]); update the defaults block to either just include _self_ or to set split: kfold_split as the default, then verify startup by running the module (e.g., python -m split.kfold_split --cfg job).
🤖 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 `@split/kfold_split.py`:
- Around line 47-58: build_stratification_labels can return a strat array where
the collapsed "background" group still has fewer than n_folds members, causing a
cryptic StratifiedKFold error in assign_folds; after computing rare and updating
strat (and before returning), compute the merged background count (sum of counts
for classes in rare plus any existing "background" count), and if that merged
count < n_folds raise a clear ValueError (or custom exception) that lists
n_folds and the affected classes/counts so callers (and assign_folds) get a
meaningful diagnostic instead of the sklearn message; reference the variables
unique, counts, rare, strat and the function build_stratification_labels when
implementing this validation.
---
Duplicate comments:
In `@split/kfold_split.py`:
- Around line 125-127: The Hydra runtime fails because
`@hydra.main`(config_name="split") requires a primary configs/split.yaml which is
missing; add a minimal configs/split.yaml (with an empty defaults or a defaults
list selecting split: kfold_split) so Hydra can load the primary config before
group overrides applied by with_cli_args(defaults=["+split=kfold_split"]);
update the defaults block to either just include _self_ or to set split:
kfold_split as the default, then verify startup by running the module (e.g.,
python -m split.kfold_split --cfg job).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
After collapsing rare classes into 'background', the merged group can still have fewer than n_folds tiles, which would surface as a cryptic sklearn error from StratifiedKFold. Raise a ValueError with the n_folds value and the list of collapsed classes so the failure is actionable.
…ists Dataset[col] materialized 80M strings as a Python list before np.array walked it — multi-GB allocation, no progress, easy to mistake for a hang. Pull the columns straight from the underlying pyarrow.Table which is near-zero-copy for fixed-width and chunked-fast for strings.
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
split/kfold_split.py (1)
137-140:⚠️ Potential issue | 🟠 Major | ⚡ Quick win
configs/split.yamlprimary config is still missing — Hydra will fail at startup.
@hydra.main(config_path="../configs", config_name="split", version_base=None)requires aconfigs/split.yamlfile to exist as the composition entry point. This was flagged in a prior review and remains unresolved.🤖 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 `@split/kfold_split.py` around lines 137 - 140, Hydra is configured in the kfold_split module via the `@hydra.main` decorator on main (decorator: `@hydra.main`(config_path="../configs", config_name="split", version_base=None)) but the required composition entry file configs/split.yaml is missing; add a top-level configs/split.yaml (the Hydra primary config) that composes or includes the split config values expected by main (or alternatively change the decorator to point to an existing config_name/path), ensuring the config keys referenced by main and downstream code match the new split.yaml entries so Hydra can start successfully.
🤖 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 `@split/kfold_split.py`:
- Around line 58-70: The guard that checks background_count < n_folds is firing
even when no classes were collapsed; move the computation and check for
background_count inside the existing if len(rare) > 0 block so it only runs when
collapse happened (i.e., after strat[np.isin(strat, rare)] = "background").
Specifically, in the function that contains strat and rare (e.g.,
build_stratification_labels), stop computing background_count and raising the
ValueError unconditionally—compute background_count = int((strat ==
"background").sum()) and run the if background_count < n_folds: raise
ValueError(...) only inside the if len(rare) > 0: branch, and remove the
unconditional check after that branch.
---
Duplicate comments:
In `@split/kfold_split.py`:
- Around line 137-140: Hydra is configured in the kfold_split module via the
`@hydra.main` decorator on main (decorator: `@hydra.main`(config_path="../configs",
config_name="split", version_base=None)) but the required composition entry file
configs/split.yaml is missing; add a top-level configs/split.yaml (the Hydra
primary config) that composes or includes the split config values expected by
main (or alternatively change the decorator to point to an existing
config_name/path), ensuring the config keys referenced by main and downstream
code match the new split.yaml entries so Hydra can start successfully.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
The background-count guard was running unconditionally, so an input with no rare classes and no pre-existing 'background' class would raise a spurious ValueError despite being valid for StratifiedKFold. Move the check inside the 'len(rare) > 0' branch — it only matters after a collapse actually merged classes into 'background'.
Resolved conflict in configs/data/dataset.yaml by keeping both additions: train_tiles_filename from this branch and tissue_masks_run_id from master. Regenerated uv.lock.
…-fold-train-split
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
There was a problem hiding this comment.
♻️ Duplicate comments (1)
split/kfold_split.py (1)
137-138:⚠️ Potential issue | 🟠 Major | ⚡ Quick win
configs/split.yamlprimary config appears to be missing — Hydra will fail at startup.
@hydra.main(config_name="split")requires a primary config file atconfigs/split.yaml. Hydra does not support group-only composition without a primary config file; the+split=kfold_splitCLI override alone cannot satisfy this. Onlyconfigs/split/kfold_split.yamlis visible in this PR — noconfigs/split.yaml.#!/bin/bash # Description: Check if configs/split.yaml exists and look at how other # Hydra entrypoints in the repo pair their config_name with a primary config file. echo "=== Check for configs/split.yaml ===" fd -t f 'split.ya?ml$' configs/ echo "" echo "=== All YAML files under configs/ top level (primary config candidates) ===" fd -t f -d 1 '\.ya?ml$' configs/ echo "" echo "=== `@hydra.main` declarations and their config_name values ===" rg -nP --type=py '@hydra\.main\s*\(' -A2🤖 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 `@split/kfold_split.py` around lines 137 - 138, Hydra is failing because `@hydra.main`(config_name="split", ...) in split/kfold_split.py expects a primary config file configs/split.yaml which is missing; fix by either (A) adding a primary config file configs/split.yaml that composes the group entry (e.g., defaults: ["split: kfold_split"] or defaults: ["+split=kfold_split"]) so the group-only override works, or (B) change the decorator in split/kfold_split.py to point directly at the existing config (e.g., set config_name to the existing file such as "split/kfold_split" or the exact primary config name you have), ensuring the config_name used by `@hydra.main` matches an actual YAML file under configs.
🤖 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.
Duplicate comments:
In `@split/kfold_split.py`:
- Around line 137-138: Hydra is failing because `@hydra.main`(config_name="split",
...) in split/kfold_split.py expects a primary config file configs/split.yaml
which is missing; fix by either (A) adding a primary config file
configs/split.yaml that composes the group entry (e.g., defaults: ["split:
kfold_split"] or defaults: ["+split=kfold_split"]) so the group-only override
works, or (B) change the decorator in split/kfold_split.py to point directly at
the existing config (e.g., set config_name to the existing file such as
"split/kfold_split" or the exact primary config name you have), ensuring the
config_name used by `@hydra.main` matches an actual YAML file under configs.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 6e20c773-ecac-4c5b-9c3c-8416bc2173b6
📒 Files selected for processing (2)
configs/data/dataset.yamlsplit/kfold_split.py
…-fold-train-split
3 participants
Summary
n_foldsvalidation folds usingStratifiedKFoldon the dominant tissue class label(
roi_coverage_*columns), so no code changes are needed when tissue classes changen_foldstiles are collapsed into background beforesplitting to prevent
StratifiedKFoldfrom crashing, with a printed warning listingaffected classes and their counts
distribution table are all logged to MLflow
label,tissue_prop, andfoldcolumnsadded, then logged as an MLflow artifact
Summary by CodeRabbit
New Features
Configuration
Chores