blob: 4c7d0911b918bfbd525b3475930549afc74c8128 (
plain) (
blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
|
from __future__ import annotations
import gc
from contextlib import contextmanager
from datetime import datetime
from functools import wraps
from logging import Logger
import psutil
@contextmanager
def log_stage(logger: Logger, msg: str, collect_garbage: bool = False):
"""Contextmanager that logs a message to the logger when it is entered and exited.
The message has debug information about memory use. Optionally, it can
run garbage collection when exiting.
"""
meminfo = psutil.Process().memory_full_info()
logger.info(f">> Starting >> {msg}")
logger.debug(str(meminfo))
start_time = datetime.now()
try:
yield
finally:
end_time = datetime.now()
duration = end_time - start_time
meminfo = psutil.Process().memory_full_info()
logger.info(f"<< Finished << {msg} ({duration})")
logger.debug(str(meminfo))
if collect_garbage:
gc.collect()
def staged(logger: Logger, log_message: str, collect_garbage: bool = False):
"""Like log_stage but a decorator."""
def deco(func):
@wraps(func)
def wrapper(*args, **kwargs):
with log_stage(logger, log_message, collect_garbage):
return func(*args, **kwargs)
return wrapper
return deco
|
