DI-engine/ding/utils/profiler_helper.py

import atexit
import pstats
import io
import cProfile
import os


def register_profiler(write_profile, pr, folder_path):
    atexit.register(write_profile, pr, folder_path)


class Profiler:
    """
    Overview:
        A class for profiling code execution. It can be used as a context manager or a decorator.

    Interfaces:
        ``__init__``, ``mkdir``, ``write_profile``, ``profile``.
    """

    def __init__(self):
        """
        Overview:
            Initialize the Profiler object.
        """

        self.pr = cProfile.Profile()

    def mkdir(self, directory: str):
        """
        OverView:
            Create a directory if it doesn't exist.

        Arguments:
            - directory (:obj:`str`): The path of the directory to be created.
        """
        if not os.path.exists(directory):
            os.makedirs(directory)

    def write_profile(self, pr: cProfile.Profile, folder_path: str):
        """
        OverView:
            Write the profiling results to files.

        Arguments:
            - pr (:obj:`cProfile.Profile`): The profiler object containing the profiling results.
            - folder_path (:obj:`str`): The path of the folder where the profiling files will be saved.
        """
        pr.disable()
        s_tottime = io.StringIO()
        s_cumtime = io.StringIO()

        ps = pstats.Stats(pr, stream=s_tottime).sort_stats('tottime')
        ps.print_stats()
        with open(folder_path + "/profile_tottime.txt", 'w+') as f:
            f.write(s_tottime.getvalue())

        ps = pstats.Stats(pr, stream=s_cumtime).sort_stats('cumtime')
        ps.print_stats()
        with open(folder_path + "/profile_cumtime.txt", 'w+') as f:
            f.write(s_cumtime.getvalue())

        pr.dump_stats(folder_path + "/profile.prof")

    def profile(self, folder_path="./tmp"):
        """
        OverView:
            Enable profiling and save the results to files.

        Arguments:
            - folder_path (:obj:`str`): The path of the folder where the profiling files will be saved. \
                Defaults to "./tmp".
        """
        self.mkdir(folder_path)
        self.pr.enable()
        register_profiler(self.write_profile, self.pr, folder_path)