|
|
import time |
|
|
from dataclasses import dataclass |
|
|
from enum import Enum |
|
|
|
|
|
import torch |
|
|
|
|
|
from finetrainers.constants import FINETRAINERS_ENABLE_TIMING |
|
|
from finetrainers.logging import get_logger |
|
|
|
|
|
|
|
|
logger = get_logger() |
|
|
|
|
|
|
|
|
class TimerDevice(str, Enum): |
|
|
CPU = "cpu" |
|
|
CUDA = "cuda" |
|
|
|
|
|
|
|
|
@dataclass |
|
|
class TimerData: |
|
|
name: str |
|
|
device: TimerDevice |
|
|
start_time: float = 0.0 |
|
|
end_time: float = 0.0 |
|
|
|
|
|
|
|
|
class Timer: |
|
|
def __init__(self, name: str, device: TimerDevice, device_sync: bool = False): |
|
|
self.data = TimerData(name=name, device=device) |
|
|
|
|
|
self._device_sync = device_sync |
|
|
self._start_event = None |
|
|
self._end_event = None |
|
|
self._active = False |
|
|
self._enabled = FINETRAINERS_ENABLE_TIMING |
|
|
|
|
|
def __enter__(self): |
|
|
self.start() |
|
|
return self |
|
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb): |
|
|
self.end() |
|
|
return False |
|
|
|
|
|
def start(self): |
|
|
if self._active: |
|
|
logger.warning(f"Timer {self.data.name} is already running. Please stop it before starting again.") |
|
|
return |
|
|
self._active = True |
|
|
if not self._enabled: |
|
|
return |
|
|
if self.data.device == TimerDevice.CUDA and torch.cuda.is_available(): |
|
|
self._start_cuda() |
|
|
else: |
|
|
self._start_cpu() |
|
|
if not self.data.device == TimerDevice.CPU: |
|
|
logger.warning( |
|
|
f"Timer device {self.data.device} is either not supported or incorrect device selected. Falling back to CPU." |
|
|
) |
|
|
|
|
|
def end(self): |
|
|
if not self._active: |
|
|
logger.warning(f"Timer {self.data.name} is not running. Please start it before stopping.") |
|
|
return |
|
|
self._active = False |
|
|
if not self._enabled: |
|
|
return |
|
|
if self.data.device == TimerDevice.CUDA and torch.cuda.is_available(): |
|
|
self._end_cuda() |
|
|
else: |
|
|
self._end_cpu() |
|
|
if not self.data.device == TimerDevice.CPU: |
|
|
logger.warning( |
|
|
f"Timer device {self.data.device} is either not supported or incorrect device selected. Falling back to CPU." |
|
|
) |
|
|
|
|
|
@property |
|
|
def elapsed_time(self) -> float: |
|
|
if self._active: |
|
|
if self.data.device == TimerDevice.CUDA and torch.cuda.is_available(): |
|
|
premature_end_event = torch.cuda.Event(enable_timing=True) |
|
|
premature_end_event.record() |
|
|
premature_end_event.synchronize() |
|
|
return self._start_event.elapsed_time(premature_end_event) / 1000.0 |
|
|
else: |
|
|
return time.time() - self.data.start_time |
|
|
else: |
|
|
if self.data.device == TimerDevice.CUDA and torch.cuda.is_available(): |
|
|
return self._start_event.elapsed_time(self._end_event) / 1000.0 |
|
|
else: |
|
|
return self.data.end_time - self.data.start_time |
|
|
|
|
|
def _start_cpu(self): |
|
|
self.data.start_time = time.time() |
|
|
|
|
|
def _start_cuda(self): |
|
|
torch.cuda.synchronize() |
|
|
self._start_event = torch.cuda.Event(enable_timing=True) |
|
|
self._end_event = torch.cuda.Event(enable_timing=True) |
|
|
self._start_event.record() |
|
|
|
|
|
def _end_cpu(self): |
|
|
self.data.end_time = time.time() |
|
|
|
|
|
def _end_cuda(self): |
|
|
if self._device_sync: |
|
|
torch.cuda.synchronize() |
|
|
self._end_event.record() |
|
|
|
