def end()

in src/nanotron/logging/timers.py [0:0]

• 24 lines of code
• 11 McCabe index (conditional complexity)

    def end(self) -> None:
        """End the timer, but don't compute elapsed time yet."""

        if self.name == "dummy" or not self.enabled:  # disabled
            return

        if not self.running:
            logger.warning(f"Timer '{self.name}' was not running. Ignoring end call.")
            return

        if self.timer_type == TimerType.CUDA:
            if torch.cuda.is_available() and self._current_start_event is not None:
                # Synchronize before ending timing if requested
                if self.cuda_sync:
                    torch.cuda.synchronize()
                # Create and record an end event
                end_event = torch.cuda.Event(enable_timing=True)
                end_event.record()

                # Store the start/end event pair for later querying
                self._cuda_events.append((self._current_start_event, end_event))
                self._current_start_event = None
            else:
                logger.warning("CUDA timer end called but CUDA events are not available.")
                self.timer_type = TimerType.CPU
                self.end_time = time.time()
                self._cpu_total_time += self.end_time - self.start_time
        else:
            self.end_time = time.time()
            self._cpu_total_time += self.end_time - self.start_time

        self.call_count += 1
        self.running = False