Coverage for src / tracekit / core / memory_check.py: 94%
42 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-11 23:04 +0000
1"""Pre-flight memory checking for TraceKit operations.
3This module provides automatic memory verification before executing
4memory-intensive operations to prevent OOM crashes.
7Example:
8 >>> from tracekit.core.memory_check import check_operation_memory, require_memory
9 >>> check = check_operation_memory('spectrogram', samples=1e9, nperseg=4096)
10 >>> if not check.sufficient:
11 ... print(check.recommendation)
13References:
14 See tracekit.utils.memory for memory estimation functions.
15"""
17from __future__ import annotations
19from typing import TYPE_CHECKING, Any
21from tracekit.utils.memory import (
22 MemoryCheck,
23 MemoryCheckError,
24 check_memory_available,
25 require_memory,
26)
28if TYPE_CHECKING:
29 from collections.abc import Callable
32# Operations that automatically perform memory checks
33_AUTO_CHECK_OPERATIONS = {
34 "fft",
35 "psd",
36 "spectrogram",
37 "eye_diagram",
38 "correlate",
39 "filter",
40 "stft",
41 "cwt",
42 "dwt",
43}
45# Global flag to bypass memory checks (use with caution)
46_force_memory = False
49def set_force_memory(enabled: bool) -> None:
50 """Enable or disable forced memory bypass.
53 Args:
54 enabled: If True, bypass memory checks (dangerous).
56 Warning:
57 Bypassing memory checks can lead to system crashes.
58 Only use when you are certain the operation will succeed.
60 Example:
61 >>> set_force_memory(True) # Bypass all memory checks
62 >>> # ... perform operation ...
63 >>> set_force_memory(False) # Re-enable checks
64 """
65 global _force_memory
66 _force_memory = enabled
69def is_force_memory() -> bool:
70 """Check if memory checks are bypassed.
72 Returns:
73 True if memory checks are disabled.
74 """
75 return _force_memory
78def check_operation_memory(
79 operation: str,
80 samples: int | float | None = None,
81 **kwargs: Any,
82) -> MemoryCheck:
83 """Check if sufficient memory is available for an operation.
86 This is a convenience wrapper around utils.memory.check_memory_available
87 with automatic bypass support.
89 Args:
90 operation: Operation name (fft, psd, spectrogram, etc.).
91 samples: Number of samples to process.
92 **kwargs: Additional operation-specific parameters.
94 Returns:
95 MemoryCheck with sufficiency status and recommendations.
97 Example:
98 >>> check = check_operation_memory('fft', samples=1e9, nfft=8192)
99 >>> if not check.sufficient:
100 ... print(f"Insufficient memory: {check.recommendation}")
101 """
102 # Bypass check if forced
103 if _force_memory:
104 return MemoryCheck(
105 sufficient=True,
106 available=0,
107 required=0,
108 recommendation="Memory check bypassed (--force-memory enabled)",
109 )
111 return check_memory_available(operation, samples, **kwargs)
114def auto_check_memory(
115 operation: str,
116 samples: int | float | None = None,
117 **kwargs: Any,
118) -> None:
119 """Automatically check memory and raise error if insufficient.
122 This function is called automatically by operations that support
123 memory checking (fft, psd, spectrogram, etc.).
125 Args:
126 operation: Operation name.
127 samples: Number of samples.
128 **kwargs: Additional parameters.
130 Example:
131 >>> try:
132 ... auto_check_memory('spectrogram', samples=1e9, nperseg=4096)
133 ... except MemoryCheckError as e:
134 ... print(f"Memory check failed: {e}")
135 ... print(f"Suggestion: {e.recommendation}")
137 Note:
138 May raise MemoryCheckError if insufficient memory and not forced.
139 """
140 # Skip check if operation doesn't require it
141 if operation not in _AUTO_CHECK_OPERATIONS:
142 return
144 # Bypass if forced
145 if _force_memory:
146 return
148 # Perform check
149 require_memory(operation, samples, **kwargs)
152def with_memory_check(func: Callable) -> Callable: # type: ignore[type-arg]
153 """Decorator to add automatic memory checking to a function.
156 The decorated function must accept 'samples' as a keyword argument
157 and should have an 'operation' attribute or name that matches
158 a supported operation type.
160 Args:
161 func: Function to decorate.
163 Returns:
164 Decorated function with memory checking.
166 Example:
167 >>> @with_memory_check
168 ... def my_fft(signal, samples=None, **kwargs):
169 ... # ... FFT implementation ...
170 ... pass
171 >>> my_fft.operation = 'fft' # Specify operation type
172 """
174 def wrapper(*args: Any, **kwargs: Any) -> Any:
175 # Extract operation name
176 operation = getattr(func, "operation", func.__name__)
178 # Extract samples if available
179 samples = kwargs.get("samples")
180 if samples is None and len(args) > 0:
181 # Try to infer from first argument
182 try:
183 import numpy as np
185 if isinstance(args[0], np.ndarray): 185 ↛ 191line 185 didn't jump to line 191 because the condition on line 185 was always true
186 samples = len(args[0])
187 except (ImportError, TypeError):
188 pass
190 # Perform check
191 if operation in _AUTO_CHECK_OPERATIONS and not _force_memory:
192 auto_check_memory(operation, samples, **kwargs)
194 # Call original function
195 return func(*args, **kwargs)
197 # Preserve function metadata
198 wrapper.__name__ = func.__name__
199 wrapper.__doc__ = func.__doc__
200 wrapper.__module__ = func.__module__
202 return wrapper
205def register_auto_check_operation(operation: str) -> None:
206 """Register an operation for automatic memory checking.
208 Args:
209 operation: Operation name to register.
211 Example:
212 >>> register_auto_check_operation('custom_transform')
213 >>> # Now custom_transform will automatically check memory
214 """
215 _AUTO_CHECK_OPERATIONS.add(operation)
218def get_auto_check_operations() -> set[str]:
219 """Get set of operations that automatically check memory.
221 Returns:
222 Set of operation names.
224 Example:
225 >>> ops = get_auto_check_operations()
226 >>> print(f"Auto-checked operations: {', '.join(sorted(ops))}")
227 """
228 return _AUTO_CHECK_OPERATIONS.copy()
231__all__ = [
232 "MemoryCheck",
233 "MemoryCheckError",
234 "auto_check_memory",
235 "check_operation_memory",
236 "get_auto_check_operations",
237 "is_force_memory",
238 "register_auto_check_operation",
239 "set_force_memory",
240 "with_memory_check",
241]