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

134 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-03-20 10:50 +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 sys 

24import time 

25from collections.abc import Callable, Collection, Iterable, Iterator, MutableMapping 

26from contextlib import contextmanager, suppress 

27from typing import TYPE_CHECKING, Any 

28 

29from astropy import units as u 

30 

31from .introspection import find_outside_stacklevel 

32from .logging import LsstLoggers 

33from .usage import _get_current_rusage, get_current_mem_usage, get_peak_mem_usage 

34 

35if TYPE_CHECKING: 

36 import cProfile 

37 

38 

39_LOG = logging.getLogger(__name__) 

40 

41 

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

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

44 

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

46 

47 Parameters 

48 ---------- 

49 metadata : `dict`-like, optional 

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

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

52 name : `str` 

53 The key to use in the metadata dictionary. 

54 value : Any 

55 Value to store in the list. 

56 """ 

57 try: 

58 try: 

59 # PropertySet should always prefer LongLong for integers 

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

61 except (TypeError, AttributeError): 

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

63 except AttributeError: 

64 pass 

65 else: 

66 return 

67 

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

69 if name not in metadata: 

70 metadata[name] = [] 

71 metadata[name].append(value) 

72 

73 

74def logPairs( 

75 obj: Any, 

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

77 logLevel: int = logging.DEBUG, 

78 metadata: MutableMapping | None = None, 

79 logger: logging.Logger | None = None, 

80 stacklevel: int | None = None, 

81) -> None: 

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

83 

84 Parameters 

85 ---------- 

86 obj : `object` 

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

88 

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

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

91 entries to a list. 

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

93 

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

95 or this function will do nothing. 

96 pairs : sequence 

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

98 logLevel : `int, optional 

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

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

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

102 logger : `logging.Logger` 

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

104 stacklevel : `int`, optional 

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

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

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

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

109 1 to account for this caller. 

110 """ 

111 if obj is not None: 

112 if metadata is None: 

113 with suppress(AttributeError): 

114 metadata = obj.metadata 

115 

116 if logger is None: 

117 with suppress(AttributeError): 

118 logger = obj.log 

119 

120 strList = [] 

121 for name, value in pairs: 

122 if metadata is not None: 

123 _add_to_metadata(metadata, name, value) 

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

125 if logger is not None: 

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

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

128 # message will be issued. 

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

130 if timer_logger.isEnabledFor(logLevel): 

131 if stacklevel is None: 

132 stacklevel = find_outside_stacklevel("lsst.utils") 

133 else: 

134 # Account for the caller stack. 

135 stacklevel += 1 

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

137 

138 

139def logInfo( 

140 obj: Any, 

141 prefix: str, 

142 logLevel: int = logging.DEBUG, 

143 metadata: MutableMapping | None = None, 

144 logger: logging.Logger | None = None, 

145 stacklevel: int | None = None, 

146) -> None: 

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

148 

149 Parameters 

150 ---------- 

151 obj : `object` 

152 An object with both or one these two attributes: 

153 

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

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

156 entries to a list. 

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

158 

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

160 or this function will do nothing. 

161 prefix : `str` 

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

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

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

165 logLevel : optional 

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

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

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

169 logger : `logging.Logger` 

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

171 stacklevel : `int`, optional 

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

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

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

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

176 1 to account for this caller. 

177 

178 Notes 

179 ----- 

180 Logged items include: 

181 

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

183 timestamps). 

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

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

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

187 

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

189 processes are excluded. 

190 

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

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

193 to be version 0. 

194 

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

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

197 """ 

198 if metadata is None and obj is not None: 

199 with suppress(AttributeError): 

200 metadata = obj.metadata 

201 

202 if metadata is not None: 

203 # Log messages already have timestamps. 

204 if sys.version_info < (3, 11, 0): 

205 now = datetime.datetime.utcnow() 

206 else: 

207 now = datetime.datetime.now(datetime.UTC) 

208 utcStr = now.isoformat() 

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

210 

211 # Force a version number into the metadata. 

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

213 metadata["__version__"] = 1 

214 if stacklevel is not None: 

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

216 # are going one down in the stack. 

217 stacklevel += 1 

218 

219 usage = _get_current_rusage() 

220 logPairs( 

221 obj=obj, 

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

223 logLevel=logLevel, 

224 metadata=metadata, 

225 logger=logger, 

226 stacklevel=stacklevel, 

227 ) 

228 

229 

230def timeMethod( 

231 _func: Any | None = None, 

232 *, 

233 metadata: MutableMapping | None = None, 

234 logger: logging.Logger | None = None, 

235 logLevel: int = logging.DEBUG, 

236) -> Callable: 

237 """Measure duration of a method. 

238 

239 Parameters 

240 ---------- 

241 _func : `~collections.abc.Callable` or `None` 

242 The method to wrap. 

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

244 Metadata to use as override if the instance object attached 

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

246 logger : `logging.Logger`, optional 

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

248 have a ``log`` property. 

249 logLevel : `int`, optional 

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

251 

252 Notes 

253 ----- 

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

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

256 

257 .. warning:: 

258 

259 This decorator only works with instance methods of any class 

260 with these attributes: 

261 

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

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

264 list. 

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

266 

267 Examples 

268 -------- 

269 To use: 

270 

271 .. code-block:: python 

272 

273 import lsst.utils as utils 

274 import lsst.pipe.base as pipeBase 

275 class FooTask(pipeBase.Task): 

276 pass 

277 

278 @utils.timeMethod 

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

280 pass 

281 """ 

282 

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

284 @functools.wraps(func) 

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

286 # Adjust the stacklevel to account for the wrappers. 

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

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

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

290 stacklevel = 2 

291 logInfo( 

292 obj=self, 

293 prefix=func.__name__ + "Start", 

294 metadata=metadata, 

295 logger=logger, 

296 logLevel=logLevel, 

297 stacklevel=stacklevel, 

298 ) 

299 try: 

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

301 finally: 

302 logInfo( 

303 obj=self, 

304 prefix=func.__name__ + "End", 

305 metadata=metadata, 

306 logger=logger, 

307 logLevel=logLevel, 

308 stacklevel=stacklevel, 

309 ) 

310 return res 

311 

312 return timeMethod_wrapper 

313 

314 if _func is None: 

315 return decorator_timer 

316 else: 

317 return decorator_timer(_func) 

318 

319 

320@contextmanager 

321def time_this( 

322 log: LsstLoggers | None = None, 

323 msg: str | None = None, 

324 level: int = logging.DEBUG, 

325 prefix: str | None = "timer", 

326 args: Iterable[Any] = (), 

327 mem_usage: bool = False, 

328 mem_child: bool = False, 

329 mem_unit: u.Quantity = u.byte, 

330 mem_fmt: str = ".0f", 

331) -> Iterator[None]: 

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

333 

334 Parameters 

335 ---------- 

336 log : `logging.Logger`, optional 

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

338 be used if none is given. 

339 msg : `str`, optional 

340 Context to include in log message. 

341 level : `int`, optional 

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

343 code block raises an exception the log message will automatically 

344 switch to level ERROR. 

345 prefix : `str`, optional 

346 Prefix to use to prepend to the supplied logger to 

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

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

349 args : iterable of any 

350 Additional parameters passed to the log command that should be 

351 written to ``msg``. 

352 mem_usage : `bool`, optional 

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

354 Defaults, to False. 

355 mem_child : `bool`, optional 

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

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

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

359 mem_fmt : `str`, optional 

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

361 Defaults to '.0f'. 

362 """ 

363 if log is None: 

364 log = logging.getLogger() 

365 if prefix: 

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

367 log = logging.getLogger(log_name) 

368 

369 success = False 

370 start = time.time() 

371 

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

373 mem_usage = False 

374 

375 if mem_usage: 

376 current_usages_start = get_current_mem_usage() 

377 peak_usages_start = get_peak_mem_usage() 

378 

379 try: 

380 yield 

381 success = True 

382 finally: 

383 end = time.time() 

384 

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

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

387 # the None message to empty string. 

388 if msg is None: 

389 msg = "" 

390 

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

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

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

394 

395 if not success: 

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

397 # this. 

398 level = logging.ERROR 

399 

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

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

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

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

404 if mem_usage: 

405 current_usages_end = get_current_mem_usage() 

406 peak_usages_end = get_peak_mem_usage() 

407 

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

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

410 

411 current_usage = current_usages_end[0] 

412 current_delta = current_deltas[0] 

413 peak_delta = peak_deltas[0] 

414 if mem_child: 

415 current_usage += current_usages_end[1] 

416 current_delta += current_deltas[1] 

417 peak_delta += peak_deltas[1] 

418 

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

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

421 mem_unit = u.byte 

422 

423 msg += ( 

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

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

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

427 ) 

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

429 

430 

431@contextmanager 

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

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

434 

435 Parameters 

436 ---------- 

437 filename : `str` or `None` 

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

439 empty string). 

440 log : `logging.Logger`, optional 

441 Log object for logging the profile operations. 

442 

443 Yields 

444 ------ 

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

446 If profiling is enabled, the context manager returns the 

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

448 which allows additional control over profiling. 

449 

450 Examples 

451 -------- 

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

453 

454 .. code-block:: python 

455 

456 with profile(filename) as prof: 

457 runYourCodeHere() 

458 

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

460 

461 .. code-block:: bash 

462 

463 python -c 'import pstats; \ 

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

465 """ 

466 if not filename: 

467 # Nothing to do 

468 yield None 

469 return 

470 from cProfile import Profile 

471 

472 profile = Profile() 

473 if log is not None: 

474 log.info("Enabling cProfile profiling") 

475 profile.enable() 

476 yield profile 

477 profile.disable() 

478 profile.dump_stats(filename) 

479 if log is not None: 

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