Limit GPU memory pool per worker to prevent OOM with multi-worker
This commit is contained in:
@ -91,6 +91,14 @@ def _init_gpu():
|
||||
_xp = _real_cupy
|
||||
_cp = _real_cupy
|
||||
_cp_ndimage = _real_cupy_ndimage
|
||||
# Limit GPU memory pool per worker to avoid OOM when multiple
|
||||
# workers share one GPU. Each worker gets at most 3.5 GB (or
|
||||
# 50 % of total VRAM on smaller cards).
|
||||
props = _real_cupy.cuda.runtime.getDeviceProperties(0)
|
||||
total_mem = props['totalGlobalMem']
|
||||
max_worker_mem = min(int(total_mem * 0.5), 3.5 * 1024**3)
|
||||
_real_cupy.cuda.set_memory_pool(0, max_worker_mem)
|
||||
logger.info(f" Pool mémoire GPU limité à {max_worker_mem // (1024**3) * 1000 // 1024} MB")
|
||||
except (ImportError, Exception) as e:
|
||||
logger.warning(f"GPU non disponible — mode CPU: {e}")
|
||||
_xp = np
|
||||
|
||||
Reference in New Issue
Block a user