Source code for dubfi.core.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
import typing as T
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()
        name = getattr(f, "__name__", "unknown function")
        if t2 - t1 >= 1e9:
            log.info(f"{name:s} took {1e-9 * (t2 - t1):.2f} s")
        else:
            log.info(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

        name = getattr(fnc, "__name__", "unknown_function")
        target = f"pstat_{name}_{mp.current_process().ident}.prof"

    @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, count: int = -1) -> int:
    """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.
    count: int, optional
        stop after this number of log message
    """
    from threading import Thread

    t = Thread(
        target=_monitor_memory_thread, args=(interval, log_fn, count), daemon=True
    )
    t.start()
    return t.ident  # type: ignore





def _monitor_memory_thread(interval: float, log_fn: T.Callable, count: int) -> 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}"
        )
        count -= 1
        if count == 0:
            break
        sleep(interval)