utils/profiling.py

from contextlib import contextmanager
import time
import mlflow


@contextmanager
def profile_step(expr_name: str, prometheus_logger, camera_id, frame_count=None):
    """With statement utility to time block of code"""
    start_time = time.time()

    try:
        # Code inside with statement
        yield
    finally:
        duration = round(time.time() - start_time, 4)
        prometheus_logger.labels(camera_id).observe(duration)
        mlflow.log_metric(
            expr_name,
            duration,
            frame_count,
        )