datashare-python 0.7.1__py3-none-any.whl → 0.7.3__py3-none-any.whl

datashare_python/config.py

@@ -1,9 +1,9 @@
 from pathlib import Path
-from typing import Annotated, Literal
+from typing import Literal
 
 from icij_common.es import ESClient
 from icij_common.pydantic_utils import ICIJSettings
-from pydantic import Field, PrivateAttr
+from pydantic import PrivateAttr
 from pydantic_settings import SettingsConfigDict
 from temporalio.contrib.pydantic import PydanticJSONPlainPayloadConverter, ToJsonOptions
 from temporalio.converter import (
@@ -83,12 +83,14 @@ class LoggingConfig(BaseModel):
     loggers: dict[str, LogLevel]
 
 
+_DEFAULT_LOGGERS = {datashare_python.__name__: "INFO"}
+_DEFAULT_LOGGING_CONFIG = LoggingConfig(log_in_json=True, loggers=_DEFAULT_LOGGERS)
+
+
 class WorkerConfig(ICIJSettings, BaseModel):
     model_config = DS_WORKER_SETTINGS_CONFIG
 
-    logging: Annotated[LoggingConfig, Field(frozen=True)] = {
-        datashare_python.__name__: "INFO"
-    }
+    logging: LoggingConfig = _DEFAULT_LOGGING_CONFIG
 
     datashare: DatashareClientConfig = DatashareClientConfig()
     elasticsearch: ESClientConfig = ESClientConfig()

datashare_python/logging_.py

@@ -2,11 +2,7 @@ import logging
 import sys
 from copy import copy
 
-from icij_common.logging_utils import (
-    DATE_FMT,
-    STREAM_HANDLER_FMT,
-    STREAM_HANDLER_FMT_WITH_WORKER_ID,
-)
+from icij_common.logging_utils import DATE_FMT, STREAM_HANDLER_FMT
 from pythonjsonlogger.core import RESERVED_ATTRS, BaseJsonFormatter
 from pythonjsonlogger.orjson import OrjsonFormatter
 from temporalio import activity, workflow
@@ -26,6 +22,11 @@ _LOGGED_ATTRIBUTES = (
 )
 
 
+_STREAM_HANDLER_FMT_WITH_WORKER_ID = (
+    "[%(levelname)s][%(asctime)s.%(msecs)03d][%(worker_id)s][%(name)s]: %(message)s"
+)
+
+
 def setup_worker_loggers(
     loggers: dict[str, LogLevel], *, worker_id: str | None, in_json: bool
 ) -> None:
@@ -35,35 +36,18 @@ def setup_worker_loggers(
         logger = logging.getLogger(logger_name)
         logger.setLevel(level)
         logger.handlers = []
-        for handler in _get_worker_handlers(level, worker_id, in_json=in_json):
+        for handler in _get_worker_handlers(level, worker_filter, in_json=in_json):
             logger.addHandler(handler)
-        logger.addFilter(worker_filter)
-
-
-def _get_worker_handlers(
-    level: int, worker_id: str | None, *, in_json: bool
-) -> list[logging.Handler]:
-    stream_handler = logging.StreamHandler(sys.stderr)
-    if in_json:
-        fmt = _json_formatter(datefmt=DATE_FMT)
-    else:
-        if worker_id is not None:
-            fmt = STREAM_HANDLER_FMT_WITH_WORKER_ID
-        else:
-            fmt = STREAM_HANDLER_FMT
-        fmt = logging.Formatter(fmt, DATE_FMT)
-    stream_handler.setFormatter(fmt)
-    stream_handler.setLevel(level)
-    return [stream_handler]
 
 
 class WorkerFilter(logging.Filter):
-    def __init__(self, worker_id: str) -> None:
+    def __init__(self, worker_id: str | None) -> None:
         super().__init__()
-        self._worker_id = worker_id
+        self.worker_id = worker_id
 
     def filter(self, record: logging.LogRecord) -> bool:
-        record.worker_id = self._worker_id
+        if self.worker_id is not None:
+            record.worker_id = self.worker_id
         if workflow.in_workflow():
             wf_info = workflow.info()
             for attr in _WF_LOGGED_ATTRS:
@@ -79,6 +63,24 @@ class WorkerFilter(logging.Filter):
         return True
 
 
+def _get_worker_handlers(
+    level: int, worker_filter: WorkerFilter, *, in_json: bool
+) -> list[logging.Handler]:
+    stream_handler = logging.StreamHandler(sys.stderr)
+    if in_json:
+        fmt = _json_formatter(datefmt=DATE_FMT)
+    else:
+        if worker_filter.worker_id is not None:
+            fmt = _STREAM_HANDLER_FMT_WITH_WORKER_ID
+        else:
+            fmt = STREAM_HANDLER_FMT
+        fmt = logging.Formatter(fmt, DATE_FMT)
+    stream_handler.setFormatter(fmt)
+    stream_handler.setLevel(level)
+    stream_handler.addFilter(worker_filter)
+    return [stream_handler]
+
+
 def _json_formatter(datefmt: str) -> BaseJsonFormatter:
     fmt = OrjsonFormatter(  # let's keep logging as fast as possible
         _LOGGED_ATTRIBUTES, datefmt=datefmt

datashare_python/utils.py

@@ -1,12 +1,10 @@
 import asyncio
+import contextlib
+import contextvars
 import inspect
 import json
-import logging
-import sys
-from collections.abc import (
-    Callable,
-    Coroutine,
-)
+import threading
+from collections.abc import Awaitable, Callable, Coroutine
 from copy import deepcopy
 from dataclasses import dataclass
 from datetime import timedelta
@@ -20,15 +18,6 @@ from uuid import uuid4
 
 import nest_asyncio
 import temporalio
-from icij_common.logging_utils import (
-    DATE_FMT,
-    STREAM_HANDLER_FMT,
-    STREAM_HANDLER_FMT_WITH_WORKER_ID,
-    WorkerIdFilter,
-)
-from icij_common.pydantic_utils import get_field_default_value
-from pydantic.fields import FieldInfo
-from pythonjsonlogger.json import JsonFormatter
 from temporalio import activity, workflow
 from temporalio.client import Client, WorkflowHandle
 from temporalio.common import RetryPolicy, SearchAttributeKey
@@ -123,6 +112,7 @@ async def execute_activity(
     *,
     args: list | None = None,
     start_to_close_timeout: timedelta | None = None,
+    heartbeat_timeout: timedelta = timedelta(minutes=1),
     retry_policy: temporalio.common.RetryPolicy | None = None,
 ) -> Any:
     if args is None:
@@ -135,6 +125,7 @@ async def execute_activity(
         start_to_close_timeout=start_to_close_timeout,
         task_queue=task_queue,
         retry_policy=retry_policy,
+        heartbeat_timeout=heartbeat_timeout,
     )
 
 
@@ -150,6 +141,8 @@ async def progress_handler(
         activity_id=activity_id, run_id=run_id, progress=progress, weight=weight
     )
     await handle.signal("update_progress", signal)
+    with contextlib.suppress(RuntimeError, asyncio.TimeoutError):
+        activity.heartbeat()
 
 
 def get_activity_progress_handler_async(
@@ -229,6 +222,74 @@ def with_progress(weight: float = 1.0) -> Callable[P, T]:
     return decorator
 
 
+def with_async_heartbeat(
+    activity_fn: Callable[P, Awaitable[T]], n_missed_before_timeout: int
+) -> Callable[P, Awaitable[T]]:
+    # Copied from
+    # https://github.com/temporalio/samples-python/blob/main/custom_decorator/activity_utils.py
+    @wraps(activity_fn)
+    async def wrapper(*args, **kwargs) -> T:
+        heartbeat_timeout = activity.info().heartbeat_timeout
+        heartbeat_task = None
+        if heartbeat_timeout:
+            period = heartbeat_timeout.total_seconds() / n_missed_before_timeout
+            heartbeat_task = asyncio.create_task(_async_heartbeat_every(period))
+        try:
+            activity.heartbeat()
+            return await activity_fn(*args, **kwargs)
+        finally:
+            if heartbeat_task:
+                heartbeat_task.cancel()
+                await asyncio.wait([heartbeat_task])
+
+    return wrapper
+
+
+async def _async_heartbeat_every(period: float, *details: Any) -> None:
+    with contextlib.suppress(RuntimeError, asyncio.TimeoutError):
+        activity.heartbeat(*details)
+    while True:
+        await asyncio.sleep(period)
+        with contextlib.suppress(RuntimeError, asyncio.TimeoutError):
+            activity.heartbeat(*details)
+
+
+def with_sync_heartbeat(
+    activity_fn: Callable[P, T], n_missed_before_timeout: int
+) -> Callable[P, T]:
+    @wraps(activity_fn)
+    def wrapper(*args, **kwargs) -> T:
+        heartbeat_timeout = activity.info().heartbeat_timeout
+        heartbeat_thread, stop_event = None, None
+        if heartbeat_timeout:
+            period = heartbeat_timeout.total_seconds() / n_missed_before_timeout
+            ctx = contextvars.copy_context()
+            run_args = (_sync_heartbeat_every, period, threading.Event())
+            heartbeat_thread, stop_event = (
+                threading.Thread(target=ctx.run, args=run_args),
+                run_args[-1],
+            )
+            heartbeat_thread.start()
+        try:
+            return activity_fn(*args, **kwargs)
+        finally:
+            if heartbeat_thread:
+                stop_event.set()
+                heartbeat_thread.join()
+
+    return wrapper
+
+
+def _sync_heartbeat_every(
+    period: float, stop_event: threading.Event, *details: Any
+) -> None:
+    with contextlib.suppress(RuntimeError, asyncio.TimeoutError):
+        activity.heartbeat(*details)
+    while not stop_event.wait(period):
+        with contextlib.suppress(RuntimeError, asyncio.TimeoutError):
+            activity.heartbeat(*details)
+
+
 def positional_args_only(activity_fn: Callable[P, T]) -> Callable[P, T]:
     sig = inspect.signature(activity_fn)
 
@@ -336,6 +397,7 @@ def activity_defn(
     name: str,
     progress_weight: float = 1.0,
     retriables: set[type[Exception]] = None,
+    n_missed_heartbeats_before_timeout: int = 5,
 ) -> Callable[[Callable[P, T]], Callable[P, T]]:
     def decorator(activity_fn: Callable[P, T]) -> Callable[P, T]:
         # TODO: some of these could probably be reimplemented more elegantly using
@@ -344,6 +406,15 @@ def activity_defn(
         activity_fn = with_retriables(retriables)(activity_fn)
         if supports_progress(activity_fn):
             activity_fn = with_progress(progress_weight)(activity_fn)
+        is_async = asyncio.iscoroutinefunction(activity_fn)
+        if is_async:
+            activity_fn = with_async_heartbeat(
+                activity_fn, n_missed_heartbeats_before_timeout
+            )
+        else:
+            activity_fn = with_sync_heartbeat(
+                activity_fn, n_missed_heartbeats_before_timeout
+            )
         activity_fn = activity.defn(activity_fn, name=name)
         return activity_fn
 
@@ -382,59 +453,6 @@ def to_scaled_progress(
     return _scaled
 
 
-class LogWithWorkerIDMixin:
-    def setup_loggers(self, worker_id: str | None = None) -> None:
-        # Ugly work around the Pydantic V1 limitations...
-        all_loggers = self.loggers
-        if isinstance(all_loggers, FieldInfo):
-            all_loggers = get_field_default_value(all_loggers)
-        all_loggers.append(__name__)
-        loggers = sorted(set(all_loggers))
-        log_level = self.log_level
-        if isinstance(log_level, FieldInfo):
-            log_level = get_field_default_value(log_level)
-        force_warning = getattr(self, "force_warning_loggers", [])
-        if isinstance(force_warning, FieldInfo):
-            force_warning = get_field_default_value(force_warning)
-        force_warning = set(force_warning)
-        worker_id_filter = None
-        if worker_id is not None:
-            worker_id_filter = WorkerIdFilter(worker_id)
-        handlers = self._handlers(worker_id_filter, log_level)
-        for logger_ in loggers:
-            logger_ = logging.getLogger(logger_)  # noqa: PLW2901
-            level = getattr(logging, log_level)
-            if logger_.name in force_warning:
-                level = max(logging.WARNING, level)
-            logger_.setLevel(level)
-            logger_.handlers = []
-            for handler in handlers:
-                logger_.addHandler(handler)
-
-    def _handlers(
-        self, worker_id_filter: logging.Filter | None, log_level: int
-    ) -> list[logging.Handler]:
-        stream_handler = logging.StreamHandler(sys.stderr)
-        if worker_id_filter is not None:
-            fmt = STREAM_HANDLER_FMT_WITH_WORKER_ID
-        else:
-            fmt = STREAM_HANDLER_FMT
-        log_in_json = getattr(self, "log_in_json", False)
-        if isinstance(log_in_json, FieldInfo):
-            log_in_json = get_field_default_value(log_in_json)
-        if log_in_json:
-            fmt = JsonFormatter(fmt, DATE_FMT)
-        else:
-            fmt = logging.Formatter(fmt, DATE_FMT)
-        stream_handler.setFormatter(fmt)
-        handlers = [stream_handler]
-        for handler in handlers:
-            if worker_id_filter is not None:
-                handler.addFilter(worker_id_filter)
-            handler.setLevel(log_level)
-        return handlers
-
-
 def safe_dir(doc_id: str) -> Path:
     if len(doc_id) < 4:
         raise ValueError(f"expected doc_id to be at least 4, found {doc_id}")

datashare_python/worker.py

@@ -11,7 +11,12 @@ from contextlib import asynccontextmanager
 from copy import copy
 from typing import Any
 
-from temporalio.worker import PollerBehaviorSimpleMaximum, Worker
+from temporalio.worker import (
+    PollerBehaviorSimpleMaximum,
+    UnsandboxedWorkflowRunner,
+    Worker,
+)
+from temporalio.worker.workflow_sandbox import SandboxedWorkflowRunner
 
 from .config import WorkerConfig
 from .dependencies import with_dependencies
@@ -62,6 +67,7 @@ def datashare_worker(
     # Scale horizontally be default for activities, each worker processes one activity
     # at a time
     max_concurrent_io_activities: int = 10,
+    sandboxed: bool = True,
 ) -> DatashareWorker:
     if workflows is None:
         workflows = []
@@ -86,6 +92,7 @@ def datashare_worker(
         if workflows:
             logger.warning(_SEPARATE_IO_AND_CPU_WORKERS)
     interceptors = [TraceContextInterceptor()]
+    wf_runner = SandboxedWorkflowRunner() if sandboxed else UnsandboxedWorkflowRunner()
     return DatashareWorker(
         client,
         interceptors=interceptors,
@@ -101,6 +108,7 @@ def datashare_worker(
         # Workflow tasks are assumed to be very lightweight and fast we can reserve
         # several of them
         workflow_task_poller_behavior=PollerBehaviorSimpleMaximum(5),
+        workflow_runner=wf_runner,
     )
 
 
@@ -144,6 +152,7 @@ async def worker_context(
     event_loop: AbstractEventLoop,
     task_queue: str,
     dependencies: list[ContextManagerFactory] | None = None,
+    sandboxed: bool = True,
 ) -> AsyncGenerator[DatashareWorker, None]:
     discovered = []
     if activities is not None:
@@ -185,6 +194,7 @@ async def worker_context(
             activities=acts,
             task_queue=task_queue,
             max_concurrent_io_activities=worker_config.max_concurrent_io_activities,
+            sandboxed=sandboxed,
         )
         async with worker:
             yield worker

{datashare_python-0.7.1.dist-info → datashare_python-0.7.3.dist-info}/METADATA

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: datashare-python
-Version: 0.7.1
-Summary: Manage Pythoœn tasks and local resources in Datashare
+Version: 0.7.3
+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/
 Project-URL: Repository, https://github.com/ICIJ/datashare-python

{datashare_python-0.7.1.dist-info → datashare_python-0.7.3.dist-info}/RECORD

@@ -1,27 +1,27 @@
 datashare_python/.gitignore,sha256=e-SRgnvGGdsjRrqgKsTzALz6Obx8IYiOjr0yaAxT6v8,22
 datashare_python/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datashare_python/__main__.py,sha256=g-fvS46zl9umKmGrSpl-OG-8PSuZgjqvTCqjpsZtSps,101
-datashare_python/config.py,sha256=Q4iu3ZGaQB7npaHJWclxPIfgzZTf_-8VxxhXrB9nlpE,3928
+datashare_python/config.py,sha256=_Cx4EB1yHXXcLUtw1OBlMk2SKcJRwxqJwgRu6klbxNg,3994
 datashare_python/conftest.py,sha256=MrmQKFcUipm_qn-cHsLovZMwMMtVxyK0s1lmKEx54bc,8651
 datashare_python/constants.py,sha256=a8-ceZKBVMXydcoNQ35fSjFjxeJ7dt-N6eAvqtPpf9g,320
 datashare_python/dependencies.py,sha256=KJuAp6Dmv8DQuFnGjbWiHu7StzZj97eBPDyZ_RfCQRc,4141
 datashare_python/discovery.py,sha256=BPB_Ak6d1-vcf9vAQA63IRb2U8h83_mIIi8MbKbFzQ0,7020
 datashare_python/exceptions.py,sha256=bVHEAXxDPKfxeeMC0hJXEsrJkgsKO2ESAhxWU96GA4M,496
 datashare_python/interceptors.py,sha256=Pl7GodPO4KbfflmacpW-vOUgLazjlXSlDNENbpOUt1c,6725
-datashare_python/logging_.py,sha256=-qHz4ztKz4mOCO2z4wunQ4M3xoVhztNvxaiozuMLFRM,2815
+datashare_python/logging_.py,sha256=XUhZTtofbOqJi1gwytYpUVqvoGPhoz5p2orsXs2FaWs,2968
 datashare_python/objects.py,sha256=pE0DGNNkl1etxz5ed7T-EaGo1o9TONjH2Lg9u1qdAWU,7571
 datashare_python/task_client.py,sha256=oTmP8bvZW0UyhLNMi1AV3XIAx7hrdbxNRss2Mw2azEc,8435
 datashare_python/template.py,sha256=RxKTYLXoS_EQ8Jc41JkBXppPdbCFqDWfP3BmC0gvB5o,4024
 datashare_python/types_.py,sha256=9Hk1XqpdXbM1TnEzwvJ5G9ABbaCZW9KgBTtiPBVn_7k,649
-datashare_python/utils.py,sha256=inVjtlBbgL88mN0UM73SSzW76koTW5MGC0NlyopqRW4,17412
-datashare_python/worker-template.tar.gz,sha256=gNSDvn4Lh8iFpFk6j8nlu7pJoIDLh7SVa6EHcjXrj54,286805
-datashare_python/worker.py,sha256=1FdmwYKWKYUKteTM3RC6kFQHR02q8NUDe91hv68QPEo,7207
+datashare_python/utils.py,sha256=gX3_RJEJS0sAYBNVfBLoWJu7_hIANhylSAohzXVW-yQ,17982
+datashare_python/worker-template.tar.gz,sha256=bOqoF6xVJRyFQaRYHIXXre31WYdmEqDLGeiXRr4Inqg,287091
+datashare_python/worker.py,sha256=czrN9Z0fPFX-6KHinX8Orx4vb9tpta2e7Qs6H0NiYyE,7534
 datashare_python/cli/__init__.py,sha256=9BPWtssDgsVfWMsZ1TtZCla0EC_kai4RHttr8oNLYOE,1401
 datashare_python/cli/project.py,sha256=w32Gy9AOL5B00uDT4in7YUCt2g68FnNbvwg2M3a8G6o,946
 datashare_python/cli/task.py,sha256=8mvKGS21bZ14BgZ0Uo-dfameljkaI2ZBha80ywCy-E8,5822
 datashare_python/cli/utils.py,sha256=p69CQb0zfixuyBkiZprhdMCc_NuYwXyAn6vC9H1UzAw,911
 datashare_python/cli/worker.py,sha256=I4KTpFIpXFowioFn72Rm6LBCYlY-Dhp4NBIPvtRgUXE,5283
-datashare_python-0.7.1.dist-info/METADATA,sha256=FF5E62UhxGi0zBseyb4b_6TQRl2JtDZlA1wLECoTblg,923
-datashare_python-0.7.1.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
-datashare_python-0.7.1.dist-info/entry_points.txt,sha256=ILE7auxabHWiu3GC-AunWnzjhOI_SbZp7D4GqZHlLw4,68
-datashare_python-0.7.1.dist-info/RECORD,,
+datashare_python-0.7.3.dist-info/METADATA,sha256=sZEiq4mFYgmvkYuJOO6KWePleBYWDex722rVR3FSm3I,921
+datashare_python-0.7.3.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
+datashare_python-0.7.3.dist-info/entry_points.txt,sha256=ILE7auxabHWiu3GC-AunWnzjhOI_SbZp7D4GqZHlLw4,68
+datashare_python-0.7.3.dist-info/RECORD,,