VLA/wall_x/utils/timers.py

import logging
import os
import time
from abc import ABC, abstractmethod
from contextlib import nullcontext
from functools import wraps
from typing import List

import torch
from torch.cuda import nvtx

logger = logging.getLogger(__name__)

ENABLE_PERFORMANCE_TIMING = (
    os.environ.get("ENABLE_PERFORMANCE_TIMING", "True").lower() == "true"
)

ENABLE_CUDA_SYNC_IN_TIMER = (
    os.environ.get("ENABLE_CUDA_SYNC_IN_TIMER", "False").lower() == "true"
)


class ScopeTimerContext:
    def __init__(self, msg):
        self.msg = msg

    def __enter__(self):
        if ENABLE_CUDA_SYNC_IN_TIMER and torch.cuda.is_available():
            torch.cuda.synchronize()
        self.start_time = time.perf_counter()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        if ENABLE_CUDA_SYNC_IN_TIMER and torch.cuda.is_available():
            torch.cuda.synchronize()
        end_time = time.perf_counter()
        cost_ms = (end_time - self.start_time) * 1e3
        logger.info("%s took %.3f ms to execute", self.msg, cost_ms)


ScopeTimer = ScopeTimerContext if ENABLE_PERFORMANCE_TIMING else nullcontext


def timer(func, msg=None):
    """Decorator to measure function execution time."""
    if msg is None:
        msg = func.__name__
    else:
        msg = f"{func.__name__:} {msg}"

    @wraps(func)
    def wrapper(*args, **kwargs):
        with ScopeTimer(msg):
            result = func(*args, **kwargs)
        return result

    return wrapper


def _is_distributed():
    return torch.distributed.is_available() and torch.distributed.is_initialized()


def _get_world_size():
    if _is_distributed():
        return torch.distributed.get_world_size()
    return 1


def _get_rank():
    if _is_distributed():
        return torch.distributed.get_rank()
    return 0


def _barrier(group=None):
    if _is_distributed():
        torch.distributed.barrier(group=group)


if torch.distributed.is_available():
    try:
        dist_all_gather_func = torch.distributed.all_gather_into_tensor
    except AttributeError:
        dist_all_gather_func = torch.distributed.all_gather
else:
    dist_all_gather_func = None


class TimerBase(ABC):
    """Timer base class."""

    def __init__(self, name):
        self.name = name

    @abstractmethod
    def start(self, barrier=False):
        """Start the timer, optionally syncing all ranks with a barrier first."""
        pass

    @abstractmethod
    def stop(self, barrier=False):
        """Stop the timer, optionally syncing all ranks with a barrier first."""
        pass

    @abstractmethod
    def reset(self):
        """Reset accumulated elapsed time to zero."""
        pass

    @abstractmethod
    def elapsed(self, reset=True, barrier=False):
        """Return accumulated elapsed time in seconds; reset if reset=True."""
        pass


class DummyTimer(TimerBase):
    """Dummy Timer - no-op placeholder used when log level exceeds threshold."""

    def __init__(self):
        super().__init__("dummy timer")

    def start(self, barrier=False, nvtx_push=False, sync=False, **kwargs):
        return

    def stop(self, barrier=False, sync=False, **kwargs):
        return

    def reset(self):
        return

    def elapsed(self, reset=True, barrier=False):
        raise Exception(
            "dummy timer should not be used to calculate elapsed time, "
            "check if timer's log_level <= self._log_level."
        )

    def active_time(self):
        raise Exception(
            "active timer should not be used to calculate elapsed time, "
            "check if timer's log_level <= self._log_level."
        )


class Timer(TimerBase):
    """
    Timer class with ability to start/stop.

    Comment on using `barrier`: If this flag is passed, then all
    the caller processes will wait till all reach the timing routine.
    It is up to the user to make sure all the ranks in `barrier_group`
    call it otherwise, it will result in a hang.
    Comment on `barrier_group`: By default it is set to None which
    in torch distributed land, it will result in the global communicator.
    """

    def __init__(self, name):
        super().__init__(name)
        self._elapsed = 0.0
        self._active_time = 0.0
        self._started = False
        self._barrier_group = None
        self._start_time = time.time()
        self.nvtx = False

    def set_barrier_group(self, barrier_group):
        self._barrier_group = barrier_group

    def start(self, barrier=False, nvtx_push=False, sync=False):
        assert not self._started, "timer has already been started"
        if barrier:
            _barrier(group=self._barrier_group)
        if torch.cuda.is_available() and sync:
            torch.cuda.synchronize()
        self._start_time = time.time()
        self._started = True
        if nvtx_push:
            nvtx.range_push("{}".format(self.name))
            self.nvtx = True

    def stop(self, barrier=False, sync=False):
        if self.nvtx:
            nvtx.range_pop()
        assert self._started, "timer is not started"
        if barrier:
            _barrier(group=self._barrier_group)
        if torch.cuda.is_available() and sync:
            torch.cuda.synchronize()
        elapsed = time.time() - self._start_time
        self._elapsed += elapsed
        self._active_time += elapsed
        self._started = False

    def reset(self):
        self._elapsed = 0.0
        self._started = False

    def elapsed(self, reset=True, barrier=False):
        _started = self._started
        if self._started:
            self.stop(barrier=barrier)
        _elapsed = self._elapsed
        if reset:
            self.reset()
        if _started:
            self.start(barrier=barrier)
        return _elapsed

    def active_time(self):
        return self._active_time


class Timers:
    """Class for a group of Timers."""

    def __init__(self, log_level, log_option):
        self._log_level = log_level
        allowed_log_options = set(["max", "minmax", "all"])
        assert (
            log_option in allowed_log_options
        ), "input log option {} is invalid. It must be one of {}".format(
            log_option, allowed_log_options
        )
        self._log_option = log_option
        self._timers = {}
        self._log_levels = {}
        self._dummy_timer = DummyTimer()
        self._max_log_level = 2

    def __call__(self, name, log_level=None):
        if name in self._timers:
            if log_level is not None:
                assert log_level == self._log_levels[name], (
                    "input log level {} does not match already existing "
                    "log level {} for {} timer".format(
                        log_level, self._log_levels[name], name
                    )
                )
            return self._timers[name]
        if log_level is None:
            log_level = self._max_log_level
        assert (
            log_level <= self._max_log_level
        ), "log level {} is larger than max supported log level {}".format(
            log_level, self._max_log_level
        )
        if log_level > self._log_level:
            return self._dummy_timer
        self._timers[name] = Timer(name)
        self._log_levels[name] = log_level
        return self._timers[name]

    def _get_elapsed_time_all_ranks(self, names, reset, barrier):
        if barrier:
            _barrier()

        world_size = _get_world_size()
        rank = _get_rank()

        if torch.cuda.is_available():
            device = torch.cuda.current_device()
        else:
            device = torch.device("cpu")

        rank_name_to_time = torch.zeros(
            (world_size, len(names)), dtype=torch.float, device=device
        )

        for i, name in enumerate(names):
            if name in self._timers:
                rank_name_to_time[rank, i] = self._timers[name].elapsed(reset=reset)

        if world_size > 1 and _is_distributed() and dist_all_gather_func is not None:
            try:
                dist_all_gather_func(
                    rank_name_to_time.view(-1), rank_name_to_time[rank, :].view(-1)
                )
            except Exception as e:
                logger.warning("all_gather failed: %s. Using single rank timing.", e)

        return rank_name_to_time

    def _get_global_min_max_time(self, names, reset, barrier, normalizer):
        rank_name_to_time = self._get_elapsed_time_all_ranks(names, reset, barrier)
        name_to_min_max_time = {}
        for i, name in enumerate(names):
            rank_to_time = rank_name_to_time[:, i]
            rank_to_time = rank_to_time[rank_to_time > 0.0]
            if rank_to_time.numel() > 0:
                name_to_min_max_time[name] = (
                    rank_to_time.min().item() / normalizer,
                    rank_to_time.max().item() / normalizer,
                )
        return name_to_min_max_time

    def _get_global_min_max_time_string(
        self, names, reset, barrier, normalizer, max_only
    ):
        name_to_min_max_time = self._get_global_min_max_time(
            names, reset, barrier, normalizer
        )
        if not name_to_min_max_time:
            return None

        world_size = _get_world_size()
        if world_size == 1:
            output_string = "time (ms):"
            for name in name_to_min_max_time:
                _, max_time = name_to_min_max_time[name]
                output_string += "\n    {}: {:.2f}".format(
                    (name + " ").ljust(48, "."), max_time
                )
        else:
            if max_only:
                output_string = "max time across ranks (ms):"
            else:
                output_string = "(min, max) time across ranks (ms):"
            for name in name_to_min_max_time:
                min_time, max_time = name_to_min_max_time[name]
                if max_only:
                    output_string += "\n    {}: {:.2f}".format(
                        (name + " ").ljust(48, "."), max_time
                    )
                else:
                    output_string += "\n    {}: ({:.2f}, {:.2f})".format(
                        (name + " ").ljust(48, "."), min_time, max_time
                    )
        return output_string

    def _get_all_ranks_time_string(self, names, reset, barrier, normalizer):
        rank_name_to_time = self._get_elapsed_time_all_ranks(names, reset, barrier)
        world_size = _get_world_size()

        output_string = "times across ranks (ms):"
        no_reported_timing = True
        for i, name in enumerate(names):
            not_yet_found = True
            for rank in range(world_size):
                if rank_name_to_time[rank, i] > 0:
                    no_reported_timing = False
                    if not_yet_found:
                        not_yet_found = False
                        output_string += "\n  {}:".format(name)
                    if world_size == 1:
                        output_string += "\n     {:.2f}".format(
                            rank_name_to_time[rank, i] / normalizer
                        )
                    else:
                        output_string += "\n     rank {:2d}: {:.2f}".format(
                            rank, rank_name_to_time[rank, i] / normalizer
                        )
        if no_reported_timing:
            return None
        return output_string

    def get_all_timers_string(
        self,
        names: List[str] = None,
        normalizer: float = 1.0,
        reset: bool = True,
        barrier: bool = False,
    ):
        """Return a formatted timing string for the given timer names.

        Args:
            names: Timers to include; defaults to all registered timers.
            normalizer: Divide raw seconds by this value (e.g. 1000 for ms output).
            reset: Reset each timer after reading its elapsed time.
            barrier: Synchronize across ranks before gathering times.
        """
        if names is None:
            names = list(self._timers.keys())

        assert normalizer > 0.0
        if self._log_option in ["max", "minmax"]:
            max_only = self._log_option == "max"
            output_string = self._get_global_min_max_time_string(
                names, reset, barrier, normalizer / 1000.0, max_only
            )
        elif self._log_option == "all":
            output_string = self._get_all_ranks_time_string(
                names, reset, barrier, normalizer / 1000.0
            )
        else:
            raise Exception("unknown timing log option {}".format(self._log_option))
        return output_string

    def log(
        self,
        names: List[str],
        rank: int = None,
        normalizer: float = 1.0,
        reset: bool = True,
        barrier: bool = False,
    ):
        """Print timing results for the given names to stdout on one rank.

        Args:
            names: Timer names to log.
            rank: Rank that prints; defaults to the last rank (world_size - 1).
            normalizer: Divide raw seconds by this value before printing.
            reset: Reset each timer after reading.
            barrier: Synchronize across ranks first.
        """
        output_string = self.get_all_timers_string(names, normalizer, reset, barrier)
        world_size = _get_world_size()
        current_rank = _get_rank()

        if rank is None:
            rank = world_size - 1
        if rank == current_rank and output_string is not None:
            logger.info("%s", output_string)

    def write(
        self,
        names: List[str],
        writer,
        iteration: int,
        normalizer: float = 1.0,
        reset: bool = True,
        barrier: bool = False,
    ):
        """Write per-timer max times as TensorBoard scalars.

        Args:
            names: Timer names to write.
            writer: TensorBoard SummaryWriter instance.
            iteration: Global step value for the scalar.
            normalizer: Divide raw seconds by this value.
            reset: Reset each timer after reading.
            barrier: Synchronize across ranks first.
        """
        assert normalizer > 0.0
        name_to_min_max_time = self._get_global_min_max_time(
            names, reset, barrier, normalizer
        )
        if writer is not None:
            for name in name_to_min_max_time:
                _, max_time = name_to_min_max_time[name]
                writer.add_scalar(name + "-time", max_time, iteration)