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

147 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2023-02-02 06:11 -0800

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 inspect 

23import logging 

24import time 

25from contextlib import contextmanager 

26from typing import ( 

27 TYPE_CHECKING, 

28 Any, 

29 Callable, 

30 Collection, 

31 Iterable, 

32 Iterator, 

33 MutableMapping, 

34 Optional, 

35 Tuple, 

36) 

37 

38from astropy import units as u 

39 

40from .usage import _get_current_rusage, get_current_mem_usage, get_peak_mem_usage 

41 

42if TYPE_CHECKING: 42 ↛ 43line 42 didn't jump to line 43, because the condition on line 42 was never true

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 _find_outside_stacklevel() -> int: 

84 """Find the stack level corresponding to caller code outside of this 

85 module. 

86 

87 This can be passed directly to `logging.Logger.log()` to ensure 

88 that log messages are issued as if they are coming from caller code. 

89 

90 Returns 

91 ------- 

92 stacklevel : `int` 

93 The stack level to use to refer to a caller outside of this module. 

94 A ``stacklevel`` of ``1`` corresponds to the caller of this internal 

95 function and that is the default expected by `logging.Logger.log()`. 

96 

97 Notes 

98 ----- 

99 Intended to be called from the function that is going to issue a log 

100 message. The result should be passed into `~logging.Logger.log` via the 

101 keyword parameter ``stacklevel``. 

102 """ 

103 stacklevel = 1 # the default for `Logger.log` 

104 for i, s in enumerate(inspect.stack()): 

105 module = inspect.getmodule(s.frame) 

106 if module is None: 

107 # Stack frames sometimes hang around so explicilty delete. 

108 del s 

109 continue 

110 if not module.__name__.startswith("lsst.utils"): 

111 # 0 will be this function. 

112 # 1 will be the caller which will be the default for `Logger.log` 

113 # and so does not need adjustment. 

114 stacklevel = i 

115 break 

116 # Stack frames sometimes hang around so explicilty delete. 

117 del s 

118 

119 return stacklevel 

120 

121 

122def logPairs( 

123 obj: Any, 

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

125 logLevel: int = logging.DEBUG, 

126 metadata: Optional[MutableMapping] = None, 

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

128 stacklevel: Optional[int] = None, 

129) -> None: 

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

131 

132 Parameters 

133 ---------- 

134 obj : `object` 

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

136 

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

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

139 entries to a list. 

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

141 

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

143 or this function will do nothing. 

144 pairs : sequence 

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

146 logLevel : `int, optional 

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

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

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

150 logger : `logging.Logger` 

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

152 stacklevel : `int`, optional 

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

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

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

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

157 1 to account for this caller. 

158 """ 

159 if obj is not None: 

160 if metadata is None: 

161 try: 

162 metadata = obj.metadata 

163 except AttributeError: 

164 pass 

165 if logger is None: 

166 try: 

167 logger = obj.log 

168 except AttributeError: 

169 pass 

170 strList = [] 

171 for name, value in pairs: 

172 if metadata is not None: 

173 _add_to_metadata(metadata, name, value) 

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

175 if logger is not None: 

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

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

178 # message will be issued. 

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

180 if timer_logger.isEnabledFor(logLevel): 

181 if stacklevel is None: 

182 stacklevel = _find_outside_stacklevel() 

183 else: 

184 # Account for the caller stack. 

185 stacklevel += 1 

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

187 

188 

189def logInfo( 

190 obj: Any, 

191 prefix: str, 

192 logLevel: int = logging.DEBUG, 

193 metadata: Optional[MutableMapping] = None, 

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

195 stacklevel: Optional[int] = None, 

196) -> None: 

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

198 

199 Parameters 

200 ---------- 

201 obj : `object` 

202 An object with both or one these two attributes: 

203 

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

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

206 entries to a list. 

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

208 

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

210 or this function will do nothing. 

211 prefix : `str` 

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

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

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

215 logLevel : optional 

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

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

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

219 logger : `logging.Logger` 

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

221 stacklevel : `int`, optional 

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

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

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

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

226 1 to account for this caller. 

227 

228 Notes 

229 ----- 

230 Logged items include: 

231 

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

233 timestamps). 

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

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

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

237 

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

239 processes are excluded. 

240 

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

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

243 to be version 0. 

244 

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

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

247 """ 

248 if metadata is None and obj is not None: 

249 try: 

250 metadata = obj.metadata 

251 except AttributeError: 

252 pass 

253 if metadata is not None: 

254 # Log messages already have timestamps. 

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

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

257 

258 # Force a version number into the metadata. 

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

260 metadata["__version__"] = 1 

261 if stacklevel is not None: 

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

263 # are going one down in the stack. 

264 stacklevel += 1 

265 

266 usage = _get_current_rusage() 

267 logPairs( 

268 obj=obj, 

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

270 logLevel=logLevel, 

271 metadata=metadata, 

272 logger=logger, 

273 stacklevel=stacklevel, 

274 ) 

275 

276 

277def timeMethod( 

278 _func: Optional[Any] = None, 

279 *, 

280 metadata: Optional[MutableMapping] = None, 

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

282 logLevel: int = logging.DEBUG, 

283) -> Callable: 

284 """Measure duration of a method. 

285 

286 Parameters 

287 ---------- 

288 func 

289 The method to wrap. 

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

291 Metadata to use as override if the instance object attached 

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

293 logger : `logging.Logger`, optional 

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

295 have a ``log`` property. 

296 logLevel : `int`, optional 

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

298 

299 Notes 

300 ----- 

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

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

303 

304 .. warning:: 

305 

306 This decorator only works with instance methods of any class 

307 with these attributes: 

308 

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

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

311 list. 

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

313 

314 Examples 

315 -------- 

316 To use: 

317 

318 .. code-block:: python 

319 

320 import lsst.utils as utils 

321 import lsst.pipe.base as pipeBase 

322 class FooTask(pipeBase.Task): 

323 pass 

324 

325 @utils.timeMethod 

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

327 pass 

328 """ 

329 

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

331 @functools.wraps(func) 

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

333 # Adjust the stacklevel to account for the wrappers. 

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

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

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

337 stacklevel = 2 

338 logInfo( 

339 obj=self, 

340 prefix=func.__name__ + "Start", 

341 metadata=metadata, 

342 logger=logger, 

343 logLevel=logLevel, 

344 stacklevel=stacklevel, 

345 ) 

346 try: 

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

348 finally: 

349 logInfo( 

350 obj=self, 

351 prefix=func.__name__ + "End", 

352 metadata=metadata, 

353 logger=logger, 

354 logLevel=logLevel, 

355 stacklevel=stacklevel, 

356 ) 

357 return res 

358 

359 return timeMethod_wrapper 

360 

361 if _func is None: 

362 return decorator_timer 

363 else: 

364 return decorator_timer(_func) 

365 

366 

367@contextmanager 

368def time_this( 

369 log: Optional[LsstLoggers] = None, 

370 msg: Optional[str] = None, 

371 level: int = logging.DEBUG, 

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

373 args: Iterable[Any] = (), 

374 mem_usage: bool = False, 

375 mem_child: bool = False, 

376 mem_unit: u.Quantity = u.byte, 

377 mem_fmt: str = ".0f", 

378) -> Iterator[None]: 

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

380 

381 Parameters 

382 ---------- 

383 log : `logging.Logger`, optional 

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

385 be used if none is given. 

386 msg : `str`, optional 

387 Context to include in log message. 

388 level : `int`, optional 

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

390 code block raises an exception the log message will automatically 

391 switch to level ERROR. 

392 prefix : `str`, optional 

393 Prefix to use to prepend to the supplied logger to 

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

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

396 args : iterable of any 

397 Additional parameters passed to the log command that should be 

398 written to ``msg``. 

399 mem_usage : `bool`, optional 

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

401 Defaults, to False. 

402 mem_child : `bool`, optional 

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

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

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

406 mem_fmt : `str`, optional 

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

408 Defaults to '.0f'. 

409 """ 

410 if log is None: 

411 log = logging.getLogger() 

412 if prefix: 

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

414 log = logging.getLogger(log_name) 

415 

416 success = False 

417 start = time.time() 

418 if mem_usage: 

419 current_usages_start = get_current_mem_usage() 

420 peak_usages_start = get_peak_mem_usage() 

421 

422 try: 

423 yield 

424 success = True 

425 finally: 

426 end = time.time() 

427 

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

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

430 # the None message to empty string. 

431 if msg is None: 

432 msg = "" 

433 

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

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

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

437 

438 if not success: 

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

440 # this. 

441 level = logging.ERROR 

442 

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

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

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

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

447 if mem_usage and log.isEnabledFor(level): 

448 current_usages_end = get_current_mem_usage() 

449 peak_usages_end = get_peak_mem_usage() 

450 

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

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

453 

454 current_usage = current_usages_end[0] 

455 current_delta = current_deltas[0] 

456 peak_delta = peak_deltas[0] 

457 if mem_child: 

458 current_usage += current_usages_end[1] 

459 current_delta += current_deltas[1] 

460 peak_delta += peak_deltas[1] 

461 

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

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

464 mem_unit = u.byte 

465 

466 msg += ( 

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

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

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

470 ) 

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

472 

473 

474@contextmanager 

475def profile( 

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

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

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

479 

480 Parameters 

481 ---------- 

482 filename : `str` or `None` 

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

484 empty string). 

485 log : `logging.Logger`, optional 

486 Log object for logging the profile operations. 

487 

488 Yields 

489 ------ 

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

491 If profiling is enabled, the context manager returns the 

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

493 which allows additional control over profiling. 

494 

495 Examples 

496 -------- 

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

498 

499 .. code-block:: python 

500 

501 with profile(filename) as prof: 

502 runYourCodeHere() 

503 

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

505 

506 .. code-block:: bash 

507 

508 python -c 'import pstats; \ 

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

510 """ 

511 if not filename: 

512 # Nothing to do 

513 yield None 

514 return 

515 from cProfile import Profile 

516 

517 profile = Profile() 

518 if log is not None: 

519 log.info("Enabling cProfile profiling") 

520 profile.enable() 

521 yield profile 

522 profile.disable() 

523 profile.dump_stats(filename) 

524 if log is not None: 

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