Coverage for orchestrant/pipeline/monitoring/system.py: 31%

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

2 

3from __future__ import annotations 

4 

5import psutil 

6from loguru import logger 

7 

8from orchestrant.monitoring.gpu import GPUProbe 

9from orchestrant.monitoring.snapshot import read_raw_snapshot 

10from orchestrant.pipeline.types import SystemStats 

11 

12 

13class SystemMonitor: 

14 """Live per-frame system snapshotter for pipeline UIs; keeps no history.""" 

15 

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

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

18 self._gpu = GPUProbe(gpu_device_id) 

19 

20 # Expose convenience attributes for backward compatibility 

21 self.gpu_available = self._gpu.available 

22 self.gpu_name = self._gpu.gpu_name 

23 

24 psutil.cpu_percent(interval=None) 

25 self.process = psutil.Process() 

26 self.process.cpu_percent() 

27 

28 @property 

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

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

31 return self._gpu._handle # noqa: SLF001 

32 

33 def get_stats(self) -> SystemStats: 

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

35 stats = SystemStats() 

36 

37 raw = read_raw_snapshot(self._gpu) 

38 stats.cpu_percent = raw.cpu_percent 

39 stats.ram_percent = raw.ram.percent 

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

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

42 

43 snapshot = raw.gpu 

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 # Relative to one logical core, so it can exceed 100% on multi-core hosts. 

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

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

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

67 } 

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

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

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

71 

72 def shutdown(self) -> None: 

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

74 self._gpu.shutdown()