raylim commited on
Commit
4fda083
·
unverified ·
1 Parent(s): 5549e2b

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

Files changed (1) hide show
  1. 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
- all_slide_masks,
277
  combined_aeon_results,
278
  aeon_output,
279
- combined_paladin_results if len(combined_paladin_results) > 0 else None,
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
  )