From 9119d63bc39a11f285dba13ca605f69fa623fb9c Mon Sep 17 00:00:00 2001 From: Antoine Jacquin Date: Sun, 31 May 2026 20:55:06 +0200 Subject: [PATCH] Auto-detect best GPU (RTX 5060 preferred) + build CuPy from source for sm_120 --- Dockerfile | 20 +-- lidar_pipeline/gpu.py | 345 +++++++++++++++--------------------------- 2 files changed, 135 insertions(+), 230 deletions(-) diff --git a/Dockerfile b/Dockerfile index dc4db24..48ec06e 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,4 +1,4 @@ -FROM nvidia/cuda:11.8.0-devel-ubuntu22.04 +FROM nvidia/cuda:12.4.0-devel-ubuntu22.04 ENV DEBIAN_FRONTEND=noninteractive ENV TZ=Europe/Paris @@ -45,15 +45,15 @@ RUN pip3 install --no-cache-dir \ pillow-avif-plugin \ cmcrameri -# Install CuPy for GPU acceleration (optional - will fallback to numpy if not available) -# We use CuPy 13.4 (CUDA 11.x wheel) because CuPy 14.x dropped JIT compilation -# support. With CUPY_CUDA_COMPILE_WITH_CACHE=1, CuPy 13.4 compiles kernels at -# runtime for GPU architectures not in the pre-built wheel (e.g. sm_89 / RTX 4060 Ti). -# The devel image includes nvcc, required for JIT compilation. -# nvcc stays in PATH for the 'lidar' user. -ENV CUPY_CUDA_COMPILE_WITH_CACHE=1 -ENV PATH=/usr/local/cuda/bin:${PATH} -RUN pip3 install --no-cache-dir cupy-cuda11x==13.4.0 || echo "CuPy not available - GPU acceleration disabled" +# Build CuPy from source with nvcc, targeting sm_120 (RTX 5060) and nearby archs. +# Pre-built wheels (cupy-cuda12x 14.x) don't include sm_120, so we compile. +# This step takes ~30 min the first time; the image is cached after that. +RUN apt-get update && apt-get install -y --no-install-recommends git && \ + git clone --depth 1 --branch v14.0.0 https://github.com/cupy/cupy.git /tmp/cupy-src && \ + cd /tmp/cupy-src && \ + CUPY_NVCC_GENERATE_CODE='sm_89;sm_90;sm_120' \ + pip3 install --no-cache-dir -e . && \ + rm -rf /tmp/cupy-src # Copy and install the pipeline package COPY setup.py . diff --git a/lidar_pipeline/gpu.py b/lidar_pipeline/gpu.py index 5477566..36297e5 100644 --- a/lidar_pipeline/gpu.py +++ b/lidar_pipeline/gpu.py @@ -1,15 +1,8 @@ """GPU acceleration helpers for LiDAR pipeline. -Provides CuPy/numpy abstraction layer. If CuPy is available and a CUDA GPU -is detected, array operations are accelerated on the GPU. Otherwise, all -operations fall back to numpy/scipy on CPU. - -GPU errors (e.g. in forked subprocesses) are caught gracefully and -cause an automatic fallback to CPU for the current operation. - -Multi-GPU support: each worker process sets CUDA_VISIBLE_DEVICES before -CuPy is imported, so CuPy only sees its assigned GPU. This avoids kernel -cache incompatibilities that occur with Device.use() switching. +Auto-selects the best NVIDIA GPU (RTX 50xx preferred) and restricts +CUDA_VISIBLE_DEVICES so all workers share a single GPU. Falls back +to CPU if no GPU is available or usable. """ import logging @@ -19,125 +12,113 @@ from scipy import ndimage logger = logging.getLogger("lidar") -# Detect total GPU count via nvidia-smi (no CUDA context created). -# This must happen before any CUDA_VISIBLE_DEVICES manipulation. +# --------------------------------------------------------------------------- +# GPU auto-detection via nvidia-smi (no CUDA context created) +# --------------------------------------------------------------------------- _NUM_GPUS = 0 HAS_GPU = False _gpu_name = None _gpu_mem_gb = 0 -# System-level GPU IDs that are currently visible (after restrict_gpus). -# Populated by restrict_gpus() or auto-detected at import time. -_available_gpu_ids: list[int] = [] +_best_gpu_id: int | None = None -try: - import subprocess - _result = subprocess.run( - ['nvidia-smi', '--query-gpu=count,name,memory.total', '--format=csv,noheader,nounits'], - capture_output=True, text=True, timeout=5 - ) - if _result.returncode == 0: - _lines = _result.stdout.strip().split('\n') - _NUM_GPUS = len(_lines) - # Parse first GPU info for logging - _parts = _lines[0].split(',') - if len(_parts) >= 3: - _gpu_name = _parts[1].strip() - try: - _gpu_mem_gb = int(float(_parts[2].strip())) // 1024 - except (ValueError, IndexError): - pass + +def _pick_gpu() -> int | None: + """Pick the best GPU from the system. + + Preference order: + 1. RTX 50xx (Blackwell, compute >= 12.0) + 2. RTX 40xx (Ada Lovelace, compute >= 8.9) + 3. Any NVIDIA GPU with highest compute capability + """ + try: + import subprocess + result = subprocess.run( + ['nvidia-smi', '--query-gpu=index,name,compute_cap,memory.total', + '--format=csv,noheader,nounits'], + capture_output=True, text=True, timeout=5, + ) + if result.returncode != 0: + return None + + global _NUM_GPUS, _gpu_name, _gpu_mem_gb, HAS_GPU, _best_gpu_id + + gpus = [] + for line in result.stdout.strip().split('\n'): + parts = [p.strip() for p in line.split(',')] + if len(parts) < 4: + continue + idx = int(parts[0]) + name = parts[1] + cap_str = parts[2] + mem_mi = int(parts[3]) + major, minor = (int(x) for x in cap_str.split('.')) + # Score: higher compute capability first, then more VRAM + score = major * 1000 + minor * 100 + mem_mi + gpus.append((idx, name, cap_str, mem_mi, score)) + + if not gpus: + return None + + _NUM_GPUS = len(gpus) + gpus.sort(key=lambda g: g[4], reverse=True) + best = gpus[0] + _best_gpu_id = best[0] + _gpu_name = best[1] + _gpu_mem_gb = best[3] // 1024 HAS_GPU = True - # All GPUs are visible by default - _available_gpu_ids = list(range(_NUM_GPUS)) -except (FileNotFoundError, subprocess.TimeoutExpired, Exception): - pass + return _best_gpu_id -# Lazy CuPy initialization — imported only when first needed. -# This allows CUDA_VISIBLE_DEVICES to be set before CuPy creates -# a CUDA context, enabling per-process GPU assignment. -_xp = np # Default: CPU -_cp = None # cupy module (or None) -_cp_ndimage = None # cupyx.scipy.ndimage (or None) + except (FileNotFoundError, subprocess.TimeoutExpired, Exception): + return None + + +_best_gpu_id = _pick_gpu() + +# --------------------------------------------------------------------------- +# Lazy CuPy initialization +# --------------------------------------------------------------------------- +_xp = np +_cp = None +_cp_ndimage = None _gpu_initialized = False def _init_gpu(): - """Lazily initialize CuPy on first GPU use. - - Import CuPy only when needed, so CUDA_VISIBLE_DEVICES can be - set before the CUDA context is created. - - With CUPY_CUDA_COMPILE_WITH_CACHE=1, JIT-compiled kernels are cached - on disk and shared across processes. We warm up on ALL visible GPUs - so workers that later restrict to a single GPU find pre-compiled kernels. - - Per-GPU warm-up errors are caught individually: if GPU 1 fails to - compile a kernel (e.g. CUDA_ERROR_NO_BINARY_FOR_GPU on sm_89), we - record it and continue warming up GPU 0. Workers assigned to a - failed GPU fall back to CPU automatically. - """ - global _xp, _cp, _cp_ndimage, _gpu_initialized, HAS_GPU, _gpu_failed_ids + """Lazily initialize CuPy on first GPU use.""" + global _xp, _cp, _cp_ndimage, _gpu_initialized, HAS_GPU if _gpu_initialized: return _gpu_initialized = True - _gpu_failed_ids = set() + + if not HAS_GPU or _best_gpu_id is None: + _xp = np + _cp = None + _cp_ndimage = None + HAS_GPU = False + return try: + # Restrict to the selected GPU before CuPy imports + os.environ['CUDA_VISIBLE_DEVICES'] = str(_best_gpu_id) + import cupy as _real_cupy import cupyx.scipy.ndimage as _real_cupy_ndimage - n_devs = _real_cupy.cuda.runtime.getDeviceCount() - all_failed = False + # Warm-up: verify kernel execution works on this GPU + _test = _real_cupy.array([1.0, 2.0, 3.0], dtype=_real_cupy.float32) + _result = _real_cupy.sum(_test * _test) + _ = _result.get() + del _test, _result - # Warm up each GPU independently — one failure doesn't kill the others. - for dev_id in range(n_devs): - try: - with _real_cupy.cuda.Device(dev_id): - props = _real_cupy.cuda.runtime.getDeviceProperties(dev_id) - name = props['name'].decode() if isinstance(props['name'], bytes) else props['name'] - _test = _real_cupy.array([1.0, 2.0, 3.0], dtype=_real_cupy.float32) - _result = _real_cupy.sum(_test * _test) - _ = _result.get() - del _test, _result - logger.info(f" GPU {dev_id}: {name} — kernels pré-compilés") - except Exception as dev_err: - _gpu_failed_ids.add(dev_id) - logger.warning( - f" GPU {dev_id}: échec warm-up ({dev_err.__class__.__name__}: " - f"{dev_err}) — workers sur ce GPU passeront en CPU" - ) + _xp = _real_cupy + _cp = _real_cupy + _cp_ndimage = _real_cupy_ndimage - # If ALL visible GPUs failed, disable GPU entirely. - # This is critical for worker subprocesses that restricted - # CUDA_VISIBLE_DEVICES to a single GPU before importing CuPy. - if len(_gpu_failed_ids) >= n_devs: - all_failed = True - logger.warning("Tous les GPUs échoués au warm-up — passage en mode CPU") - - if not all_failed: - _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). - try: - 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é à " - f"{max_worker_mem // (1024**3) * 1000 // 1024} MB" - ) - except Exception: - pass # pool config is best-effort - else: - _xp = np - _cp = None - _cp_ndimage = None - HAS_GPU = False + props = _real_cupy.cuda.runtime.getDeviceProperties(0) + total_mem = props['totalGlobalMem'] + # Memory pool: up to 90% of VRAM (single GPU shared by workers) + pool_size = int(total_mem * 0.9) + _real_cupy.cuda.set_memory_pool(0, pool_size) except (ImportError, Exception) as e: logger.warning(f"GPU non disponible — mode CPU: {e}") @@ -147,94 +128,32 @@ def _init_gpu(): HAS_GPU = False -def restrict_gpus(gpu_ids: list[int], set_env_var: bool = False): - """Restrict which GPUs are visible to the process. - - By default (set_env_var=False), only records the GPU IDs for - num_gpus() and logging. Does NOT touch CUDA_VISIBLE_DEVICES - in the main process because CuPy 13.x JIT compilation (sm_89) - needs ALL GPUs visible during warm-up to pre-compile kernels. - - Workers call this with set_env_var=True (via assign_gpu_to_worker) - after forking, when CuPy hasn't been imported yet in the child. - - Args: - gpu_ids: List of system-level GPU indices to make visible. - set_env_var: If True, actually set CUDA_VISIBLE_DEVICES. - Default False (safe for main process). - """ - global _NUM_GPUS, HAS_GPU, _available_gpu_ids - if not gpu_ids or not HAS_GPU: - return - - # Validate IDs against total GPU count from nvidia-smi - total_count = _NUM_GPUS or 1 - valid_ids = [gid % total_count for gid in gpu_ids] - _available_gpu_ids = valid_ids - _NUM_GPUS = len(valid_ids) - - if set_env_var: - os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(str(g) for g in valid_ids) - logger.info(f"GPU visibles: {_available_gpu_ids}") - +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- def num_gpus(): - """Return the number of visible (available) GPUs.""" - return _NUM_GPUS + """Return 1 if GPU is active, 0 otherwise.""" + return 1 if HAS_GPU else 0 + + +def restrict_gpus(gpu_ids: list[int], set_env_var: bool = False): + """No-op — GPU is auto-selected at import time.""" + pass def set_active_gpu(gpu_id): - """Set the active GPU for the current process via CUDA_VISIBLE_DEVICES. - - gpu_id is an index into the currently visible GPU list - (_available_gpu_ids), not a system-level ID. - - MUST be called before any GPU operation (to_gpu, etc.) to ensure - CuPy creates its CUDA context on the correct device. With lazy - initialization, CuPy is imported AFTER this call, so it only - sees the assigned GPU. - - If the assigned GPU failed warm-up (e.g. NO_BINARY_FOR_GPU), this - function disables GPU for this worker so all operations fall back - to CPU without crashing. - - Args: - gpu_id: 0-based index into the visible GPU list. - """ - if not HAS_GPU or _NUM_GPUS <= 1: - return # Nothing to do for single GPU or no GPU - - gpu_id = gpu_id % _NUM_GPUS - - # Map visible-GPU index back to the real system GPU ID - system_gpu_id = _available_gpu_ids[gpu_id] - - # If this GPU failed warm-up, disable GPU for this worker - if system_gpu_id in _gpu_failed_ids: - logger.warning( - f" GPU {system_gpu_id} échouée au warm-up — " - f"worker passe en mode CPU" - ) - disable_gpu() - return - - # Set CUDA_VISIBLE_DEVICES to isolate this worker to one GPU. - # This MUST happen before CuPy is imported (lazy init). - # The JIT kernels were already pre-compiled by the main process - # on all GPUs, so the worker finds them in the cache. - os.environ['CUDA_VISIBLE_DEVICES'] = str(system_gpu_id) - - logger.info(f" GPU {system_gpu_id} sélectionnée pour ce worker") + """No-op — GPU is auto-selected at import time.""" + pass def _gpu_available(): - """Check if GPU is usable right now (may fail in forked subprocesses).""" + """Check if GPU is usable right now.""" if not HAS_GPU: return False try: _init_gpu() - _cp.cuda.runtime.getDevice() - return True + return _cp is not None except Exception: return False @@ -242,34 +161,30 @@ def _gpu_available(): def log_gpu_status(): """Log GPU detection result. Called after logging is configured.""" if _gpu_available(): - # Get actual device name from CuPy (after init) try: - dev = _cp.cuda.Device() name = _cp.cuda.runtime.getDeviceProperties(0)['name'] if isinstance(name, bytes): name = name.decode() - mem_gb = _cp.cuda.runtime.getDeviceProperties(0)['totalGlobalMem'] // (1024 ** 3) - gpu_info = f"GPU: {name} ({mem_gb} Go VRAM)" + mem_gb = _cp.cuda.runtime.getDeviceProperties(0)['totalGlobalMem'] // (1024**3) + gpu_info = f"GPU: {name} ({mem_gb} Go VRAM) — ID {_best_gpu_id}" except Exception: gpu_info = f"GPU: {_gpu_name} ({_gpu_mem_gb} Go VRAM)" - if _NUM_GPUS > 1: - gpu_info += f" × {_NUM_GPUS}" logger.info(gpu_info) else: logger.info("Pas de GPU — mode CPU uniquement") -def to_gpu(arr): - """Send array to GPU if available, otherwise return as float32 numpy. +# --------------------------------------------------------------------------- +# Array transfer +# --------------------------------------------------------------------------- - Uses float32 to reduce GPU memory usage. Falls back to CPU if GPU - is unavailable (e.g. in forked subprocess). - """ +def to_gpu(arr): + """Send array to GPU if available, otherwise return as float32 numpy.""" if _gpu_available(): try: return _cp.asarray(arr.astype(np.float32)) except Exception: - pass # Fall back to CPU + pass return arr.astype(np.float32) @@ -279,12 +194,15 @@ def to_cpu(arr): try: return _cp.asnumpy(arr) except Exception: - pass # Already on CPU or GPU error + pass return arr +# --------------------------------------------------------------------------- +# Filters — GPU if array is on GPU, CPU otherwise +# --------------------------------------------------------------------------- + def xp_gaussian_filter(arr, sigma): - """Gaussian filter — uses GPU if array is on GPU, CPU otherwise.""" if _cp is not None and isinstance(arr, _cp.ndarray): try: return _cp_ndimage.gaussian_filter(arr, sigma) @@ -294,7 +212,6 @@ def xp_gaussian_filter(arr, sigma): def xp_uniform_filter(arr, size): - """Uniform filter — uses GPU if array is on GPU, CPU otherwise.""" if _cp is not None and isinstance(arr, _cp.ndarray): try: return _cp_ndimage.uniform_filter(arr, size) @@ -304,7 +221,6 @@ def xp_uniform_filter(arr, size): def xp_minimum_filter(arr, footprint=None, size=None): - """Minimum filter — uses GPU if array is on GPU, CPU otherwise.""" if _cp is not None and isinstance(arr, _cp.ndarray): try: return _cp_ndimage.minimum_filter(arr, footprint=footprint, size=size) @@ -314,7 +230,6 @@ def xp_minimum_filter(arr, footprint=None, size=None): def xp_maximum_filter(arr, footprint=None, size=None): - """Maximum filter — uses GPU if array is on GPU, CPU otherwise.""" if _cp is not None and isinstance(arr, _cp.ndarray): try: return _cp_ndimage.maximum_filter(arr, footprint=footprint, size=size) @@ -323,6 +238,10 @@ def xp_maximum_filter(arr, footprint=None, size=None): return ndimage.maximum_filter(arr, footprint=footprint, size=size) +# --------------------------------------------------------------------------- +# Misc +# --------------------------------------------------------------------------- + def gpu_cleanup(): """Free GPU memory. Call between visualizations to prevent OOM.""" if _cp is not None: @@ -333,15 +252,10 @@ def gpu_cleanup(): def disable_gpu(): - """Disable GPU acceleration for the rest of this process. - - Called when a CUDA error indicates the GPU is unusable (e.g. - CUDA_ERROR_NO_BINARY_FOR_GPU). Falls back to numpy for all - subsequent operations. - """ + """Disable GPU acceleration for the rest of this process.""" global HAS_GPU, _xp, _cp, _cp_ndimage if not HAS_GPU: - return # Already disabled + return logger.warning("GPU désactivé — passage en mode CPU pour la suite du processus") HAS_GPU = False _xp = np @@ -350,21 +264,12 @@ def disable_gpu(): def is_gpu_active(): - """Check if GPU acceleration is currently active. - - Unlike the HAS_GPU module-level variable (which can go stale if - imported directly), this always reflects the current runtime state. - Use this in logging tags and conditional GPU paths. - """ + """Check if GPU acceleration is currently active.""" return HAS_GPU def safe_gpu_call(func, *args, **kwargs): - """Call a function with GPU arrays, retrying on CPU if GPU fails. - - Usage: - result = safe_gpu_call(generate_svf, dem_file, basename, vis_dir, resolution, shared=shared) - """ + """Call a function with GPU arrays, retrying on CPU if GPU fails.""" try: return func(*args, **kwargs) except Exception as e: @@ -373,4 +278,4 @@ def safe_gpu_call(func, *args, **kwargs): logger.warning(f"Erreur GPU ({e.__class__.__name__}), retry en CPU...") disable_gpu() return func(*args, **kwargs) - raise \ No newline at end of file + raise