diff --git a/docs/examples/state/progress_tracker.py b/docs/examples/state/progress_tracker.py new file mode 100644 index 00000000..d1204c30 --- /dev/null +++ b/docs/examples/state/progress_tracker.py @@ -0,0 +1,11 @@ +import asyncio + +from taskiq_dependencies import Depends + +from taskiq.depends.progress_tracker import ProgressTracker, TaskState + + +async def my_task(progress: ProgressTracker[str] = Depends()) -> None: + for i in range(10): + await asyncio.sleep(1) + await progress.set_progress(TaskState.STARTED, meta=f"{(i + 1) * 10}%") diff --git a/docs/examples/state/progress_tracker_annot.py b/docs/examples/state/progress_tracker_annot.py new file mode 100644 index 00000000..12d095ea --- /dev/null +++ b/docs/examples/state/progress_tracker_annot.py @@ -0,0 +1,12 @@ +import asyncio +from typing import Annotated + +from taskiq_dependencies import Depends + +from taskiq.depends.progress_tracker import ProgressTracker, TaskState + + +async def my_task(progress: Annotated[ProgressTracker[str], Depends()]) -> None: + for i in range(10): + await asyncio.sleep(1) + await progress.set_progress(TaskState.STARTED, meta=f"{(i + 1) * 10}%") diff --git a/docs/guide/state-and-deps.md b/docs/guide/state-and-deps.md index aa79b322..ef8eda6e 100644 --- a/docs/guide/state-and-deps.md +++ b/docs/guide/state-and-deps.md @@ -319,6 +319,48 @@ taskiq worker my_file:broker --no-propagate-errors In this case, no exception will ever going to be propagated to any dependency. +## Progress tracking + +Sometimes a task runs for a long time and you want to report how far it has progressed, so other parts of your system +(e.g. a web handler polling for status) can display it. + +Taskiq provides a `ProgressTracker` dependency for this. It's not a method on `Context`, it's a dependency, just like +anything else built with `taskiq-dependencies`. It grabs the current `task_id` from the context for you and stores +progress using your broker's result backend. + +::: tabs + +@tab Annotated 3.10+ + +@[code python](../examples/state/progress_tracker_annot.py) + +@tab default values + +@[code python](../examples/state/progress_tracker.py) + +::: + +You can read the progress back from anywhere that has access to the `AsyncTaskiqTask` returned by `kiq`: + +```python +task = await my_task.kiq() + +progress = await task.get_progress() +if progress is not None: + print(progress.state, progress.meta) +``` + +`state` can be one of the `TaskState` enum values (`STARTED`, `SUCCESS`, `FAILURE`, `RETRY`) or any custom string. +`meta` is generic (`ProgressTracker[MetaType]`) and can be any value your result backend can serialize, such as a plain +string, a `dict`, or a pydantic model. If you call `set_progress` without `meta`, the previously stored `meta` value is +preserved, which is handy when you only want to update `state`. + +::: warning important note + +`set_progress`/`get_progress` are no-ops by default on `AsyncResultBackend`. Make sure the result backend you use +actually implements progress storage (`InMemoryBroker`'s built-in backend does), before relying on this in production. + +::: ## Generics diff --git a/taskiq/cli/scheduler/run.py b/taskiq/cli/scheduler/run.py index 854162ea..06382c43 100644 --- a/taskiq/cli/scheduler/run.py +++ b/taskiq/cli/scheduler/run.py @@ -313,7 +313,7 @@ def _is_schedule_ready_to_send( return is_ready_to_send - async def run( + async def run( # noqa: C901 self, *, update_interval: timedelta | None = None, diff --git a/tests/abc/test_broker.py b/tests/abc/test_broker.py index 636f9576..7f36eb9e 100644 --- a/tests/abc/test_broker.py +++ b/tests/abc/test_broker.py @@ -82,7 +82,6 @@ async def test_task() -> None: ... assert test_task.labels == old_labels -@pytest.mark.anyio @pytest.mark.parametrize( ("is_worker_process", "startup", "shutdown"), [ @@ -120,7 +119,6 @@ async def track_shutdown(state: TaskiqState) -> None: assert shutdown_called is True -@pytest.mark.anyio @pytest.mark.parametrize( ("is_worker_process", "startup", "shutdown"), [ diff --git a/tests/cli/scheduler/test_send_with_timeout.py b/tests/cli/scheduler/test_send_with_timeout.py index f97c94bc..151a58eb 100644 --- a/tests/cli/scheduler/test_send_with_timeout.py +++ b/tests/cli/scheduler/test_send_with_timeout.py @@ -3,6 +3,7 @@ from typing import Any import pytest +from typing_extensions import override from taskiq.abc.schedule_source import ScheduleSource from taskiq.brokers.inmemory_broker import InMemoryBroker @@ -30,8 +31,11 @@ def __init__(self, hang_seconds: float = 60.0) -> None: self.hang_seconds = hang_seconds self.on_ready_calls = 0 - async def on_ready( # type: ignore[override] - self, source: ScheduleSource, task: ScheduledTask + @override + async def on_ready( + self, + source: ScheduleSource, + task: ScheduledTask, ) -> None: self.on_ready_calls += 1 await asyncio.sleep(self.hang_seconds) @@ -48,7 +52,6 @@ def _task() -> ScheduledTask: ) -@pytest.mark.anyio async def test_send_with_timeout_returns_on_timeout_without_raising( caplog: pytest.LogCaptureFixture, ) -> None: @@ -85,7 +88,6 @@ async def test_send_with_timeout_returns_on_timeout_without_raising( assert "dummy" in msg -@pytest.mark.anyio async def test_send_with_timeout_does_not_log_on_success( caplog: pytest.LogCaptureFixture, ) -> None: @@ -97,8 +99,11 @@ def __init__(self) -> None: super().__init__(broker=InMemoryBroker(), sources=[_StubSource()]) self.calls = 0 - async def on_ready( # type: ignore[override] - self, source: ScheduleSource, task: ScheduledTask + @override + async def on_ready( + self, + source: ScheduleSource, + task: ScheduledTask, ) -> None: self.calls += 1 @@ -117,7 +122,6 @@ async def on_ready( # type: ignore[override] assert warnings == [] -@pytest.mark.anyio async def test_send_with_timeout_propagates_non_timeout_exceptions() -> None: """Errors inside on_ready that AREN'T a timeout must still propagate. @@ -130,8 +134,11 @@ class _BoomScheduler(TaskiqScheduler): def __init__(self) -> None: super().__init__(broker=InMemoryBroker(), sources=[_StubSource()]) - async def on_ready( # type: ignore[override] - self, source: ScheduleSource, task: ScheduledTask + @override + async def on_ready( + self, + source: ScheduleSource, + task: ScheduledTask, ) -> None: raise RuntimeError("boom") @@ -139,7 +146,6 @@ async def on_ready( # type: ignore[override] await send_with_timeout(_BoomScheduler(), _StubSource(), _task(), timeout=5.0) -@pytest.mark.anyio async def test_send_with_timeout_cancels_inner_send() -> None: """The inner ``send`` coroutine must actually be cancelled on timeout. @@ -152,8 +158,11 @@ class _CancelObservingScheduler(TaskiqScheduler): def __init__(self) -> None: super().__init__(broker=InMemoryBroker(), sources=[_StubSource()]) - async def on_ready( # type: ignore[override] - self, source: ScheduleSource, task: ScheduledTask + @override + async def on_ready( + self, + source: ScheduleSource, + task: ScheduledTask, ) -> None: try: await asyncio.sleep(30.0) @@ -171,7 +180,6 @@ async def on_ready( # type: ignore[override] await asyncio.wait_for(cancelled.wait(), timeout=2.0) -@pytest.mark.anyio async def test_plain_send_still_works_unchanged() -> None: """The original ``send`` function must remain unchanged in behavior. @@ -183,8 +191,11 @@ def __init__(self) -> None: super().__init__(broker=InMemoryBroker(), sources=[_StubSource()]) self.calls: list[Any] = [] - async def on_ready( # type: ignore[override] - self, source: ScheduleSource, task: ScheduledTask + @override + async def on_ready( + self, + source: ScheduleSource, + task: ScheduledTask, ) -> None: self.calls.append((source, task))