Coverage for orchestr_ant_ion / pipeline / monitoring / system.py: 29%

1"""System resource monitoring helpers.""" 

2 

3from __future__ import annotations 

4 

5import psutil 

6from loguru import logger 

7 

8from orchestr_ant_ion.monitoring.gpu import GPUProbe 

9from orchestr_ant_ion.pipeline.types import SystemStats 

10 

11 

12class SystemMonitor: 

13 """Monitor system resources including CPU, RAM, and GPU.""" 

14 

15 def __init__(self, gpu_device_id: int = 0) -> None: 

16 """Initialize system monitoring and optional GPU probing.""" 

17 self._gpu = GPUProbe(gpu_device_id) 

18 

19 # Expose convenience attributes for backward compatibility 

20 self.gpu_available = self._gpu.available 

21 self.gpu_name = self._gpu.gpu_name 

22 

23 psutil.cpu_percent(interval=None) 

24 self.process = psutil.Process() 

25 self.process.cpu_percent() 

26 

27 @property 

28 def gpu_handle(self) -> object | None: 

29 """Return the GPU handle if initialized.""" 

30 return self._gpu._handle # noqa: SLF001 

31 

32 def get_stats(self) -> SystemStats: 

33 """Collect current system-wide CPU, RAM, and GPU statistics.""" 

34 stats = SystemStats() 

35 

36 stats.cpu_percent = psutil.cpu_percent(interval=None) 

37 

38 ram = psutil.virtual_memory() 

39 stats.ram_percent = ram.percent 

40 stats.ram_used_gb = ram.used / (1024**3) 

41 stats.ram_total_gb = ram.total / (1024**3) 

42 

43 snapshot = self._gpu.read() 

44 if snapshot is not None: 

45 stats.gpu_percent = snapshot.utilization 

46 stats.gpu_memory_used_gb = snapshot.memory_used_bytes / (1024**3) 

47 stats.gpu_memory_total_gb = snapshot.memory_total_bytes / (1024**3) 

48 stats.gpu_memory_percent = ( 

49 (snapshot.memory_used_bytes / snapshot.memory_total_bytes) * 100 

50 if snapshot.memory_total_bytes > 0 

51 else 0 

52 ) 

53 stats.gpu_temp_celsius = snapshot.temperature_celsius 

54 stats.gpu_power_watts = snapshot.power_watts 

55 stats.gpu_name = self._gpu.gpu_name 

56 

57 return stats 

58 

59 def get_process_stats(self) -> dict: 

60 """Collect current process CPU/RAM/thread statistics.""" 

61 try: 

62 return { 

63 # psutil returns process CPU% relative to a single logical core; 

64 # values can exceed 100% on multi-core systems. 

65 "cpu_percent": self.process.cpu_percent(), 

66 "memory_mb": self.process.memory_info().rss / (1024**2), 

67 "threads": self.process.num_threads(), 

68 } 

69 except (psutil.NoSuchProcess, psutil.AccessDenied) as exc: 

70 logger.debug("Failed to read process stats: {}", exc) 

71 return {"cpu_percent": 0, "memory_mb": 0, "threads": 0} 

72 

73 def shutdown(self) -> None: 

74 """Shutdown GPU monitoring resources, if available.""" 

75 self._gpu.shutdown()