Skip to content

Commit

Permalink
Add timeout() class (ultralytics#3460)
Browse files Browse the repository at this point in the history
* Add `timeout()` class

* rearrange order
  • Loading branch information
glenn-jocher committed Jun 4, 2021
1 parent fd86784 commit 1678c00
Showing 1 changed file with 23 additions and 1 deletion.
24 changes: 23 additions & 1 deletion utils/general.py
Original file line number Diff line number Diff line change
@@ -1,12 +1,14 @@
# YOLOv5 general utils

import contextlib
import glob
import logging
import math
import os
import platform
import random
import re
import signal
import subprocess
import time
import urllib
Expand Down Expand Up @@ -34,6 +36,26 @@
os.environ['NUMEXPR_MAX_THREADS'] = str(min(os.cpu_count(), 8)) # NumExpr max threads


class timeout(contextlib.ContextDecorator):
# Usage: @timeout(seconds) decorator or 'with timeout(seconds):' context manager
def __init__(self, seconds, *, timeout_message="", suppress_timeout_errors=True):
self.seconds = int(seconds)
self.timeout_message = timeout_message
self.suppress = bool(suppress_timeout_errors)

def _timeout_handler(self, signum, frame):
raise TimeoutError(self.timeout_message)

def __enter__(self):
signal.signal(signal.SIGALRM, self._timeout_handler) # Set handler for SIGALRM
signal.alarm(self.seconds) # start countdown for SIGALRM to be raised

def __exit__(self, exc_type, exc_val, exc_tb):
signal.alarm(0) # Cancel SIGALRM if it's scheduled
if self.suppress and exc_type is TimeoutError: # Suppress TimeoutError
return True


def set_logging(rank=-1, verbose=True):
logging.basicConfig(
format="%(message)s",
Expand Down Expand Up @@ -86,7 +108,7 @@ def check_online():
# Check internet connectivity
import socket
try:
socket.create_connection(("1.1.1.1", 443), 5) # check host accesability
socket.create_connection(("1.1.1.1", 443), 5) # check host accessibility
return True
except OSError:
return False
Expand Down

0 comments on commit 1678c00

Please sign in to comment.