mirror of
https://github.com/qdrant/fastembed.git
synced 2026-10-03 19:37:46 -05:00
e572d004f7
* fix: terminate interrupted parallel workers * fix: reap terminated parallel workers * test: cover worker cleanup after pool reuse * fix: bound parallel worker shutdown and free unsent input batches * fix: reset emergency shutdown on start --------- Co-authored-by: FU-max-boop <214359569+FU-max-boop@users.noreply.github.com> Co-authored-by: George Panchuk <george.panchuk@qdrant.tech>
41 lines
1.4 KiB
Python
41 lines
1.4 KiB
Python
import threading
|
|
from itertools import count
|
|
from multiprocessing import get_all_start_methods
|
|
|
|
import pytest
|
|
|
|
from fastembed.parallel_processor import ParallelWorkerPool, Worker
|
|
|
|
|
|
class EchoWorker(Worker):
|
|
@classmethod
|
|
def start(cls, **kwargs):
|
|
return cls()
|
|
|
|
def process(self, items):
|
|
yield from items
|
|
|
|
|
|
# semi_ordered_map is closed by garbage collection, so an error in its cleanup is only reported as unraisable
|
|
@pytest.mark.filterwarnings("error::pytest.PytestUnraisableExceptionWarning")
|
|
def test_closing_partially_consumed_iterator_stops_workers():
|
|
start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
|
|
pool = ParallelWorkerPool(2, EchoWorker, start_method=start_method)
|
|
# the stream never ends, so the workers never get their stop signals
|
|
results = pool.ordered_map(count())
|
|
assert next(results) == 0
|
|
workers = list(pool.processes)
|
|
|
|
# close() used to wait in join() forever, run it in a thread so a regression fails instead of
|
|
# hanging the test session
|
|
closer = threading.Thread(target=results.close, daemon=True)
|
|
closer.start()
|
|
closer.join(timeout=30)
|
|
try:
|
|
assert not closer.is_alive(), "closing the iterator hung"
|
|
assert not any(worker.is_alive() for worker in workers)
|
|
finally:
|
|
for worker in workers:
|
|
if worker.is_alive():
|
|
worker.kill()
|