From 6b418c1386678267c7b4493cee0c998a195bf7a1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cle=CC=81ment=20Doumouro?= Date: Fri, 7 Aug 2026 16:37:40 +0200 Subject: [PATCH 1/2] fix(passport-worker): reduce progress updates --- workers/passport-worker/passport_worker/preprocessing.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/workers/passport-worker/passport_worker/preprocessing.py b/workers/passport-worker/passport_worker/preprocessing.py index 955c8d5..6290cca 100644 --- a/workers/passport-worker/passport_worker/preprocessing.py +++ b/workers/passport-worker/passport_worker/preprocessing.py @@ -204,16 +204,17 @@ async def convert_to_pdfs_act( res_i = 0 successes = [] errors = [] + progress_modulo = n_docs // 5 async for res in run_with_concurrency(aws, max_concurrency): if isinstance(res, FileProcessingError): errors.append(res) else: successes.append(res) - if progress is not None and res_i % 10 == 0: + if progress is not None and res_i % progress_modulo == 0: await progress(res_i) res_i += 1 logger.info( - "done converting docs to PDfs: %s success, %s errors", + "done converting docs to PDFs: %s success, %s errors", len(successes), len(errors), ) From 987d73884fb977e52c290323b96b4390be137442 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cle=CC=81ment=20Doumouro?= Date: Fri, 7 Aug 2026 17:20:49 +0200 Subject: [PATCH 2/2] feature(datashare-python): progress update throttling --- datashare-python/datashare_python/config.py | 1 + .../datashare_python/interceptors.py | 91 ++++++++++++++----- datashare-python/datashare_python/types_.py | 20 ++-- datashare-python/datashare_python/utils.py | 18 ++-- datashare-python/datashare_python/worker.py | 4 +- datashare-python/tests/test_interceptors.py | 91 +++++++++++++++++++ .../passport_worker/preprocessing.py | 2 +- 7 files changed, 185 insertions(+), 42 deletions(-) diff --git a/datashare-python/datashare_python/config.py b/datashare-python/datashare_python/config.py index 8a43656..70698eb 100644 --- a/datashare-python/datashare_python/config.py +++ b/datashare-python/datashare_python/config.py @@ -122,6 +122,7 @@ class WorkerConfig(ICIJSettings, BaseModel): temporal: TemporalClientConfig = TemporalClientConfig() max_concurrent_activities: int = 5 + min_progress_interval_s: float = 30.0 paths: WorkerPaths | None = None diff --git a/datashare-python/datashare_python/interceptors.py b/datashare-python/datashare_python/interceptors.py index 3f99d6c..41192ee 100644 --- a/datashare-python/datashare_python/interceptors.py +++ b/datashare-python/datashare_python/interceptors.py @@ -1,11 +1,12 @@ import asyncio import dataclasses +import datetime import secrets from collections.abc import Callable, Generator, Mapping from contextlib import contextmanager from contextvars import ContextVar from copy import deepcopy -from functools import partial, wraps +from functools import wraps from inspect import signature from types import UnionType from typing import ( @@ -234,11 +235,14 @@ def _with_trace_context_header[InputWithHeaders]( class ProgressInterceptor(Interceptor): + def __init__(self, min_progress_interval_s: float = 30.0): + self._min_progress_interval_s: float = min_progress_interval_s + def intercept_activity( self, next: ActivityInboundInterceptor, # noqa: A002 ) -> ActivityInboundInterceptor: - return _ProgressInboundInterceptor(next) + return _ProgressInboundInterceptor(next, self._min_progress_interval_s) def _parse_progress_weight(act_fn: Callable) -> float: @@ -254,18 +258,41 @@ def _parse_progress_weight(act_fn: Callable) -> float: return 1.0 -async def progress_handler( - progress: float, - handle: WorkflowHandle, - *, - activity_id: str, - run_id: str, - weight: float = 1.0, -) -> None: - signal = ProgressSignal( - activity_id=activity_id, run_id=run_id, progress=progress, weight=weight - ) - await handle.signal("update_progress", signal) +class TemporalProgressHandler: + def __init__( + self, + handle: WorkflowHandle, + activity_id: str, + *, + run_id: str, + min_progress_interval_s: float = 30.0, + weight: float = 1.0, + ) -> None: + self._handle = handle + self._activity_id = activity_id + self._run_id = run_id + self._weight = weight + self._min_progress_interval_s = min_progress_interval_s + self._last: datetime.datetime | None = None + + async def progress(self, progress: float, *, force: bool = False) -> None: + # TODO: we could lock here to avoid race conditions, it's not critical though + now = datetime.datetime.now(datetime.UTC) + report_progress = ( + force + or self._last is None + or (now - self._last).total_seconds() >= self._min_progress_interval_s + ) + if not report_progress: + return + self._last = now + signal = ProgressSignal( + activity_id=self._activity_id, + run_id=self._run_id, + progress=progress, + weight=self._weight, + ) + await self._handle.signal("update_progress", signal) def supports_progress(task_fn: Callable) -> bool: @@ -275,7 +302,9 @@ def supports_progress(task_fn: Callable) -> bool: ) -def _get_progress_handler(act_fn: Callable) -> ProgressRateHandler: +def _get_progress_handler( + act_fn: Callable, min_progress_interval_s: float +) -> ProgressRateHandler: act = getattr(act_fn, "__self__", None) # Weirdly isinstance doesn't work here if act is None or not isinstance(act, ActivityWithProgress): @@ -291,14 +320,14 @@ def _get_progress_handler(act_fn: Callable) -> ProgressRateHandler: activity_id = activity.info().activity_id client = act._temporal_client workflow_handle = client.get_workflow_handle(workflow_id, run_id=run_id) - handler = partial( - progress_handler, - handle=workflow_handle, + handler = TemporalProgressHandler( + workflow_handle, + activity_id, run_id=run_id, - activity_id=activity_id, weight=weight, + min_progress_interval_s=min_progress_interval_s, ) - return handler + return handler.progress def _is_progress(t: type) -> bool: @@ -320,13 +349,23 @@ def _without_progress(arg_types: list[type] | None) -> list[type] | None: class _ProgressInboundInterceptor(ActivityInboundInterceptor): + def __init__( + self, + next: ActivityInboundInterceptor, # noqa: A002 + min_progress_interval_s: float, + ) -> None: + super().__init__(next) + self._min_progress_interval_s = min_progress_interval_s + async def execute_activity(self, input: ExecuteActivityInput) -> Any: # noqa: A002 if not supports_progress(input.fn): return await super().execute_activity(input) # The progress args breaks trigger a bypass of the dataloader: # https://github.com/temporalio/sdk-python/blob/631ebaf0e20fb214b16589b45627b358048a5d77/temporalio/worker/_activity.py#L600 # we have to force it here again - progress_handler = _get_progress_handler(input.fn) + progress_handler = _get_progress_handler( + input.fn, self._min_progress_interval_s + ) new_args = [] act_definition = _Definition.must_from_callable(input.fn) if input.args: @@ -343,9 +382,9 @@ async def execute_activity(self, input: ExecuteActivityInput) -> Any: # noqa: A ) new_args.append(injected_progress) new_input = dataclasses.replace(input, args=new_args) - await progress_handler(0.0) + await progress_handler(0.0, force=True) res = await super().execute_activity(new_input) - await progress_handler(1.0) + await progress_handler(1.0, force=True) return res @@ -408,9 +447,11 @@ def _sync_progress( progress_handler: AsyncProgressRateHandler, ) -> SyncProgressRateHandler: @wraps(progress_handler) - def p(progress: float, event_loop: asyncio.AbstractEventLoop) -> None: + def p( + progress: float, event_loop: asyncio.AbstractEventLoop, *, force: bool = False + ) -> None: asyncio.run_coroutine_threadsafe( - progress_handler(progress), event_loop + progress_handler(progress, force=force), event_loop ).result() return p diff --git a/datashare-python/datashare_python/types_.py b/datashare-python/datashare_python/types_.py index 32af2e8..f6311c1 100644 --- a/datashare-python/datashare_python/types_.py +++ b/datashare-python/datashare_python/types_.py @@ -10,15 +10,17 @@ class AsyncProgressRateHandler(Protocol): - async def __call__(self, progress_rate: float) -> None: - pass + async def __call__(self, progress_rate: float, *, force: bool = False) -> None: ... class SyncProgressRateHandler(Protocol): def __call__( - self, progress_rate: float, event_loop: asyncio.AbstractEventLoop - ) -> None: - pass + self, + progress_rate: float, + event_loop: asyncio.AbstractEventLoop, + *, + force: bool = False, + ) -> None: ... ProgressRateHandler = SyncProgressRateHandler | AsyncProgressRateHandler @@ -30,12 +32,16 @@ class Weight: class RawAsyncProgressHandler(Protocol): - async def __call__(self, iteration: int) -> None: ... + async def __call__(self, iteration: int, *, force: bool = False) -> None: ... class RawSyncProgressHandler(Protocol): async def __call__( - self, iteration: int, event_loop: asyncio.AbstractEventLoop + self, + iteration: int, + event_loop: asyncio.AbstractEventLoop, + *, + force: bool = False, ) -> None: ... diff --git a/datashare-python/datashare_python/utils.py b/datashare-python/datashare_python/utils.py index daf5cd9..8e96452 100644 --- a/datashare-python/datashare_python/utils.py +++ b/datashare-python/datashare_python/utils.py @@ -317,8 +317,8 @@ def to_raw_async_progress( if not max_progress > 0: raise ValueError("max_progress must be > 0") - async def raw(p: int) -> None: - await progress(p / max_progress) + async def raw(p: int, *, force: bool = False) -> None: + await progress(p / max_progress, force=force) return raw @@ -329,10 +329,10 @@ def to_incremental_async_progress( offset = 0 - async def incremental(p: int) -> None: + async def incremental(p: int, *, force: bool = False) -> None: nonlocal offset offset += p - await progress(offset) + await progress(offset, force=force) return incremental @@ -343,8 +343,10 @@ def to_raw_sync_progress( if not max_progress > 0: raise ValueError("max_progress must be > 0") - def raw(iteration: int, event_loop: asyncio.AbstractEventLoop) -> None: - progress(iteration / max_progress, event_loop) + def raw( + iteration: int, event_loop: asyncio.AbstractEventLoop, *, force: bool = False + ) -> None: + progress(iteration / max_progress, event_loop, force=force) return raw @@ -357,8 +359,8 @@ def to_scaled_async_progress( if not start < end <= 1.0: raise ValueError("end must be ]start, 1.0]") - async def _scaled(p: float) -> None: - await progress(start + p * (end - start)) + async def _scaled(p: float, *, force: bool = False) -> None: + await progress(start + p * (end - start), force=force) return _scaled diff --git a/datashare-python/datashare_python/worker.py b/datashare-python/datashare_python/worker.py index e73e9dd..6965557 100644 --- a/datashare-python/datashare_python/worker.py +++ b/datashare-python/datashare_python/worker.py @@ -83,6 +83,7 @@ def datashare_worker( # at a time max_concurrent_activities: int = 1, max_activities_per_second: float = 20.0, + min_progress_interval_s: float = 30.0, sandboxed: bool = True, ) -> DatashareWorker: if workflows is None: @@ -108,7 +109,7 @@ def datashare_worker( logger.warning(_SEPARATE_IO_AND_CPU_WORKERS) interceptors = [ TraceContextInterceptor(), - ProgressInterceptor(), + ProgressInterceptor(min_progress_interval_s=min_progress_interval_s), HeartbeatInterceptor(), ] wf_runner = SandboxedWorkflowRunner() if sandboxed else UnsandboxedWorkflowRunner() @@ -216,6 +217,7 @@ async def worker_context( activities=acts, task_queue=task_queue, max_concurrent_activities=worker_config.max_concurrent_activities, + min_progress_interval_s=worker_config.min_progress_interval_s, sandboxed=sandboxed, ) async with worker: diff --git a/datashare-python/tests/test_interceptors.py b/datashare-python/tests/test_interceptors.py index 0558652..06c17f1 100644 --- a/datashare-python/tests/test_interceptors.py +++ b/datashare-python/tests/test_interceptors.py @@ -5,12 +5,14 @@ from concurrent.futures import ThreadPoolExecutor from datetime import timedelta from typing import Annotated, Any +from unittest.mock import AsyncMock, call import pytest import temporalio from datashare_python.objects import DatashareModel from datashare_python.utils import ( ActivityWithProgress, + ProgressSignal, WorkflowWithProgress, activity_defn, execute_activity, @@ -23,6 +25,7 @@ from datashare_python.interceptors import ( HeartbeatInterceptor, ProgressInterceptor, + TemporalProgressHandler, TraceContext, TraceContextInterceptor, get_trace_context, @@ -409,3 +412,91 @@ async def test_heartbeat_interceptor_should_fail_when_no_heartbeat( cause = ctx.value.cause.__cause__ assert isinstance(cause, temporalio_exceptions.TimeoutError) assert "Heartbeat timeout" in cause.args[0] + + +async def test_should_progress_handler_should_not_report_progress() -> None: + # Given + mocked_wf_handle = AsyncMock() + activity_id = "activity-id" + run_id = "run-id" + min_progress_interval_s = float("inf") + handler = TemporalProgressHandler( + mocked_wf_handle, + activity_id, + run_id=run_id, + min_progress_interval_s=min_progress_interval_s, + ) + # When + await handler.progress(0.1) + await handler.progress(0.2) + # Then + expected_signal = ProgressSignal( + activity_id=activity_id, run_id=run_id, progress=0.1, weight=1.0 + ) + mocked_wf_handle.signal.assert_called_once_with("update_progress", expected_signal) + + +async def test_should_progress_handler_should_report_progress() -> None: + # Given + mocked_wf_handle = AsyncMock() + activity_id = "activity-id" + run_id = "run-id" + min_progress_interval_s = 0.0 + handler = TemporalProgressHandler( + mocked_wf_handle, + activity_id, + run_id=run_id, + min_progress_interval_s=min_progress_interval_s, + ) + # When + await handler.progress(0.1) + await handler.progress(0.2) + # Then + expected_calls = [ + call( + "update_progress", + ProgressSignal( + activity_id=activity_id, run_id=run_id, progress=0.1, weight=1.0 + ), + ), + call( + "update_progress", + ProgressSignal( + activity_id=activity_id, run_id=run_id, progress=0.2, weight=1.0 + ), + ), + ] + mocked_wf_handle.signal.assert_has_calls(expected_calls) + + +async def test_should_progress_handler_should_report_progress_on_force() -> None: + # Given + mocked_wf_handle = AsyncMock() + activity_id = "activity-id" + run_id = "run-id" + min_progress_interval_s = float("inf") + handler = TemporalProgressHandler( + mocked_wf_handle, + activity_id, + run_id=run_id, + min_progress_interval_s=min_progress_interval_s, + ) + # When + await handler.progress(0.1) + await handler.progress(0.2, force=True) + # Then + expected_calls = [ + call( + "update_progress", + ProgressSignal( + activity_id=activity_id, run_id=run_id, progress=0.1, weight=1.0 + ), + ), + call( + "update_progress", + ProgressSignal( + activity_id=activity_id, run_id=run_id, progress=0.2, weight=1.0 + ), + ), + ] + mocked_wf_handle.signal.assert_has_calls(expected_calls) diff --git a/workers/passport-worker/passport_worker/preprocessing.py b/workers/passport-worker/passport_worker/preprocessing.py index 6290cca..c6117a3 100644 --- a/workers/passport-worker/passport_worker/preprocessing.py +++ b/workers/passport-worker/passport_worker/preprocessing.py @@ -204,7 +204,7 @@ async def convert_to_pdfs_act( res_i = 0 successes = [] errors = [] - progress_modulo = n_docs // 5 + progress_modulo = max(n_docs // 5, 1) async for res in run_with_concurrency(aws, max_concurrency): if isinstance(res, FileProcessingError): errors.append(res)