feat(cli): add event loop factory option

Author: Davixk

docs/guide/cli.md

@@ -167,6 +167,7 @@ The number of signals before a hard kill can be configured with the `--hardkill-
 * `--no-propagate-errors` - if this parameter is enabled, exceptions won't be thrown in generator dependencies.
 * `--receiver` - python path to custom receiver class.
 * `--receiver_arg` - custom args for receiver.
+* `--loop-factory` - python path to an event loop factory in `module:variable` format. When set, this overrides automatic uvloop selection.
 * `--ack-type` - Type of acknowledgement. This parameter is used to set when to acknowledge the task. Possible values are `when_received`, `when_executed`, `when_saved`, `manual`. Default is `when_saved`.
 * `--max-tasks-per-child` - maximum number of tasks to be executed by a single worker process before restart.
 * `--max-fails` - Maximum number of child process exits.
@@ -200,5 +201,6 @@ Path to scheduler is the only required argument.
 - `--fs-discover` or `-fsd`. This option enables search of task files in current directory recursively, using the given pattern.
 - `--no-configure-logging` - use this parameter if your application configures custom logging.
 - `--log-level` is used to set a log level (default `INFO`).
+- `--loop-factory` - python path to an event loop factory in `module:variable` format.
 - `--skip-first-run` - skip first run of scheduler. This option skips running tasks immediately after scheduler start.
 - `--update-interval` - interval in seconds to check for new tasks. By default scheduler will check for new scheduled tasks every first second of the minute.

taskiq/cli/scheduler/args.py

@@ -20,6 +20,7 @@ class SchedulerArgs:
     skip_first_run: bool = False
     update_interval: int | None = None
     loop_interval: int | None = None
+    loop_factory: str | None = None
 
     @classmethod
     def from_cli(cls, args: Sequence[str] | None = None) -> "SchedulerArgs":
@@ -111,6 +112,15 @@ def from_cli(cls, args: Sequence[str] | None = None) -> "SchedulerArgs":
                 "If not specified, scheduler will run once a second."
             ),
         )
+        parser.add_argument(
+            "--loop-factory",
+            default=None,
+            help=(
+                "Where to search for an event loop factory. "
+                "This string must be specified in "
+                "'module.module:variable' format."
+            ),
+        )
 
         namespace = parser.parse_args(args)
         # If there are any patterns specified, remove default.

taskiq/cli/scheduler/cmd.py

@@ -1,9 +1,13 @@
 import asyncio
 from collections.abc import Sequence
+from functools import partial
+
+import anyio
 
 from taskiq.abc.cmd import TaskiqCMD
 from taskiq.cli.scheduler.args import SchedulerArgs
 from taskiq.cli.scheduler.run import run_scheduler
+from taskiq.cli.utils import create_event_loop, resolve_loop_factory
 
 
 class SchedulerCMD(TaskiqCMD):
@@ -23,4 +27,17 @@ def exec(self, args: Sequence[str]) -> None:
         :param args: CLI arguments.
         """
         parsed = SchedulerArgs.from_cli(args)
-        asyncio.run(run_scheduler(parsed))
+        if parsed.loop_factory is None:
+            asyncio.run(run_scheduler(parsed))
+            return
+        loop_factory = resolve_loop_factory(
+            parsed.loop_factory,
+            app_dir=parsed.app_dir,
+        )
+        anyio.run(
+            run_scheduler,
+            parsed,
+            backend_options={
+                "loop_factory": partial(create_event_loop, loop_factory),
+            },
+        )

taskiq/cli/utils.py

@@ -1,6 +1,7 @@
+import asyncio
 import os
 import sys
-from collections.abc import Generator, Sequence
+from collections.abc import Callable, Generator, Sequence
 from contextlib import contextmanager
 from importlib import import_module
 from logging import getLogger
@@ -9,6 +10,8 @@
 
 logger = getLogger("taskiq.worker")
 
+LoopFactory = Callable[[], asyncio.AbstractEventLoop]
+
 
 @contextmanager
 def add_cwd_in_path() -> Generator[None, None, None]:
@@ -55,6 +58,38 @@ def import_object(object_spec: str, app_dir: str | None = None) -> Any:
     return getattr(module, import_spec[1])
 
 
+def resolve_loop_factory(
+    loop_factory: str,
+    app_dir: str | None = None,
+) -> LoopFactory:
+    """
+    Resolve an event loop factory from a callable or import string.
+
+    :param loop_factory: path in `module:variable` format.
+    :param app_dir: directory to add in sys.path for importing.
+    :raises ValueError: if the resolved object is not callable.
+    :return: event loop factory.
+    """
+    factory = import_object(loop_factory, app_dir=app_dir)
+    if not callable(factory):
+        raise ValueError("Event loop factory must be callable.")
+    return factory
+
+
+def create_event_loop(loop_factory: LoopFactory) -> asyncio.AbstractEventLoop:
+    """
+    Create and validate an event loop from a factory.
+
+    :param loop_factory: event loop factory.
+    :raises ValueError: if the factory does not return an event loop.
+    :return: created event loop.
+    """
+    loop = loop_factory()
+    if not isinstance(loop, asyncio.AbstractEventLoop):
+        raise ValueError("Event loop factory must return an event loop.")
+    return loop
+
+
 def import_from_modules(modules: list[str]) -> None:
     """
     Import all modules from modules variable.

taskiq/cli/worker/args.py

@@ -55,6 +55,7 @@ class WorkerArgs:
     wait_tasks_timeout: float | None = None
     hardkill_count: int = 3
     use_process_pool: bool = False
+    loop_factory: str | None = None
 
     @classmethod
     def from_cli(
@@ -281,6 +282,15 @@ def from_cli(
             default=None,
             help="Maximum number of processes in process pool.",
         )
+        parser.add_argument(
+            "--loop-factory",
+            default=None,
+            help=(
+                "Where to search for an event loop factory. "
+                "This string must be specified in "
+                "'module.module:variable' format."
+            ),
+        )
 
         namespace = parser.parse_args(
             args,

taskiq/cli/worker/run.py

@@ -10,7 +10,12 @@
 from typing import Any
 
 from taskiq.abc.broker import AsyncBroker
-from taskiq.cli.utils import import_object, import_tasks
+from taskiq.cli.utils import (
+    create_event_loop,
+    import_object,
+    import_tasks,
+    resolve_loop_factory,
+)
 from taskiq.cli.worker.args import WorkerArgs
 from taskiq.cli.worker.process_manager import ProcessManager
 from taskiq.receiver import Receiver
@@ -29,6 +34,16 @@
 logger = logging.getLogger("taskiq.worker")
 
 
+def _create_worker_event_loop(args: WorkerArgs) -> asyncio.AbstractEventLoop:
+    if args.loop_factory is not None:
+        loop_factory = resolve_loop_factory(args.loop_factory, app_dir=args.app_dir)
+        return create_event_loop(loop_factory)
+    if uvloop is not None:
+        logger.debug("UVLOOP found. Using it as async runner")
+        return uvloop.new_event_loop()  # type: ignore
+    return asyncio.new_event_loop()
+
+
 async def shutdown_broker(broker: AsyncBroker, timeout: float) -> None:
     """
     This function used to shutdown broker.
@@ -120,11 +135,7 @@ def interrupt_handler(signum: int, _frame: Any) -> None:
     if sys.platform != "win32":
         signal.signal(signal.SIGHUP, interrupt_handler)
 
-    if uvloop is not None:
-        logger.debug("UVLOOP found. Using it as async runner")
-        loop = uvloop.new_event_loop()  # type: ignore
-    else:
-        loop = asyncio.new_event_loop()
+    loop = _create_worker_event_loop(args)
 
     asyncio.set_event_loop(loop)
 

tests/cli/scheduler/test_cmd.py

@@ -0,0 +1,48 @@
+import asyncio
+from unittest.mock import patch
+
+from taskiq.cli.scheduler.args import SchedulerArgs
+from taskiq.cli.scheduler.cmd import SchedulerCMD
+
+
+def test_scheduler_runs_on_configured_event_loop() -> None:
+    parsed = SchedulerArgs(
+        scheduler="example:scheduler",
+        modules=[],
+        loop_factory="asyncio:SelectorEventLoop",
+    )
+    running_loop: asyncio.AbstractEventLoop | None = None
+
+    async def run_scheduler(_args: SchedulerArgs) -> None:
+        nonlocal running_loop
+        running_loop = asyncio.get_running_loop()
+
+    with (
+        patch.object(SchedulerArgs, "from_cli", return_value=parsed),
+        patch("taskiq.cli.scheduler.cmd.run_scheduler", new=run_scheduler),
+    ):
+        SchedulerCMD().exec([])
+
+    assert isinstance(running_loop, asyncio.SelectorEventLoop)
+
+
+def test_scheduler_uses_default_event_loop_without_factory() -> None:
+    parsed = SchedulerArgs(
+        scheduler="example:scheduler",
+        modules=[],
+    )
+    running_loop: asyncio.AbstractEventLoop | None = None
+
+    async def run_scheduler(_args: SchedulerArgs) -> None:
+        nonlocal running_loop
+        running_loop = asyncio.get_running_loop()
+
+    with (
+        patch.object(SchedulerArgs, "from_cli", return_value=parsed),
+        patch("taskiq.cli.scheduler.cmd.run_scheduler", new=run_scheduler),
+        patch("taskiq.cli.scheduler.cmd.anyio.run") as anyio_run,
+    ):
+        SchedulerCMD().exec([])
+
+    assert running_loop is not None
+    anyio_run.assert_not_called()

tests/cli/scheduler/test_scheduler_args.py

@@ -0,0 +1,9 @@
+from taskiq.cli.scheduler.args import SchedulerArgs
+
+
+def test_loop_factory_accepts_import_string() -> None:
+    args = SchedulerArgs.from_cli(
+        ["example:scheduler", "--loop-factory", "asyncio:SelectorEventLoop"],
+    )
+
+    assert args.loop_factory == "asyncio:SelectorEventLoop"

tests/cli/test_utils.py

@@ -1,8 +1,27 @@
+import asyncio
 from contextlib import suppress
 from pathlib import Path
 from unittest.mock import patch
 
-from taskiq.cli.utils import import_tasks
+import pytest
+
+from taskiq.cli.utils import create_event_loop, import_tasks, resolve_loop_factory
+
+
+def test_resolve_loop_factory_from_import_string() -> None:
+    assert resolve_loop_factory("asyncio:new_event_loop") is asyncio.new_event_loop
+
+
+def test_resolve_loop_factory_rejects_non_callable() -> None:
+    with pytest.raises(ValueError, match="must be callable"):
+        resolve_loop_factory("asyncio:ALL_COMPLETED")
+
+
+def test_create_event_loop_rejects_invalid_result() -> None:
+    factory = resolve_loop_factory("builtins:object")
+
+    with pytest.raises(ValueError, match="must return an event loop"):
+        create_event_loop(factory)
 
 
 def test_import_tasks_list_pattern() -> None:

tests/cli/worker/test_args.py

@@ -35,3 +35,11 @@ def test_max_prefetch_rejects_negative_default(
 
     assert exc_info.value.code == 2
     assert "max_prefetch cannot be negative" in capsys.readouterr().err
+
+
+def test_loop_factory_accepts_import_string() -> None:
+    args = WorkerArgs.from_cli(
+        ["example:broker", "--loop-factory", "asyncio:SelectorEventLoop"],
+    )
+
+    assert args.loop_factory == "asyncio:SelectorEventLoop"