Repository navigation
Parallelize iceberg image crop export across worker processes - #495
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (5)
🚧 Files skipped from review as they are similar to previous changes (3)
📝 WalkthroughWalkthroughThe PR adds ChangesIceberg image export parallelism
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant convert
participant write_iceberg_warehouse
participant image_export
participant ProcessPoolExecutor
participant ArrowTables
convert->>write_iceberg_warehouse: pass image_crop_workers
write_iceberg_warehouse->>image_export: submit crop or source-image task
image_export->>ProcessPoolExecutor: process shards when threshold is met
ProcessPoolExecutor->>ArrowTables: build shard tables
ArrowTables-->>image_export: merge ordered and deduplicated results
image_export-->>write_iceberg_warehouse: return image export table
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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 |
|
@coderabbitai review |
✅ Action performedReview finished.
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
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)
cytotable/warehouse/images.py (1)
778-805: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winThicken empty shard outputs to the expected schema.
Both table builders build
key_field_namesfrom the received rows. A shard that skips all rows returns only fixed columns, while sibling shards include key columns, and_image_crop_table_parallel/_source_image_table_parallelconcatenate without schema promotion.Compute the expected rows from chunk data, then pass those key names into each shard worker so empty shards emit the same Arrow schema, or normalize the empty shard output before
pa.concatenate.🤖 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 `@cytotable/warehouse/images.py` around lines 778 - 805, Update _image_crop_table_parallel and _source_image_table_parallel to derive expected key-field names from the full chunk data and pass them to _rows_to_crop_table and the corresponding source-table builder, so shards with no rows still emit the complete schema before concatenation. Apply this at cytotable/warehouse/images.py:778-805 and cytotable/warehouse/images.py:1176-1183; preserve existing fixed-column behavior while ensuring empty shard outputs match non-empty siblings.
🧹 Nitpick comments (4)
cytotable/warehouse/images.py (1)
907-917: 🚀 Performance & Scalability | 🔵 Trivial | 🏗️ Heavy liftIndex rebuild cost is not cheap for cloud image directories.
Each shard worker rebuilds the image, mask, and outline indexes.
_build_file_indexcallscloud_glob(root, "**/*"), which performs a full recursive listing. With 8 shards this becomes 24 recursive listings per chunk instead of 3. On a local filesystem this is minor. On ans3://,gs://, oraz://prefix each listing is a paginated remote LIST, so the parallel path can add latency and request cost that partly offsets the crop speedup.The docstring calls these "a few cheap dir walks", which does not hold for the cloud case. Consider building the index once in the parent and passing the resolved path strings to the workers, or restricting parallel mode to local
image_dirvalues.🤖 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 `@cytotable/warehouse/images.py` around lines 907 - 917, Update the multiprocessing path around the worker that calls _build_file_index so cloud-backed image, mask, and outline indexes are built once in the parent and passed to workers as resolved path strings, avoiding repeated recursive cloud listings per shard. Preserve local behavior and ensure the worker reuses the supplied indexes rather than rebuilding them.cytotable/warehouse/iceberg.py (1)
896-925: 🚀 Performance & Scalability | 🔵 Trivial | 🏗️ Heavy liftNote the multiplicative process count across chunks.
When the Parsl thread executor is present, this loop submits one task per chunk concurrently. Each task then creates its own
ProcessPoolExecutorwith up toimage_crop_workersprocesses. The total process count is therefore the thread-pool width multiplied byimage_crop_workers, which can oversubscribe cores and multiply peak memory on large plates.The docstring at lines 657-659 already tells users to lower the value. Consider deriving a per-chunk budget from the thread-executor width instead, so the default stays safe without manual tuning.
🤖 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 `@cytotable/warehouse/iceberg.py` around lines 896 - 925, Adjust the concurrent chunk submission in the use_threads path so each _image_crop_table_app and _source_image_table_app receives a per-chunk image_crop_workers budget derived from the Parsl thread-executor width, rather than the full configured value. Keep the total process usage bounded across concurrently submitted chunks while preserving the existing behavior for non-threaded execution and the include_source_images condition.tests/test_iceberg.py (1)
1226-1260: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for shards that produce no rows.
This test uses 30 rows with 2 workers, so
_write_row_shardsproduces two evenly sized, fully populated shards. Every existing parallel test has the same shape. Two failure modes therefore stay uncovered:
- A shard whose input rows are all skipped, for example an unresolvable image filename or an invalid bounding box across a contiguous row block. That shard returns a table without key columns, and
pa.concat_tablesthen sees mismatched schemas.- A chunk with fewer rows than workers, which makes
_write_row_shardsemit empty trailing shards.Add a case that mixes resolvable and unresolvable image filenames across the shard boundary, and a case with
crop_workersgreater than the row count.🤖 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 `@tests/test_iceberg.py` around lines 1226 - 1260, Extend test_source_image_parallel_dedup_matches_serial with coverage for parallel shards that yield no rows: include resolvable and unresolvable filenames arranged across a shard boundary, and add a scenario where crop_workers exceeds the chunk’s row count to create empty trailing shards. Compare serial and parallel outputs and verify schemas, row counts, and deduplicated Metadata_ImageID values remain consistent.docs/source/overview.md (1)
181-181: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse source-level links instead of hard-coded HTML paths.
docs/source/tutorials/iceberg_warehouse_with_images.mdis in thedocs/source/tutorials/tree, whileoverview.mdis underdocs/source/. The link should betutorials/iceberg_warehouse_with_images.mdfor the Markdown/MyST source; keep the generated.htmlpath separate or derive it with a Sphinx role.🤖 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/source/overview.md` at line 181, Update the image crops tutorial link in the overview documentation to use the source-level relative path tutorials/iceberg_warehouse_with_images.md instead of the hard-coded generated HTML path, preserving the existing link text and surrounding guidance.
🤖 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 `@cytotable/warehouse/images.py`:
- Around line 1305-1337: Update _source_image_table_parallel to shard work by
unique source_image_file values rather than rows, ensuring each distinct image
name is assigned to exactly one shard and all rows referencing it remain
together. Replace or adapt _write_row_shards usage accordingly, preserve the
existing worker execution and result merge behavior, and retain
_dedup_table_by_first_occurrence only as a safety net.
- Around line 942-953: Update _write_row_shards to clamp each computed end
offset to n before the start >= end check, ensuring no shard is created when
start is beyond the table’s row count. Preserve the existing shard sizing and
return behavior for non-empty ranges, including the final shard ending at n.
- Around line 983-984: Update both ProcessPoolExecutor constructions in the
shard-processing flows, including the one using worker and the corresponding
_crop_shard_worker/_source_shard_worker flow, to pass an explicit
multiprocessing context obtained with the "spawn" start method. Keep the
existing worker mapping and max_workers behavior unchanged.
---
Outside diff comments:
In `@cytotable/warehouse/images.py`:
- Around line 778-805: Update _image_crop_table_parallel and
_source_image_table_parallel to derive expected key-field names from the full
chunk data and pass them to _rows_to_crop_table and the corresponding
source-table builder, so shards with no rows still emit the complete schema
before concatenation. Apply this at cytotable/warehouse/images.py:778-805 and
cytotable/warehouse/images.py:1176-1183; preserve existing fixed-column behavior
while ensuring empty shard outputs match non-empty siblings.
---
Nitpick comments:
In `@cytotable/warehouse/iceberg.py`:
- Around line 896-925: Adjust the concurrent chunk submission in the use_threads
path so each _image_crop_table_app and _source_image_table_app receives a
per-chunk image_crop_workers budget derived from the Parsl thread-executor
width, rather than the full configured value. Keep the total process usage
bounded across concurrently submitted chunks while preserving the existing
behavior for non-threaded execution and the include_source_images condition.
In `@cytotable/warehouse/images.py`:
- Around line 907-917: Update the multiprocessing path around the worker that
calls _build_file_index so cloud-backed image, mask, and outline indexes are
built once in the parent and passed to workers as resolved path strings,
avoiding repeated recursive cloud listings per shard. Preserve local behavior
and ensure the worker reuses the supplied indexes rather than rebuilding them.
In `@docs/source/overview.md`:
- Line 181: Update the image crops tutorial link in the overview documentation
to use the source-level relative path tutorials/iceberg_warehouse_with_images.md
instead of the hard-coded generated HTML path, preserving the existing link text
and surrounding guidance.
In `@tests/test_iceberg.py`:
- Around line 1226-1260: Extend test_source_image_parallel_dedup_matches_serial
with coverage for parallel shards that yield no rows: include resolvable and
unresolvable filenames arranged across a shard boundary, and add a scenario
where crop_workers exceeds the chunk’s row count to create empty trailing
shards. Compare serial and parallel outputs and verify schemas, row counts, and
deduplicated Metadata_ImageID values remain consistent.
🪄 Autofix
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 Plus
Run ID: 9f0c2154-9a99-4791-9e56-f9741477fdc7
📒 Files selected for processing (6)
cytotable/convert.pycytotable/warehouse/iceberg.pycytotable/warehouse/images.pydocs/source/overview.mddocs/source/tutorials/iceberg_warehouse_with_images.mdtests/test_iceberg.py
Co-Authored-By: Gregory Way <gregory.way@gmail.com>
jenna-tomkinson
left a comment
jenna-tomkinson
left a comment
There was a problem hiding this comment.
I have one general comment regarding the workers. Glad to see we found a bottleneck!
|
@d33bs In your test, can you confirm if the number of cells in the warehouse joined_profiles matches the number that we expect from the SQLite (based on number of cytoplasm in the table)? |
|
Thanks @gwaybio and @jenna-tomkinson .
To my knowledge, it does match. The source of how I compared can be found here: https://github.com/d33bs/SK-N-AS_cytotable_pilot . |
I confirmed from the SQLite that this looks like it is correct! I asked as we had an old bug in CT v1 where there was a mismatch so making sure this did not resurrect like a bug zombie of sort. Thanks, Dave!
|
|
Thanks @jenna-tomkinson and @gwaybio ! |
Description
Image export is much faster now. Crop work runs across multiple worker
processes automatically (~8 by default), cutting a ~17 min runs from ExampleHuman to a few
minutes. Tunable via
image_crop_workers; small runs are unchanged.What is the nature of your change?
Checklist
Please ensure that all boxes are checked before indicating that a pull request is ready for review.
Summary by CodeRabbit
New Features
image_crop_workersoption to conversion and Iceberg export workflows.Documentation