Coverage for python/lsst/utils/timer.py: 16%

134 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2023-06-21 09:53 +0000

1# This file is part of utils. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (https://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# Use of this source code is governed by a 3-clause BSD-style 

10# license that can be found in the LICENSE file. 

11# 

12 

13"""Utilities for measuring execution time. 

14""" 

15 

16from __future__ import annotations 

17 

18__all__ = ["profile", "logInfo", "timeMethod", "time_this"] 

19 

20import datetime 

21import functools 

22import logging 

23import time 

24from contextlib import contextmanager 

25from typing import ( 

26 TYPE_CHECKING, 

27 Any, 

28 Callable, 

29 Collection, 

30 Iterable, 

31 Iterator, 

32 MutableMapping, 

33 Optional, 

34 Tuple, 

35) 

36 

37from astropy import units as u 

38 

39from .introspection import find_outside_stacklevel 

40from .usage import _get_current_rusage, get_current_mem_usage, get_peak_mem_usage 

41 

42if TYPE_CHECKING: 

43 import cProfile 

44 

45 from .logging import LsstLoggers 

46 

47 

48_LOG = logging.getLogger(__name__) 

49 

50 

51def _add_to_metadata(metadata: MutableMapping, name: str, value: Any) -> None: 

52 """Add a value to dict-like object, creating list as needed. 

53 

54 The list grows as more values are added for that key. 

55 

56 Parameters 

57 ---------- 

58 metadata : `dict`-like, optional 

59 `dict`-like object that can store keys. Uses `add()` method if 

60 one is available, else creates list and appends value if needed. 

61 name : `str` 

62 The key to use in the metadata dictionary. 

63 value : Any 

64 Value to store in the list. 

65 """ 

66 try: 

67 try: 

68 # PropertySet should always prefer LongLong for integers 

69 metadata.addLongLong(name, value) # type: ignore 

70 except (TypeError, AttributeError): 

71 metadata.add(name, value) # type: ignore 

72 except AttributeError: 

73 pass 

74 else: 

75 return 

76 

77 # Fallback code where `add` is not implemented. 

78 if name not in metadata: 

79 metadata[name] = [] 

80 metadata[name].append(value) 

81 

82 

83def logPairs( 

84 obj: Any, 

85 pairs: Collection[Tuple[str, Any]], 

86 logLevel: int = logging.DEBUG, 

87 metadata: Optional[MutableMapping] = None, 

88 logger: Optional[logging.Logger] = None, 

89 stacklevel: Optional[int] = None, 

90) -> None: 

91 """Log ``(name, value)`` pairs to ``obj.metadata`` and ``obj.log`` 

92 

93 Parameters 

94 ---------- 

95 obj : `object` 

96 An object with one or both of these two attributes: 

97 

98 - ``metadata`` a `dict`-like container for storing metadata. 

99 Can use the ``add(name, value)`` method if found, else will append 

100 entries to a list. 

101 - ``log`` an instance of `logging.Logger` or subclass. 

102 

103 If `None`, at least one of ``metadata`` or ``logger`` should be passed 

104 or this function will do nothing. 

105 pairs : sequence 

106 A sequence of ``(name, value)`` pairs, with value typically numeric. 

107 logLevel : `int, optional 

108 Log level (an `logging` level constant, such as `logging.DEBUG`). 

109 metadata : `collections.abc.MutableMapping`, optional 

110 Metadata object to write entries to. Ignored if `None`. 

111 logger : `logging.Logger` 

112 Log object to write entries to. Ignored if `None`. 

113 stacklevel : `int`, optional 

114 The stack level to pass to the logger to adjust which stack frame 

115 is used to report the file information. If `None` the stack level 

116 is computed such that it is reported as the first package outside 

117 of the utils package. If a value is given here it is adjusted by 

118 1 to account for this caller. 

119 """ 

120 if obj is not None: 

121 if metadata is None: 

122 try: 

123 metadata = obj.metadata 

124 except AttributeError: 

125 pass 

126 if logger is None: 

127 try: 

128 logger = obj.log 

129 except AttributeError: 

130 pass 

131 strList = [] 

132 for name, value in pairs: 

133 if metadata is not None: 

134 _add_to_metadata(metadata, name, value) 

135 strList.append(f"{name}={value}") 

136 if logger is not None: 

137 # Want the file associated with this log message to be that 

138 # of the caller. This is expensive so only do it if we know the 

139 # message will be issued. 

140 timer_logger = logging.getLogger("timer." + logger.name) 

141 if timer_logger.isEnabledFor(logLevel): 

142 if stacklevel is None: 

143 stacklevel = find_outside_stacklevel("lsst.utils") 

144 else: 

145 # Account for the caller stack. 

146 stacklevel += 1 

147 timer_logger.log(logLevel, "; ".join(strList), stacklevel=stacklevel) 

148 

149 

150def logInfo( 

151 obj: Any, 

152 prefix: str, 

153 logLevel: int = logging.DEBUG, 

154 metadata: Optional[MutableMapping] = None, 

155 logger: Optional[logging.Logger] = None, 

156 stacklevel: Optional[int] = None, 

157) -> None: 

158 """Log timer information to ``obj.metadata`` and ``obj.log``. 

159 

160 Parameters 

161 ---------- 

162 obj : `object` 

163 An object with both or one these two attributes: 

164 

165 - ``metadata`` a `dict`-like container for storing metadata. 

166 Can use the ``add(name, value)`` method if found, else will append 

167 entries to a list. 

168 - ``log`` an instance of `logging.Logger` or subclass. 

169 

170 If `None`, at least one of ``metadata`` or ``logger`` should be passed 

171 or this function will do nothing. 

172 prefix : `str` 

173 Name prefix, the resulting entries are ``CpuTime``, etc.. For example 

174 `timeMethod` uses ``prefix = Start`` when the method begins and 

175 ``prefix = End`` when the method ends. 

176 logLevel : optional 

177 Log level (a `logging` level constant, such as `logging.DEBUG`). 

178 metadata : `collections.abc.MutableMapping`, optional 

179 Metadata object to write entries to, overriding ``obj.metadata``. 

180 logger : `logging.Logger` 

181 Log object to write entries to, overriding ``obj.log``. 

182 stacklevel : `int`, optional 

183 The stack level to pass to the logger to adjust which stack frame 

184 is used to report the file information. If `None` the stack level 

185 is computed such that it is reported as the first package outside 

186 of the utils package. If a value is given here it is adjusted by 

187 1 to account for this caller. 

188 

189 Notes 

190 ----- 

191 Logged items include: 

192 

193 - ``Utc``: UTC date in ISO format (only in metadata since log entries have 

194 timestamps). 

195 - ``CpuTime``: System + User CPU time (seconds). This should only be used 

196 in differential measurements; the time reference point is undefined. 

197 - ``MaxRss``: maximum resident set size. Always in bytes. 

198 

199 All logged resource information is only for the current process; child 

200 processes are excluded. 

201 

202 The metadata will be updated with a ``__version__`` field to indicate the 

203 version of the items stored. If there is no version number it is assumed 

204 to be version 0. 

205 

206 * Version 0: ``MaxResidentSetSize`` units are platform-dependent. 

207 * Version 1: ``MaxResidentSetSize`` will be stored in bytes. 

208 """ 

209 if metadata is None and obj is not None: 

210 try: 

211 metadata = obj.metadata 

212 except AttributeError: 

213 pass 

214 if metadata is not None: 

215 # Log messages already have timestamps. 

216 utcStr = datetime.datetime.utcnow().isoformat() 

217 _add_to_metadata(metadata, name=prefix + "Utc", value=utcStr) 

218 

219 # Force a version number into the metadata. 

220 # v1: Ensure that max_rss field is always bytes. 

221 metadata["__version__"] = 1 

222 if stacklevel is not None: 

223 # Account for the caller of this routine not knowing that we 

224 # are going one down in the stack. 

225 stacklevel += 1 

226 

227 usage = _get_current_rusage() 

228 logPairs( 

229 obj=obj, 

230 pairs=[(prefix + k[0].upper() + k[1:], v) for k, v in usage.dict().items()], 

231 logLevel=logLevel, 

232 metadata=metadata, 

233 logger=logger, 

234 stacklevel=stacklevel, 

235 ) 

236 

237 

238def timeMethod( 

239 _func: Optional[Any] = None, 

240 *, 

241 metadata: Optional[MutableMapping] = None, 

242 logger: Optional[logging.Logger] = None, 

243 logLevel: int = logging.DEBUG, 

244) -> Callable: 

245 """Measure duration of a method. 

246 

247 Parameters 

248 ---------- 

249 func 

250 The method to wrap. 

251 metadata : `collections.abc.MutableMapping`, optional 

252 Metadata to use as override if the instance object attached 

253 to this timer does not support a ``metadata`` property. 

254 logger : `logging.Logger`, optional 

255 Logger to use when the class containing the decorated method does not 

256 have a ``log`` property. 

257 logLevel : `int`, optional 

258 Log level to use when logging messages. Default is `~logging.DEBUG`. 

259 

260 Notes 

261 ----- 

262 Writes various measures of time and possibly memory usage to the 

263 metadata; all items are prefixed with the function name. 

264 

265 .. warning:: 

266 

267 This decorator only works with instance methods of any class 

268 with these attributes: 

269 

270 - ``metadata``: an instance of `collections.abc.Mapping`. The ``add`` 

271 method will be used if available, else entries will be added to a 

272 list. 

273 - ``log``: an instance of `logging.Logger` or subclass. 

274 

275 Examples 

276 -------- 

277 To use: 

278 

279 .. code-block:: python 

280 

281 import lsst.utils as utils 

282 import lsst.pipe.base as pipeBase 

283 class FooTask(pipeBase.Task): 

284 pass 

285 

286 @utils.timeMethod 

287 def run(self, ...): # or any other instance method you want to time 

288 pass 

289 """ 

290 

291 def decorator_timer(func: Callable) -> Callable: 

292 @functools.wraps(func) 

293 def timeMethod_wrapper(self: Any, *args: Any, **keyArgs: Any) -> Any: 

294 # Adjust the stacklevel to account for the wrappers. 

295 # stacklevel 1 would make the log message come from this function 

296 # but we want it to come from the file that defined the method 

297 # so need to increment by 1 to get to the caller. 

298 stacklevel = 2 

299 logInfo( 

300 obj=self, 

301 prefix=func.__name__ + "Start", 

302 metadata=metadata, 

303 logger=logger, 

304 logLevel=logLevel, 

305 stacklevel=stacklevel, 

306 ) 

307 try: 

308 res = func(self, *args, **keyArgs) 

309 finally: 

310 logInfo( 

311 obj=self, 

312 prefix=func.__name__ + "End", 

313 metadata=metadata, 

314 logger=logger, 

315 logLevel=logLevel, 

316 stacklevel=stacklevel, 

317 ) 

318 return res 

319 

320 return timeMethod_wrapper 

321 

322 if _func is None: 

323 return decorator_timer 

324 else: 

325 return decorator_timer(_func) 

326 

327 

328@contextmanager 

329def time_this( 

330 log: Optional[LsstLoggers] = None, 

331 msg: Optional[str] = None, 

332 level: int = logging.DEBUG, 

333 prefix: Optional[str] = "timer", 

334 args: Iterable[Any] = (), 

335 mem_usage: bool = False, 

336 mem_child: bool = False, 

337 mem_unit: u.Quantity = u.byte, 

338 mem_fmt: str = ".0f", 

339) -> Iterator[None]: 

340 """Time the enclosed block and issue a log message. 

341 

342 Parameters 

343 ---------- 

344 log : `logging.Logger`, optional 

345 Logger to use to report the timer message. The root logger will 

346 be used if none is given. 

347 msg : `str`, optional 

348 Context to include in log message. 

349 level : `int`, optional 

350 Python logging level to use to issue the log message. If the 

351 code block raises an exception the log message will automatically 

352 switch to level ERROR. 

353 prefix : `str`, optional 

354 Prefix to use to prepend to the supplied logger to 

355 create a new logger to use instead. No prefix is used if the value 

356 is set to `None`. Defaults to "timer". 

357 args : iterable of any 

358 Additional parameters passed to the log command that should be 

359 written to ``msg``. 

360 mem_usage : `bool`, optional 

361 Flag indicating whether to include the memory usage in the report. 

362 Defaults, to False. 

363 mem_child : `bool`, optional 

364 Flag indication whether to include memory usage of the child processes. 

365 mem_unit : `astropy.units.Unit`, optional 

366 Unit to use when reporting the memory usage. Defaults to bytes. 

367 mem_fmt : `str`, optional 

368 Format specifier to use when displaying values related to memory usage. 

369 Defaults to '.0f'. 

370 """ 

371 if log is None: 

372 log = logging.getLogger() 

373 if prefix: 

374 log_name = f"{prefix}.{log.name}" if not isinstance(log, logging.RootLogger) else prefix 

375 log = logging.getLogger(log_name) 

376 

377 success = False 

378 start = time.time() 

379 

380 if mem_usage and not log.isEnabledFor(level): 

381 mem_usage = False 

382 

383 if mem_usage: 

384 current_usages_start = get_current_mem_usage() 

385 peak_usages_start = get_peak_mem_usage() 

386 

387 try: 

388 yield 

389 success = True 

390 finally: 

391 end = time.time() 

392 

393 # The message is pre-inserted to allow the logger to expand 

394 # the additional args provided. Make that easier by converting 

395 # the None message to empty string. 

396 if msg is None: 

397 msg = "" 

398 

399 # Convert user provided parameters (if any) to mutable sequence to make 

400 # mypy stop complaining when additional parameters will be added below. 

401 params = list(args) if args else [] 

402 

403 if not success: 

404 # Something went wrong so change the log level to indicate 

405 # this. 

406 level = logging.ERROR 

407 

408 # Specify stacklevel to ensure the message is reported from the 

409 # caller (1 is this file, 2 is contextlib, 3 is user) 

410 params += (": " if msg else "", end - start) 

411 msg += "%sTook %.4f seconds" 

412 if mem_usage: 

413 current_usages_end = get_current_mem_usage() 

414 peak_usages_end = get_peak_mem_usage() 

415 

416 current_deltas = [end - start for end, start in zip(current_usages_end, current_usages_start)] 

417 peak_deltas = [end - start for end, start in zip(peak_usages_end, peak_usages_start)] 

418 

419 current_usage = current_usages_end[0] 

420 current_delta = current_deltas[0] 

421 peak_delta = peak_deltas[0] 

422 if mem_child: 

423 current_usage += current_usages_end[1] 

424 current_delta += current_deltas[1] 

425 peak_delta += peak_deltas[1] 

426 

427 if not mem_unit.is_equivalent(u.byte): 

428 _LOG.warning("Invalid memory unit '%s', using '%s' instead", mem_unit, u.byte) 

429 mem_unit = u.byte 

430 

431 msg += ( 

432 f"; current memory usage: {current_usage.to(mem_unit):{mem_fmt}}" 

433 f", delta: {current_delta.to(mem_unit):{mem_fmt}}" 

434 f", peak delta: {peak_delta.to(mem_unit):{mem_fmt}}" 

435 ) 

436 log.log(level, msg, *params, stacklevel=3) 

437 

438 

439@contextmanager 

440def profile( 

441 filename: Optional[str], log: Optional[logging.Logger] = None 

442) -> Iterator[Optional[cProfile.Profile]]: 

443 """Profile the enclosed code block and save the result to a file. 

444 

445 Parameters 

446 ---------- 

447 filename : `str` or `None` 

448 Filename to which to write profile (profiling disabled if `None` or 

449 empty string). 

450 log : `logging.Logger`, optional 

451 Log object for logging the profile operations. 

452 

453 Yields 

454 ------ 

455 prof : `cProfile.Profile` or `None` 

456 If profiling is enabled, the context manager returns the 

457 `cProfile.Profile` object (otherwise it returns `None`), 

458 which allows additional control over profiling. 

459 

460 Examples 

461 -------- 

462 You can obtain the `cProfile.Profile` object using the "as" clause, e.g.: 

463 

464 .. code-block:: python 

465 

466 with profile(filename) as prof: 

467 runYourCodeHere() 

468 

469 The output cumulative profile can be printed with a command-line like: 

470 

471 .. code-block:: bash 

472 

473 python -c 'import pstats; \ 

474 pstats.Stats("<filename>").sort_stats("cumtime").print_stats(30)' 

475 """ 

476 if not filename: 

477 # Nothing to do 

478 yield None 

479 return 

480 from cProfile import Profile 

481 

482 profile = Profile() 

483 if log is not None: 

484 log.info("Enabling cProfile profiling") 

485 profile.enable() 

486 yield profile 

487 profile.disable() 

488 profile.dump_stats(filename) 

489 if log is not None: 

490 log.info("cProfile stats written to %s", filename)