From a8fd8addb7429a444a3b0c8054adaa6212467ad3 Mon Sep 17 00:00:00 2001 From: Antoine Jacquin Date: Sun, 31 May 2026 19:51:32 +0200 Subject: [PATCH] Limit GPU memory pool per worker to prevent OOM with multi-worker --- lidar_pipeline/gpu.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/lidar_pipeline/gpu.py b/lidar_pipeline/gpu.py index ad9e72c..5b3650b 100644 --- a/lidar_pipeline/gpu.py +++ b/lidar_pipeline/gpu.py @@ -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