Class _Timeout

bench.py:43–82

Context-manager wall-clock timeout. Works on both Unix (SIGALRM) and Windows (thread-based fallback).

class _Timeout:
    """Context-manager wall-clock timeout. Works on both Unix (SIGALRM) and
    Windows (thread-based fallback)."""

    def __init__(self, seconds: int):
        self.seconds = seconds
        self._use_signal = hasattr(signal, "SIGALRM")

    # --- signal-based (Unix) -------------------------------------------
    def _handler(self, signum, frame):
        raise BenchTimeoutError(f"Timed out after {self.seconds}s")

    def __enter__(self):
        if self._use_signal:
            self._old = signal.signal(signal.SIGALRM, self._handler)
            signal.alarm(self.seconds)
        else:
            import threading
            self._timer = threading.Timer(self.seconds, self._timeout_thread)
            self._timer.daemon = True
            self._timed_out = False
            self._timer.start()
        return self

    def __exit__(self, *exc):
        if self._use_signal:
            signal.alarm(0)
            signal.signal(signal.SIGALRM, self._old)
        else:
            self._timer.cancel()
        return False

    def _timeout_thread(self):
        self._timed_out = True
        # On Windows we cannot forcefully interrupt the main thread the same
        # way SIGALRM does.  We set a flag; callers that iterate can check it.
        # For truly blocking GPU calls, this will not help -- but at least
        # the outer try/except will catch it after the call returns.
        import _thread
        _thread.interrupt_main()


# =========================================================================