Timeout handling while using run_in_executor and asyncio
Question:
I’m using asyncio to run a piece of blocking code like this:
result = await loop.run_in_executor(None, long_running_function)
My question is: Can I impose a timeout for the execution of long_running_function
?
Basically I don’t want long_running_function
to last more than 2 seconds and I can’t do proper timeout handling within it because that function comes from a third-party library.
Answers:
You could use asyncio.wait_for:
future = loop.run_in_executor(None, long_running_function)
result = await asyncio.wait_for(future, timeout, loop=loop)
A warning about cancelling long running functions:
Although wrapping the Future
returned by loop.run_in_executor
with an asyncio.wait_for
call will allow the event loop to stop waiting for long_running_function
after some x
seconds, it won’t necessarily stop the underlying long_running_function
. This is one of the shortcomings of concurrent.futures
and to the best of my knowledge there is no simple way to just cancel a concurrent.futures.Future
.
although not using run_in_executor
, i have some workaround about "wrap a block function async with timeout handling"
import asyncio
import threading
import time
import ctypes
def terminate_thread(t: threading.Thread, exc_type=SystemExit):
if not t.is_alive(): return
try:
tid = next(tid for tid, tobj in threading._active.items() if tobj is t)
except StopIteration:
raise ValueError("tid not found")
if ctypes.pythonapi.PyThreadState_SetAsyncExc(tid, ctypes.py_object(exc_type)) != 1:
raise SystemError("PyThreadState_SetAsyncExc failed")
class AsyncResEvent(asyncio.Event):
def __init__(self):
super().__init__()
self.res = None
self.is_exc = False
self._loop = asyncio.get_event_loop()
def set(self, data=None) -> None:
self.res = data
self.is_exc = False
self._loop.call_soon_threadsafe(super().set)
def set_exception(self, exc) -> None:
self.res = exc
self.is_exc = True
self._loop.call_soon_threadsafe(super().set)
async def wait(self, timeout: float | None = None):
await asyncio.wait_for(super().wait(), timeout)
if self.is_exc:
raise self.res
else:
return self.res
async def sub_thread_async(func, *args, _timeout: float | None = None, **kwargs):
res = AsyncResEvent()
def f():
try:
res.set(func(*args, **kwargs))
except Exception as e:
res.set_exception(e)
except SystemExit:
res.set_exception(TimeoutError)
(t := threading.Thread(target=f)).start()
try:
return await res.wait(_timeout)
except TimeoutError:
raise TimeoutError
finally:
if not res.is_set():
terminate_thread(t)
_lock = threading.Lock()
def test(n):
_tid = threading.get_ident()
for i in range(n):
with _lock:
print(f'print from thread {_tid} ({i})')
time.sleep(1)
return n
async def main():
res_normal = await asyncio.gather(*(sub_thread_async(test, 5) for _ in range(2)))
print(res_normal) # [5,5]
res_normal_2 = await asyncio.gather(*(sub_thread_async(test, 2, _timeout=3) for _ in range(2)))
print(res_normal_2) # [2,2]
res_should_not_get = await asyncio.gather(*(sub_thread_async(test, 5, _timeout=3) for _ in range(2)))
print(res_should_not_get) # timeout error
if __name__ == '__main__':
asyncio.new_event_loop().run_until_complete(main())
I’m using asyncio to run a piece of blocking code like this:
result = await loop.run_in_executor(None, long_running_function)
My question is: Can I impose a timeout for the execution of long_running_function
?
Basically I don’t want long_running_function
to last more than 2 seconds and I can’t do proper timeout handling within it because that function comes from a third-party library.
You could use asyncio.wait_for:
future = loop.run_in_executor(None, long_running_function)
result = await asyncio.wait_for(future, timeout, loop=loop)
A warning about cancelling long running functions:
Although wrapping the Future
returned by loop.run_in_executor
with an asyncio.wait_for
call will allow the event loop to stop waiting for long_running_function
after some x
seconds, it won’t necessarily stop the underlying long_running_function
. This is one of the shortcomings of concurrent.futures
and to the best of my knowledge there is no simple way to just cancel a concurrent.futures.Future
.
although not using run_in_executor
, i have some workaround about "wrap a block function async with timeout handling"
import asyncio
import threading
import time
import ctypes
def terminate_thread(t: threading.Thread, exc_type=SystemExit):
if not t.is_alive(): return
try:
tid = next(tid for tid, tobj in threading._active.items() if tobj is t)
except StopIteration:
raise ValueError("tid not found")
if ctypes.pythonapi.PyThreadState_SetAsyncExc(tid, ctypes.py_object(exc_type)) != 1:
raise SystemError("PyThreadState_SetAsyncExc failed")
class AsyncResEvent(asyncio.Event):
def __init__(self):
super().__init__()
self.res = None
self.is_exc = False
self._loop = asyncio.get_event_loop()
def set(self, data=None) -> None:
self.res = data
self.is_exc = False
self._loop.call_soon_threadsafe(super().set)
def set_exception(self, exc) -> None:
self.res = exc
self.is_exc = True
self._loop.call_soon_threadsafe(super().set)
async def wait(self, timeout: float | None = None):
await asyncio.wait_for(super().wait(), timeout)
if self.is_exc:
raise self.res
else:
return self.res
async def sub_thread_async(func, *args, _timeout: float | None = None, **kwargs):
res = AsyncResEvent()
def f():
try:
res.set(func(*args, **kwargs))
except Exception as e:
res.set_exception(e)
except SystemExit:
res.set_exception(TimeoutError)
(t := threading.Thread(target=f)).start()
try:
return await res.wait(_timeout)
except TimeoutError:
raise TimeoutError
finally:
if not res.is_set():
terminate_thread(t)
_lock = threading.Lock()
def test(n):
_tid = threading.get_ident()
for i in range(n):
with _lock:
print(f'print from thread {_tid} ({i})')
time.sleep(1)
return n
async def main():
res_normal = await asyncio.gather(*(sub_thread_async(test, 5) for _ in range(2)))
print(res_normal) # [5,5]
res_normal_2 = await asyncio.gather(*(sub_thread_async(test, 2, _timeout=3) for _ in range(2)))
print(res_normal_2) # [2,2]
res_should_not_get = await asyncio.gather(*(sub_thread_async(test, 5, _timeout=3) for _ in range(2)))
print(res_should_not_get) # timeout error
if __name__ == '__main__':
asyncio.new_event_loop().run_until_complete(main())