GPU Memory Cleanup - Research Notes

Experiment Overview

Context

Jupyter notebook cells using CuPy for GPU-accelerated processing can hang after the main computation completes. The cell shows "complete" in output but the execution indicator keeps running. This is caused by CuPy's lazy memory cleanup blocking the Python interpreter.

Symptoms

• Cell output shows processing is "complete" or prints final summary
• Cell execution indicator (star

  [*]

  or spinning) remains active
• Kernel appears responsive but cell never finishes
• Interrupting the cell doesn't lose any work (files already saved)

Verified Workflow

Add Explicit GPU Cleanup at End of Processing Cells

# After your main processing loop completes...

# Explicit GPU cleanup to prevent hang
try:
    import cupy as cp
    cp.get_default_memory_pool().free_all_blocks()
    cp.get_default_pinned_memory_pool().free_all_blocks()
    print("GPU memory cleared")
except Exception:
    pass

import gc
gc.collect()
print("Stage complete")

Full Pattern for Processing Cells

import gc
import time

# Your processing code here
for item in items:
    process_item(item)
    gc.collect()  # Incremental cleanup during processing

# Final timing
end_time = time.time()
print(f"Processing complete: {end_time - start_time:.1f}s")

# CRITICAL: Explicit GPU cleanup before cell ends
try:
    import cupy as cp
    cp.get_default_memory_pool().free_all_blocks()
    cp.get_default_pinned_memory_pool().free_all_blocks()
    print("GPU memory cleared")
except ImportError:
    pass  # CuPy not available (CPU mode)
except Exception as e:
    print(f"GPU cleanup warning: {e}")

gc.collect()
print("Stage complete")  # This line confirms cell finished cleanly

Failed Attempts (Critical)

gc.collect()

Final Parameters

Minimal Cleanup Block

try:
    import cupy as cp
    cp.get_default_memory_pool().free_all_blocks()
    cp.get_default_pinned_memory_pool().free_all_blocks()
except Exception:
    pass
gc.collect()

With Logging

try:
    import cupy as cp
    mempool = cp.get_default_memory_pool()
    pinned_mempool = cp.get_default_pinned_memory_pool()

    used_before = mempool.used_bytes()
    mempool.free_all_blocks()
    pinned_mempool.free_all_blocks()

    print(f"GPU memory freed: {used_before / 1e9:.2f} GB")
except Exception:
    pass
gc.collect()

Key Insights

• CuPy uses lazy memory management for performance
• When cell ends, Python may wait for CuPy cleanup indefinitely
• Explicit cleanup forces immediate memory release
• This is especially important after long-running processing loops
• Interrupting a hung cell is safe if all files have been saved
• The cleanup block should be the LAST code in the cell

When to Apply This Pattern

• After BaSiC illumination correction loops
• After deconvolution processing
• After any multi-GPU parallel processing
• After EDF (Extended Depth of Focus) processing
• Any cell with

  >30 seconds

  of GPU-intensive work

References

• CuPy Memory Management: https://docs.cupy.dev/en/stable/user_guide/memory.html
• KINTSUGI Notebook 2 (Cycle Processing)