What is gpu-memory-cleanup?

This Claude skill provides a specialized solution for preventing Jupyter notebook cells from hanging after GPU-accelerated tasks. By implementing explicit CuPy memory pool cleanup and garbage collection, it resolves issues where the Python interpreter remains blocked by lazy memory management, ensuring smooth execution transitions and stable kernel performance in NVIDIA GPU environments.

When should I use gpu-memory-cleanup?

gpu-memory-cleanup is useful in the following scenarios: • Eliminating persistent execution hangs (spinning indicators) in Jupyter notebooks after completing CuPy-based computations. • Managing GPU memory overhead in long-running research pipelines such as image deconvolution, illumination correction, or EDF processing. • Ensuring immediate resource release in multi-GPU parallel processing workflows to prevent the Python interpreter from blocking indefinitely. • Improving notebook reliability for data scientists working with large-scale GPU-intensive loops where automatic garbage collection fails to release memory pools.

name	gpu-memory-cleanup
description	"Preventing Jupyter cell hangs with explicit CuPy GPU memory cleanup"
author	KINTSUGI Team
date	2024-12-15

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

When & Why to Use This Skill

Use Cases