Coverage for src / tracekit / core / memory_warnings.py: 100%

1"""Memory warning threshold system for TraceKit. 

2 

3This module provides configurable memory pressure warnings with 

4automatic alerts and optional operation cancellation. 

5 

6 

7Example: 

8 >>> from tracekit.core.memory_warnings import check_memory_warnings, MemoryWarningLevel 

9 >>> level = check_memory_warnings() 

10 >>> if level == MemoryWarningLevel.CRITICAL: 

11 ... print("Critical memory pressure!") 

12 

13References: 

14 See tracekit.config.memory for threshold configuration. 

15""" 

16 

17from __future__ import annotations 

18 

19import warnings 

20from enum import Enum 

21from typing import TYPE_CHECKING 

22 

23from tracekit.config.memory import get_memory_config 

24from tracekit.utils.memory import get_available_memory, get_memory_pressure 

25 

26if TYPE_CHECKING: 

27 from collections.abc import Callable 

28 

29 

30class MemoryWarningLevel(Enum): 

31 """Memory warning levels based on pressure thresholds. 

32 

33 Attributes: 

34 OK: Memory usage is normal (below warn threshold). 

35 WARNING: Memory usage is elevated (above warn threshold). 

36 CRITICAL: Memory usage is critical (above critical threshold). 

37 """ 

38 

39 OK = "ok" 

40 WARNING = "warning" 

41 CRITICAL = "critical" 

42 

43 

44def check_memory_warnings() -> MemoryWarningLevel: 

45 """Check current memory pressure against configured thresholds. 

46 

47 

48 Returns: 

49 Current memory warning level. 

50 

51 Example: 

52 >>> level = check_memory_warnings() 

53 >>> if level != MemoryWarningLevel.OK: 

54 ... print(f"Memory pressure: {level.value}") 

55 """ 

56 pressure = get_memory_pressure() 

57 config = get_memory_config() 

58 

59 if pressure >= config.critical_threshold: 

60 return MemoryWarningLevel.CRITICAL 

61 elif pressure >= config.warn_threshold: 

62 return MemoryWarningLevel.WARNING 

63 else: 

64 return MemoryWarningLevel.OK 

65 

66 

67def emit_memory_warning(force: bool = False) -> None: 

68 """Emit warning if memory pressure exceeds thresholds. 

69 

70 

71 Args: 

72 force: If True, always emit warning regardless of level. 

73 

74 Example: 

75 >>> emit_memory_warning() # Emits warning if pressure high 

76 """ 

77 level = check_memory_warnings() 

78 pressure = get_memory_pressure() 

79 available = get_available_memory() 

80 

81 if force or level != MemoryWarningLevel.OK: 

82 if level == MemoryWarningLevel.CRITICAL: 

83 warnings.warn( 

84 f"CRITICAL memory pressure: {pressure * 100:.1f}% utilized. " 

85 f"Only {available / 1e9:.2f} GB available. " 

86 "Consider closing applications or canceling operations.", 

87 ResourceWarning, 

88 stacklevel=2, 

89 ) 

90 elif level == MemoryWarningLevel.WARNING: 

91 warnings.warn( 

92 f"High memory pressure: {pressure * 100:.1f}% utilized. " 

93 f"{available / 1e9:.2f} GB available. " 

94 "Monitor memory usage closely.", 

95 ResourceWarning, 

96 stacklevel=2, 

97 ) 

98 

99 

100def check_and_abort_if_critical(operation: str = "operation") -> None: 

101 """Check memory and abort if critical threshold exceeded. 

102 

103 

104 Args: 

105 operation: Name of operation to abort. 

106 

107 Raises: 

108 MemoryError: If memory pressure is critical. 

109 

110 Example: 

111 >>> try: 

112 ... check_and_abort_if_critical('spectrogram') 

113 ... # ... perform operation ... 

114 ... except MemoryError: 

115 ... print("Operation aborted due to critical memory") 

116 """ 

117 level = check_memory_warnings() 

118 

119 if level == MemoryWarningLevel.CRITICAL: 

120 pressure = get_memory_pressure() 

121 available = get_available_memory() 

122 config = get_memory_config() 

123 

124 raise MemoryError( 

125 f"Critical memory pressure ({pressure * 100:.1f}% > {config.critical_threshold * 100:.0f}%). " 

126 f"Only {available / 1e9:.2f} GB available. " 

127 f"Operation '{operation}' aborted to prevent system crash. " 

128 "Free memory or increase critical threshold to proceed." 

129 ) 

130 

131 

132class MemoryWarningMonitor: 

133 """Context manager for monitoring memory warnings during operations. 

134 

135 

136 Monitors memory pressure during an operation and emits warnings 

137 when thresholds are crossed. 

138 

139 Example: 

140 >>> with MemoryWarningMonitor('spectrogram', check_interval=100): 

141 ... for i in range(1000): 

142 ... # ... perform work ... 

143 ... pass 

144 """ 

145 

146 def __init__( 

147 self, 

148 operation: str, 

149 *, 

150 check_interval: int = 100, 

151 abort_on_critical: bool = True, 

152 ): 

153 """Initialize memory warning monitor. 

154 

155 Args: 

156 operation: Name of operation being monitored. 

157 check_interval: Check memory every N iterations. 

158 abort_on_critical: Abort if critical threshold reached. 

159 """ 

160 self.operation = operation 

161 self.check_interval = check_interval 

162 self.abort_on_critical = abort_on_critical 

163 self._iteration = 0 

164 self._warned = False 

165 self._critical_warned = False 

166 

167 def __enter__(self) -> MemoryWarningMonitor: 

168 """Enter context and perform initial check.""" 

169 # Initial check 

170 emit_memory_warning() 

171 return self 

172 

173 def __exit__(self, exc_type: type, exc_val: Exception, exc_tb: object) -> None: 

174 """Exit context.""" 

175 # Note: exc_val and exc_tb intentionally unused but required for Python 3.11+ compatibility 

176 pass 

177 

178 def check(self, iteration: int | None = None) -> None: 

179 """Check memory pressure and emit warnings. 

180 

181 Args: 

182 iteration: Current iteration number (for periodic checking). 

183 

184 Raises: 

185 MemoryError: If critical threshold exceeded and abort_on_critical=True. 

186 """ 

187 self._iteration += 1 

188 

189 # Only check periodically 

190 if iteration is not None and iteration % self.check_interval != 0: 

191 return 

192 

193 level = check_memory_warnings() 

194 pressure = get_memory_pressure() 

195 available = get_available_memory() 

196 

197 if level == MemoryWarningLevel.CRITICAL: 

198 if not self._critical_warned: 

199 warnings.warn( 

200 f"Critical memory pressure during {self.operation}: " 

201 f"{pressure * 100:.1f}% utilized, {available / 1e9:.2f} GB available.", 

202 ResourceWarning, 

203 stacklevel=2, 

204 ) 

205 self._critical_warned = True 

206 

207 if self.abort_on_critical: 

208 raise MemoryError( 

209 f"Critical memory pressure during {self.operation}. " 

210 f"Operation aborted to prevent system crash. " 

211 f"Pressure: {pressure * 100:.1f}%, Available: {available / 1e9:.2f} GB" 

212 ) 

213 

214 elif level == MemoryWarningLevel.WARNING: 

215 if not self._warned: 

216 warnings.warn( 

217 f"High memory pressure during {self.operation}: " 

218 f"{pressure * 100:.1f}% utilized, {available / 1e9:.2f} GB available.", 

219 ResourceWarning, 

220 stacklevel=2, 

221 ) 

222 self._warned = True 

223 

224 

225def register_memory_warning_callback(callback: Callable[[MemoryWarningLevel], None]) -> None: 

226 """Register a callback for memory warnings. 

227 

228 The callback will be invoked whenever memory warnings are checked 

229 via emit_memory_warning() or MemoryWarningMonitor. 

230 

231 Args: 

232 callback: Function that accepts MemoryWarningLevel. 

233 

234 Example: 

235 >>> def my_callback(level): 

236 ... if level == MemoryWarningLevel.CRITICAL: 

237 ... print("CRITICAL MEMORY!") 

238 >>> register_memory_warning_callback(my_callback) 

239 """ 

240 _warning_callbacks.append(callback) 

241 

242 

243def clear_memory_warning_callbacks() -> None: 

244 """Clear all registered memory warning callbacks. 

245 

246 Example: 

247 >>> clear_memory_warning_callbacks() 

248 """ 

249 _warning_callbacks.clear() 

250 

251 

252# Global list of warning callbacks 

253_warning_callbacks: list[Callable[[MemoryWarningLevel], None]] = [] 

254 

255 

256def _invoke_callbacks(level: MemoryWarningLevel) -> None: 

257 """Invoke all registered callbacks with current warning level. 

258 

259 Args: 

260 level: Current memory warning level. 

261 """ 

262 for callback in _warning_callbacks: 

263 try: 

264 callback(level) 

265 except Exception as e: 

266 warnings.warn( 

267 f"Memory warning callback failed: {e}", 

268 RuntimeWarning, 

269 stacklevel=2, 

270 ) 

271 

272 

273def format_memory_warning(level: MemoryWarningLevel) -> str: 

274 """Format memory warning message. 

275 

276 Args: 

277 level: Memory warning level. 

278 

279 Returns: 

280 Formatted warning message. 

281 

282 Example: 

283 >>> msg = format_memory_warning(MemoryWarningLevel.WARNING) 

284 >>> print(msg) 

285 High memory pressure: 75.0% utilized... 

286 """ 

287 pressure = get_memory_pressure() 

288 available = get_available_memory() 

289 config = get_memory_config() 

290 

291 if level == MemoryWarningLevel.CRITICAL: 

292 return ( 

293 f"CRITICAL memory pressure: {pressure * 100:.1f}% utilized " 

294 f"(threshold: {config.critical_threshold * 100:.0f}%). " 

295 f"Only {available / 1e9:.2f} GB available." 

296 ) 

297 elif level == MemoryWarningLevel.WARNING: 

298 return ( 

299 f"High memory pressure: {pressure * 100:.1f}% utilized " 

300 f"(threshold: {config.warn_threshold * 100:.0f}%). " 

301 f"{available / 1e9:.2f} GB available." 

302 ) 

303 else: 

304 return f"Memory OK: {pressure * 100:.1f}% utilized, {available / 1e9:.2f} GB available." 

305 

306 

307__all__ = [ 

308 "MemoryWarningLevel", 

309 "MemoryWarningMonitor", 

310 "check_and_abort_if_critical", 

311 "check_memory_warnings", 

312 "clear_memory_warning_callbacks", 

313 "emit_memory_warning", 

314 "format_memory_warning", 

315 "register_memory_warning_callback", 

316]