datashare-python 0.10.0rc1__tar.gz → 0.10.4__tar.gz

{datashare_python-0.10.0rc1 → datashare_python-0.10.4}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.5
 Name: datashare-python
-Version: 0.10.0rc1
+Version: 0.10.4
 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.0rc1 → datashare_python-0.10.4}/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.0rc1 → datashare_python-0.10.4}/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.0rc1 → datashare_python-0.10.4}/datashare_python/interceptors.py

@@ -1,12 +1,13 @@
 import asyncio
-import contextlib
 import dataclasses
+import datetime
+import logging
 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 +30,7 @@ from temporalio.client import WorkflowHandle
 from temporalio.converter import DataConverter
 from temporalio.worker import (
     ActivityInboundInterceptor,
+    ActivityOutboundInterceptor,
     ContinueAsNewInput,
     ExecuteActivityInput,
     ExecuteWorkflowInput,
@@ -65,6 +67,8 @@ from .utils import (
     ProgressSignal,
 )
 
+logger = logging.getLogger(__name__)
+
 _TRACEPARENT = "traceparent"
 _DEFAULT_PAYLOAD_CONVERTER = DataConverter.default.payload_converter
 _PROGRESS_TYPES = {
@@ -189,6 +193,8 @@ class _TraceContextWorkflowOutboundInterceptor(WorkflowOutboundInterceptor):
 
 class _TraceContextActivityInboundInterceptor(ActivityInboundInterceptor):
     async def execute_activity(self, input: ExecuteActivityInput) -> Any:  # noqa: A002
+        info = activity.info()
+        logger.info("start executing activity: %s", info.activity_id)
         with _trace_context(input.headers):
             return await super().execute_activity(input)
 
@@ -219,7 +225,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 +240,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 +263,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 +307,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 +325,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 +354,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 +387,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 +404,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 +421,25 @@ 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))
+            # We don't want a failing hearbeat to fail the worker task so we just
+            # ignore any exception
+            heartbeat_fn = _fail_safe(self._outbound.heartbeat)
+            heartbeat_task = asyncio.create_task(_heartbeat_every(heartbeat_fn, period))
         try:
             activity.heartbeat()
             return await super().execute_activity(input)
@@ -393,13 +449,32 @@ class _HeartbeatInboundInterceptor(ActivityInboundInterceptor):
                 await asyncio.wait([heartbeat_task])
 
 
+def _fail_safe(
+    fn: Callable[[Any], None], excs: tuple[type[Exception]] | None = None
+) -> Callable[[Any], None]:
+    if excs is None:
+        excs = (Exception,)
+
+    @wraps(fn)
+    def wrapper(*args, **kwargs) -> None:
+        try:
+            fn(*args, **kwargs)
+        except excs as exc:
+            msg = f"failed to heartbeat due to {exc}"
+            logger.exception(msg)
+
+    return wrapper
+
+
 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.0rc1 → datashare_python-0.10.4}/datashare_python/objects.py

@@ -7,14 +7,14 @@ 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
 from pydantic_core import PydanticCustomError, ValidationError, core_schema
 from pydantic_core.core_schema import PlainValidatorFunctionSchema
 from pydantic_extra_types.language_code import LanguageName
-from temporalio import workflow
+from temporalio import activity, workflow
 
 from .constants import TIKA_METADATA_RESOURCENAME
 
@@ -323,8 +323,10 @@ class TaskArgs(DatashareModel, ABC):
 A = TypeVar("A", bound=TaskArgs)
 
 
-class ManifestEntry(DatashareModel, Generic[A], ABC):
+class ManifestEntry[A](DatashareModel, ABC):
     status: ManifestEntryStatus
+    # TODO: make this one non optional in the next major !
+    task_id: str | None
     label: str | None = None
     input: Annotated[
         dict[str, Any] | None,
@@ -336,7 +338,11 @@ class ManifestEntry(DatashareModel, Generic[A], ABC):
 
     @classmethod
     def complete(cls, args: A, label: str | None = None, **kwargs) -> Self:
+        task_id = None
+        if activity.in_activity():
+            task_id = activity.info().workflow_id
         return cls(
+            task_id=task_id,
             input=args.as_manifest_task_input(),
             label=label,
             status=ManifestEntryStatus.COMPLETE,
@@ -345,7 +351,11 @@ class ManifestEntry(DatashareModel, Generic[A], ABC):
 
     @classmethod
     def partial(cls, args: A, label: str | None = None, **kwargs) -> Self:
+        task_id = None
+        if activity.in_activity():
+            task_id = activity.info().workflow_id
         return cls(
+            task_id=task_id,
             input=args.as_manifest_task_input(),
             label=label,
             status=ManifestEntryStatus.PARTIAL,

{datashare_python-0.10.0rc1 → datashare_python-0.10.4}/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.0rc1 → datashare_python-0.10.4}/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.0rc1 → datashare_python-0.10.4}/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.0rc1 → datashare_python-0.10.4}/pyproject.toml

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