This repository was archived by the owner on Feb 24, 2022. It is now read-only.
forked from facebookresearch/TransCoder
-
Notifications
You must be signed in to change notification settings - Fork 1
Expand file tree
/
Copy pathtimeout.py
More file actions
executable file
·51 lines (42 loc) · 1.6 KB
/
Copy pathtimeout.py
File metadata and controls
executable file
·51 lines (42 loc) · 1.6 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
# Copyright (c) 2019-present, Facebook, Inc.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
#
import errno
import math
import os
import signal
import time
from functools import partial, wraps
class TimeoutError(BaseException):
pass
def timeout(seconds=10, error_message=os.strerror(errno.ETIME)):
def decorator(func):
def _handle_timeout(repeat_id, signum, frame):
# logger.warning(f"Catched the signal ({repeat_id}) Setting signal handler {repeat_id + 1}")
signal.signal(signal.SIGALRM, partial(
_handle_timeout, repeat_id + 1))
signal.alarm(seconds)
raise TimeoutError(error_message)
def wrapper(*args, **kwargs):
old_signal = signal.signal(
signal.SIGALRM, partial(_handle_timeout, 0))
old_time_left = signal.alarm(seconds)
assert type(old_time_left) is int and old_time_left >= 0
if 0 < old_time_left < seconds: # do not exceed previous timer
signal.alarm(old_time_left)
start_time = time.time()
try:
result = func(*args, **kwargs)
finally:
if old_time_left == 0:
signal.alarm(0)
else:
sub = time.time() - start_time
signal.signal(signal.SIGALRM, old_signal)
signal.alarm(max(0, math.ceil(old_time_left - sub)))
return result
return wraps(func)(wrapper)
return decorator