Coverage for presidio_analyzer / nlp_engine / device_detector.py: 100%

1"""GPU/CPU device detection for Presidio NLP engines. 

2 

3This module provides lazy device detection using module-level __getattr__. 

4Access via `device_detector` attribute which is lazily initialized on first use. 

5 

6The device can be explicitly set via the PRESIDIO_DEVICE environment variable. 

7Any valid PyTorch device string is accepted (e.g., 'cpu', 'cuda', 'cuda:0', 'cuda:1'). 

8 

9If PRESIDIO_DEVICE is not set, automatic detection is performed. 

10 

11Usage: 

12 from presidio_analyzer.nlp_engine.device_detector import device_detector 

13 device = device_detector.get_device() 

14""" 

15 

16import logging 

17import os 

18import threading 

19 

20logger = logging.getLogger("presidio-analyzer") 

21 

22PRESIDIO_DEVICE_ENV_VAR = "PRESIDIO_DEVICE" 

23 

24_detector_instance = None 

25_lock = threading.Lock() 

26 

27 

28class DeviceDetector: 

29 """Detect and expose PyTorch device availability. 

30 

31 Device selection priority: 

32 1. PRESIDIO_DEVICE environment variable (any valid device string) 

33 2. Automatic detection: Prefer CUDA when available, otherwise CPU 

34 

35 Note: MPS (Apple Silicon/Metal) is currently not supported. 

36 """ 

37 

38 def __init__(self) -> None: 

39 self._device = self._get_device_from_env() or self._detect() 

40 logger.info(f"Using device of type: {self._device}") 

41 

42 def _get_device_from_env(self) -> str | None: 

43 """Get device from PRESIDIO_DEVICE environment variable. 

44 

45 Returns 

46 Device string if env var is set, None otherwise. 

47 """ 

48 env_device = os.environ.get(PRESIDIO_DEVICE_ENV_VAR, "").strip() 

49 if not env_device: 

50 return None 

51 

52 logger.info(f"Device set to '{env_device}' via {PRESIDIO_DEVICE_ENV_VAR}") 

53 return env_device 

54 

55 def _detect(self) -> str: 

56 """Auto-detect PyTorch CUDA support. 

57 

58 Returns 

59 'cuda' if available, 'cpu' otherwise. 

60 """ 

61 try: 

62 import torch 

63 except ImportError: 

64 # torch not installed - this is expected, silently fall back to CPU 

65 return "cpu" 

66 

67 try: 

68 if torch.cuda.is_available(): 

69 _ = str(torch.tensor([1.0], device="cuda")) 

70 _ = torch.cuda.get_device_name(0) 

71 torch.cuda.get_device_capability(0) 

72 torch.cuda.empty_cache() 

73 return "cuda" 

74 except Exception as e: 

75 logger.warning(f"CUDA device detection failed, falling back to CPU: {e}") 

76 

77 return "cpu" 

78 

79 def get_device(self) -> str: 

80 """Return device string ('cuda' or 'cpu').""" 

81 return self._device 

82 

83 

84def __getattr__(name: str): 

85 """Lazily initialize device_detector on first access.""" 

86 if name == "device_detector": 

87 global _detector_instance 

88 if _detector_instance is None: 

89 with _lock: 

90 _detector_instance = _detector_instance or DeviceDetector() 

91 return _detector_instance 

92 raise AttributeError(f"module {__name__} has no attribute {name}")