Coverage for src / quber / utils / cuda_runtime.py: 85%

1""" 

2Eagerly preload the CUDA 12 runtime shared libraries bundled by the 

3`nvidia-*-cu12` wheels so that `onnxruntime-gpu`'s `CUDAExecutionProvider` 

4can resolve them via `dlopen`. 

5 

6Why this exists 

7--------------- 

8`onnxruntime-gpu` (>=1.26.0) is built against CUDA 12. On hosts running 

9CUDA 13 (driver 580+) the system loader cannot find `libcublasLt.so.12`, 

10`libcudnn.so.9`, etc. — and `CUDAExecutionProvider` silently falls back 

11to CPU. The `nvidia-*-cu12` pip wheels install the libraries under 

12`<site-packages>/nvidia/<pkg>/lib/`, but the dynamic linker does not 

13search those paths by default and the wheels do not register them. 

14 

15`ctypes.CDLL(path, RTLD_GLOBAL)` at import time makes the libraries' 

16symbols available to subsequent `dlopen` calls from onnxruntime, without 

17requiring callers to export `LD_LIBRARY_PATH`. 

18 

19This is best-effort: a CPU-only environment without the `nvidia-*-cu12` 

20wheels installed will skip the preload silently. RapidOCR will then run 

21on CPU as before — extraction still works, just slower. 

22""" 

23 

24from __future__ import annotations 

25 

26import ctypes 

27import sys 

28from pathlib import Path 

29 

30_CUDA_LIB_SONAMES = ( 

31 "libcublasLt.so.12", 

32 "libcublas.so.12", 

33 "libcurand.so.10", 

34 "libcufft.so.11", 

35 "libcudart.so.12", 

36 "libcudnn.so.9", 

37) 

38 

39 

40def preload_cuda_runtime_libs() -> list[str]: 

41 """Preload CUDA 12 sonames from the `nvidia-*-cu12` pip wheels. 

42 

43 Returns the list of soname strings that were successfully loaded. 

44 Silent on missing wheels (CPU-only environments) and on individual 

45 `OSError`s during `dlopen`. 

46 """ 

47 site_packages = ( 

48 Path(sys.prefix) 

49 / "lib" 

50 / f"python{sys.version_info.major}.{sys.version_info.minor}" 

51 / "site-packages" 

52 ) 

53 nvidia_root = site_packages / "nvidia" 

54 if not nvidia_root.is_dir(): 

55 return [] 

56 

57 loaded: list[str] = [] 

58 for soname in _CUDA_LIB_SONAMES: 

59 for path in nvidia_root.rglob(soname): 

60 try: 

61 ctypes.CDLL(str(path), mode=ctypes.RTLD_GLOBAL) 

62 loaded.append(soname) 

63 break 

64 except OSError: 

65 continue 

66 return loaded