Spaces:
Sleeping
Sleeping
Add aggressive GPU memory cleanup after UI analysis
Browse files- Clear intermediate data structures (all_slide_masks, all_aeon_results, etc)
- Force garbage collection and CUDA cache clearing
- Log both allocated and reserved GPU memory after cleanup
- Should reduce memory footprint from 4GB to near zero after analysis
- Helps prevent OOM on consecutive T4 runs
- src/mosaic/ui/app.py +28 -2
src/mosaic/ui/app.py
CHANGED
|
@@ -271,12 +271,38 @@ def analyze_slides(
|
|
| 271 |
# Final yield with complete results
|
| 272 |
# Hide settings table if only one slide, keep visible for multiple slides
|
| 273 |
settings_visible = len(slides) > 1
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 274 |
yield (
|
| 275 |
gr.Dataframe(value=settings_input, visible=settings_visible), # Hide if single slide
|
| 276 |
-
|
| 277 |
combined_aeon_results,
|
| 278 |
aeon_output,
|
| 279 |
-
|
| 280 |
paladin_output,
|
| 281 |
user_dir,
|
| 282 |
)
|
|
|
|
| 271 |
# Final yield with complete results
|
| 272 |
# Hide settings table if only one slide, keep visible for multiple slides
|
| 273 |
settings_visible = len(slides) > 1
|
| 274 |
+
|
| 275 |
+
# Store final results before cleanup
|
| 276 |
+
final_slide_masks = all_slide_masks
|
| 277 |
+
final_combined_paladin = combined_paladin_results if len(combined_paladin_results) > 0 else None
|
| 278 |
+
|
| 279 |
+
# Aggressive memory cleanup after storing final results
|
| 280 |
+
import gc
|
| 281 |
+
import torch
|
| 282 |
+
|
| 283 |
+
# Clear intermediate data structures
|
| 284 |
+
all_slide_masks = None
|
| 285 |
+
all_aeon_results = None
|
| 286 |
+
all_paladin_results = None
|
| 287 |
+
combined_paladin_results = None
|
| 288 |
+
|
| 289 |
+
# Force garbage collection
|
| 290 |
+
gc.collect()
|
| 291 |
+
|
| 292 |
+
# Clear GPU cache if available
|
| 293 |
+
if torch.cuda.is_available():
|
| 294 |
+
torch.cuda.synchronize()
|
| 295 |
+
torch.cuda.empty_cache()
|
| 296 |
+
mem_allocated = torch.cuda.memory_allocated() / (1024**3)
|
| 297 |
+
mem_reserved = torch.cuda.memory_reserved() / (1024**3)
|
| 298 |
+
logger.info(f"GPU memory after final cleanup: {mem_allocated:.2f} GB allocated, {mem_reserved:.2f} GB reserved")
|
| 299 |
+
|
| 300 |
yield (
|
| 301 |
gr.Dataframe(value=settings_input, visible=settings_visible), # Hide if single slide
|
| 302 |
+
final_slide_masks,
|
| 303 |
combined_aeon_results,
|
| 304 |
aeon_output,
|
| 305 |
+
final_combined_paladin,
|
| 306 |
paladin_output,
|
| 307 |
user_dir,
|
| 308 |
)
|