Source code for dubfi.profile_util

# Copyright 2025 DWD
# SPDX-License-Identifier: BSD-3-Clause

"""Profiling and performance monitoring utilities.

.. codeauthor:: Valentin Bruch, DWD <valentin.bruch@dwd.de>
.. versionadded:: 0.1.0 (initial release)
"""

import functools
from time import time_ns

from . import log




def timer(f, log=log):
    """Log duration of wall clock time of function call."""

    @functools.wraps(f)
    def wrap(*args, **kwargs):
        t1 = time_ns()
        res = f(*args, **kwargs)
        t2 = time_ns()
        if t2 - t1 >= 1e9:
            log.info(f"{f.__name__:s} took {1e-9 * (t2 - t1):.2f} s")
        else:
            log.info(f"{f.__name__:s} took {1e-6 * (t2 - t1):.2f} ms")
        return res

    return wrap





def profile_fnc(fnc, target=None):
    """Profile function and save result to file (decorator)."""
    import cProfile

    if target is None:
        import multiprocessing as mp

        target = f"pstat_{fnc.__name__}_{mp.current_process().ident}.dat"

    @functools.wraps(fnc)
    def wrapper(*args, **kwargs):
        log.info(f"Start profiling, output will be written to {target}")
        with cProfile.Profile() as pr:
            try:
                result = fnc(*args, **kwargs)
            finally:
                pr.dump_stats(target)
                log.info(f"Completed profiling, output written to {target}")
        return result

    return wrapper





def monitor_memory(interval: float = 10, log_fn=log.info) -> None:
    """Create daemon thread that logs memory usage.

    Parameters
    ----------
    interval: float
        logging interval in seconds
    log_fn: callable
        output function, can be print or log.info or similar. This
        function should take one string as argument.
    """
    from threading import Thread

    Thread(target=_monitor_memory_thread, args=(interval, log_fn), daemon=True).start()





def _monitor_memory_thread(interval: float, log_fn) -> None:
    """Log memory usage to log_fn in infinite loop."""
    from time import sleep

    from psutil import Process

    p = Process()
    scale = 2**-30
    while True:
        rss, vms, shared, text, lib, data, dirty = p.memory_info()
        log_fn(
            f"{p.pid: 6d}: res={rss * scale:5.2f}GiB, virt={vms * scale:5.2f}GiB, data={data * scale:5.2f}GiB, shared={shared * scale:5.2f}GiB, dirty:{dirty: 2d}"
        )
        sleep(interval)