feature: worker side timeouts

datashare-python/datashare_python/config.py

@@ -1,5 +1,6 @@
+from datetime import timedelta
 from enum import StrEnum
-from typing import Literal
+from typing import TYPE_CHECKING, Literal
 
 from icij_common.es import ESClient
 from icij_common.pydantic_utils import ICIJSettings
@@ -12,7 +13,9 @@
 from .objects import BaseModel, WorkerPaths
 from .task_client import DatashareTaskClient
 from .types_ import TemporalClient
-from .utils import PYDANTIC_DATA_CONVERTER, SharedResources, close_cm_callback
+
+if TYPE_CHECKING:
+    from .utils import SharedResources
 
 _ALL_LOGGERS = [datashare_python.__name__]
 
@@ -63,13 +66,16 @@ class TemporalClientConfig(BaseModel):
     _client: TemporalClient | None = PrivateAttr(default=None)
 
     async def to_client(self) -> TemporalClient:
+        from .utils import PYDANTIC_DATA_CONVERTER  # noqa: PLC0415
+
         if self._client is None:
             runtime = None
             if self.prometheus_host is not None:
                 telemetry_config = TelemetryConfig(
                     metrics=PrometheusConfig(bind_address="0.0.0.0:9000")
                 )
                 runtime = Runtime(telemetry=telemetry_config)
+
             self._client = await TemporalClient.connect(
                 target_host=self.host,
                 namespace=self.namespace,
@@ -97,7 +103,9 @@ class ResourceCacheConfig(BaseModel):
     size: int = 1
     exit_context_managers: bool = True
 
-    def to_resource_cache(self) -> SharedResources:
+    def to_resource_cache(self) -> "SharedResources":
+        from .utils import SharedResources, close_cm_callback  # noqa: PLC0415
+
         eviction_callback = None
         if self.exit_context_managers:
             eviction_callback = close_cm_callback
@@ -134,3 +142,8 @@ def to_task_client(self) -> DatashareTaskClient:
 
     async def to_temporal_client(self) -> TemporalClient:
         return await self.temporal.to_client()
+
+
+class ActivityTimeouts(BaseModel):
+    start_to_close: timedelta | None = None
+    heartbeat: timedelta = timedelta(minutes=1)

datashare-python/datashare_python/utils.py

@@ -20,7 +20,6 @@
 )
 from copy import deepcopy
 from dataclasses import dataclass
-from datetime import timedelta
 from functools import cache, wraps
 from hashlib import sha256
 from io import BytesIO
@@ -46,6 +45,7 @@
 )
 from temporalio.exceptions import ApplicationError
 
+from datashare_python.config import ActivityTimeouts
 from datashare_python.types_ import (
     AsyncProgressRateHandler,
     RawSyncProgressHandler,
@@ -165,21 +165,22 @@ async def execute_activity(
     arg: Any = temporalio.common._arg_unset,
     *,
     args: list | None = None,
-    start_to_close_timeout: timedelta | None = None,
-    heartbeat_timeout: timedelta = timedelta(minutes=1),
+    timeouts: ActivityTimeouts | None = None,
     retry_policy: temporalio.common.RetryPolicy | None = None,
 ) -> Any:
+    if timeouts is None:
+        timeouts = ActivityTimeouts()
     if args is None:
         args = []
     retry_policy = _retry_policy_with_default(retry_policy)
     return await workflow.execute_activity(
         activity,
         arg=arg,
         args=args,
-        start_to_close_timeout=start_to_close_timeout,
         task_queue=task_queue,
         retry_policy=retry_policy,
-        heartbeat_timeout=heartbeat_timeout,
+        start_to_close_timeout=timeouts.start_to_close,
+        heartbeat_timeout=timeouts.heartbeat,
     )
 
 

datashare-python/tests/test_interceptors.py

@@ -21,7 +21,7 @@
 from temporalio.common import RetryPolicy
 
 with temporalio.workflow.unsafe.imports_passed_through():
-    from datashare_python.config import WorkerConfig
+    from datashare_python.config import ActivityTimeouts, WorkerConfig
     from datashare_python.interceptors import (
         HeartbeatInterceptor,
         ProgressInterceptor,
@@ -81,7 +81,7 @@ def intercept_client(self, next: OutboundInterceptor) -> OutboundInterceptor:  #
         return super().intercept_client(_MockOutboundInterceptor(next))
 
 
-_TIMEOUT = timedelta(seconds=180)
+_TIMEOUTS = ActivityTimeouts(start_to_close=timedelta(seconds=180))
 
 
 @workflow.defn
@@ -90,17 +90,11 @@ class _TestTraceContentWorkflow:
     async def run(self) -> list[TraceContext]:
         current_ctx = get_trace_context()
         ctx_log = [current_ctx]
-        ctx_log = await workflow.execute_activity(
-            ctx_test_act,
-            ctx_log,
-            task_queue=TestTaskQueue.TRACE,
-            start_to_close_timeout=_TIMEOUT,
+        ctx_log = await execute_activity(
+            ctx_test_act, TestTaskQueue.TRACE, ctx_log, timeouts=_TIMEOUTS
         )
-        ctx_log = await workflow.execute_activity(
-            ctx_test_act,
-            ctx_log,
-            task_queue=TestTaskQueue.TRACE,
-            start_to_close_timeout=_TIMEOUT,
+        ctx_log = await execute_activity(
+            ctx_test_act, TestTaskQueue.TRACE, ctx_log, timeouts=_TIMEOUTS
         )
         return ctx_log
 
@@ -163,13 +157,13 @@ async def run(self, args: ProgressArg) -> None:
             _ProgressAct.hello_sync_act,
             args=[args],
             task_queue=TestTaskQueue.PROGRESS_SYNC,
-            start_to_close_timeout=_TIMEOUT,
+            timeouts=_TIMEOUTS,
         )
         await execute_activity(
             _ProgressAct.hello_async_act,
             args=[args],
             task_queue=TestTaskQueue.PROGRESS_ASYNC,
-            start_to_close_timeout=_TIMEOUT,
+            timeouts=_TIMEOUTS,
         )
 
 
@@ -181,8 +175,10 @@ async def run(self) -> None:
             sleep_for_act,
             arg=1,
             task_queue=TestTaskQueue.HEARTBEAT,
-            start_to_close_timeout=_TIMEOUT,
-            heartbeat_timeout=timedelta(milliseconds=500),
+            timeouts=ActivityTimeouts(
+                start_to_close=timedelta(seconds=180),
+                heartbeat=timedelta(milliseconds=500),
+            ),
             retry_policy=RetryPolicy(maximum_attempts=1),
         )
 
@@ -195,8 +191,10 @@ async def run(self) -> None:
             sleep_for_act,
             arg=1,
             task_queue=TestTaskQueue.NO_HEARTBEAT,
-            start_to_close_timeout=_TIMEOUT,
-            heartbeat_timeout=timedelta(milliseconds=500),
+            timeouts=ActivityTimeouts(
+                start_to_close=timedelta(seconds=180),
+                heartbeat=timedelta(milliseconds=500),
+            ),
             retry_policy=RetryPolicy(maximum_attempts=1),
         )