groot/vla/utils/timer.py

import time


class ContextTimer:

    def __init__(self, trainer):
        self.last_key = None
        self.trainer = trainer
        self.start_times = {}
        self.key_stack = []

    def with_label(self, key):
        self.last_key = key
        return self

    def __enter__(self):
        self.key_stack.append(self.last_key)  # Push key to stack
        self.start_times[self.last_key] = time.time()  # Start timing for this key
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        key = self.key_stack.pop()  # Pop key from stack
        diff = time.time() - self.start_times[key]
        self.trainer.log({f"{key}_time": diff})
        # print(f"{key}: {diff:.2f} seconds")


if __name__ == "__main__":

    class MockTrainer:
        def log(self, data):
            print("Logging:", data)

    trainer = MockTrainer()
    my_timer = ContextTimer(trainer)

    with my_timer.with_label("outer"):
        time.sleep(1)
        with my_timer.with_label("inner"):
            time.sleep(2)
        with my_timer.with_label("inner"):  # Another inner block
            time.sleep(1)