feat(preprocessing): generalize qc_stats into coverage_stats for reuse with tissue masks#6
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 a new coverage-stats preprocessing module, related experiment and dataset configs, two job-submission scripts, and dependency additions; computes per-tile and per-ROI mask coverage using summed-area tables, parallelized with Ray, and logs artifacts/metrics to MLflow. Changes
Sequence Diagram(s)sequenceDiagram
participant Submit as Submit Script
participant Job as Compute Job
participant Repo as Repository
participant MLflow as MLflow Server
participant Ray as Ray Worker Pool
participant Storage as Parquet Storage
Submit->>Job: submit job (8 CPU, 64Gi)
Job->>Repo: git clone repo
Job->>Repo: uv sync (install deps)
Job->>Job: run preprocessing.coverage_stats
Job->>MLflow: download `slides.parquet` / `tiles.parquet`
MLflow-->>Job: return Parquet artifacts
Job->>Job: resolve mask_source (local dir or MLflow)
alt masks on MLflow
Job->>MLflow: download masks artifact
MLflow-->>Job: return mask images
end
Job->>Ray: dispatch slides for processing
Ray->>Ray: for each slide: load masks -> build SAT -> compute tile & ROI coverage
Ray-->>Job: return per-slide coverage results
Job->>Storage: stream-write coverage Parquet
Storage-->>Job: ack write
Job->>MLflow: log metrics and upload output artifacts
MLflow-->>Job: artifacts stored
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes 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 docstrings
🧪 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. Review rate limit: 0/1 reviews remaining, refill in 60 minutes.Comment |
There was a problem hiding this comment.
Code Review
This pull request introduces a new preprocessing script and associated configurations to calculate quality control coverage statistics for whole-slide image tiles using Ray and Summed Area Tables. The implementation includes new Hydra configuration files and a Kubernetes job submission script. A review comment identifies a potential integer overflow risk when calculating the Summed Area Table for very large images and suggests using np.int64 to ensure numerical stability.
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (6)
scripts/submit_qc_stats.py (2)
4-17: Optional: addgpu=Nonefor consistency with other submit scripts.Other submit scripts in the repo (e.g.,
scripts/submit_qc.py) explicitly passgpu=None. Omitting it here works only ifsubmit_job's default already implies no GPU; making it explicit avoids accidentally inheriting a different default ifkube_jobsever changes.Proposed alignment
cpu=8, memory="64Gi", + gpu=None, public=False,🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@scripts/submit_qc_stats.py` around lines 4 - 17, The submit_job invocation for "tissue-classification-qc-stats" should explicitly set gpu=None to match other submit scripts and avoid relying on submit_job defaults; update the call to include gpu=None alongside cpu, memory, public, script, and storage so the submit_job(...) invocation (function name submit_job and parameter job_name="tissue-classification-qc-stats") clearly declares no GPU.
14-14: Align with the existing pattern insubmit_qc.pyfor consistency and robustness.While
uv run -m module_nameis supported in uv v0.4.18+, the reference scriptsubmit_qc.pyusesuv run python -m preprocessing.qc. For consistency and to ensure compatibility with any older uv versions that may be present in the cluster environment, use the same pattern here.Proposed change
- "uv run -m preprocessing.qc_stats +experiment=...", + "uv run python -m preprocessing.qc_stats +experiment=...",🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@scripts/submit_qc_stats.py` at line 14, Replace the uv invocation string to match the pattern used in submit_qc.py: change the command "uv run -m preprocessing.qc_stats +experiment=..." to use "uv run python -m preprocessing.qc_stats +experiment=..." so it remains compatible with older uv versions; update the string wherever it's defined in scripts/submit_qc_stats.py (the command variable / argument construction that currently contains "-m preprocessing.qc_stats") to the new form.preprocessing/qc_stats.py (4)
195-214: Cancel pending Ray tasks on exception.If
ray.get(done[0])raises (e.g., a corrupted QC mask on one slide), thefinallyblock closes the writer but the remainingpendingfutures keep running on the cluster, wasting CPU/memory until the driver exits. Adding a cancel pass keeps shutdowns clean:Proposed cancellation
finally: if writer is not None: writer.close() + for fut in pending: + ray.cancel(fut, force=False)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@preprocessing/qc_stats.py` around lines 195 - 214, If ray.get(done[0]) raises, the remaining Ray tasks in pending keep running; modify the finally block in the loop that uses pending/futures (the code that calls ray.wait, ray.get and manages writer/output_path, cov_cols, cov_sums) to cancel all remaining futures before closing the writer: iterate over pending and call ray.cancel(...) (optionally with force=True) inside a safe try/except to swallow cancel errors, then close writer; this ensures any running tasks are terminated cleanly and resources freed when an exception occurs.
92-95: SAT dtype is brittle ifqc_mppis lowered.
satisint32with a max of ~2.1B. At the currentqc_mpp: 8.0a typical WSI mask is ≲10⁷ pixels, so this is safe today. However, if an experiment ever sets a finerqc_mpp(say 1.0), a 100k×100k mask easily exceeds int32 and silently wraps — coverages would become negative/garbage with no error. Usingint64(ornp.uintp) costs only memory for the SAT and removes the foot-gun.Proposed dtype hardening
- sat = np.zeros((mask_h + 1, mask_w + 1), dtype=np.int32) + sat = np.zeros((mask_h + 1, mask_w + 1), dtype=np.int64) sat[1:, 1:] = mask > 0🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@preprocessing/qc_stats.py` around lines 92 - 95, The summed-area table `sat` is created with dtype=np.int32 which can overflow for finer qc_mpp values; change its dtype to a 64-bit integer (e.g. dtype=np.int64 or np.intp) when creating `sat` so cumulative sums won't wrap: update the allocation of `sat` where `sat = np.zeros((mask_h + 1, mask_w + 1), dtype=...)` to use a 64-bit integer type (refer to the `sat` variable, and the lines that set `sat[1:, 1:] = mask > 0` and call `np.cumsum`) so cumsum outputs remain correct for large masks.
173-182: Use the configured logger instead of
logger: MLFlowLoggeris plumbed in via@autologand the rest of the module relies on MLflow for observability.logging(consistent with library norms) or surface the missing-slide count as an MLflow tag/metric so it's visible alongside the run.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@preprocessing/qc_stats.py` around lines 173 - 182, Replace the print(...) that reports the number of missing slides with the configured MLFlowLogger instead of printing to stdout: call the provided logger (variable name logger: MLFlowLogger) to emit a warning (e.g., logger.warning or logger.log) and also record the missing_slides count as an MLflow artifact/tag/metric (e.g., logger.log_metric or logger.set_tag) so the value surfaces in run telemetry; keep the subsequent RuntimeError behavior unchanged. Use the existing missing_slides variable and the same conditional that currently wraps the print to locate where to change this.
34-45: Switch from internalDataset.data.tableAPI topyarrow.parquet.read_table()
Dataset.data.tableis not part of the public datasets API and has changed shape between minor releases. Withdatasets>=4.0.0pinned with>=, an upstream refactor could silently break this loader.A more stable alternative is to read the parquet directly with
pyarrow(already a dependency), which also avoids the redundant Arrow IPC cache that HF datasets writes:-def load_tiles_columns(run_id: str, artifact_path: str, columns: list[str]) -> pa.Table: - """Load a tiles parquet column subset as a memory-mapped Arrow table. - - HF datasets caches the parquet into Arrow IPC on first load. We return the - underlying pyarrow Table without materializing 80M rows into pandas — only - small per-slide slices are converted before dispatching Ray tasks. - """ - local_path = mlflow.artifacts.download_artifacts( - run_id=run_id, artifact_path=artifact_path - ) - dataset = load_dataset("parquet", data_files=local_path, split="train") - return dataset.select_columns(columns).data.table +def load_tiles_columns(run_id: str, artifact_path: str, columns: list[str]) -> pa.Table: + """Load a tiles parquet column subset as an Arrow table (memory-mapped where possible).""" + local_path = mlflow.artifacts.download_artifacts( + run_id=run_id, artifact_path=artifact_path + ) + return pq.read_table(local_path, columns=columns, memory_map=True)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@preprocessing/qc_stats.py` around lines 34 - 45, The loader uses the internal Dataset.data.table access which is unstable; update load_tiles_columns to read the parquet directly with pyarrow.parquet.read_table by opening the downloaded local_path, pass columns to read_table to return a pyarrow Table, and remove reliance on load_dataset/dataset.select_columns; ensure the function still returns a pa.Table and that pyarrow.parquet is imported and used when reading local_path and selecting columns (refer to load_tiles_columns, local_path, and columns).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@preprocessing/qc_stats.py`:
- Around line 215-219: add a guard in the end of add_qc_coverage so you don't
divide by zero: check if count == 0 and, if so, return tile_count 0 with
mean_<col> set to a safe sentinel (e.g., float("nan") or None) instead of
computing cov_sums[col] / count; otherwise compute the means as before using
cov_sums, cov_cols and count. This uses the existing symbols cov_sums, cov_cols,
and count to locate where to add the conditional.
---
Nitpick comments:
In `@preprocessing/qc_stats.py`:
- Around line 195-214: If ray.get(done[0]) raises, the remaining Ray tasks in
pending keep running; modify the finally block in the loop that uses
pending/futures (the code that calls ray.wait, ray.get and manages
writer/output_path, cov_cols, cov_sums) to cancel all remaining futures before
closing the writer: iterate over pending and call ray.cancel(...) (optionally
with force=True) inside a safe try/except to swallow cancel errors, then close
writer; this ensures any running tasks are terminated cleanly and resources
freed when an exception occurs.
- Around line 92-95: The summed-area table `sat` is created with dtype=np.int32
which can overflow for finer qc_mpp values; change its dtype to a 64-bit integer
(e.g. dtype=np.int64 or np.intp) when creating `sat` so cumulative sums won't
wrap: update the allocation of `sat` where `sat = np.zeros((mask_h + 1, mask_w +
1), dtype=...)` to use a 64-bit integer type (refer to the `sat` variable, and
the lines that set `sat[1:, 1:] = mask > 0` and call `np.cumsum`) so cumsum
outputs remain correct for large masks.
- Around line 173-182: Replace the print(...) that reports the number of missing
slides with the configured MLFlowLogger instead of printing to stdout: call the
provided logger (variable name logger: MLFlowLogger) to emit a warning (e.g.,
logger.warning or logger.log) and also record the missing_slides count as an
MLflow artifact/tag/metric (e.g., logger.log_metric or logger.set_tag) so the
value surfaces in run telemetry; keep the subsequent RuntimeError behavior
unchanged. Use the existing missing_slides variable and the same conditional
that currently wraps the print to locate where to change this.
- Around line 34-45: The loader uses the internal Dataset.data.table access
which is unstable; update load_tiles_columns to read the parquet directly with
pyarrow.parquet.read_table by opening the downloaded local_path, pass columns to
read_table to return a pyarrow Table, and remove reliance on
load_dataset/dataset.select_columns; ensure the function still returns a
pa.Table and that pyarrow.parquet is imported and used when reading local_path
and selecting columns (refer to load_tiles_columns, local_path, and columns).
In `@scripts/submit_qc_stats.py`:
- Around line 4-17: The submit_job invocation for
"tissue-classification-qc-stats" should explicitly set gpu=None to match other
submit scripts and avoid relying on submit_job defaults; update the call to
include gpu=None alongside cpu, memory, public, script, and storage so the
submit_job(...) invocation (function name submit_job and parameter
job_name="tissue-classification-qc-stats") clearly declares no GPU.
- Line 14: Replace the uv invocation string to match the pattern used in
submit_qc.py: change the command "uv run -m preprocessing.qc_stats
+experiment=..." to use "uv run python -m preprocessing.qc_stats
+experiment=..." so it remains compatible with older uv versions; update the
string wherever it's defined in scripts/submit_qc_stats.py (the command variable
/ argument construction that currently contains "-m preprocessing.qc_stats") to
the new form.
🪄 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: 54d57456-5971-475c-91e0-0c5c0a80a802
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (6)
configs/data/dataset.yamlconfigs/experiment/preprocessing/qc_stats_standard.yamlconfigs/preprocessing/qc_stats.yamlpreprocessing/qc_stats.pypyproject.tomlscripts/submit_qc_stats.py
…ue masks Replaces qc_stats with a config-driven coverage_stats pipeline that takes a masks dict and a mask_source (local path or mlflow artifact). QC keeps its existing artifact path, column names, and metrics; tissue_stats reuses the same module via a new experiment config. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@preprocessing/coverage_stats.py`:
- Around line 145-163: The loop assumes there are tiles but
indices/sorted_indices can be empty; before indexing sorted_indices[start]
inside the for i in range(...) loop (and before using
dictionary[int(sorted_indices[start])].as_py()), add a guard to skip empty
splits — e.g., check if sorted_indices.size == 0 or if boundaries length makes
start==end and continue; ensure this check is placed alongside the existing
futures/missing_slides logic so empty splits are skipped without causing
IndexError.
🪄 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: c76eadfd-d721-4911-b050-c91052a1174e
📒 Files selected for processing (6)
configs/experiment/preprocessing/qc_stats_standard.yamlconfigs/experiment/preprocessing/tissue_stats_05mpp.yamlconfigs/preprocessing/coverage_stats.yamlpreprocessing/coverage_stats.pyscripts/submit_qc_stats.pyscripts/submit_tissue_stats.py
✅ Files skipped from review due to trivial changes (3)
- configs/preprocessing/coverage_stats.yaml
- configs/experiment/preprocessing/qc_stats_standard.yaml
- scripts/submit_tissue_stats.py
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@preprocessing/coverage_stats.py`:
- Around line 95-98: The summed-area table 'sat' is using np.int32 and can
overflow for high-resolution masks; change its dtype to a 64-bit integer (e.g.,
np.int64) when creating sat in preprocessing/coverage_stats.py (the sat =
np.zeros((mask_h + 1, mask_w + 1), dtype=...) line) so that sat[1:, 1:] = mask >
0 and the subsequent np.cumsum calls operate on a 64-bit accumulator and avoid
silent overflow.
- Around line 21-32: load_tiles_columns currently calls load_dataset without a
columns projection and then uses dataset.select_columns, which defeats parquet
projection pushdown; change the load_dataset call in load_tiles_columns to pass
the columns argument (i.e., use load_dataset(..., data_files=local_path,
split="train", columns=columns)) so the parquet reader only reads requested
fields, and remove or keep select_columns as redundant after that change.
🪄 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: 05903518-7015-4d9b-be60-06c09706507a
📒 Files selected for processing (2)
configs/data/dataset.yamlpreprocessing/coverage_stats.py
✅ Files skipped from review due to trivial changes (1)
- configs/data/dataset.yaml
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
preprocessing/coverage_stats.py (1)
21-32:⚠️ Potential issue | 🟠 Major | ⚡ Quick winPush the parquet projection into
load_dataset().
select_columns()trims the dataset after the parquet scan, so this still reads every column fromtiles.parquet. On the 80M-row path mentioned in the docstring, that is a large avoidable I/O and cache hit. Passcolumns=columnstoload_dataset(...)and keep or dropselect_columns()as a no-op safeguard.Suggested change
- dataset = load_dataset("parquet", data_files=local_path, split="train") - return dataset.select_columns(columns).data.table + dataset = load_dataset( + "parquet", + data_files=local_path, + split="train", + columns=columns, + ) + return dataset.data.tableIn Hugging Face Datasets, when loading Parquet with `load_dataset("parquet", ...)`, does passing `columns=[...]` enable projection pushdown while calling `Dataset.select_columns()` after load does not avoid reading the other columns?🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@preprocessing/coverage_stats.py` around lines 21 - 32, The function load_tiles_columns currently calls load_dataset("parquet", ...) without a columns projection so the parquet reader still scans all columns; update load_tiles_columns to pass columns=columns into load_dataset(...) to enable parquet projection pushdown (and you may keep the subsequent dataset.select_columns(columns) as a no-op safeguard). Locate load_tiles_columns and change the load_dataset invocation to include the columns argument while retaining the existing return of dataset.select_columns(columns).data.table.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@preprocessing/coverage_stats.py`:
- Around line 35-48: The resolve_mask_dir function currently prefers local_path
when both keys exist and lets partial MLflow configs cause AttributeError;
update resolve_mask_dir to explicitly validate mask_source as an exclusive
union: if "local_path" is present, ensure neither "mlflow_run_id" nor
"artifact_path" are present and return Path(mask_source.local_path); else if
both "mlflow_run_id" and "artifact_path" are present, use
mlflow.artifacts.download_artifacts(...) and return its Path; otherwise raise a
clear ValueError explaining the required shapes ("either 'local_path' OR both
'mlflow_run_id' and 'artifact_path'") so mixed or partial configs fail fast
instead of silently misbehaving (refer to resolve_mask_dir and mask_source keys
to locate the change).
---
Duplicate comments:
In `@preprocessing/coverage_stats.py`:
- Around line 21-32: The function load_tiles_columns currently calls
load_dataset("parquet", ...) without a columns projection so the parquet reader
still scans all columns; update load_tiles_columns to pass columns=columns into
load_dataset(...) to enable parquet projection pushdown (and you may keep the
subsequent dataset.select_columns(columns) as a no-op safeguard). Locate
load_tiles_columns and change the load_dataset invocation to include the columns
argument while retaining the existing return of
dataset.select_columns(columns).data.table.
🪄 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: c1343dd9-965d-493c-815e-3da588e5048a
📒 Files selected for processing (1)
preprocessing/coverage_stats.py
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@preprocessing/coverage_stats.py`:
- Around line 84-85: The ROI offset is currently computed from tm_extent
(roi_offset = tm_extent // 4) which mis-centers when tm_extent is odd; instead
compute roi_offset from roi_extent so the ROI is centered: set roi_offset =
(tm_extent - roi_extent) // 2 after you compute roi_extent = max(1, tm_extent //
2). Update the assignment in preprocessing/coverage_stats.py where roi_extent
and roi_offset are defined so downstream uses (e.g., roi_*_coverage
calculations) use the correctly centered ROI.
🪄 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: dc5ee7c4-ff40-476a-bb3d-89333ed4874c
📒 Files selected for processing (3)
configs/experiment/preprocessing/qc_standard.yamlconfigs/experiment/preprocessing/qc_stats_standard.yamlpreprocessing/coverage_stats.py
✅ Files skipped from review due to trivial changes (2)
- configs/experiment/preprocessing/qc_standard.yaml
- configs/experiment/preprocessing/qc_stats_standard.yaml
Summary
Generalizes the per-tile QC coverage pipeline into a config-driven
coverage_statsmodule that any mask-based coverage stats can reuse. Tissue stats (previously a near-duplicate script onfeature/tile-statistics) now drops in as a single experiment config, so that branch can be closed.What changed
preprocessing/coverage_stats.py— generic per-tile coverage computation. Driven by:masks: dict[str, str]—{name: filename_template}where{stem}expands to the WSI file stemmask_source— either{local_path: ...}or{mlflow_run_id: ..., artifact_path: ...}mpp— MPP at which the masks were renderedconfigs/preprocessing/coverage_stats.yaml— defaults shared by all consumersconfigs/experiment/preprocessing/qc_stats_standard.yaml— updated to the new schema; preserves existingmlflow_artifact_path: qc_stats, output column names (tile_residual_coverage, etc.), and metric names. No behavior change.configs/experiment/preprocessing/tissue_stats_05mpp.yaml— single masktissue: "{stem}.tiff",mask_sourcepulls fromthe
tissue_masks_run_idMLflow artifactpreprocessing.coverage_statstissue_masks_run_idadded (was previously only onfeature/tile-statistics)preprocessing/qc_stats.py,configs/preprocessing/qc_stats.yaml(rename detected by git)Summary by CodeRabbit
New Features
Chores
Chores