diff --git a/README.rst b/README.rst index 2d92bb3f..7abb1e82 100644 --- a/README.rst +++ b/README.rst @@ -141,6 +141,42 @@ Run a worker (make sure the task code can be found, e.g. using ``PYTHONPATH``). {"task_id": "6fa07a91642363593cddef7a9e0c70ae3480921231710aa7648b467e637baa79", "level": "debug", "timestamp": "2015-08-27T21:03:56.732457Z", "pid": 69840, "queue": "default", "event": "done"} +Named tasks +----------- + +Tasks can be queued by a stable name instead of a Python import path. Install +a dispatch callback on each TaskTiger instance that processes these tasks. The +callback returns a callable for a known name and ``None`` for an unknown name. +Unknown names retain the existing import lookup, so ordinary ``delay`` tasks +do not need to change. + +.. code:: python + + handlers = {"mail.send": send_mail} + tiger.set_dispatch(handlers.get) + tiger.enqueue("mail.send", kwargs={"email_id": "123"}, queue="mail") + +Names passed to ``enqueue`` share the same namespace as the function paths +stored by ``delay``. For example, ``delay(tasks.my_task)`` stores +``"tasks:my_task"``; ``enqueue("tasks:my_task")`` queues that existing task +by name. If the dispatch callback returns ``None`` for that name, TaskTiger +imports ``tasks.my_task`` as usual. A matching route can instead select a +different callable for either form of queued task. + +For batch tasks, the callback can return ``TaskDispatch(handler, batch=True)`` +instead of a plain callable. The handler receives the same list of per-task +``args`` and ``kwargs`` dictionaries as a function decorated with +``@tiger.task(batch=True)``. Plain callable results continue to work as before. + +``enqueue`` accepts the same task options as ``delay``. Its ``args`` and +``kwargs`` must be JSON-serializable. The name is stored in the existing +``func`` field and is used for unique task IDs and locks. Configure the +dispatch callback in workers before queueing names that cannot be imported. +The callback selects a handler and may declare it as a batch handler; TaskTiger +executes it using the normal runner, timeout, and retry behavior. If the +callback raises an exception, the error is not treated as a missing route. + + Configuration ------------- diff --git a/tasktiger/__init__.py b/tasktiger/__init__.py index 2580bdb0..029b9cd8 100644 --- a/tasktiger/__init__.py +++ b/tasktiger/__init__.py @@ -1,3 +1,4 @@ +from .dispatch import TaskDispatch from .exceptions import ( JobTimeoutException, QueueFullException, @@ -17,6 +18,7 @@ "TaskTiger", "Worker", "Task", + "TaskDispatch", # Exceptions "JobTimeoutException", "RetryException", diff --git a/tasktiger/dispatch.py b/tasktiger/dispatch.py new file mode 100644 index 00000000..65eee7c6 --- /dev/null +++ b/tasktiger/dispatch.py @@ -0,0 +1,13 @@ +from dataclasses import dataclass +from typing import Any, Callable + + +@dataclass(frozen=True) +class TaskDispatch: + """A resolved task handler with execution metadata.""" + + handler: Callable[..., Any] + batch: bool = False + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + return self.handler(*args, **kwargs) diff --git a/tasktiger/executor.py b/tasktiger/executor.py index c42653cb..f7d6fd2d 100644 --- a/tasktiger/executor.py +++ b/tasktiger/executor.py @@ -129,7 +129,7 @@ def execute_tasks(self, tasks: List[Task], log: BoundLogger) -> bool: runner_class = get_runner_class(log, tasks) runner = runner_class(self.tiger) - is_batch_func = getattr(func, "_task_batch", False) + is_batch_func = tasks[0].is_batch g["tiger"] = self.tiger g["current_task_is_batch"] = is_batch_func diff --git a/tasktiger/runner.py b/tasktiger/runner.py index 36b8b4d1..823463af 100644 --- a/tasktiger/runner.py +++ b/tasktiger/runner.py @@ -72,7 +72,7 @@ def run_batch_tasks(self, tasks: List["Task"], hard_timeout: float) -> None: def run_eager_task(self, task: "Task") -> None: func = task.func - is_batch_func = getattr(func, "_task_batch", False) + is_batch_func = task.is_batch if is_batch_func: return func([{"args": task.args, "kwargs": task.kwargs}]) diff --git a/tasktiger/task.py b/tasktiger/task.py index 2343ba60..c46257eb 100644 --- a/tasktiger/task.py +++ b/tasktiger/task.py @@ -30,6 +30,7 @@ serialize_func_name, serialize_retry_method, ) +from .dispatch import TaskDispatch from .exceptions import QueueFullException, TaskImportError, TaskNotFound from .runner import BaseRunner, get_runner_class from .types import RetryStrategy @@ -66,16 +67,19 @@ def __init__( _state: Any = None, _ts: Any = None, _executions: Optional[List[Dict[str, Any]]] = None, + name: Optional[str] = None, ): """ Queues a task. See README.rst for an explanation of the options. """ - if func and queue is None: + if queue is None and (func or name is not None): queue = Task.queue_from_function(func, tiger) self.tiger = tiger self._func = func + self._func_resolved = False + self._dispatch_batch: Optional[bool] = None self._queue = queue self._state = _state self._ts = _ts @@ -86,9 +90,16 @@ def __init__( self._data = _data return - assert func + if (func is None) == (name is None): + raise ValueError("Exactly one of func or name is required") + if name is not None and (not isinstance(name, str) or not name): + raise ValueError("Task name must be a non-empty string") - serialized_name = serialize_func_name(func) + if name is not None: + serialized_name = name + else: + assert func is not None + serialized_name = serialize_func_name(func) if unique is None: unique = getattr(func, "_task_unique", False) @@ -270,10 +281,28 @@ def should_retry_on( @property def func(self) -> Callable: - if not self._func: - self._func = import_attribute(self.serialized_func) + if not self._func_resolved: + dispatched = self.tiger.dispatch(self.serialized_func) + if dispatched is not None: + if isinstance(dispatched, TaskDispatch): + self._dispatch_batch = dispatched.batch + dispatched = dispatched.handler + if not callable(dispatched): + raise TypeError("Task dispatcher must return a callable or None") + self._func = dispatched + elif self._func is None: + self._func = import_attribute(self.serialized_func) + self._func_resolved = True + assert self._func is not None return self._func + @property + def is_batch(self) -> bool: + func = self.func + if self._dispatch_batch is not None: + return self._dispatch_batch + return getattr(func, "_task_batch", False) + @property def max_stored_executions(self) -> Optional[int]: return self._data.get("max_stored_executions") @@ -347,8 +376,7 @@ def _move( self._state = to_state def execute(self) -> None: - func = self.func - is_batch_func = getattr(func, "_task_batch", False) + is_batch_func = self.is_batch g["current_task_is_batch"] = is_batch_func g["current_tasks"] = [self] diff --git a/tasktiger/tasktiger.py b/tasktiger/tasktiger.py index aeb611b9..545dd0d3 100644 --- a/tasktiger/tasktiger.py +++ b/tasktiger/tasktiger.py @@ -30,6 +30,7 @@ queue_matches, serialize_func_name, ) +from .dispatch import TaskDispatch from .executor import Executor, ForkExecutor, SyncExecutor from .redis_scripts import RedisScripts from .redis_semaphore import Semaphore @@ -113,6 +114,9 @@ def __init__( """ self.config: Dict[str, Any] = None # type: ignore[assignment] + self._dispatch: Optional[ + Callable[[str], Optional[Union[Callable, TaskDispatch]]] + ] = None # List of task functions that are executed periodically. self.periodic_task_funcs: Dict[str, Callable] = {} @@ -301,6 +305,23 @@ def _key(self, *parts: str) -> str: """ return ":".join([self.config["REDIS_PREFIX"]] + list(parts)) + def set_dispatch( + self, + dispatch: Optional[Callable[[str], Optional[Union[Callable, TaskDispatch]]]], + ) -> None: + """Set a callback that maps serialized task names to callables. + + Return None for names that should use the legacy import lookup. The + callback is consulted before importing a task, including eager tasks. + """ + if dispatch is not None and not callable(dispatch): + raise TypeError("dispatch must be callable or None") + self._dispatch = dispatch + + def dispatch(self, name: str) -> Optional[Union[Callable, TaskDispatch]]: + """Look up a task by its serialized name without invoking it.""" + return self._dispatch(name) if self._dispatch is not None else None + def task( self, _fn: Optional[Callable] = None, @@ -474,6 +495,52 @@ def delay( return task + def enqueue( + self, + name: str, + args: Any = None, + kwargs: Any = None, + queue: Optional[str] = None, + hard_timeout: Optional[float] = None, + unique: Optional[bool] = None, + unique_key: Optional[Collection[str]] = None, + lock: Optional[bool] = None, + lock_key: Optional[Collection[str]] = None, + when: Optional[Union[datetime.datetime, datetime.timedelta]] = None, + retry: Optional[bool] = None, + retry_on: Optional[Collection[Type[BaseException]]] = None, + retry_method: Optional[ + Union[Callable[[int], float], Tuple[Callable[..., float], Tuple]] + ] = None, + max_queue_size: Optional[int] = None, + max_stored_executions: Optional[int] = None, + runner_class: Optional[Type["BaseRunner"]] = None, + ) -> Task: + """Queue a task by serialized name instead of a Python function. + + Arguments and task options use the same storage and execution path as + delay(). The name is also used for unique task IDs and locks. + """ + task = Task( + self, + name=name, + args=args, + kwargs=kwargs, + queue=queue, + hard_timeout=hard_timeout, + unique=unique, + unique_key=unique_key, + lock=lock, + lock_key=lock_key, + retry=retry, + retry_on=retry_on, + retry_method=retry_method, + max_stored_executions=max_stored_executions, + runner_class=runner_class, + ) + task.delay(when=when, max_queue_size=max_queue_size) + return task + def get_queue_sizes(self, queue: str) -> Dict[str, int]: """ Get the queue's number of tasks in each state. diff --git a/tasktiger/worker.py b/tasktiger/worker.py index 38d4fd1c..65ce278f 100644 --- a/tasktiger/worker.py +++ b/tasktiger/worker.py @@ -356,7 +356,7 @@ def _worker_queue_expired_tasks(self) -> None: ) def get_hard_timeouts(self, func: Any, tasks: List[Task]) -> List[float]: - is_batch_func = getattr(func, "_task_batch", False) + is_batch_func = tasks[0].is_batch if is_batch_func: task_timeouts = [ task.hard_timeout for task in tasks if task.hard_timeout is not None diff --git a/tests/test_named_tasks.py b/tests/test_named_tasks.py new file mode 100644 index 00000000..a7af9201 --- /dev/null +++ b/tests/test_named_tasks.py @@ -0,0 +1,154 @@ +import pytest + +from tasktiger import Task, TaskDispatch, TaskImportError, TaskTiger, Worker +from tasktiger.executor import SyncExecutor +from tasktiger.runner import DefaultRunner + +from .tasks import simple_task + + +def test_enqueue_dispatches_named_task(tiger): + def handler(value, *, suffix): + tiger.connection.set("named_task_result", value + suffix) + + tiger.set_dispatch(lambda name: handler if name == "example.named" else None) + task = tiger.enqueue( + "example.named", args=["hello"], kwargs={"suffix": " world"}, queue="named" + ) + + assert task.serialized_func == "example.named" + assert task.data["args"] == ["hello"] + assert task.data["kwargs"] == {"suffix": " world"} + assert Task.from_id(tiger, "named", "queued", task.id).func is handler + + Worker(tiger, executor_class=SyncExecutor).run(once=True) + assert tiger.connection.get("named_task_result") == "hello world" + + +def test_named_task_runs_in_default_worker(tiger): + def handler(value): + tiger.connection.set("forked_named_task_result", value) + + tiger.set_dispatch(lambda name: handler if name == "example.forked" else None) + tiger.enqueue("example.forked", args=["done"]) + + Worker(tiger).run(once=True) + assert tiger.connection.get("forked_named_task_result") == "done" + + +def test_named_task_runs_eagerly(tiger): + called = [] + tiger.config["ALWAYS_EAGER"] = True + tiger.set_dispatch(lambda name: called.append if name == "example.eager" else None) + + tiger.enqueue("example.eager", args=["done"]) + + assert called == ["done"] + + +def test_named_batch_dispatch(tiger): + calls = [] + tiger.config["BATCH_QUEUES"]["named"] = 2 + + def handler(params): + calls.append(params) + + tiger.set_dispatch( + lambda name: ( + TaskDispatch(handler, batch=True) if name == "example.batch" else None + ) + ) + tiger.enqueue("example.batch", kwargs={"value": "one"}, queue="named") + tiger.enqueue("example.batch", kwargs={"value": "two"}, queue="named") + + Worker(tiger, executor_class=SyncExecutor).run(once=True) + + assert len(calls) == 1 + assert sorted(call["kwargs"]["value"] for call in calls[0]) == ["one", "two"] + + +def test_named_batch_runs_eagerly_without_redis(): + calls = [] + tiger = TaskTiger(lazy_init=True) + tiger.set_dispatch( + lambda name: ( + TaskDispatch(calls.append, batch=True) + if name == "example.eager_batch" + else None + ) + ) + + task = Task( + tiger, + name="example.eager_batch", + kwargs={"value": "one"}, + queue="named", + ) + assert task.is_batch is True + DefaultRunner(tiger).run_eager_task(task) + + assert calls == [[{"args": [], "kwargs": {"value": "one"}}]] + + +def test_dispatch_can_override_legacy_name(tiger): + called = [] + + def handler(): + called.append(True) + + tiger.set_dispatch( + lambda name: handler if name == "tests.tasks:simple_task" else None + ) + tiger.config["ALWAYS_EAGER"] = True + + tiger.delay(simple_task) + + assert called == [True] + + +def test_dispatch_miss_uses_legacy_import(tiger): + looked_up = [] + + def dispatch(name): + looked_up.append(name) + return None + + tiger.set_dispatch(dispatch) + task = tiger.delay(simple_task) + + assert Task.from_id(tiger, "default", "queued", task.id).func is simple_task + assert looked_up == ["tests.tasks:simple_task"] + + +def test_enqueue_importable_name_uses_legacy_import(tiger): + tiger.set_dispatch(lambda name: None) + task = tiger.enqueue("tests.tasks:simple_task") + + assert Task.from_id(tiger, "default", "queued", task.id).func is simple_task + + +def test_unrouted_name_uses_legacy_import(tiger): + task = tiger.enqueue("not.an.import.path") + + with pytest.raises(TaskImportError): + Task.from_id(tiger, "default", "queued", task.id).func + + +def test_named_task_unique_ids_use_name(tiger): + first = tiger.enqueue("example.first", args=[1], unique=True) + duplicate = tiger.enqueue("example.first", args=[1], unique=True) + other = tiger.enqueue("example.other", args=[1], unique=True) + + assert first.id == duplicate.id + assert first.id != other.id + + +def test_dispatch_errors_do_not_fall_back_to_import(tiger): + def dispatch(name): + raise RuntimeError("route error") + + tiger.set_dispatch(dispatch) + task = tiger.enqueue("tests.tasks:simple_task") + + with pytest.raises(RuntimeError, match="route error"): + Task.from_id(tiger, "default", "queued", task.id).func