From 3b7d32ea4caaa4f8d27ab0a097a16dbd0119c9fe Mon Sep 17 00:00:00 2001 From: Dave Date: Tue, 18 Aug 2026 21:49:53 +0200 Subject: [PATCH] feat(cli): add event loop factory option --- docs/guide/cli.md | 2 + taskiq/cli/scheduler/args.py | 10 +++ taskiq/cli/scheduler/cmd.py | 19 +++++- taskiq/cli/utils.py | 37 +++++++++- taskiq/cli/worker/args.py | 10 +++ taskiq/cli/worker/run.py | 23 +++++-- tests/cli/scheduler/test_cmd.py | 48 +++++++++++++ tests/cli/scheduler/test_scheduler_args.py | 9 +++ tests/cli/test_utils.py | 21 +++++- tests/cli/worker/test_args.py | 8 +++ tests/cli/worker/test_run.py | 78 ++++++++++++++++++++++ 11 files changed, 256 insertions(+), 9 deletions(-) create mode 100644 tests/cli/scheduler/test_cmd.py create mode 100644 tests/cli/scheduler/test_scheduler_args.py create mode 100644 tests/cli/worker/test_run.py diff --git a/docs/guide/cli.md b/docs/guide/cli.md index 0c8c4d81..7b3c4bd0 100644 --- a/docs/guide/cli.md +++ b/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. diff --git a/taskiq/cli/scheduler/args.py b/taskiq/cli/scheduler/args.py index 8f82b89f..bfb9dd30 100644 --- a/taskiq/cli/scheduler/args.py +++ b/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. diff --git a/taskiq/cli/scheduler/cmd.py b/taskiq/cli/scheduler/cmd.py index 48c81e30..06d8dd51 100644 --- a/taskiq/cli/scheduler/cmd.py +++ b/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), + }, + ) diff --git a/taskiq/cli/utils.py b/taskiq/cli/utils.py index a8554d37..51bc296a 100644 --- a/taskiq/cli/utils.py +++ b/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. diff --git a/taskiq/cli/worker/args.py b/taskiq/cli/worker/args.py index 8187113c..3f96d2b3 100644 --- a/taskiq/cli/worker/args.py +++ b/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, diff --git a/taskiq/cli/worker/run.py b/taskiq/cli/worker/run.py index 24d8f8db..4727f468 100644 --- a/taskiq/cli/worker/run.py +++ b/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) diff --git a/tests/cli/scheduler/test_cmd.py b/tests/cli/scheduler/test_cmd.py new file mode 100644 index 00000000..e7a464e7 --- /dev/null +++ b/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() diff --git a/tests/cli/scheduler/test_scheduler_args.py b/tests/cli/scheduler/test_scheduler_args.py new file mode 100644 index 00000000..fd7bbb3c --- /dev/null +++ b/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" diff --git a/tests/cli/test_utils.py b/tests/cli/test_utils.py index 0ee165dd..2492868c 100644 --- a/tests/cli/test_utils.py +++ b/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: diff --git a/tests/cli/worker/test_args.py b/tests/cli/worker/test_args.py index 258d7e1c..5e2b5849 100644 --- a/tests/cli/worker/test_args.py +++ b/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" diff --git a/tests/cli/worker/test_run.py b/tests/cli/worker/test_run.py new file mode 100644 index 00000000..d864af7b --- /dev/null +++ b/tests/cli/worker/test_run.py @@ -0,0 +1,78 @@ +import asyncio +from unittest.mock import patch + +from taskiq.brokers.inmemory_broker import InMemoryBroker +from taskiq.cli.worker.args import WorkerArgs +from taskiq.cli.worker.run import _create_worker_event_loop, start_listen + + +def test_create_worker_event_loop_uses_configured_factory() -> None: + args = WorkerArgs( + broker="example:broker", + modules=[], + loop_factory="asyncio:SelectorEventLoop", + ) + + with patch("taskiq.cli.worker.run.uvloop") as uvloop: + loop = _create_worker_event_loop(args) + + try: + assert isinstance(loop, asyncio.SelectorEventLoop) + uvloop.new_event_loop.assert_not_called() + finally: + loop.close() + + +def test_create_worker_event_loop_uses_uvloop_by_default() -> None: + args = WorkerArgs(broker="example:broker", modules=[]) + expected_loop = asyncio.new_event_loop() + + with patch("taskiq.cli.worker.run.uvloop") as uvloop: + uvloop.new_event_loop.return_value = expected_loop + loop = _create_worker_event_loop(args) + + try: + assert loop is expected_loop + uvloop.new_event_loop.assert_called_once_with() + finally: + loop.close() + + +def test_create_worker_event_loop_uses_asyncio_without_uvloop() -> None: + args = WorkerArgs(broker="example:broker", modules=[]) + + with patch("taskiq.cli.worker.run.uvloop", new=None): + loop = _create_worker_event_loop(args) + + try: + assert isinstance(loop, asyncio.AbstractEventLoop) + finally: + loop.close() + + +def test_start_listen_uses_created_event_loop() -> None: + args = WorkerArgs(broker="example:broker", modules=[]) + broker = InMemoryBroker() + loop = asyncio.new_event_loop() + + class Receiver: + def __init__(self, **_kwargs: object) -> None: + pass + + async def listen(self, _shutdown_event: asyncio.Event) -> None: + pass + + with ( + patch("taskiq.cli.worker.run.signal.signal"), + patch("taskiq.cli.worker.run._create_worker_event_loop", return_value=loop), + patch("taskiq.cli.worker.run.import_object", return_value=broker), + patch("taskiq.cli.worker.run.import_tasks"), + patch("taskiq.cli.worker.run.get_receiver_type", return_value=Receiver), + ): + start_listen(args) + + try: + assert loop.is_closed() is False + finally: + asyncio.set_event_loop(None) + loop.close()