feature: progress update throttling

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
 

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

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: ...
 
 

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
 

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: