Coverage for src / tracekit / config / memory.py: 100%

1"""Memory configuration module for TraceKit. 

2 

3This module provides global memory limit configuration and settings. 

4 

5 

6Example: 

7 >>> from tracekit.config.memory import set_memory_limit, get_memory_config 

8 >>> set_memory_limit("4GB") 

9 >>> config = get_memory_config() 

10 >>> print(f"Max memory: {config.max_memory / 1e9:.1f} GB") 

11 

12References: 

13 See tracekit.utils.memory for memory estimation and checking functions. 

14""" 

15 

16from __future__ import annotations 

17 

18import contextlib 

19import os 

20from dataclasses import dataclass 

21 

22 

23@dataclass 

24class MemoryConfiguration: 

25 """Global memory configuration for TraceKit operations. 

26 

27 

28 Attributes: 

29 max_memory: Global memory limit in bytes (None = auto-detect 80% available). 

30 warn_threshold: Memory pressure warning threshold (0.0-1.0). 

31 critical_threshold: Memory pressure critical threshold (0.0-1.0). 

32 auto_degrade: Automatically downsample if memory exceeded. 

33 memory_reserve: Reserved memory headroom in bytes. 

34 """ 

35 

36 max_memory: int | None = None 

37 warn_threshold: float = 0.7 

38 critical_threshold: float = 0.9 

39 auto_degrade: bool = False 

40 memory_reserve: int = 0 

41 

42 def __post_init__(self) -> None: 

43 """Validate configuration on initialization.""" 

44 if not 0.0 <= self.warn_threshold <= 1.0: 

45 raise ValueError( 

46 f"warn_threshold must be in range [0.0, 1.0], got {self.warn_threshold}" 

47 ) 

48 if not 0.0 <= self.critical_threshold <= 1.0: 

49 raise ValueError( 

50 f"critical_threshold must be in range [0.0, 1.0], got {self.critical_threshold}" 

51 ) 

52 if self.warn_threshold >= self.critical_threshold: 

53 raise ValueError( 

54 f"warn_threshold ({self.warn_threshold}) must be less than " 

55 f"critical_threshold ({self.critical_threshold})" 

56 ) 

57 if self.memory_reserve < 0: 

58 raise ValueError(f"memory_reserve must be non-negative, got {self.memory_reserve}") 

59 

60 

61# Global memory configuration instance 

62_global_config = MemoryConfiguration() 

63 

64 

65def get_memory_config() -> MemoryConfiguration: 

66 """Get the current global memory configuration. 

67 

68 Returns: 

69 Current MemoryConfiguration instance. 

70 

71 Example: 

72 >>> config = get_memory_config() 

73 >>> print(f"Max memory: {config.max_memory or 'auto'}") 

74 """ 

75 return _global_config 

76 

77 

78def set_memory_limit(limit: int | str | None) -> None: 

79 """Set global memory limit for all TraceKit operations. 

80 

81 

82 Args: 

83 limit: Memory limit as bytes (int), string ("4GB", "512MB"), or None for auto. 

84 

85 Example: 

86 >>> set_memory_limit("4GB") 

87 >>> set_memory_limit(4 * 1024**3) # 4 GB in bytes 

88 >>> set_memory_limit(None) # Auto (80% of available) 

89 

90 Environment: 

91 Can also be set via TK_MAX_MEMORY environment variable. 

92 """ 

93 global _global_config # noqa: PLW0602 

94 

95 if limit is None: 

96 _global_config.max_memory = None 

97 return 

98 

99 if isinstance(limit, str): 

100 _global_config.max_memory = _parse_memory_string(limit) 

101 else: 

102 _global_config.max_memory = int(limit) 

103 

104 

105def set_memory_thresholds( 

106 warn_threshold: float | None = None, 

107 critical_threshold: float | None = None, 

108) -> None: 

109 """Set memory pressure warning thresholds. 

110 

111 

112 Args: 

113 warn_threshold: Warning threshold (0.0-1.0, e.g., 0.7 for 70% utilization). 

114 critical_threshold: Critical threshold (0.0-1.0, e.g., 0.9 for 90% utilization). 

115 

116 Example: 

117 >>> set_memory_thresholds(warn_threshold=0.7, critical_threshold=0.9) 

118 >>> set_memory_thresholds(critical_threshold=0.95) # Keep warn unchanged 

119 """ 

120 global _global_config # noqa: PLW0602 

121 

122 if warn_threshold is not None: 

123 _global_config.warn_threshold = warn_threshold 

124 if critical_threshold is not None: 

125 _global_config.critical_threshold = critical_threshold 

126 

127 # Re-validate 

128 _global_config.__post_init__() 

129 

130 

131def enable_auto_degrade(enabled: bool = True) -> None: 

132 """Enable or disable automatic downsampling when memory limits exceeded. 

133 

134 

135 Args: 

136 enabled: Whether to enable automatic downsampling. 

137 

138 Example: 

139 >>> enable_auto_degrade(True) 

140 >>> # Operations will now auto-downsample if memory insufficient 

141 """ 

142 global _global_config # noqa: PLW0602 

143 _global_config.auto_degrade = enabled 

144 

145 

146def set_memory_reserve(reserve: int | str) -> None: 

147 """Set memory headroom to reserve (not use for operations). 

148 

149 

150 Args: 

151 reserve: Reserved memory as bytes (int) or string ("1GB", "512MB"). 

152 

153 Example: 

154 >>> set_memory_reserve("1GB") # Reserve 1 GB for system 

155 >>> set_memory_reserve(512 * 1024**2) # 512 MB 

156 """ 

157 global _global_config # noqa: PLW0602 

158 

159 if isinstance(reserve, str): 

160 _global_config.memory_reserve = _parse_memory_string(reserve) 

161 else: 

162 _global_config.memory_reserve = int(reserve) 

163 

164 

165def configure_from_environment() -> None: 

166 """Configure memory settings from environment variables. 

167 

168 

169 Environment Variables: 

170 TK_MAX_MEMORY: Maximum memory limit (e.g., "4GB", "512MB"). 

171 TK_MEMORY_RESERVE: Reserved memory headroom (e.g., "1GB"). 

172 TK_MEMORY_WARN_THRESHOLD: Warning threshold (e.g., "0.7"). 

173 TK_MEMORY_CRITICAL_THRESHOLD: Critical threshold (e.g., "0.9"). 

174 TK_AUTO_DEGRADE: Enable auto downsampling ("1", "true", "yes"). 

175 

176 Example: 

177 >>> import os 

178 >>> os.environ['TK_MAX_MEMORY'] = '4GB' 

179 >>> configure_from_environment() 

180 """ 

181 # Max memory 

182 if max_mem := os.environ.get("TK_MAX_MEMORY"): 

183 set_memory_limit(max_mem) 

184 

185 # Memory reserve 

186 if reserve := os.environ.get("TK_MEMORY_RESERVE"): 

187 set_memory_reserve(reserve) 

188 

189 # Thresholds 

190 warn = None 

191 critical = None 

192 if warn_str := os.environ.get("TK_MEMORY_WARN_THRESHOLD"): 

193 with contextlib.suppress(ValueError): 

194 warn = float(warn_str) 

195 if crit_str := os.environ.get("TK_MEMORY_CRITICAL_THRESHOLD"): 

196 with contextlib.suppress(ValueError): 

197 critical = float(crit_str) 

198 if warn is not None or critical is not None: 

199 set_memory_thresholds(warn, critical) 

200 

201 # Auto degrade 

202 if auto_str := os.environ.get("TK_AUTO_DEGRADE"): 

203 enable_auto_degrade(auto_str.lower() in ("1", "true", "yes", "on")) 

204 

205 

206def reset_to_defaults() -> None: 

207 """Reset memory configuration to default values. 

208 

209 Example: 

210 >>> reset_to_defaults() 

211 >>> config = get_memory_config() 

212 >>> assert config.max_memory is None # Auto-detect 

213 """ 

214 global _global_config 

215 _global_config = MemoryConfiguration() 

216 

217 

218def _parse_memory_string(limit_str: str) -> int: 

219 """Parse memory limit string to bytes. 

220 

221 Args: 

222 limit_str: Memory string like "4GB", "512MB", "1024KB". 

223 

224 Returns: 

225 Memory limit in bytes. 

226 

227 Raises: 

228 ValueError: If format is invalid. 

229 

230 Example: 

231 >>> _parse_memory_string("4GB") 

232 4000000000 

233 >>> _parse_memory_string("512MB") 

234 512000000 

235 """ 

236 limit_upper = limit_str.upper().strip() 

237 

238 try: 

239 if limit_upper.endswith("GB"): 

240 return int(float(limit_upper[:-2]) * 1e9) 

241 elif limit_upper.endswith("MB"): 

242 return int(float(limit_upper[:-2]) * 1e6) 

243 elif limit_upper.endswith("KB"): 

244 return int(float(limit_upper[:-2]) * 1e3) 

245 elif limit_upper.endswith("GIB"): 

246 return int(float(limit_upper[:-3]) * 1024**3) 

247 elif limit_upper.endswith("MIB"): 

248 return int(float(limit_upper[:-3]) * 1024**2) 

249 elif limit_upper.endswith("KIB"): 

250 return int(float(limit_upper[:-3]) * 1024) 

251 else: 

252 # Assume bytes 

253 return int(float(limit_upper)) 

254 except ValueError as e: 

255 raise ValueError(f"Invalid memory limit format: {limit_str}") from e 

256 

257 

258# Auto-configure from environment on import 

259configure_from_environment() 

260 

261 

262__all__ = [ 

263 "MemoryConfiguration", 

264 "configure_from_environment", 

265 "enable_auto_degrade", 

266 "get_memory_config", 

267 "reset_to_defaults", 

268 "set_memory_limit", 

269 "set_memory_reserve", 

270 "set_memory_thresholds", 

271]