diff --git a/ms_agent/utils/thread_util.py b/ms_agent/utils/thread_util.py index 380dd99b3..ed2922901 100644 --- a/ms_agent/utils/thread_util.py +++ b/ms_agent/utils/thread_util.py @@ -102,8 +102,7 @@ def weakref_cb(_, q=self._work_queue): thread_name = '%s_%d' % (self._thread_name_prefix or self, num_threads) # Import internal helpers from stdlib to keep behavior consistent. - from concurrent.futures.thread import ( # type: ignore - _threads_queues, _worker) + from concurrent.futures.thread import _worker # type: ignore t = threading.Thread( name=thread_name, @@ -114,4 +113,7 @@ def weakref_cb(_, q=self._work_queue): t.daemon = True t.start() self._threads.add(t) - _threads_queues[t] = self._work_queue + # Do not register daemon workers in _threads_queues: CPython + # joins every registered worker at interpreter exit, even after + # shutdown(wait=False). Explicit shutdown(wait=True) still joins + # these workers through self._threads. diff --git a/tests/utils/test_thread_util.py b/tests/utils/test_thread_util.py new file mode 100644 index 000000000..29d6a8735 --- /dev/null +++ b/tests/utils/test_thread_util.py @@ -0,0 +1,72 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Daemon workers must not turn non-waiting shutdown into an exit-time join.""" +import subprocess +import sys +import textwrap +import threading + +import pytest +from ms_agent.utils.thread_util import DaemonThreadPoolExecutor + + +def test_nonwaiting_shutdown_does_not_block_process_exit(): + script = textwrap.dedent(''' + import threading + from ms_agent.utils.thread_util import DaemonThreadPoolExecutor + + started = threading.Event() + blocked = threading.Event() + + def work(): + started.set() + blocked.wait() + + executor = DaemonThreadPoolExecutor(max_workers=1) + executor.submit(work) + assert started.wait(5) + executor.shutdown(wait=False, cancel_futures=True) + print('shutdown returned', flush=True) + ''') + result = subprocess.run( + [sys.executable, '-c', script], capture_output=True, text=True, + timeout=15) + assert result.returncode == 0, result.stderr + assert 'shutdown returned' in result.stdout + + +def test_results_and_exceptions_are_preserved(): + with DaemonThreadPoolExecutor(max_workers=1) as executor: + assert executor.submit(lambda: 42).result(timeout=5) == 42 + future = executor.submit(lambda: 1 / 0) + with pytest.raises(ZeroDivisionError): + future.result(timeout=5) + + +def test_waiting_shutdown_joins_workers(): + executor = DaemonThreadPoolExecutor(max_workers=1) + future = executor.submit(lambda: threading.current_thread()) + worker = future.result(timeout=5) + assert worker.daemon + executor.shutdown(wait=True) + assert not worker.is_alive() + + +def test_nonwaiting_shutdown_cancels_queued_work(): + started = threading.Event() + release = threading.Event() + + def work(): + started.set() + release.wait(5) + + executor = DaemonThreadPoolExecutor(max_workers=1) + running = executor.submit(work) + try: + assert started.wait(5) + queued = executor.submit(lambda: 'should not run') + executor.shutdown(wait=False, cancel_futures=True) + assert queued.cancelled() + finally: + release.set() + running.result(timeout=5) + executor.shutdown(wait=True)