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

136 statements  

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

25from contextlib import contextmanager 

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 try: 

113 metadata = obj.metadata 

114 except AttributeError: 

115 pass 

116 if logger is None: 

117 try: 

118 logger = obj.log 

119 except AttributeError: 

120 pass 

121 strList = [] 

122 for name, value in pairs: 

123 if metadata is not None: 

124 _add_to_metadata(metadata, name, value) 

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

126 if logger is not None: 

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

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

129 # message will be issued. 

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

131 if timer_logger.isEnabledFor(logLevel): 

132 if stacklevel is None: 

133 stacklevel = find_outside_stacklevel("lsst.utils") 

134 else: 

135 # Account for the caller stack. 

136 stacklevel += 1 

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

138 

139 

140def logInfo( 

141 obj: Any, 

142 prefix: str, 

143 logLevel: int = logging.DEBUG, 

144 metadata: MutableMapping | None = None, 

145 logger: logging.Logger | None = None, 

146 stacklevel: int | None = None, 

147) -> None: 

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

149 

150 Parameters 

151 ---------- 

152 obj : `object` 

153 An object with both or one these two attributes: 

154 

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

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

157 entries to a list. 

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

159 

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

161 or this function will do nothing. 

162 prefix : `str` 

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

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

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

166 logLevel : optional 

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

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

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

170 logger : `logging.Logger` 

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

172 stacklevel : `int`, optional 

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

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

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

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

177 1 to account for this caller. 

178 

179 Notes 

180 ----- 

181 Logged items include: 

182 

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

184 timestamps). 

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

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

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

188 

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

190 processes are excluded. 

191 

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

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

194 to be version 0. 

195 

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

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

198 """ 

199 if metadata is None and obj is not None: 

200 try: 

201 metadata = obj.metadata 

202 except AttributeError: 

203 pass 

204 if metadata is not None: 

205 # Log messages already have timestamps. 

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

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

208 

209 # Force a version number into the metadata. 

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

211 metadata["__version__"] = 1 

212 if stacklevel is not None: 

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

214 # are going one down in the stack. 

215 stacklevel += 1 

216 

217 usage = _get_current_rusage() 

218 logPairs( 

219 obj=obj, 

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

221 logLevel=logLevel, 

222 metadata=metadata, 

223 logger=logger, 

224 stacklevel=stacklevel, 

225 ) 

226 

227 

228def timeMethod( 

229 _func: Any | None = None, 

230 *, 

231 metadata: MutableMapping | None = None, 

232 logger: logging.Logger | None = None, 

233 logLevel: int = logging.DEBUG, 

234) -> Callable: 

235 """Measure duration of a method. 

236 

237 Parameters 

238 ---------- 

239 func 

240 The method to wrap. 

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

242 Metadata to use as override if the instance object attached 

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

244 logger : `logging.Logger`, optional 

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

246 have a ``log`` property. 

247 logLevel : `int`, optional 

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

249 

250 Notes 

251 ----- 

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

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

254 

255 .. warning:: 

256 

257 This decorator only works with instance methods of any class 

258 with these attributes: 

259 

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

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

262 list. 

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

264 

265 Examples 

266 -------- 

267 To use: 

268 

269 .. code-block:: python 

270 

271 import lsst.utils as utils 

272 import lsst.pipe.base as pipeBase 

273 class FooTask(pipeBase.Task): 

274 pass 

275 

276 @utils.timeMethod 

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

278 pass 

279 """ 

280 

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

282 @functools.wraps(func) 

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

284 # Adjust the stacklevel to account for the wrappers. 

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

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

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

288 stacklevel = 2 

289 logInfo( 

290 obj=self, 

291 prefix=func.__name__ + "Start", 

292 metadata=metadata, 

293 logger=logger, 

294 logLevel=logLevel, 

295 stacklevel=stacklevel, 

296 ) 

297 try: 

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

299 finally: 

300 logInfo( 

301 obj=self, 

302 prefix=func.__name__ + "End", 

303 metadata=metadata, 

304 logger=logger, 

305 logLevel=logLevel, 

306 stacklevel=stacklevel, 

307 ) 

308 return res 

309 

310 return timeMethod_wrapper 

311 

312 if _func is None: 

313 return decorator_timer 

314 else: 

315 return decorator_timer(_func) 

316 

317 

318@contextmanager 

319def time_this( 

320 log: LsstLoggers | None = None, 

321 msg: str | None = None, 

322 level: int = logging.DEBUG, 

323 prefix: str | None = "timer", 

324 args: Iterable[Any] = (), 

325 mem_usage: bool = False, 

326 mem_child: bool = False, 

327 mem_unit: u.Quantity = u.byte, 

328 mem_fmt: str = ".0f", 

329) -> Iterator[None]: 

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

331 

332 Parameters 

333 ---------- 

334 log : `logging.Logger`, optional 

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

336 be used if none is given. 

337 msg : `str`, optional 

338 Context to include in log message. 

339 level : `int`, optional 

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

341 code block raises an exception the log message will automatically 

342 switch to level ERROR. 

343 prefix : `str`, optional 

344 Prefix to use to prepend to the supplied logger to 

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

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

347 args : iterable of any 

348 Additional parameters passed to the log command that should be 

349 written to ``msg``. 

350 mem_usage : `bool`, optional 

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

352 Defaults, to False. 

353 mem_child : `bool`, optional 

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

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

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

357 mem_fmt : `str`, optional 

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

359 Defaults to '.0f'. 

360 """ 

361 if log is None: 

362 log = logging.getLogger() 

363 if prefix: 

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

365 log = logging.getLogger(log_name) 

366 

367 success = False 

368 start = time.time() 

369 

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

371 mem_usage = False 

372 

373 if mem_usage: 

374 current_usages_start = get_current_mem_usage() 

375 peak_usages_start = get_peak_mem_usage() 

376 

377 try: 

378 yield 

379 success = True 

380 finally: 

381 end = time.time() 

382 

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

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

385 # the None message to empty string. 

386 if msg is None: 

387 msg = "" 

388 

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

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

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

392 

393 if not success: 

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

395 # this. 

396 level = logging.ERROR 

397 

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

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

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

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

402 if mem_usage: 

403 current_usages_end = get_current_mem_usage() 

404 peak_usages_end = get_peak_mem_usage() 

405 

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

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

408 

409 current_usage = current_usages_end[0] 

410 current_delta = current_deltas[0] 

411 peak_delta = peak_deltas[0] 

412 if mem_child: 

413 current_usage += current_usages_end[1] 

414 current_delta += current_deltas[1] 

415 peak_delta += peak_deltas[1] 

416 

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

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

419 mem_unit = u.byte 

420 

421 msg += ( 

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

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

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

425 ) 

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

427 

428 

429@contextmanager 

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

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

432 

433 Parameters 

434 ---------- 

435 filename : `str` or `None` 

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

437 empty string). 

438 log : `logging.Logger`, optional 

439 Log object for logging the profile operations. 

440 

441 Yields 

442 ------ 

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

444 If profiling is enabled, the context manager returns the 

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

446 which allows additional control over profiling. 

447 

448 Examples 

449 -------- 

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

451 

452 .. code-block:: python 

453 

454 with profile(filename) as prof: 

455 runYourCodeHere() 

456 

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

458 

459 .. code-block:: bash 

460 

461 python -c 'import pstats; \ 

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

463 """ 

464 if not filename: 

465 # Nothing to do 

466 yield None 

467 return 

468 from cProfile import Profile 

469 

470 profile = Profile() 

471 if log is not None: 

472 log.info("Enabling cProfile profiling") 

473 profile.enable() 

474 yield profile 

475 profile.disable() 

476 profile.dump_stats(filename) 

477 if log is not None: 

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