name	gpu-memory-cleanup
description	Preventing Jupyter cell hangs with explicit CuPy GPU memory cleanup
author	KINTSUGI Team
date	Sun Dec 15 2024 00:00:00 GMT+0000 (Coordinated Universal Time)

GPU Memory Cleanup - Research Notes

Experiment Overview

Item	Details
Date	2024-12-15
Goal	Prevent Jupyter notebook cells from hanging after GPU processing completes
Environment	KINTSUGI pipeline, CuPy, Jupyter, NVIDIA GPU
Status	Success

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)

Attempt	Why it Failed	Lesson Learned
Relying on automatic cleanup	CuPy lazy cleanup can block indefinitely	Always add explicit cleanup at end of GPU-heavy cells
Only using `gc.collect()`	Python GC doesn't release GPU memory pools	Must use CuPy's memory pool methods
Adding cleanup inside loops	Slows processing, doesn't fix end-of-cell hang	Cleanup belongs at the END of the cell, after all work

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)

gpu-memory-cleanup

Install Skill

SKILL.md