Metaflow `@time_step` Decorator — Track and store execution time of each `@step` in your flow using `flow.step_timings`. Handy for lightweight performance profiling.

time_step_decorator.py

import time
import functools


def time_step(func):
    """
    A Metaflow-specific decorator to measure and log the execution time of a @step.

    This decorator is intended to be used on Metaflow step functions only.
    It stores the step's execution duration (in seconds) in the `step_timings`
    dictionary of the flow object.

    - Automatically records the time taken by a step.
    - Saves timing under flow.step_timings[<step_name>].
    - Requires the flow object (self) to be the first argument.

    Example:
        @step
        @time_step
        def preprocess(self):
            # step logic
            pass
    """
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        flow, *_ = args

        try:
            start_time = time.time()
            result = func(*args, **kwargs)
            end_time = time.time()
            duration = round(end_time - start_time, 3)
            step_name = func.__name__

            if not hasattr(flow, "step_timings"):
                flow.step_timings = {}
            flow.step_timings[step_name] = duration
            return result
        except Exception as e:
            raise e("failed")

    return wrapper