Fix memory leak from cached xarray wrappers in CellMapImage

[Copilot]

• Analyze memory leak issue and identify root causes
• Fix cached_property memory accumulation in CellMapImage.array
• Add explicit cleanup method (_clear_array_cache) to clear cached array after data retrieval
• Invalidate array cache after each getitem call to prevent accumulation
• Add comprehensive tests for memory management including simulated training loop
• Document the memory management approach and clarify TensorStore cache behavior
• Verify minimal performance impact (TensorStore chunk cache is preserved)
• Create detailed documentation of the fix (MEMORY_LEAK_FIX.md)
• Verify EmptyImage doesn't have the same issue
• Address PR review feedback:
  • Wrap cache clearing in try/finally to ensure cleanup on exceptions
  • Short-circuit empty classes in multidataset class_counts
  • Use self.executor instead of creating new ThreadPoolExecutor
  • Handle Windows+TensorStore in multidataset class_counts computation
  • Fix UserWarning/Warning construction to properly emit warnings using logger

Original prompt

---

This section details on the original issue you should resolve

<issue_title>Potential memory leak, or unintentional accumulation</issue_title>
<issue_description>
With a batch size that should be ~350MB each, num_workers=11, persistent_workers=False, prefetch_factor=1, and CELLMAP_TENSORSTORE_CACHE_BYTES=1, a training loop blows past the 480GB memory limit for the job. Could this be due to upsampling happening during array retrieval not getting cleared properly?</issue_description>

<agent_instructions>The training loop mentioned can be found here if you need it: https://github.com/janelia-cellmap/cellmap-segmentation-challenge/blob/main/src/cellmap_segmentation_challenge/train.py</agent_instructions>

Comments on the Issue (you are @copilot in this section)

• Fixes Potential memory leak, or unintentional accumulation #63

---

✨ Let Copilot coding agent set things up for you — coding agent works faster and does higher quality work when set up for your repo.

[Copilot]

Pull request overview

This pull request fixes a critical memory leak in CellMapImage where cached xarray wrappers accumulated during training loops, causing memory consumption to grow from 0 to nearly 500GB over 20 minutes despite 350MB batches. The fix adds explicit cache clearing after each data retrieval operation while preserving TensorStore's chunk cache performance.

Changes:

• Adds _clear_array_cache() method to remove cached xarray wrapper from __dict__ after each __getitem__ call
• Preserves TensorStore connection and chunk cache for performance
• Includes comprehensive test suite and documentation

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.

File	Description
src/cellmap_data/image.py	Implements cache clearing mechanism by adding _clear_array_cache() method and calling it after data retrieval in __getitem__(), plus updated docstrings
tests/test_memory_management.py	New comprehensive test suite verifying cache clearing behavior across different scenarios including spatial transforms, value transforms, and interpolation
MEMORY_LEAK_FIX.md	Detailed documentation explaining the root cause, solution, and architectural considerations
tests/demo_memory_fix.py	Demo script showing the fix (though it only prints information rather than executing actual demonstrations)

[Copilot]

Pull request overview

Copilot reviewed 8 out of 10 changed files in this pull request and generated 3 comments.

[Copilot]

Pull request overview

Copilot reviewed 10 out of 12 changed files in this pull request and generated 7 comments.

[rhoadesScholar]

@copilot apply changes based on the comments in this thread

[Copilot]

Pull request overview

Copilot reviewed 10 out of 12 changed files in this pull request and generated 2 comments.

[Copilot]

Pull request overview

Copilot reviewed 11 out of 13 changed files in this pull request and generated 2 comments.

[rhoadesScholar]

@copilot apply changes based on the comments in this thread

[Copilot]

Pull request overview

Copilot reviewed 11 out of 13 changed files in this pull request and generated 2 comments.

[codecov]

Codecov Report

❌ Patch coverage is 78.98551% with 29 lines in your changes missing coverage. Please review.
✅ Project coverage is 71.83%. Comparing base (2c9a16e) to head (907b28a).
⚠️ Report is 22 commits behind head on main.

Files with missing lines	Patch %	Lines
src/cellmap_data/image.py	66.66%	12 Missing ⚠️
src/cellmap_data/image_writer.py	75.67%	9 Missing ⚠️
src/cellmap_data/multidataset.py	83.33%	7 Missing ⚠️
src/cellmap_data/dataset_writer.py	50.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #64      +/-   ##
==========================================
+ Coverage   70.48%   71.83%   +1.34%     
==========================================
  Files          28       28              
  Lines        2450     2464      +14     
==========================================
+ Hits         1727     1770      +43     
+ Misses        723      694      -29     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.