A universal Python decorator function to record the runtime of arbitrary functions (Copied from MegatronLM)

time_decorator.py

import time
import logging
from typing import Callable
from functools import wraps


# Reference: https://github.com/NVIDIA/Megatron-LM/blob/9a496c976e12a62ce8e39e14496e52a985588730/megatron/core/dist_checkpointing/strategies/two_stage.py#L35
def timed(verbose=True) -> Callable:
    def timed_dec(fn):
        name = fn.__name__

        @wraps(fn)
        def wrapped(*args, **kwargs):
            if verbose:
                logging.debug(f'{name} init')
            start = time.time()
            ret = fn(*args, **kwargs)
            took = time.time() - start
            if verbose:
                logging.debug(f'{name} took {took}s')
            return ret

        return wrapped

    return timed_dec

  
@timed(verbose=False)
def func(n_sec: int) -> None:
  time.sleep(n_sec)