Added option to run in processpool.

Author: s3rius

taskiq/api/receiver.py

@@ -1,5 +1,5 @@
 import asyncio
-from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import Executor, ProcessPoolExecutor, ThreadPoolExecutor
 from logging import getLogger
 from typing import Optional, Type
 
@@ -13,13 +13,14 @@
 async def run_receiver_task(
     broker: AsyncBroker,
     receiver_cls: Type[Receiver] = Receiver,
-    sync_workers: int = 4,
+    sync_workers: Optional[int] = None,
     validate_params: bool = True,
     max_async_tasks: int = 100,
     max_prefetch: int = 0,
     propagate_exceptions: bool = True,
     run_startup: bool = False,
     ack_time: Optional[AcknowledgeType] = None,
+    use_process_pool: bool = False,
 ) -> None:
     """
     Function to run receiver programmatically.
@@ -46,6 +47,7 @@ async def run_receiver_task(
     :param propagate_exceptions: whether to propagate exceptions in generators or not.
     :param run_startup: whether to run startup function or not.
     :param ack_time: acknowledge type to use.
+    :param use_process_pool: whether to use process pool or threadpool.
     :raises asyncio.CancelledError: if the task was cancelled.
     """
     finish_event = asyncio.Event()
@@ -62,7 +64,12 @@ def on_exit(_: Receiver) -> None:
         finish_event.set()
         raise asyncio.CancelledError
 
-    with ThreadPoolExecutor(max_workers=sync_workers) as executor:
+    executor: Executor
+    if use_process_pool:
+        executor = ProcessPoolExecutor(max_workers=sync_workers)
+    else:
+        executor = ThreadPoolExecutor(max_workers=sync_workers)
+    with executor as executor:
         broker.is_worker_process = True
         while True:
             try:

taskiq/cli/worker/args.py

@@ -46,6 +46,7 @@ class WorkerArgs:
     max_tasks_per_child: Optional[int] = None
     wait_tasks_timeout: Optional[float] = None
     hardkill_count: int = 3
+    use_process_pool: bool = False
 
     @classmethod
     def from_cli(
@@ -210,8 +211,7 @@ def from_cli(
             "--wait-tasks-timeout",
             type=float,
             default=None,
-            help="Maximum time to wait for all current tasks "
-            "to finish before exiting.",
+            help="Maximum time to wait for all current tasks to finish before exiting.",
         )
         parser.add_argument(
             "--hardkill-count",
@@ -220,6 +220,11 @@ def from_cli(
             help="Number of termination signals to the main "
             "process before performing a hardkill.",
         )
+        parser.add_argument(
+            "--use-process-pool",
+            action="store_true",
+            help="Use process pool instead of thread pool for sync tasks.",
+        )
 
         namespace = parser.parse_args(args)
         # If there are any patterns specified, remove default.

taskiq/cli/worker/run.py

@@ -3,7 +3,7 @@
 import os
 import signal
 import sys
-from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import Executor, ProcessPoolExecutor, ThreadPoolExecutor
 from multiprocessing import set_start_method
 from sys import platform
 from typing import Any, Optional, Type
@@ -135,9 +135,15 @@ def interrupt_handler(signum: int, _frame: Any) -> None:
     receiver_type = get_receiver_type(args)
     receiver_kwargs = dict(args.receiver_arg)
 
+    executor: Executor
+    if args.use_process_pool:
+        executor = ProcessPoolExecutor(max_workers=args.max_threadpool_threads)
+    else:
+        executor = ThreadPoolExecutor(max_workers=args.max_threadpool_threads)
+
     try:
         logger.debug("Initialize receiver.")
-        with ThreadPoolExecutor(args.max_threadpool_threads) as pool:
+        with executor as pool:
             receiver = receiver_type(
                 broker=broker,
                 executor=pool,