Fix: issues#9
Conversation
Co-authored-by: Copilot <copilot@github.com>
|
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:
📝 WalkthroughWalkthroughRefactors HeatmapBuilder to run concurrent per-tile predictions into a shared MaskBuilder and write a BigTIFF pyramid; extends SemanticSegmentation config and Helm user_config; adds Prov-GigaPath docs, model integration/get_config guidance, README updates, and Helm registration instructions. ChangesHeatmap Building Refactor
Models & Helm
Foundation Models & WSI Documentation
Sequence DiagramsequenceDiagram
participant Root as HeatmapBuilder.root
participant Worker as per-tile worker
participant Model as Ray Model (predict.remote)
participant Mask as MaskBuilder
participant Writer as pyvips Image
Root->>Worker: enumerate tiles & schedule bounded tasks
Worker->>Worker: fetch tissue tile (thread executor)
alt tile is empty
Worker-->>Root: skip
else
Worker->>Model: model.predict.remote(tile batch)
Model-->>Worker: prediction array
Worker->>Worker: normalize to batched tensor shape
alt first initialization
Worker->>Mask: create MaskBuilder(extents,tile,stride,channels)
end
Worker->>Mask: update_batch(pred, coords=[[y,x]])
Worker-->>Root: task complete
end
Root->>Root: await all tasks
Root->>Mask: finalize (NaNs→0, resize)
Root->>Writer: convert to 8-bit & write BigTIFF pyramid (DEFLATE, tiles)
Root->>Mask: cleanup()
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Possibly related issues
Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (3 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 transitions the HeatmapBuilder to use the MaskBuilder from the ratiopath library, enabling support for diverse prediction shapes and multi-resolution BigTIFF generation via pyvips. It also introduces documentation for the Prov-GigaPath model and provides architectural guidance for foundation model integration and whole-slide inference. The review feedback highlights critical performance improvements, recommending that CPU and I/O-bound operations—such as batch updates and TIFF saving—be offloaded to an executor to avoid blocking the asyncio event loop. A memory optimization was also suggested to handle large arrays in-place during the finalization step.
![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
🤖 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 `@builders/heatmap_builder.py`:
- Around line 136-156: The code may leak resources because
mask_builder.cleanup() is only called after the work and will be skipped if
finalize(), resize_to_source(), or vips_image.tiffsave(...) raise; wrap the
post-processing block that calls mask_builder.finalize(),
mask_builder.resize_to_source(...), the vips_image manipulation and tiffsave()
in a try/finally where mask_builder.cleanup() is invoked in the finally to
guarantee cleanup; ensure you still return or re-raise the exception after
cleanup as appropriate and keep existing behavior when mask_builder is None.
🪄 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: 5ff453ed-90c5-4434-8653-b0e23c0f9c1f
📒 Files selected for processing (6)
README.mdbuilders/heatmap_builder.pydocs/available-models.mddocs/guides/adding-models.mddocs/guides/deployment-guide.mdmisc/tile_heatmap_builder.py
💤 Files with no reviewable changes (1)
- misc/tile_heatmap_builder.py
There was a problem hiding this comment.
Pull request overview
This PR addresses the linked issues by replacing the local heatmap aggregation helper with the ratiopath mask builder and clarifying deployment documentation.
Changes:
- Replaces
TileHeatmapBuilderusage withratiopath.masks.mask_builders.MaskBuilder. - Deletes the old local tile heatmap builder implementation.
- Updates documentation for Helm application registration, model listings, and WSI/foundation-model usage.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
README.md |
Refreshes repository structure and key documentation links. |
misc/tile_heatmap_builder.py |
Removes the old local heatmap builder implementation. |
docs/guides/deployment-guide.md |
Adds the missing Helm values.yaml application registration step. |
docs/guides/adding-models.md |
Documents foundation-model reuse and WSI builder guidance. |
docs/available-models.md |
Adds Prov-GigaPath to available model documentation and renumbers Heatmap Builder. |
builders/heatmap_builder.py |
Migrates heatmap aggregation to ratiopath MaskBuilder. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Co-authored-by: Copilot <copilot@github.com>
Co-authored-by: Matěj Pekár <matej.pekar120@gmail.com>
Co-authored-by: Copilot <copilot@github.com>
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
builders/heatmap_builder.py (1)
71-147:⚠️ Potential issue | 🟠 Major | ⚡ Quick winCleanup gap remains for exceptions during tile processing.
The
try/finallyonly wraps the finalization phase.mask_builderis allocated at lines 71-78 inside thewithblock, but if any worker task raises andtask.result()at line 122 or 128 re-raises, control leaves thewithblock entirely and thetryat line 130 is never entered —mask_builder.cleanup()won't run and the memmap files leak.Widen the
try/finallyto cover the entire lifetime ofmask_builder.🛡️ Proposed fix to widen the cleanup scope
loop = asyncio.get_running_loop() tasks: set[asyncio.Task[Any]] = set() - with ( - OpenSlide(slide_path) as slide, - OpenSlide(tissue_mask_path) as tissue_slide, - ThreadPoolExecutor(max_workers=self.num_threads) as executor, - ): - level = slide.closest_level(model_config["mpp"]) - ... - mask_builder = MaskBuilder( - ... - ) - ... - - try: - result = np.nan_to_num(mask_builder.finalize(), nan=0.0, copy=False) - ... - finally: - mask_builder.cleanup() + mask_builder: MaskBuilder | None = None + try: + with ( + OpenSlide(slide_path) as slide, + OpenSlide(tissue_mask_path) as tissue_slide, + ThreadPoolExecutor(max_workers=self.num_threads) as executor, + ): + level = slide.closest_level(model_config["mpp"]) + ... + mask_builder = MaskBuilder( + ... + ) + ... # tile processing + + result = np.nan_to_num(mask_builder.finalize(), nan=0.0, copy=False) + ... + finally: + if mask_builder is not None: + mask_builder.cleanup()🤖 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 `@builders/heatmap_builder.py` around lines 71 - 147, The cleanup for MaskBuilder is not guaranteed if worker tasks raise because the try/finally only wraps finalization; expand the try/finally to begin immediately after constructing mask_builder (the MaskBuilder(...) call) and encompass the tile processing loop (process_tile, task scheduling/awaiting, and the finalize/resizing/saving block) so that mask_builder.cleanup() is always called in the finally; also ensure any pending asyncio Tasks are cancelled/awaited inside that finally before cleanup to avoid dangling tasks referencing the memmap.
🧹 Nitpick comments (2)
docs/guides/adding-models.md (2)
191-214: ⚡ Quick winClarify return type when calling batched methods through handles.
The example shows:
embedding = await self.foundation_model.predict.remote(image)When calling a method decorated with
@serve.batchthrough a Ray Serve handle, the return type matches the method's signature. If the foundation model'spredictreturnslist[float](as shown earlier in line 130), thenembeddinghere would be a list, not a single array. Consider adding a brief note about the expected return type to avoid confusion during implementation.📝 Suggested clarification
Add a comment or note after line 210:
embedding = await self.foundation_model.predict.remote(image) + # Note: embedding type matches the foundation model's predict return type🤖 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 `@docs/guides/adding-models.md` around lines 191 - 214, The example should clarify that calling a batched Serve method via a handle returns whatever the method's signature declares (e.g., a list), so update the example around self.foundation_model.predict.remote to note that await self.foundation_model.predict.remote(image) will yield the same return type as the foundation model's predict (for instance list[float] if predict is annotated to return list[float]) and recommend extracting or converting that list into the expected format before passing to your own predict; reference the foundation model's predict method, the `@serve.batch` decorator, and the await self.foundation_model.predict.remote call so readers know where to add the explanatory note.
216-223: ⚡ Quick winConsider adding reference to HeatmapBuilder implementation.
The section mentions creating "a custom Application (similar to
HeatmapBuilder)" but doesn't link to the HeatmapBuilder code. Consider adding a reference to help developers understand the pattern.📝 Suggested addition
2. **Custom WSI Aggregations (Non-Heatmap Outputs):** If your model generates something else across the entire slide (for example, a single slide-level scalar score, diagnostic classification, custom tabular statistics, embedded feature bags), you must **implement your own WSI aggregator service**. You should create a custom Application (similar to `HeatmapBuilder`) that takes paths to WSI files, iterates through the WSI tiles querying your base model for each tile, and correctly aggregates the results into your desired slide-level output format. + + See [`builders/heatmap_builder.py`](https://github.com/RationAI/model-service/blob/main/builders/heatmap_builder.py) for reference on implementing a WSI aggregation service.🤖 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 `@docs/guides/adding-models.md` around lines 216 - 223, Add a concrete reference to the HeatmapBuilder implementation so readers can inspect the pattern; update the Whole-Slide (WSI) section to link or point to the HeatmapBuilder service (HeatmapBuilder, running under /heatmap-builder) and the repository or file that contains its implementation (e.g., the HeatmapBuilder application/module), and suggest looking at its tile-iteration, model-querying, and aggregation logic as an example for building custom WSI aggregators.
🤖 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.
Outside diff comments:
In `@builders/heatmap_builder.py`:
- Around line 71-147: The cleanup for MaskBuilder is not guaranteed if worker
tasks raise because the try/finally only wraps finalization; expand the
try/finally to begin immediately after constructing mask_builder (the
MaskBuilder(...) call) and encompass the tile processing loop (process_tile,
task scheduling/awaiting, and the finalize/resizing/saving block) so that
mask_builder.cleanup() is always called in the finally; also ensure any pending
asyncio Tasks are cancelled/awaited inside that finally before cleanup to avoid
dangling tasks referencing the memmap.
---
Nitpick comments:
In `@docs/guides/adding-models.md`:
- Around line 191-214: The example should clarify that calling a batched Serve
method via a handle returns whatever the method's signature declares (e.g., a
list), so update the example around self.foundation_model.predict.remote to note
that await self.foundation_model.predict.remote(image) will yield the same
return type as the foundation model's predict (for instance list[float] if
predict is annotated to return list[float]) and recommend extracting or
converting that list into the expected format before passing to your own
predict; reference the foundation model's predict method, the `@serve.batch`
decorator, and the await self.foundation_model.predict.remote call so readers
know where to add the explanatory note.
- Around line 216-223: Add a concrete reference to the HeatmapBuilder
implementation so readers can inspect the pattern; update the Whole-Slide (WSI)
section to link or point to the HeatmapBuilder service (HeatmapBuilder, running
under /heatmap-builder) and the repository or file that contains its
implementation (e.g., the HeatmapBuilder application/module), and suggest
looking at its tile-iteration, model-querying, and aggregation logic as an
example for building custom WSI aggregators.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 1223a048-8f36-4c03-bb78-a219f8504305
📒 Files selected for processing (4)
builders/heatmap_builder.pydocs/guides/adding-models.mdhelm/rayservice/applications/episeg-1.yamlmodels/semantic_segmentation.py
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <copilot@github.com>
…into fix/issues
matejpekar
left a comment
matejpekar
left a comment
There was a problem hiding this comment.
@Jurgee make sure to test it before merging!
Co-authored-by: Copilot <copilot@github.com>
Co-authored-by: Copilot <copilot@github.com>
Co-authored-by: Copilot <copilot@github.com>
Co-authored-by: Copilot <copilot@github.com>
Co-authored-by: Matěj Pekár <matej.pekar120@gmail.com>
Mentioned issues in this PR: #7 and #8
Summary by CodeRabbit
New Features
Documentation
Refactor
Chores