profiling

optimus_dl.modules.data.profiling

PipelineTracer

Convenience context manager for enabling profiling during build phase.

Source code in optimus_dl/modules/data/profiling.py

class PipelineTracer:
    """Convenience context manager for enabling profiling during build phase."""

    def __init__(self, name: str = "manual", report_freq: int | None = None):
        self.profiler = PipelineProfiler(name, report_freq)
        self._scope = None

    def __enter__(self):
        self._scope = scope_profiler(self.profiler)
        self._scope.__enter__()
        return self.profiler

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self._scope:
            self._scope.__exit__(exc_type, exc_val, exc_tb)

scope_profiler(profiler)

Context manager to set the active profiler for the current context.

Source code in optimus_dl/modules/data/profiling.py

@contextmanager
def scope_profiler(profiler: PipelineProfiler | None):
    """Context manager to set the active profiler for the current context."""
    token = _active_profiler.set(profiler)
    try:
        yield
    finally:
        _active_profiler.reset(token)