Fix multi-GPU with lazy CuPy init + rendering improvements

GPU fix:
- Revert to CUDA_VISIBLE_DEVICES approach but with lazy CuPy init
- gpu.py: CuPy is no longer imported at module level; _init_gpu()
  imports it lazily on first to_gpu() call. This allows workers to
  set CUDA_VISIBLE_DEVICES before CuPy creates a CUDA context.
- gpu.py: detect GPU count via nvidia-smi (no CUDA context needed)
- pipeline.py: each worker sets CUDA_VISIBLE_DEVICES=N before CuPy
  init, so each process uses only its assigned GPU

Rendering improvements:
- Title: split into bold title (14pt) + italic description (10pt)
  instead of single 15pt bold block
- North arrow: moved inside data area (top-right corner) with
  semi-transparent white background for readability over data
- Colorbar: full height (no gap for compass rose), added
  ScalarFormatter(useOffset=False) to avoid scientific notation
- Colorbar compass rose gap removed since north arrow is now
  inside the data area
This commit is contained in:
Antoine Jacquin
2026-05-15 12:24:57 +02:00
parent b4a0e384c9
commit a3f7b44874
3 changed files with 58 additions and 48 deletions

View File

@ -7,8 +7,9 @@ operations fall back to numpy/scipy on CPU.
GPU errors (e.g. in forked subprocesses) are caught gracefully and GPU errors (e.g. in forked subprocesses) are caught gracefully and
cause an automatic fallback to CPU for the current operation. cause an automatic fallback to CPU for the current operation.
Multi-GPU support: when multiple GPUs are available, each worker process Multi-GPU support: each worker process sets CUDA_VISIBLE_DEVICES before
can be assigned a different GPU via set_active_gpu() for balanced load. CuPy is imported, so CuPy only sees its assigned GPU. This avoids kernel
cache incompatibilities that occur with Device.use() switching.
""" """
import logging import logging
@ -18,16 +19,13 @@ from scipy import ndimage
logger = logging.getLogger("lidar") logger = logging.getLogger("lidar")
# Detect GPU count at import time WITHOUT importing CuPy. # Detect total GPU count via nvidia-smi (no CUDA context created).
# We use nvidia-smi or CUDA_VISIBLE_DEVICES to count GPUs, # This must happen before any CUDA_VISIBLE_DEVICES manipulation.
# so that CUDA_VISIBLE_DEVICES can be set BEFORE CuPy context creation
# in worker processes.
_NUM_GPUS = 0 _NUM_GPUS = 0
HAS_GPU = False HAS_GPU = False
_gpu_name = None _gpu_name = None
_gpu_mem_gb = 0 _gpu_mem_gb = 0
# Check if GPUs are available via nvidia-smi (no CUDA context created)
try: try:
import subprocess import subprocess
_result = subprocess.run( _result = subprocess.run(
@ -49,9 +47,9 @@ try:
except (FileNotFoundError, subprocess.TimeoutExpired, Exception): except (FileNotFoundError, subprocess.TimeoutExpired, Exception):
pass pass
# Lazy-initialized GPU module references # Lazy CuPy initialization — imported only when first needed.
# CuPy is imported only when first needed, allowing CUDA_VISIBLE_DEVICES # This allows CUDA_VISIBLE_DEVICES to be set before CuPy creates
# to be set before CuPy context creation in worker processes. # a CUDA context, enabling per-process GPU assignment.
_xp = np # Default: CPU _xp = np # Default: CPU
_cp = None # cupy module (or None) _cp = None # cupy module (or None)
_cp_ndimage = None # cupyx.scipy.ndimage (or None) _cp_ndimage = None # cupyx.scipy.ndimage (or None)
@ -61,8 +59,8 @@ _gpu_initialized = False
def _init_gpu(): def _init_gpu():
"""Lazily initialize CuPy on first GPU use. """Lazily initialize CuPy on first GPU use.
This allows CUDA_VISIBLE_DEVICES to take effect in worker processes Import CuPy only when needed, so CUDA_VISIBLE_DEVICES can be
before CuPy creates a CUDA context. set before the CUDA context is created.
""" """
global _xp, _cp, _cp_ndimage, _gpu_initialized global _xp, _cp, _cp_ndimage, _gpu_initialized
if _gpu_initialized: if _gpu_initialized:
@ -76,39 +74,36 @@ def _init_gpu():
_xp = _real_cupy _xp = _real_cupy
_cp = _real_cupy _cp = _real_cupy
_cp_ndimage = _real_cupy_ndimage _cp_ndimage = _real_cupy_ndimage
except (ImportError, Exception): except (ImportError, Exception) as e:
logger.debug(f"CuPy non disponible: {e}")
_xp = np _xp = np
_cp = None _cp = None
_cp_ndimage = None _cp_ndimage = None
def num_gpus(): def num_gpus():
"""Return the number of available CUDA GPUs.""" """Return the total number of CUDA GPUs in the system."""
return _NUM_GPUS return _NUM_GPUS
def set_active_gpu(gpu_id): def set_active_gpu(gpu_id):
"""Set the active GPU for the current process. """Set the active GPU for the current process via CUDA_VISIBLE_DEVICES.
Must be called BEFORE any GPU operation (to_gpu, etc.) to ensure MUST be called before any GPU operation (to_gpu, etc.) to ensure
the CUDA context is created on the correct device. 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.
Args: Args:
gpu_id: 0-based GPU index. Clamped to valid range. gpu_id: 0-based GPU index (referring to the system GPU numbering).
""" """
if not HAS_GPU or _NUM_GPUS <= 1: if not HAS_GPU or _NUM_GPUS <= 1:
return # Nothing to do for single GPU or no GPU return # Nothing to do for single GPU or no GPU
gpu_id = gpu_id % _NUM_GPUS gpu_id = gpu_id % _NUM_GPUS
# Set CUDA_VISIBLE_DEVICES before CuPy context is created
# This is the most reliable way in spawn processes # Set CUDA_VISIBLE_DEVICES before CuPy context creation
os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_id) os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_id)
# Reset lazy init so CuPy re-detects with the new env
global _gpu_initialized, _cp, _cp_ndimage, _xp
_gpu_initialized = False
_cp = None
_cp_ndimage = None
_xp = np
logger.info(f" GPU {gpu_id} sélectionnée pour ce worker") logger.info(f" GPU {gpu_id} sélectionnée pour ce worker")
@ -127,8 +122,17 @@ def _gpu_available():
def log_gpu_status(): def log_gpu_status():
"""Log GPU detection result. Called after logging is configured.""" """Log GPU detection result. Called after logging is configured."""
if HAS_GPU: if _gpu_available():
gpu_info = f"GPU détectée: {_gpu_name} ({_gpu_mem_gb} Go VRAM)" # 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)"
except Exception:
gpu_info = f"GPU: {_gpu_name} ({_gpu_mem_gb} Go VRAM)"
if _NUM_GPUS > 1: if _NUM_GPUS > 1:
gpu_info += f" × {_NUM_GPUS}" gpu_info += f" × {_NUM_GPUS}"
logger.info(gpu_info) logger.info(gpu_info)

View File

@ -537,10 +537,9 @@ def _process_file_standalone(laz_file_str, input_dir, output_dir, resolution, fo
Each worker gets its own temp directory to avoid file conflicts. Each worker gets its own temp directory to avoid file conflicts.
When multiple GPUs are available, each worker is assigned a GPU via When multiple GPUs are available, each worker is assigned a GPU via
CUDA_VISIBLE_DEVICES to balance load across GPUs. CuPy's Device API to balance load across GPUs.
""" """
# Assign GPU FIRST — before any CuPy import happens # Assign GPU to this worker using CuPy's Device API
# This sets CUDA_VISIBLE_DEVICES and resets lazy CuPy init
if gpu_id is not None and gpu_id >= 0: if gpu_id is not None and gpu_id >= 0:
from .gpu import set_active_gpu from .gpu import set_active_gpu
set_active_gpu(gpu_id) set_active_gpu(gpu_id)

View File

@ -30,6 +30,7 @@ import matplotlib.pyplot as plt
from matplotlib import rcParams from matplotlib import rcParams
from matplotlib.patches import Polygon as MplPolygon, Rectangle as RectPatch, FancyBboxPatch from matplotlib.patches import Polygon as MplPolygon, Rectangle as RectPatch, FancyBboxPatch
from matplotlib.colors import ListedColormap from matplotlib.colors import ListedColormap
from matplotlib.ticker import ScalarFormatter
try: try:
from cmcrameri import cm as cmc from cmcrameri import cm as cmc
@ -495,15 +496,16 @@ def tif_to_png(tif_file, vis_dir, resolution, keep_tif=False, source_info=None,
im = ax.imshow(data, cmap=cmap, aspect='equal', origin='upper', im = ax.imshow(data, cmap=cmap, aspect='equal', origin='upper',
interpolation='bilinear') interpolation='bilinear')
ax.set_title(f"{title}\n{description}", fontsize=15, fontweight='bold', pad=10) ax.set_title(f"{title}", fontsize=14, fontweight='bold', pad=8)
ax.text(0.5, 1.01, description, transform=ax.transAxes,
fontsize=10, fontstyle='italic', color='#555555',
ha='center', va='bottom')
# Colorbar/legend area — reduced height to leave room for compass rose above # Colorbar/legend area — full height alongside data
cbar_left = data_left + data_width_frac + 0.02 cbar_left = data_left + data_width_frac + 0.02
cbar_width = 0.04 cbar_width = 0.04
compass_height = 0.07
compass_gap = 0.02
cbar_bottom = data_bottom cbar_bottom = data_bottom
cbar_height = data_height_frac - compass_height - compass_gap cbar_height = data_height_frac
if is_rgb: if is_rgb:
# RGB: descriptive text label instead of gradient colorbar # RGB: descriptive text label instead of gradient colorbar
cbar_ax = fig.add_axes([cbar_left, cbar_bottom, cbar_width, cbar_height]) cbar_ax = fig.add_axes([cbar_left, cbar_bottom, cbar_width, cbar_height])
@ -521,12 +523,14 @@ def tif_to_png(tif_file, vis_dir, resolution, keep_tif=False, source_info=None,
sm.set_array([]) sm.set_array([])
cbar = plt.colorbar(sm, cax=cbar_ax) cbar = plt.colorbar(sm, cax=cbar_ax)
cbar.ax.tick_params(labelsize=9, width=1.5) cbar.ax.tick_params(labelsize=9, width=1.5)
cbar.ax.yaxis.set_major_formatter(ScalarFormatter(useOffset=False))
cbar.outline.set_linewidth(1.5) cbar.outline.set_linewidth(1.5)
cbar.set_label(legend_label, fontsize=10, fontweight='bold') cbar.set_label(legend_label, fontsize=10, fontweight='bold')
else: else:
cbar_ax = fig.add_axes([cbar_left, cbar_bottom, cbar_width, cbar_height]) cbar_ax = fig.add_axes([cbar_left, cbar_bottom, cbar_width, cbar_height])
cbar = plt.colorbar(im, cax=cbar_ax) cbar = plt.colorbar(im, cax=cbar_ax)
cbar.ax.tick_params(labelsize=9, width=1.5) cbar.ax.tick_params(labelsize=9, width=1.5)
cbar.ax.yaxis.set_major_formatter(ScalarFormatter(useOffset=False))
cbar.outline.set_linewidth(1.5) cbar.outline.set_linewidth(1.5)
cbar.set_label(legend_label, fontsize=10, fontweight='bold') cbar.set_label(legend_label, fontsize=10, fontweight='bold')
@ -565,32 +569,35 @@ def tif_to_png(tif_file, vis_dir, resolution, keep_tif=False, source_info=None,
spine.set_color('black') spine.set_color('black')
spine.set_linewidth(0.8) spine.set_linewidth(0.8)
# North arrow — compass rose style, positioned above the colorbar # North arrow — compass rose style, inside the data area (top-right corner)
compass_bottom = data_bottom + data_height_frac + 0.02 # Semi-transparent background for readability over any data
compass_height = 0.07 north_ax = fig.add_axes([data_left + data_width_frac - 0.06,
compass_width = cbar_width + 0.03 data_bottom + data_height_frac - 0.12,
north_ax = fig.add_axes([cbar_left, compass_bottom, compass_width, compass_height]) 0.05, 0.10],
facecolor='none')
north_ax.set_xlim(-1.2, 1.2) north_ax.set_xlim(-1.2, 1.2)
north_ax.set_ylim(-0.5, 1.5) north_ax.set_ylim(-0.3, 1.5)
north_ax.axis('off') north_ax.axis('off')
north_ax.set_aspect('equal') north_ax.set_aspect('equal')
north_ax.set_facecolor('white') # Semi-transparent white background circle
circle_bg = plt.Circle((0, 0.5), 0.85, facecolor='white', edgecolor='#888888',
linewidth=0.5, alpha=0.7, zorder=1)
north_ax.add_patch(circle_bg)
# N arrow # N arrow
north_ax.annotate('N', xy=(0, 1.3), fontsize=11, fontweight='bold', north_ax.annotate('N', xy=(0, 1.3), fontsize=9, fontweight='bold',
ha='center', va='bottom', color='#b22222') ha='center', va='bottom', color='#b22222', zorder=10)
north_ax.plot([0, 0], [0.0, 1.0], color='#b22222', linewidth=2.0, zorder=10) north_ax.plot([0, 0], [0.0, 1.0], color='#b22222', linewidth=2.0, zorder=10)
north_ax.add_patch(MplPolygon([[0, 0.3], [-0.2, 0.7], [0, 1.0], [0.2, 0.7]], north_ax.add_patch(MplPolygon([[0, 0.3], [-0.2, 0.7], [0, 1.0], [0.2, 0.7]],
closed=True, facecolor='#b22222', edgecolor='#b22222', zorder=9)) closed=True, facecolor='#b22222', edgecolor='#b22222', zorder=9))
# Cardinal ticks # Cardinal ticks
for angle, label in [(90, ''), (0, 'E'), (180, 'O'), (270, 'S')]: for angle, label in [(90, ''), (0, 'E'), (180, 'O'), (270, 'S')]:
rad = np.radians(angle) rad = np.radians(angle)
r_text = 1.25
north_ax.plot([0.85*np.cos(rad), 1.05*np.cos(rad)], north_ax.plot([0.85*np.cos(rad), 1.05*np.cos(rad)],
[0.85*np.sin(rad), 1.05*np.sin(rad)], [0.85*np.sin(rad), 1.05*np.sin(rad)],
color='#555555', linewidth=0.8, zorder=5) color='#555555', linewidth=0.8, zorder=5)
if label: if label:
north_ax.text(r_text*np.cos(rad), r_text*np.sin(rad), label, north_ax.text(1.15*np.cos(rad), 1.15*np.sin(rad), label,
fontsize=7, ha='center', va='center', color='#555555') fontsize=6, ha='center', va='center', color='#555555', zorder=5)
# Bottom info bar — enriched with source, method, date # Bottom info bar — enriched with source, method, date
info_ax = fig.add_axes([data_left, 0.015, data_width_frac + cbar_width + 0.02, 0.09]) info_ax = fig.add_axes([data_left, 0.015, data_width_frac + cbar_width + 0.02, 0.09])