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

130 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2023-07-25 09:27 +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 collections.abc import Callable, Collection, Iterable, Iterator, MutableMapping 

25from contextlib import contextmanager, suppress 

26from typing import TYPE_CHECKING, Any 

27 

28from astropy import units as u 

29 

30from .introspection import find_outside_stacklevel 

31from .logging import LsstLoggers 

32from .usage import _get_current_rusage, get_current_mem_usage, get_peak_mem_usage 

33 

34if TYPE_CHECKING: 

35 import cProfile 

36 

37 

38_LOG = logging.getLogger(__name__) 

39 

40 

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

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

43 

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

45 

46 Parameters 

47 ---------- 

48 metadata : `dict`-like, optional 

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

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

51 name : `str` 

52 The key to use in the metadata dictionary. 

53 value : Any 

54 Value to store in the list. 

55 """ 

56 try: 

57 try: 

58 # PropertySet should always prefer LongLong for integers 

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

60 except (TypeError, AttributeError): 

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

62 except AttributeError: 

63 pass 

64 else: 

65 return 

66 

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

68 if name not in metadata: 

69 metadata[name] = [] 

70 metadata[name].append(value) 

71 

72 

73def logPairs( 

74 obj: Any, 

75 pairs: Collection[tuple[str, Any]], 

76 logLevel: int = logging.DEBUG, 

77 metadata: MutableMapping | None = None, 

78 logger: logging.Logger | None = None, 

79 stacklevel: int | None = None, 

80) -> None: 

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

82 

83 Parameters 

84 ---------- 

85 obj : `object` 

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

87 

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

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

90 entries to a list. 

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

92 

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

94 or this function will do nothing. 

95 pairs : sequence 

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

97 logLevel : `int, optional 

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

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

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

101 logger : `logging.Logger` 

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

103 stacklevel : `int`, optional 

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

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

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

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

108 1 to account for this caller. 

109 """ 

110 if obj is not None: 

111 if metadata is None: 

112 with suppress(AttributeError): 

113 metadata = obj.metadata 

114 

115 if logger is None: 

116 with suppress(AttributeError): 

117 logger = obj.log 

118 

119 strList = [] 

120 for name, value in pairs: 

121 if metadata is not None: 

122 _add_to_metadata(metadata, name, value) 

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

124 if logger is not None: 

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

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

127 # message will be issued. 

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

129 if timer_logger.isEnabledFor(logLevel): 

130 if stacklevel is None: 

131 stacklevel = find_outside_stacklevel("lsst.utils") 

132 else: 

133 # Account for the caller stack. 

134 stacklevel += 1 

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

136 

137 

138def logInfo( 

139 obj: Any, 

140 prefix: str, 

141 logLevel: int = logging.DEBUG, 

142 metadata: MutableMapping | None = None, 

143 logger: logging.Logger | None = None, 

144 stacklevel: int | None = None, 

145) -> None: 

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

147 

148 Parameters 

149 ---------- 

150 obj : `object` 

151 An object with both or one these two attributes: 

152 

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

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

155 entries to a list. 

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

157 

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

159 or this function will do nothing. 

160 prefix : `str` 

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

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

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

164 logLevel : optional 

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

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

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

168 logger : `logging.Logger` 

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

170 stacklevel : `int`, optional 

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

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

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

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

175 1 to account for this caller. 

176 

177 Notes 

178 ----- 

179 Logged items include: 

180 

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

182 timestamps). 

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

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

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

186 

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

188 processes are excluded. 

189 

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

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

192 to be version 0. 

193 

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

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

196 """ 

197 if metadata is None and obj is not None: 

198 with suppress(AttributeError): 

199 metadata = obj.metadata 

200 

201 if metadata is not None: 

202 # Log messages already have timestamps. 

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

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

205 

206 # Force a version number into the metadata. 

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

208 metadata["__version__"] = 1 

209 if stacklevel is not None: 

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

211 # are going one down in the stack. 

212 stacklevel += 1 

213 

214 usage = _get_current_rusage() 

215 logPairs( 

216 obj=obj, 

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

218 logLevel=logLevel, 

219 metadata=metadata, 

220 logger=logger, 

221 stacklevel=stacklevel, 

222 ) 

223 

224 

225def timeMethod( 

226 _func: Any | None = None, 

227 *, 

228 metadata: MutableMapping | None = None, 

229 logger: logging.Logger | None = None, 

230 logLevel: int = logging.DEBUG, 

231) -> Callable: 

232 """Measure duration of a method. 

233 

234 Parameters 

235 ---------- 

236 func 

237 The method to wrap. 

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

239 Metadata to use as override if the instance object attached 

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

241 logger : `logging.Logger`, optional 

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

243 have a ``log`` property. 

244 logLevel : `int`, optional 

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

246 

247 Notes 

248 ----- 

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

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

251 

252 .. warning:: 

253 

254 This decorator only works with instance methods of any class 

255 with these attributes: 

256 

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

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

259 list. 

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

261 

262 Examples 

263 -------- 

264 To use: 

265 

266 .. code-block:: python 

267 

268 import lsst.utils as utils 

269 import lsst.pipe.base as pipeBase 

270 class FooTask(pipeBase.Task): 

271 pass 

272 

273 @utils.timeMethod 

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

275 pass 

276 """ 

277 

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

279 @functools.wraps(func) 

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

281 # Adjust the stacklevel to account for the wrappers. 

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

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

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

285 stacklevel = 2 

286 logInfo( 

287 obj=self, 

288 prefix=func.__name__ + "Start", 

289 metadata=metadata, 

290 logger=logger, 

291 logLevel=logLevel, 

292 stacklevel=stacklevel, 

293 ) 

294 try: 

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

296 finally: 

297 logInfo( 

298 obj=self, 

299 prefix=func.__name__ + "End", 

300 metadata=metadata, 

301 logger=logger, 

302 logLevel=logLevel, 

303 stacklevel=stacklevel, 

304 ) 

305 return res 

306 

307 return timeMethod_wrapper 

308 

309 if _func is None: 

310 return decorator_timer 

311 else: 

312 return decorator_timer(_func) 

313 

314 

315@contextmanager 

316def time_this( 

317 log: LsstLoggers | None = None, 

318 msg: str | None = None, 

319 level: int = logging.DEBUG, 

320 prefix: str | None = "timer", 

321 args: Iterable[Any] = (), 

322 mem_usage: bool = False, 

323 mem_child: bool = False, 

324 mem_unit: u.Quantity = u.byte, 

325 mem_fmt: str = ".0f", 

326) -> Iterator[None]: 

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

328 

329 Parameters 

330 ---------- 

331 log : `logging.Logger`, optional 

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

333 be used if none is given. 

334 msg : `str`, optional 

335 Context to include in log message. 

336 level : `int`, optional 

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

338 code block raises an exception the log message will automatically 

339 switch to level ERROR. 

340 prefix : `str`, optional 

341 Prefix to use to prepend to the supplied logger to 

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

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

344 args : iterable of any 

345 Additional parameters passed to the log command that should be 

346 written to ``msg``. 

347 mem_usage : `bool`, optional 

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

349 Defaults, to False. 

350 mem_child : `bool`, optional 

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

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

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

354 mem_fmt : `str`, optional 

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

356 Defaults to '.0f'. 

357 """ 

358 if log is None: 

359 log = logging.getLogger() 

360 if prefix: 

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

362 log = logging.getLogger(log_name) 

363 

364 success = False 

365 start = time.time() 

366 

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

368 mem_usage = False 

369 

370 if mem_usage: 

371 current_usages_start = get_current_mem_usage() 

372 peak_usages_start = get_peak_mem_usage() 

373 

374 try: 

375 yield 

376 success = True 

377 finally: 

378 end = time.time() 

379 

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

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

382 # the None message to empty string. 

383 if msg is None: 

384 msg = "" 

385 

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

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

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

389 

390 if not success: 

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

392 # this. 

393 level = logging.ERROR 

394 

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

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

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

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

399 if mem_usage: 

400 current_usages_end = get_current_mem_usage() 

401 peak_usages_end = get_peak_mem_usage() 

402 

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

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

405 

406 current_usage = current_usages_end[0] 

407 current_delta = current_deltas[0] 

408 peak_delta = peak_deltas[0] 

409 if mem_child: 

410 current_usage += current_usages_end[1] 

411 current_delta += current_deltas[1] 

412 peak_delta += peak_deltas[1] 

413 

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

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

416 mem_unit = u.byte 

417 

418 msg += ( 

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

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

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

422 ) 

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

424 

425 

426@contextmanager 

427def profile(filename: str | None, log: logging.Logger | None = None) -> Iterator[cProfile.Profile | None]: 

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

429 

430 Parameters 

431 ---------- 

432 filename : `str` or `None` 

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

434 empty string). 

435 log : `logging.Logger`, optional 

436 Log object for logging the profile operations. 

437 

438 Yields 

439 ------ 

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

441 If profiling is enabled, the context manager returns the 

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

443 which allows additional control over profiling. 

444 

445 Examples 

446 -------- 

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

448 

449 .. code-block:: python 

450 

451 with profile(filename) as prof: 

452 runYourCodeHere() 

453 

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

455 

456 .. code-block:: bash 

457 

458 python -c 'import pstats; \ 

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

460 """ 

461 if not filename: 

462 # Nothing to do 

463 yield None 

464 return 

465 from cProfile import Profile 

466 

467 profile = Profile() 

468 if log is not None: 

469 log.info("Enabling cProfile profiling") 

470 profile.enable() 

471 yield profile 

472 profile.disable() 

473 profile.dump_stats(filename) 

474 if log is not None: 

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