diff --git a/deeplabcut/gui/window.py b/deeplabcut/gui/window.py index de40d326fa..1e967fa4a2 100644 --- a/deeplabcut/gui/window.py +++ b/deeplabcut/gui/window.py @@ -16,7 +16,6 @@ from pathlib import Path from typing import List from urllib.error import URLError -from concurrent.futures import ThreadPoolExecutor, TimeoutError import qdarkstyle import deeplabcut @@ -25,6 +24,7 @@ from deeplabcut.gui import BASE_DIR, components, utils from deeplabcut.gui.tabs import * from deeplabcut.gui.widgets import StreamReceiver, StreamWriter +from deeplabcut.utils.multiprocessing import call_with_timeout from napari_deeplabcut import misc from PySide6.QtWidgets import ( QMessageBox, @@ -41,19 +41,13 @@ from PySide6.QtCore import Qt, QTimer -def call_with_timeout(func, timeout, *args, **kwargs): - with ThreadPoolExecutor(max_workers=1) as executor: - future = executor.submit(func, *args, **kwargs) - return future.result(timeout=timeout) - - def _check_for_updates(silent=True): try: is_latest, latest_version = call_with_timeout( - utils.is_latest_deeplabcut_version, 1 + utils.is_latest_deeplabcut_version, 5 ) is_latest_plugin, latest_plugin_version = call_with_timeout( - misc.is_latest_version, 1 + misc.is_latest_version, 5 ) except (URLError, TimeoutError): # Handle internet connectivity issues is_latest = is_latest_plugin = True diff --git a/deeplabcut/utils/multiprocessing.py b/deeplabcut/utils/multiprocessing.py new file mode 100644 index 0000000000..3515b73125 --- /dev/null +++ b/deeplabcut/utils/multiprocessing.py @@ -0,0 +1,54 @@ +# +# DeepLabCut Toolbox (deeplabcut.org) +# © A. & M.W. Mathis Labs +# https://github.com/DeepLabCut/DeepLabCut +# +# Please see AUTHORS for contributors. +# https://github.com/DeepLabCut/DeepLabCut/blob/master/AUTHORS +# +# Licensed under GNU Lesser General Public License v3.0 +# +""" +DeepLabCut2.2 Toolbox (deeplabcut.org) +© A. & M. Mathis Labs +https://github.com/DeepLabCut/DeepLabCut +Please see AUTHORS for contributors. + +https://github.com/DeepLabCut/DeepLabCut/blob/master/AUTHORS +Licensed under GNU Lesser General Public License v3.0 +""" +import multiprocessing + + +def _wrapper(func, queue, *args, **kwargs): + try: + result = func(*args, **kwargs) + queue.put(result) # Pass the result back via the queue + except Exception as e: + queue.put(e) # Pass any exception back via the queue + + +def call_with_timeout(func, timeout, *args, **kwargs): + queue = multiprocessing.Queue() + process = multiprocessing.Process( + target=_wrapper, args=(func, queue, *args), kwargs=kwargs + ) + process.start() + process.join(timeout) + + if process.is_alive(): + process.terminate() # Forcefully terminate the process + process.join() + raise TimeoutError( + f"Function {func.__name__} did not complete within {timeout} seconds." + ) + + if not queue.empty(): + result = queue.get() + if isinstance(result, Exception): + raise result # Re-raise the exception if it occurred in the function + return result + else: + raise TimeoutError( + f"Function {func.__name__} completed but did not return a result." + ) diff --git a/tests/utils/test_multiprocessing.py b/tests/utils/test_multiprocessing.py new file mode 100644 index 0000000000..34333be81d --- /dev/null +++ b/tests/utils/test_multiprocessing.py @@ -0,0 +1,37 @@ +# +# DeepLabCut Toolbox (deeplabcut.org) +# © A. & M.W. Mathis Labs +# https://github.com/DeepLabCut/DeepLabCut +# +# Please see AUTHORS for contributors. +# https://github.com/DeepLabCut/DeepLabCut/blob/master/AUTHORS +# +# Licensed under GNU Lesser General Public License v3.0 +# +import pytest +import time +from deeplabcut.utils.multiprocessing import call_with_timeout + + +def _succeeding_method(parameter): + return parameter + + +def _failing_method(): + raise ValueError("Raise value error on purpose") + + +def _hanging_method(): + while True: + time.sleep(5) + + +def test_call_with_timeout(): + parameter = (10, "Hello test") + assert call_with_timeout(_succeeding_method, 30, parameter) == parameter + + with pytest.raises(ValueError): + call_with_timeout(_failing_method, timeout=30) + + with pytest.raises(TimeoutError): + call_with_timeout(_hanging_method, timeout=1)