diff --git a/comet/core/execution.py b/comet/core/execution.py index 8967430..52821f9 100644 --- a/comet/core/execution.py +++ b/comet/core/execution.py @@ -1,17 +1,42 @@ +import atexit +import multiprocessing +import os +import signal from concurrent.futures import ProcessPoolExecutor +_mp_context = None +try: + _mp_context = multiprocessing.get_context("forkserver") +except ValueError: + _mp_context = multiprocessing.get_context("spawn") + app_executor = None -def setup_executor(): +def worker_initializer(): + signal.signal(signal.SIGINT, signal.SIG_IGN) + + +def setup_executor(max_workers: int | None = None): global app_executor - app_executor = ProcessPoolExecutor() + + if max_workers is None: + cpu_count = os.cpu_count() or 1 + max_workers = min(cpu_count, 4) + + app_executor = ProcessPoolExecutor( + max_workers=max_workers, mp_context=_mp_context, initializer=worker_initializer + ) def shutdown_executor(): global app_executor if app_executor: - app_executor.shutdown() + app_executor.shutdown(wait=True, cancel_futures=True) + app_executor = None + + +atexit.register(shutdown_executor) def get_executor(): diff --git a/comet/services/orchestration.py b/comet/services/orchestration.py index bc0c1ec..9edb1fe 100644 --- a/comet/services/orchestration.py +++ b/comet/services/orchestration.py @@ -172,7 +172,7 @@ class TorrentManager: return loop = asyncio.get_running_loop() - chunk_size = 50 + chunk_size = 20 tasks = [ loop.run_in_executor( get_executor(),