Fixed docs and qualname generation.

Author: s3rius

taskiq/api/receiver.py

@@ -40,7 +40,8 @@ async def run_receiver_task(
 
     :param broker: current broker instance.
     :param receiver_cls: receiver class to use.
-    :param sync_workers: number of threads of a threadpool that runs sync tasks.
+    :param sync_workers: number of threads of a threadpool
+        or processes in processpool that runs sync tasks.
     :param validate_params: whether to validate params or not.
     :param max_async_tasks: maximum number of simultaneous async tasks.
     :param max_prefetch: maximum number of tasks to prefetch.

taskiq/cli/worker/args.py

@@ -32,6 +32,7 @@ class WorkerArgs:
     log_level: LogLevel = LogLevel.INFO
     workers: int = 2
     max_threadpool_threads: int = 10
+    max_process_pool_processes: Optional[int] = None
     no_parse: bool = False
     shutdown_timeout: float = 5
     reload: bool = False
@@ -223,8 +224,16 @@ def from_cli(
         parser.add_argument(
             "--use-process-pool",
             action="store_true",
+            dest="use_process_pool",
             help="Use process pool instead of thread pool for sync tasks.",
         )
+        parser.add_argument(
+            "--max-process-pool-processes",
+            type=int,
+            dest="max_process_pool_processes",
+            default=None,
+            help="Maximum number of processes in process pool.",
+        )
 
         namespace = parser.parse_args(args)
         # If there are any patterns specified, remove default.

taskiq/cli/worker/run.py

@@ -137,7 +137,7 @@ def interrupt_handler(signum: int, _frame: Any) -> None:
 
     executor: Executor
     if args.use_process_pool:
-        executor = ProcessPoolExecutor(max_workers=args.max_threadpool_threads)
+        executor = ProcessPoolExecutor(max_workers=args.max_process_pool_processes)
     else:
         executor = ThreadPoolExecutor(max_workers=args.max_threadpool_threads)
 

taskiq/decor.py

@@ -55,17 +55,39 @@ def __init__(
         self.broker = broker
         self.task_name = task_name
         self.original_func = original_func
+        self.labels = labels
+
+        # This is a hack to make ProcessPoolExecutor work
+        # with decorated functions.
+        #
+        # The problem is that when we decorate a function
+        # it becomes a new class. This class has the same
+        # name as the original function.
+        #
+        # When receiver sends original function to another
+        # process, it will have the same name as the decorated
+        # class. This will cause an error, because ProcessPoolExecutor
+        # uses `__name__` and `__qualname__` attributes to
+        # import functions from other processes and then it verifies
+        # that the function is the same as the original one.
+        #
+        # This hack renames the original function and injects
+        # it back to the module where it was defined.
+        # This way ProcessPoolExecutor will be able to import
+        # the function by it's name and verify its correctness.
         new_name = f"{original_func.__name__}__taskiq_original"
         self.original_func.__name__ = new_name
-        self.original_func.__qualname__ = new_name
+        if hasattr(self.original_func, "__qualname__"):
+            original_qualname = self.original_func.__qualname__.rsplit(".")
+            original_qualname[-1] = new_name
+            new_qualname = ".".join(original_qualname)
+            self.original_func.__qualname__ = new_qualname
         setattr(
             sys.modules[original_func.__module__],
             new_name,
             original_func,
         )
 
-        self.labels = labels
-
     # Docs for this method are omitted in order to help
     # your IDE resolve correct docs for it.
     def __call__(  # noqa: D102