datashare-python 0.10.0.dev0__tar.gz → 0.10.1__tar.gz

{datashare_python-0.10.0.dev0 → datashare_python-0.10.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datashare-python
-Version: 0.10.0.dev0
+Version: 0.10.1
 Summary: Manage Python tasks and local resources in Datashare
 Project-URL: Homepage, https://icij.github.io/datashare-python/
 Project-URL: Documentation, https://icij.github.io/datashare-python/

{datashare_python-0.10.0.dev0 → datashare_python-0.10.1}/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-0.10.0.dev0 → datashare_python-0.10.1}/datashare_python/conftest.py

@@ -10,7 +10,6 @@ from multiprocessing import Event
 from pathlib import Path
 
 import aiohttp
-import nest_asyncio
 import pytest
 from elasticsearch._async.helpers import async_streaming_bulk
 from icij_common.es import DOC_ROOT_ID, ES_DOCUMENT_TYPE, ID, ESClient
@@ -171,7 +170,7 @@ async def test_temporal_client(
 
 
 @pytest.fixture
-async def indexed_docs(
+async def indexed_docs(  # noqa: PLR0917
     doc_0: Document,
     doc_1: Document,
     doc_2: Document,
@@ -299,6 +298,8 @@ async def all_done(task_client: DatashareTaskClient, not_done: list[str]) -> boo
 
 @pytest.fixture  # noqa: F405
 def typer_asyncio_patch() -> None:
+    import nest_asyncio  # noqa: PLC0415
+
     nest_asyncio.apply()
 
 

{datashare_python-0.10.0.dev0 → datashare_python-0.10.1}/datashare_python/interceptors.py

@@ -1,12 +1,12 @@
 import asyncio
-import contextlib
 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 (
@@ -29,6 +29,7 @@ from temporalio.client import WorkflowHandle
 from temporalio.converter import DataConverter
 from temporalio.worker import (
     ActivityInboundInterceptor,
+    ActivityOutboundInterceptor,
     ContinueAsNewInput,
     ExecuteActivityInput,
     ExecuteWorkflowInput,
@@ -219,7 +220,7 @@ def _trace_context(headers: Mapping[str, Payload]) -> Generator[None, None, None
 InputWithHeaders = TypeVar("InputWithHeaders")
 
 
-def _with_trace_context_header(
+def _with_trace_context_header[InputWithHeaders](
     input_with_headers: InputWithHeaders,
 ) -> InputWithHeaders:
     ctx = get_trace_context()
@@ -234,11 +235,14 @@ def _with_trace_context_header(
 
 
 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 @@ class _ProgressInboundInterceptor(ActivityInboundInterceptor):
         )
         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
 
 
@@ -360,13 +399,13 @@ class HeartbeatInterceptor(Interceptor):
         return _HeartbeatInboundInterceptor(next, self._n_missed_before_timeout)
 
 
-async def _heartbeat_every(period: float, *details: Any) -> None:
-    with contextlib.suppress(RuntimeError, asyncio.TimeoutError):
-        activity.heartbeat(*details)
+async def _heartbeat_every(
+    heartbeat_fn: Callable[[Any], None], period: float, *details: Any
+) -> None:
+    heartbeat_fn(*details)
     while True:
         await asyncio.sleep(period)
-        with contextlib.suppress(RuntimeError, asyncio.TimeoutError):
-            activity.heartbeat(*details)
+        heartbeat_fn(*details)
 
 
 class _HeartbeatInboundInterceptor(ActivityInboundInterceptor):
@@ -377,13 +416,24 @@ class _HeartbeatInboundInterceptor(ActivityInboundInterceptor):
     ) -> None:
         super().__init__(next)
         self._n_missed_before_timeout = n_missed_before_timeout
+        self._outbound: ActivityOutboundInterceptor | None = None
+
+    def init(self, outbound: ActivityOutboundInterceptor) -> None:
+        super().init(outbound)
+        self._outbound = outbound
 
     async def execute_activity(self, input: ExecuteActivityInput) -> Any:  # noqa: A002
+        # We need to capture the output interceptor to access the outbound hearbeat
+        # function. This inbound interceptor runs on the temporal event loop side.
+        # We want the heartbeat to run on the worker thread, on on the temporal event
+        # loop !
         heartbeat_timeout = activity.info().heartbeat_timeout
         heartbeat_task = None
         if heartbeat_timeout:
             period = heartbeat_timeout.total_seconds() / self._n_missed_before_timeout
-            heartbeat_task = asyncio.create_task(_heartbeat_every(period))
+            heartbeat_task = asyncio.create_task(
+                _heartbeat_every(self._outbound.heartbeat, period)
+            )
         try:
             activity.heartbeat()
             return await super().execute_activity(input)
@@ -397,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-0.10.0.dev0 → datashare_python-0.10.1}/datashare_python/objects.py

@@ -7,7 +7,7 @@ from datetime import UTC, datetime
 from enum import StrEnum, unique
 from io import BytesIO
 from pathlib import Path
-from typing import Annotated, Any, ClassVar, Generic, Literal, Self, TypeVar, cast
+from typing import Annotated, Any, ClassVar, Literal, Self, TypeVar, cast
 
 import langcodes
 from icij_common.registrable import Registrable
@@ -323,7 +323,7 @@ class TaskArgs(DatashareModel, ABC):
 A = TypeVar("A", bound=TaskArgs)
 
 
-class ManifestEntry(DatashareModel, Generic[A], ABC):
+class ManifestEntry[A](DatashareModel, ABC):
     status: ManifestEntryStatus
     label: str | None = None
     input: Annotated[

{datashare_python-0.10.0.dev0 → datashare_python-0.10.1}/datashare_python/types_.py

@@ -10,15 +10,17 @@ TemporalClient = Client
 
 
 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-0.10.0.dev0 → datashare_python-0.10.1}/datashare_python/utils.py

@@ -25,7 +25,7 @@ from functools import cache, wraps
 from hashlib import sha256
 from io import BytesIO
 from pathlib import Path
-from typing import Any, ParamSpec, Self, TypeVar
+from typing import Any, Self, TypeVar
 from uuid import uuid4
 
 import temporalio
@@ -70,8 +70,6 @@ DependencyAsyncSetup = Callable[..., Coroutine[None, None, None]]
 
 PROGRESS_HANDLER_ARG = "progress"
 
-P = ParamSpec("P")
-T = TypeVar("T")
 
 _NEVER_RETRIABLES = {
     "ValidationError",
@@ -178,7 +176,7 @@ async def execute_activity(
     )
 
 
-def positional_args_only(activity_fn: Callable[P, T]) -> Callable[P, T]:
+def positional_args_only[**P, T](activity_fn: Callable[P, T]) -> Callable[P, T]:
     sig = inspect.signature(activity_fn)
 
     # Keep track of kwargs-only
@@ -226,7 +224,7 @@ def _unpack_positional_args(
     return new_args, new_kwargs
 
 
-def with_retriables(
+def with_retriables[**P, T](
     retriables: set[type[Exception]] = None,
 ) -> Callable[[Callable[P, T]], Callable[P, T]]:
     if retriables is None:
@@ -281,7 +279,7 @@ def with_retriables(
     return decorator
 
 
-def activity_defn(
+def activity_defn[**P, T](
     name: str, retriables: set[type[Exception]] = None
 ) -> Callable[[Callable[P, T]], Callable[P, T]]:
 
@@ -319,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
 
@@ -331,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
 
@@ -345,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
 
@@ -359,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
 
@@ -615,11 +615,11 @@ def _read_jsonl(path: Path) -> Iterable[dict]:
 M = TypeVar("M", bound=BaseModel)
 
 
-def read_jsonl_as(path: Path, cls: type[M]) -> Iterable[M]:
+def read_jsonl_as[M](path: Path, cls: type[M]) -> Iterable[M]:
     return (cls.model_validate(d) for d in _read_jsonl(path))
 
 
-async def async_read_jsonl_as(
+async def async_read_jsonl_as[M](
     path: Path, processed_file_cls: type[M]
 ) -> AsyncIterable[M]:
     async with async_open(path, "r") as f:

{datashare_python-0.10.0.dev0 → datashare_python-0.10.1}/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:

{datashare_python-0.10.0.dev0 → datashare_python-0.10.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "datashare-python"
-version = "0.10.0dev0"
+version = "0.10.1"
 description = "Manage Python tasks and local resources in Datashare"
 authors = [
   { name = "Clément Doumouro", email = "cdoumouro@icij.org" },