vibespatial.raster.memory

Tiered GPU memory pool manager (ADR-0040).

Provides a three-tier RMM-based memory pool with CuPy fallback:

Tier | Activation | Allocator Stack |

|------|————————————–|----------------------------------------------| | A | VIBESPATIAL_GPU_POOL_ONLY=1 | PoolMemoryResource -> CudaMemoryResource | | B | RMM installed (default) | FailureCallback -> Pool -> Cuda | | C | VIBESPATIAL_GPU_MANAGED_MEMORY=1 | Bare ManagedMemoryResource | | Fall | RMM not installed | CuPy MemoryPool (default) |

Tier B is the default because its OOM callback provides robust recovery from pool fragmentation. The callback first tries gc.collect() to free unreferenced CuPy arrays; if that fails, it tears down and rebuilds the pool to release all held blocks back to the CUDA driver, then retries.

Tier A (pool-only, no OOM callback) is opt-in via VIBESPATIAL_GPU_POOL_ONLY=1 for users who want maximum performance and are willing to manage memory pressure themselves.

All imports are lazy – this module does not import CuPy or RMM at module level. The pool is configured exactly once on the first GPU operation via _ensure_memory_pool().

Functions

configure_memory_pool(→ str)

Configure the GPU memory pool. Returns the tier name.

free_pool_memory(→ None)

Release unused memory back to the driver.

memory_pool_stats(→ dict)

Return current pool statistics.

Module Contents

vibespatial.raster.memory.configure_memory_pool() → str

Configure the GPU memory pool. Returns the tier name.

This function is idempotent: calling it multiple times returns the previously-configured tier without reconfiguring.

Tier selection priority:

  1. VIBESPATIAL_GPU_MANAGED_MEMORY=1 -> Tier C

  2. VIBESPATIAL_GPU_POOL_ONLY=1 -> Tier A (no OOM recovery)

  3. RMM importable -> Tier B (default, with OOM recovery)

  4. Fallback -> CuPy default pool

vibespatial.raster.memory.free_pool_memory() → None

Release unused memory back to the driver.

For RMM pool tiers (A/B), this tears down the current pool and rebuilds it so that all unused blocks are returned to the CUDA driver. PoolMemoryResource does not support incremental shrink, so a full rebuild is the only way to reclaim fragmented pool memory on a shared GPU.

For Tier C (managed memory) this is a no-op since the OS manages page migration.

For the CuPy fallback, calls free_all_blocks() on CuPy’s pool.

This function is safe to call between independent GPU operations to reduce memory pressure. It must NOT be called while device arrays are still live – all CuPy/RMM allocations must be freed first.

vibespatial.raster.memory.memory_pool_stats() → dict

Return current pool statistics.

Keys always present:

  • tier: active tier name (“A”, “B”, “C”, “fallback”, or “” if not yet configured)

  • configured: whether the pool has been initialized

When RMM is active (tiers A/B/C):

  • current_bytes: bytes currently allocated

  • current_count: number of live allocations

  • peak_bytes: high-water-mark bytes

  • peak_count: high-water-mark allocation count

  • total_bytes: cumulative bytes allocated (lifetime)

  • total_count: cumulative allocation count (lifetime)

  • pool_size: current RMM pool size (tier A/B only)

When using the CuPy fallback:

  • used_bytes: bytes in use by CuPy pool

  • free_bytes: bytes held by CuPy pool but currently free

  • total_bytes: used + free