Coverage for agentos/tools/startup_accelerator.py: 0%

1""" 

2Startup Acceleration Tools for AgentOS. 

3Lazy loading, module pre-compilation, and startup sequence optimization. 

4""" 

5 

6import importlib 

7import threading 

8import time 

9import types 

10from collections import OrderedDict 

11from dataclasses import dataclass, field 

12from typing import Any 

13 

14# ============================================================================ 

15# LazyLoader 

16# ============================================================================ 

17 

18 

19class _LazyModule: 

20 """Proxy that defers module import until an attribute is accessed.""" 

21 

22 def __init__(self, module_name: str): 

23 self._module_name = module_name 

24 self._module: types.ModuleType | None = None 

25 

26 def _load(self) -> types.ModuleType: 

27 if self._module is None: 

28 self._module = importlib.import_module(self._module_name) 

29 return self._module 

30 

31 def __getattr__(self, name: str) -> Any: 

32 return getattr(self._load(), name) 

33 

34 def __repr__(self) -> str: 

35 if self._module is None: 

36 return f"<LazyModule: {self._module_name} (unloaded)>" 

37 return repr(self._module) 

38 

39 

40class LazyLoader: 

41 """Registry for lazy-loaded modules with batch loading and dependency tracking.""" 

42 

43 def __init__(self): 

44 self._proxies: dict[str, _LazyModule] = {} 

45 self._load_times: dict[str, float] = {} 

46 self._lock = threading.Lock() 

47 

48 def register(self, module_name: str) -> _LazyModule: 

49 """Register a module for lazy loading.""" 

50 with self._lock: 

51 if module_name not in self._proxies: 

52 self._proxies[module_name] = _LazyModule(module_name) 

53 return self._proxies[module_name] 

54 

55 def load_now(self, module_name: str) -> types.ModuleType: 

56 """Eagerly load a registered module.""" 

57 proxy = self.register(module_name) 

58 with self._lock: 

59 start = time.perf_counter() 

60 result = proxy._load() 

61 elapsed = time.perf_counter() - start 

62 self._load_times[module_name] = elapsed 

63 return result 

64 

65 def load_all(self) -> list[tuple[str, float]]: 

66 """Eagerly load all registered modules. Returns load times.""" 

67 results: list[tuple[str, float]] = [] 

68 for name in list(self._proxies.keys()): 

69 start = time.perf_counter() 

70 self._proxies[name]._load() 

71 elapsed = time.perf_counter() - start 

72 self._load_times[name] = elapsed 

73 results.append((name, elapsed)) 

74 return results 

75 

76 def preload(self, module_names: list[str]) -> list[tuple[str, float]]: 

77 """Register and load a batch of modules.""" 

78 for name in module_names: 

79 self.register(name) 

80 results: list[tuple[str, float]] = [] 

81 for name in module_names: 

82 start = time.perf_counter() 

83 self._proxies[name]._load() 

84 elapsed = time.perf_counter() - start 

85 self._load_times[name] = elapsed 

86 results.append((name, elapsed)) 

87 return results 

88 

89 @property 

90 def stats(self) -> dict[str, Any]: 

91 with self._lock: 

92 loaded = {k: v for k, v in self._proxies.items() if v._module is not None} 

93 return { 

94 "registered": len(self._proxies), 

95 "loaded": len(loaded), 

96 "unloaded": len(self._proxies) - len(loaded), 

97 "load_times": dict(self._load_times), 

98 "total_load_time": sum(self._load_times.values()), 

99 } 

100 

101 def __getitem__(self, module_name: str) -> _LazyModule: 

102 return self.register(module_name) 

103 

104 

105# ============================================================================ 

106# ModulePreloader 

107# ============================================================================ 

108 

109 

110class ModulePreloader: 

111 """Pre-compile and cache frequently used modules for fast startup.""" 

112 

113 def __init__(self, max_concurrent: int = 4): 

114 self._cache: dict[str, types.ModuleType] = {} 

115 self._max_concurrent = max_concurrent 

116 self._lock = threading.Lock() 

117 self._preload_times: dict[str, float] = {} 

118 

119 def precompile(self, module_names: list[str], parallel: bool = True) -> dict[str, float]: 

120 """Pre-compile a list of modules, optionally in parallel.""" 

121 results: dict[str, float] = {} 

122 

123 if parallel and len(module_names) > 1: 

124 results = self._precompile_parallel(module_names) 

125 else: 

126 for name in module_names: 

127 start = time.perf_counter() 

128 self._cache[name] = importlib.import_module(name) 

129 elapsed = time.perf_counter() - start 

130 results[name] = elapsed 

131 

132 self._preload_times.update(results) 

133 return results 

134 

135 def _precompile_parallel(self, module_names: list[str]) -> dict[str, float]: 

136 results: dict[str, float] = {} 

137 errors: list[str] = [] 

138 

139 def _worker(name: str) -> None: 

140 try: 

141 start = time.perf_counter() 

142 module = importlib.import_module(name) 

143 elapsed = time.perf_counter() - start 

144 with self._lock: 

145 self._cache[name] = module 

146 results[name] = elapsed 

147 except Exception as e: 

148 errors.append(f"{name}: {e}") 

149 

150 # Batch into groups 

151 for i in range(0, len(module_names), self._max_concurrent): 

152 batch = module_names[i : i + self._max_concurrent] 

153 threads = [threading.Thread(target=_worker, args=(name,)) for name in batch] 

154 for t in threads: 

155 t.start() 

156 for t in threads: 

157 t.join() 

158 

159 return results 

160 

161 def get(self, module_name: str) -> types.ModuleType | None: 

162 return self._cache.get(module_name) 

163 

164 def warm_cache(self, hot_modules: list[str]) -> int: 

165 """Preload hot modules into cache. Returns number newly cached.""" 

166 count = 0 

167 for name in hot_modules: 

168 if name not in self._cache: 

169 try: 

170 self._cache[name] = importlib.import_module(name) 

171 count += 1 

172 except Exception: 

173 pass 

174 return count 

175 

176 def clear(self) -> None: 

177 with self._lock: 

178 self._cache.clear() 

179 

180 @property 

181 def stats(self) -> dict[str, Any]: 

182 with self._lock: 

183 return { 

184 "cached_modules": len(self._cache), 

185 "total_preload_time": sum(self._preload_times.values()), 

186 "module_times": dict(self._preload_times), 

187 } 

188 

189 

190# ============================================================================ 

191# StartupOptimizer 

192# ============================================================================ 

193 

194 

195@dataclass 

196class _StartupPhase: 

197 name: str 

198 start_time: float = 0.0 

199 end_time: float = 0.0 

200 metadata: dict[str, Any] = field(default_factory=dict) 

201 

202 @property 

203 def duration(self) -> float: 

204 return self.end_time - self.start_time 

205 

206 

207class StartupOptimizer: 

208 """Profile and optimize application startup sequence.""" 

209 

210 def __init__(self): 

211 self._phases: OrderedDict[str, _StartupPhase] = OrderedDict() 

212 self._lock = threading.Lock() 

213 self._total_start: float = 0.0 

214 self._total_end: float = 0.0 

215 

216 def start(self) -> None: 

217 """Mark the beginning of the startup sequence.""" 

218 self._total_start = time.perf_counter() 

219 

220 def begin_phase(self, name: str, **metadata) -> None: 

221 """Begin profiling a startup phase.""" 

222 with self._lock: 

223 phase = _StartupPhase(name=name, start_time=time.perf_counter(), metadata=metadata) 

224 self._phases[name] = phase 

225 

226 def end_phase(self, name: str) -> float | None: 

227 """End profiling a startup phase. Returns duration.""" 

228 with self._lock: 

229 phase = self._phases.get(name) 

230 if phase: 

231 phase.end_time = time.perf_counter() 

232 return phase.duration 

233 return None 

234 

235 def end(self) -> None: 

236 """Mark the end of the startup sequence.""" 

237 self._total_end = time.perf_counter() 

238 

239 def report(self) -> dict[str, Any]: 

240 """Generate a startup performance report.""" 

241 phases = [] 

242 for name, phase in self._phases.items(): 

243 phases.append( 

244 { 

245 "name": name, 

246 "duration_ms": round(phase.duration * 1000, 2), 

247 "pct_of_total": 0.0, 

248 **phase.metadata, 

249 } 

250 ) 

251 

252 total_duration = self._total_end - self._total_start 

253 total_ms = round(total_duration * 1000, 2) 

254 

255 for p in phases: 

256 if total_ms > 0: 

257 p["pct_of_total"] = round(p["duration_ms"] / total_ms * 100, 1) 

258 

259 sorted_phases = sorted(phases, key=lambda x: x["duration_ms"], reverse=True) 

260 return { 

261 "total_duration_ms": total_ms, 

262 "phase_count": len(phases), 

263 "phases": sorted_phases, 

264 "bottleneck": sorted_phases[0]["name"] if sorted_phases else None, 

265 } 

266 

267 def total_duration_ms(self) -> float: 

268 return round((self._total_end - self._total_start) * 1000, 2) 

269 

270 

271# ============================================================================ 

272# Convenience Functions 

273# ============================================================================ 

274 

275 

276def create_lazy_loader() -> LazyLoader: 

277 """Create a lazy module loader.""" 

278 return LazyLoader() 

279 

280 

281def create_module_preloader(max_concurrent: int = 4) -> ModulePreloader: 

282 """Create a module preloader for startup acceleration.""" 

283 return ModulePreloader(max_concurrent=max_concurrent) 

284 

285 

286def create_startup_optimizer() -> StartupOptimizer: 

287 """Create a startup sequence profiler and optimizer.""" 

288 return StartupOptimizer() 

289 

290 

291def quick_start( 

292 essential_modules: list[str], 

293 lazy_modules: list[str], 

294 hot_modules: list[str] | None = None, 

295) -> dict[str, Any]: 

296 """One-shot startup optimization: preload essentials, lazy-load the rest.""" 

297 preloader = ModulePreloader() 

298 loader = LazyLoader() 

299 

300 # Preload essential modules 

301 essential_times = preloader.precompile(essential_modules, parallel=True) 

302 

303 # Register lazy modules 

304 for name in lazy_modules: 

305 loader.register(name) 

306 

307 # Warm cache with hot modules (extras) 

308 if hot_modules: 

309 preloader.warm_cache(hot_modules) 

310 

311 return { 

312 "essential": essential_times, 

313 "lazy_count": len(lazy_modules), 

314 "cached_total": len(preloader._cache), 

315 "total_essential_ms": round(sum(essential_times.values()) * 1000, 2), 

316 }