Added multiprocess wrapper to do simple function calls in a seperate process

This commit is contained in:
Kevin F 2023-01-12 21:42:22 -06:00
parent 39c01fdbc5
commit 9eb2e5d413
1 changed files with 31 additions and 0 deletions

31
src/utils/multiproc.py Normal file
View File

@ -0,0 +1,31 @@
from typing import Callable
import multiprocessing
import time
def _compute(q: multiprocessing.Queue, func: Callable, *args, **kwargs):
q.put(func(*args, **kwargs))
def subprocess_compute(func: Callable, wallclock_timeout: int, *args, **kwargs):
"""
Call func in a subprocess, and return the result. Set wallclock_timeout to <= 0 to disable
Raises TimeoutError if the function does not return in time
Raises ChildProcessError if the subprocess dies before returning
"""
q = multiprocessing.Queue()
p = multiprocessing.Process(
target=_compute, args=(q, func, *args), kwargs=kwargs, daemon=True)
wallclock_timeout = max(wallclock_timeout, 0)
p.start()
start = time.time()
while True:
try:
return q.get(timeout=1)
except multiprocessing.queues.Empty:
if not p.is_alive():
raise ChildProcessError("Process died before returning")
if wallclock_timeout:
if time.time() - start >= wallclock_timeout:
raise TimeoutError("Process did not return in time")