datashare-python 0.2.23__py3-none-any.whl → 0.2.25__py3-none-any.whl

datashare_python/cli/worker.py

@@ -1,12 +1,15 @@
+import asyncio
 import logging
+from pathlib import Path
 from typing import Annotated
 
 import typer
+import yaml
+from icij_common.pydantic_utils import safe_copy
 
-from datashare_python.constants import DEFAULT_NAMESPACE, DEFAULT_TEMPORAL_ADDRESS
-from datashare_python.discovery import discover_activities, discover_workflows
-from datashare_python.types_ import TemporalClient
-from datashare_python.worker import datashare_worker
+from datashare_python.config import TemporalClientConfig, WorkerConfig
+from datashare_python.discovery import discover, discover_activities, discover_workflows
+from datashare_python.worker import bootstrap_worker, create_worker_id
 
 from .utils import AsyncTyper
 
@@ -20,11 +23,13 @@ _LIST_ACTIVITY_NAMES_HELP = "activity names filters (supports regexes)"
 
 _START_WORKER_WORKFLOWS_HELP = "workflow names run by the worker (supports regexes)"
 _START_WORKER_ACTIVITIES_HELP = "activity names run by the worker (supports regexes)"
-_WORKER_QUEUE_HELP = "worker task queue"
-_WORKER_MAX_ACTIVITIES_HELP = (
-    "maximum number of concurrent activities/tasks"
-    " concurrently run by the worker. Defaults to 1 to encourage horizontal scaling."
+_START_WORKER_DEPS_HELP = "worker lifetime dependencies name in the registry"
+_START_WORKER_WORKER_ID_PREFIX_HELP = "worker ID prefix"
+_START_WORKER_CONFIG_PATH_HELP = (
+    "path to a worker config YAML file,"
+    " if not provided will load worker configuration from env variables"
 )
+_WORKER_QUEUE_HELP = "worker task queue"
 _TEMPORAL_NAMESPACE_HELP = "worker temporal namespace"
 
 _TEMPORAL_URL_HELP = "address for temporal server"
@@ -73,42 +78,60 @@ async def start(
     workflows: Annotated[list[str], typer.Option(help=_START_WORKER_WORKFLOWS_HELP)],
     activities: Annotated[list[str], typer.Option(help=_START_WORKER_ACTIVITIES_HELP)],
     queue: Annotated[str, typer.Option("--queue", "-q", help=_WORKER_QUEUE_HELP)],
+    dependencies: Annotated[
+        str | None, typer.Option(help=_START_WORKER_DEPS_HELP)
+    ] = None,
+    config_path: Annotated[
+        Path | None,
+        typer.Option(
+            "--config-path", "--config", "-c", help=_START_WORKER_CONFIG_PATH_HELP
+        ),
+    ] = None,
+    worker_id_prefix: Annotated[
+        str | None, typer.Option(help=_START_WORKER_WORKER_ID_PREFIX_HELP)
+    ] = None,
     temporal_address: Annotated[
-        str, typer.Option("--temporal-address", "-a", help=_TEMPORAL_URL_HELP)
-    ] = DEFAULT_TEMPORAL_ADDRESS,
+        str | None, typer.Option("--temporal-address", "-a", help=_TEMPORAL_URL_HELP)
+    ] = None,
     namespace: Annotated[
-        str, typer.Option("--temporal-namespace", "-ns", help=_TEMPORAL_NAMESPACE_HELP)
-    ] = DEFAULT_NAMESPACE,
-    max_concurrent_activities: Annotated[
-        int, typer.Option("--max-activities", help=_WORKER_MAX_ACTIVITIES_HELP)
-    ] = 1,
+        str | None,
+        typer.Option("--temporal-namespace", "-ns", help=_TEMPORAL_NAMESPACE_HELP),
+    ] = None,
 ) -> None:
-    wf_names, wfs = zip(*discover_workflows(workflows), strict=False)
-    registered = ""
-    if wf_names:
-        n_wfs = len(wf_names)
-        registered += (
-            f"- {n_wfs} workflow{'s' if n_wfs > 1 else ''}: {','.join(wf_names)}"
-        )
-    act_names, acts = zip(*discover_activities(activities), strict=False)
-    if act_names:
-        if registered:
-            registered += "\n"
-        i = len(act_names)
-        registered += f"- {i} activit{'ies' if i > 1 else 'y'}: {','.join(act_names)}"
-    if not acts and not wfs:
-        raise ValueError("Couldn't find any registered activity or workflow.")
-    logger.info("Starting datashare worker running:\n%s", registered)
-    client = await TemporalClient.connect(temporal_address, namespace=namespace)
-    worker = datashare_worker(
-        client,
-        workflows=wfs,
-        activities=acts,
-        task_queue=queue,
-        max_concurrent_activities=max_concurrent_activities,
+    if config_path is not None:
+        with config_path.open() as f:
+            bootstrap_config = WorkerConfig.model_validate(
+                yaml.load(f, Loader=yaml.Loader)
+            )
+    else:
+        bootstrap_config = WorkerConfig()
+    temporal_override = dict()
+    if temporal_address is not None:
+        temporal_override["host"] = temporal_address
+    if namespace is not None:
+        temporal_override["namespace"] = namespace
+    if temporal_override:
+        temporal_config = TemporalClientConfig(**temporal_override)
+        update = {"temporal": temporal_config}
+        bootstrap_config = safe_copy(bootstrap_config, update=update)
+    registered_wfs, registered_acts, registered_deps = discover(
+        workflows, act_names=activities, deps_name=dependencies
     )
-    try:
-        await worker.run()
-    except Exception as e:  # noqa: BLE001
-        await worker.shutdown()
-        raise e
+    worker_id = create_worker_id(worker_id_prefix or "worker")
+    client = await bootstrap_config.to_temporal_client()
+    event_loop = asyncio.get_event_loop()
+    async with bootstrap_worker(
+        worker_id,
+        activities=registered_acts,
+        workflows=registered_wfs,
+        dependencies=registered_deps,
+        bootstrap_config=bootstrap_config,
+        client=client,
+        event_loop=event_loop,
+        task_queue=queue,
+    ) as worker:
+        try:
+            await worker.run()
+        except Exception as e:  # noqa: BLE001
+            await worker.shutdown()
+            raise e

datashare_python/config.py

@@ -4,7 +4,13 @@ from icij_common.es import ESClient
 from icij_common.pydantic_utils import ICIJSettings
 from pydantic import Field, PrivateAttr
 from pydantic_settings import SettingsConfigDict
-from temporalio.contrib.pydantic import pydantic_data_converter
+from temporalio.contrib.pydantic import PydanticJSONPlainPayloadConverter, ToJsonOptions
+from temporalio.converter import (
+    CompositePayloadConverter,
+    DataConverter,
+    DefaultPayloadConverter,
+    JSONPlainPayloadConverter,
+)
 
 import datashare_python
 
@@ -64,7 +70,7 @@ class TemporalClientConfig(BaseModel):
             self._client = await TemporalClient.connect(
                 target_host=self.host,
                 namespace=self.namespace,
-                data_converter=pydantic_data_converter,
+                data_converter=PYDANTIC_DATA_CONVERTER,
             )
         return self._client
 
@@ -83,6 +89,8 @@ class WorkerConfig(ICIJSettings, LogWithWorkerIDMixin, BaseModel):
     elasticsearch: ESClientConfig = ESClientConfig()
     temporal: TemporalClientConfig = TemporalClientConfig()
 
+    max_concurrent_io_activities: int = 5
+
     def to_es_client(self) -> ESClient:
         return self.elasticsearch.to_es_client(self.datashare.api_key)
 
@@ -91,3 +99,23 @@ class WorkerConfig(ICIJSettings, LogWithWorkerIDMixin, BaseModel):
 
     async def to_temporal_client(self) -> TemporalClient:
         return await self.temporal.to_client()
+
+
+class _PydanticPayloadConverter(CompositePayloadConverter):
+    def __init__(self) -> None:
+        json_payload_converter = PydanticJSONPlainPayloadConverter(
+            ToJsonOptions(exclude_unset=False)
+        )
+        super().__init__(
+            *(
+                c
+                if not isinstance(c, JSONPlainPayloadConverter)
+                else json_payload_converter
+                for c in DefaultPayloadConverter.default_encoding_payload_converters
+            )
+        )
+
+
+PYDANTIC_DATA_CONVERTER = DataConverter(
+    payload_converter_class=_PydanticPayloadConverter
+)

datashare_python/conftest.py

@@ -1,6 +1,7 @@
 import asyncio
 from asyncio import AbstractEventLoop
 from collections.abc import AsyncGenerator, Generator, Iterator, Sequence
+from pathlib import Path
 
 import aiohttp
 import pytest
@@ -50,6 +51,9 @@ _INDEX_BODY = {
             "language": {"type": "keyword"},
             "documentId": {"type": "keyword"},
             "join": {"type": "join", "relations": {"Document": "NamedEntity"}},
+            "contentType": {"type": "keyword"},
+            "content": {"type": "text"},
+            "contentTranslated": {"type": "text"},
         }
     }
 }
@@ -102,6 +106,13 @@ def test_worker_config() -> WorkerConfig:
     )
 
 
+@pytest.fixture
+def test_worker_config_path(test_worker_config: WorkerConfig, tmpdir: Path) -> Path:
+    config_path = Path(tmpdir) / "config.json"
+    config_path.write_text(test_worker_config.model_dump_json())
+    return config_path
+
+
 @pytest.fixture(scope="session")
 async def worker_lifetime_deps(
     event_loop: AbstractEventLoop,

datashare_python/constants.py

@@ -12,3 +12,5 @@ DEFAULT_TEMPORAL_ADDRESS = "temporal:7233"
 DEFAULT_DS_ADDRESS = "http://localhost:8080"
 
 DEFAULT_NAMESPACE = "datashare-default"
+
+METADATA_JSON = "metadata.json"

datashare_python/dependencies.py

@@ -1,8 +1,11 @@
+import inspect
 import logging
 from asyncio import AbstractEventLoop, iscoroutine
-from collections.abc import AsyncGenerator
+from collections.abc import AsyncGenerator, Callable
 from contextlib import AsyncExitStack, asynccontextmanager
 from contextvars import ContextVar
+from copy import deepcopy
+from typing import Any
 
 from icij_common.es import ESClient
 
@@ -20,7 +23,7 @@ TASK_CLIENT: ContextVar[DatashareTaskClient] = ContextVar("task_client")
 TEMPORAL_CLIENT: ContextVar[TemporalClient] = ContextVar("temporal_client")
 
 
-def set_event_loop(event_loop: AbstractEventLoop, **_) -> None:
+def set_event_loop(event_loop: AbstractEventLoop) -> None:
     EVENT_LOOP.set(event_loop)
 
 
@@ -31,13 +34,13 @@ def lifespan_event_loop() -> AbstractEventLoop:
         raise DependencyInjectionError("event loop") from e
 
 
-def set_loggers(worker_config: WorkerConfig, worker_id: str, **_) -> None:
+def set_loggers(worker_config: WorkerConfig, worker_id: str) -> None:
     worker_config.setup_loggers(worker_id=worker_id)
     logger.info("worker loggers ready to log 💬")
     logger.info("app config: %s", worker_config.model_dump_json(indent=2))
 
 
-async def set_es_client(worker_config: WorkerConfig, **_) -> ESClient:
+async def set_es_client(worker_config: WorkerConfig) -> ESClient:
     client = worker_config.to_es_client()
     ES_CLIENT.set(client)
     return client
@@ -52,7 +55,7 @@ def lifespan_es_client() -> ESClient:
 
 
 # Task client setup
-async def set_task_client(worker_config: WorkerConfig, **_) -> DatashareTaskClient:
+async def set_task_client(worker_config: WorkerConfig) -> DatashareTaskClient:
     task_client = worker_config.to_task_client()
     TASK_CLIENT.set(task_client)
     return task_client
@@ -67,7 +70,7 @@ def lifespan_task_client() -> DatashareTaskClient:
 
 
 # Temporal client setup
-async def set_temporal_client(worker_config: WorkerConfig, **_) -> None:
+async def set_temporal_client(worker_config: WorkerConfig) -> None:
     client = await worker_config.to_temporal_client()
     TEMPORAL_CLIENT.set(client)
 
@@ -86,7 +89,7 @@ async def with_dependencies(
 ) -> AsyncGenerator[None, None]:
     async with AsyncExitStack() as stack:
         for dep in dependencies:
-            cm = dep(**kwargs)
+            cm = dep(**add_missing_args(dep, dict(), **kwargs))
             if hasattr(cm, "__aenter__"):
                 await stack.enter_async_context(cm)
             elif hasattr(cm, "__enter__"):
@@ -94,3 +97,20 @@ async def with_dependencies(
             elif iscoroutine(cm):
                 await cm
         yield
+
+
+def add_missing_args(fn: Callable, args: dict[str, Any], **kwargs) -> dict[str, Any]:
+    # We make the choice not to raise in case of missing argument here, the error will
+    # be correctly raise when the function is called
+    from_kwargs = dict()
+    sig = inspect.signature(fn)
+    for param_name in sig.parameters:
+        if param_name in args:
+            continue
+        kwargs_value = kwargs.get(param_name)
+        if kwargs_value is not None:
+            from_kwargs[param_name] = kwargs_value
+    if from_kwargs:
+        args = deepcopy(args)
+        args.update(from_kwargs)
+    return args

datashare_python/discovery.py

@@ -1,18 +1,77 @@
+import logging
 import re
 from collections.abc import Callable, Iterable
 from importlib.metadata import entry_points
 
+from .types_ import ContextManagerFactory
 from .utils import ActivityWithProgress
 
+logger = logging.getLogger(__name__)
+
 Activity = ActivityWithProgress | Callable | type
 
+_DEPENDENCIES = "dependencies"
 _WORKFLOW_GROUPS = "datashare.workflows"
 _ACTIVITIES_GROUPS = "datashare.activities"
+_DEPENDENCIES_GROUPS = "datashare.dependencies"
+
+_RegisteredWorkflow = tuple[str, type]
+_RegisteredActivity = tuple[str, Activity]
+_Dependencies = list[ContextManagerFactory]
+_Discovery = tuple[
+    Iterable[_RegisteredWorkflow] | None,
+    Iterable[_RegisteredActivity] | None,
+    _Dependencies | None,
+]
 
 
-def discover_workflows(names: list[str]) -> Iterable[tuple[str, type]]:
+def discover(
+    wf_names: list[str] | None, *, act_names: list[str] | None, deps_name: str | None
+) -> _Discovery:
+    discovered = ""
+    wfs = None
+    if wf_names is not None:
+        discovered_wfs = discover_workflows(wf_names)
+        if discovered_wfs:
+            wf_names, wfs = zip(*discovered_wfs, strict=True)
+            if wf_names:
+                n_wfs = len(wf_names)
+                discovered += (
+                    f"- {n_wfs} workflow{'s' if n_wfs > 1 else ''}:"
+                    f" {', '.join(wf_names)}"
+                )
+    acts = None
+    if act_names is not None:
+        discovered_acts = discover_activities(act_names)
+        if discovered_acts:
+            act_names, acts = zip(*discovered_acts, strict=True)
+            if act_names:
+                if discovered:
+                    discovered += "\n"
+                n_acts = len(act_names)
+                discovered += (
+                    f"- {n_acts} activit{'ies' if n_acts > 1 else 'y'}:"
+                    f" {', '.join(act_names)}"
+                )
+    if not acts and not wfs:
+        raise ValueError("Couldn't find any registered activity or workflow.")
+    deps = discover_dependencies(deps_name)
+    if deps:
+        n_deps = len(deps)
+        discovered += "\n"
+        deps_names = (d.__name__ for d in deps)
+        discovered += (
+            f"- {n_deps} dependenc{'ies' if n_deps > 1 else 'y'}:"
+            f" {', '.join(deps_names)}"
+        )
+    logger.info("discovered:\n%s", discovered)
+    return wfs, acts, deps
+
+
+def discover_workflows(names: list[str]) -> list[_RegisteredWorkflow]:
     pattern = None if not names else re.compile(rf"^{'|'.join(names)}$")
     impls = entry_points(group=_WORKFLOW_GROUPS)
+    registered = []
     for wf_impls in impls:
         wf_impls = wf_impls.load()  # noqa: PLW2901
         if not isinstance(wf_impls, list | tuple | set):
@@ -21,12 +80,14 @@ def discover_workflows(names: list[str]) -> Iterable[tuple[str, type]]:
             wf_name = _parse_wf_name(wf_impl)
             if pattern and not pattern.match(wf_name):
                 continue
-            yield wf_name, wf_impl
+            registered.append((wf_name, wf_impl))
+    return registered
 
 
-def discover_activities(names: list[str]) -> Iterable[tuple[str, Activity]]:
+def discover_activities(names: list[str]) -> list[_RegisteredActivity]:
     pattern = None if not names else re.compile(rf"^{'|'.join(names)}$")
     impls = entry_points(group=_ACTIVITIES_GROUPS)
+    registered = []
     for act_impls in impls:
         act_impls = act_impls.load()  # noqa: PLW2901
         if not isinstance(act_impls, list | tuple | set):
@@ -35,7 +96,45 @@ def discover_activities(names: list[str]) -> Iterable[tuple[str, Activity]]:
             act_name = _parse_activity_name(act_impl)
             if pattern and not pattern.match(act_name):
                 continue
-            yield act_name, act_impl
+            registered.append((act_name, act_impl))
+    return registered
+
+
+def discover_dependencies(name: str | None) -> _Dependencies | None:
+    impls = entry_points(name=_DEPENDENCIES, group=_DEPENDENCIES_GROUPS)
+    if not impls:
+        if name is None:
+            return None
+        available_impls = entry_points(group=_DEPENDENCIES_GROUPS)
+        msg = (
+            f'failed to find dependency: "{name}", '
+            f"available dependencies: {available_impls}"
+        )
+        raise LookupError(msg)
+    if len(impls) > 1:
+        msg = f'found multiple dependencies for name "{name}": {impls}'
+        raise ValueError(msg)
+    deps_registry = impls[_DEPENDENCIES].load()
+    if name:
+        try:
+            return deps_registry[name]
+        except KeyError as e:
+            available = list(deps_registry)
+            msg = (
+                f'failed to find dependency for name "{name}", available dependencies: '
+                f"{available}"
+            )
+            raise LookupError(msg) from e
+    if not deps_registry:
+        raise ValueError("empty dependency registry !")
+    if len(deps_registry) > 1:
+        available = ", ".join('"' + d + '"' for d in deps_registry)
+        msg = (
+            f"dependency registry contains multiples entries {available},"
+            f" please select one by providing a name"
+        )
+        raise ValueError(msg)
+    return next(iter(deps_registry.values()))
 
 
 def _parse_wf_name(wf_type: type) -> str:

datashare_python/objects.py

@@ -1,9 +1,9 @@
 import logging
-from collections.abc import Callable
+from collections.abc import Awaitable, Callable
 from dataclasses import dataclass
 from datetime import UTC, datetime
 from enum import StrEnum, unique
-from typing import Any, Literal, Self
+from typing import Any, Literal, Self, TypeVar
 
 from temporalio import workflow
 
@@ -23,14 +23,26 @@ from pydantic.main import IncEx
 logger = logging.getLogger(__name__)
 
 
+T = TypeVar("T")
+Predicate = Callable[[T], bool] | Callable[[T], Awaitable[bool]]
+
+
 class BaseModel(_BaseModel):
     model_config = merge_configs(icij_config(), no_enum_values_config())
 
 
+class BasePayload(_BaseModel):
+    model_config = icij_config()
+
+
 class DatashareModel(BaseModel):
     model_config = merge_configs(BaseModel.model_config, lowercamel_case_config())
 
 
+class LowerCamelCaseModel(_BaseModel):
+    model_config = merge_configs(icij_config(), lowercamel_case_config())
+
+
 @unique
 class TaskState(StrEnum):
     CREATED = "CREATED"
@@ -135,12 +147,14 @@ class TaskGroup:
 class Document(DatashareModel):
     id: str
     root_document: str
-    content: str
     language: str
+    content: str | None = None
+    content_type: str | None = None
     tags: list[str] = Field(default_factory=list)
     content_translated: dict[str, str] = Field(
         default_factory=dict, alias="content_translated"
     )
+    type: str = Field(default="Document", frozen=True)
 
     @classmethod
     def from_es(cls, es_doc: dict) -> Self:

datashare_python/template.py

@@ -71,12 +71,12 @@ def _update_pyproject_toml(
 ) -> dict[str, Any]:
     pyproject_toml = deepcopy(pyproject_toml)
 
-    pyproject_toml["tool"]["uv"].pop("sources")
-    pyproject_toml["tool"]["uv"].pop("index")
+    pyproject_toml["tool"]["uv"].pop("sources", None)
+    pyproject_toml["tool"]["uv"].pop("index", None)
 
     project = pyproject_toml["project"]
     project["authors"] = []
-    project.pop("urls")
+    project.pop("urls", None)
     project["dependencies"] = sorted(
         d
         for d in project["dependencies"]
@@ -87,7 +87,7 @@ def _update_pyproject_toml(
         for d in project["dependencies"]
         if any(d.startswith(base) for base in _BASE_DEPS)
     )
-    project.pop("optional-dependencies")
+    project.pop("optional-dependencies", None)
 
     entry_points = project["entry-points"]
 
@@ -105,4 +105,11 @@ def _update_pyproject_toml(
         i if i != "worker_template" else package_name for i in hatch_sdist["packages"]
     ]
 
+    hatch_sdist = pyproject_toml["tool"]["hatch"]["build"]["targets"]["sdist"]
+    if "only-include" in hatch_sdist:
+        hatch_sdist["only-include"] = [
+            i if i != "worker_template" else package_name
+            for i in hatch_sdist["only-include"]
+        ]
+
     return pyproject_toml

datashare_python/utils.py

@@ -1,14 +1,24 @@
 import asyncio
 import inspect
+import json
 import logging
 import sys
-from collections.abc import Callable, Coroutine
+from collections.abc import (
+    Callable,
+    Coroutine,
+)
+from copy import deepcopy
 from dataclasses import dataclass
+from datetime import timedelta
 from functools import partial, wraps
+from hashlib import sha256
 from inspect import signature
-from typing import ParamSpec, TypeVar
+from pathlib import Path
+from typing import Any, ParamSpec, TypeVar
+from uuid import uuid4
 
 import nest_asyncio
+import temporalio
 from icij_common.logging_utils import (
     DATE_FMT,
     STREAM_HANDLER_FMT,
@@ -20,9 +30,10 @@ 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 SearchAttributeKey
+from temporalio.common import RetryPolicy, SearchAttributeKey
 from temporalio.exceptions import ApplicationError
 
+from .constants import METADATA_JSON
 from .types_ import ProgressRateHandler, RawProgressHandler
 
 DependencyLabel = str | None
@@ -34,6 +45,13 @@ PROGRESS_HANDLER_ARG = "progress"
 P = ParamSpec("P")
 T = TypeVar("T")
 
+_NEVER_RETRIABLES = {
+    "ValidationError",
+    "TypeError",
+    "PydanticSchemaGenerationError",
+    "PydanticSerializationError",
+}
+
 
 @dataclass(frozen=True)
 class Progress:
@@ -86,6 +104,38 @@ class WorkflowWithProgress:
             workflow.upsert_search_attributes(attributes)
 
 
+def _retry_policy_with_default(retry_policy: RetryPolicy | None) -> RetryPolicy:
+    if retry_policy is None:
+        retry_policy = RetryPolicy(non_retryable_error_types=[])
+    retry_policy = deepcopy(retry_policy)
+    non_retryable_error_types = set(retry_policy.non_retryable_error_types)
+    non_retryable_error_types.update(_NEVER_RETRIABLES)
+    retry_policy.non_retryable_error_types = list(non_retryable_error_types)
+    return retry_policy
+
+
+async def execute_activity(
+    activity: Callable,
+    task_queue: str,
+    arg: Any = temporalio.common._arg_unset,
+    *,
+    args: list | None = None,
+    start_to_close_timeout: timedelta | None = None,
+    retry_policy: temporalio.common.RetryPolicy | None = None,
+) -> Any:
+    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,
+    )
+
+
 async def progress_handler(
     progress: float,
     handle: WorkflowHandle,
@@ -147,10 +197,7 @@ def with_progress(weight: float = 1.0) -> Callable[P, T]:
                     client=self._temporal_client, weight=weight
                 )
                 await handler(0.0)
-                if supports_progress(activity_fn):
-                    res = await activity_fn(self, *args, progress=handler)
-                else:
-                    res = await activity_fn(self, *args)
+                res = await activity_fn(self, *args, progress=handler)
                 await handler(1.0)
                 return res
 
@@ -170,12 +217,9 @@ def with_progress(weight: float = 1.0) -> Callable[P, T]:
                     client=self._temporal_client, weight=weight
                 )
                 event_loop = self._event_loop
-                event_loop.run_until_complete(handler(0.0))
-                if supports_progress(activity_fn):
-                    res = activity_fn(self, *args, progress=handler)
-                else:
-                    res = activity_fn(self, *args)
-                event_loop.run_until_complete(handler(1.0))
+                asyncio.run_coroutine_threadsafe(handler(0.0), event_loop).result()
+                res = activity_fn(self, *args, progress=handler)
+                asyncio.run_coroutine_threadsafe(handler(1.0), event_loop).result()
                 return res
 
         return wrapper
@@ -197,6 +241,7 @@ def positional_args_only(activity_fn: Callable[P, T]) -> Callable[P, T]:
             # recreate kwargs from pargs
             new_args, new_kwargs = _unpack_positional_args(args, keyword_only, params)
             return await activity_fn(*new_args, **new_kwargs, **kwargs)
+
     else:
 
         @wraps(activity_fn)
@@ -234,8 +279,28 @@ def with_retriables(
     retriables: set[type[Exception]] = None,
 ) -> Callable[[Callable[P, T]], Callable[P, T]]:
     if retriables is None:
-        retriables = set()
-    retriables = tuple(retriables)
+
+        def decorator(activity_fn: Callable[P, T]) -> Callable[P, T]:
+            if asyncio.iscoroutinefunction(activity_fn):
+
+                @wraps(activity_fn)
+                async def wrapper(*args, **kwargs) -> T:
+                    try:
+                        return await activity_fn(*args, **kwargs)
+                    except Exception as e:
+                        raise fatal_error_from_exception(e) from e
+            else:
+
+                @wraps(activity_fn)
+                def wrapper(*args, **kwargs) -> T:
+                    try:
+                        return activity_fn(*args, **kwargs)
+                    except Exception as e:
+                        raise fatal_error_from_exception(e) from e
+
+            return wrapper
+
+        return decorator
 
     def decorator(activity_fn: Callable[P, T]) -> Callable[P, T]:
         if asyncio.iscoroutinefunction(activity_fn):
@@ -248,6 +313,7 @@ def with_retriables(
                     raise
                 except Exception as e:
                     raise fatal_error_from_exception(e) from e
+
         else:
 
             @wraps(activity_fn)
@@ -268,20 +334,13 @@ def activity_defn(
     name: str,
     progress_weight: float = 1.0,
     retriables: set[type[Exception]] = None,
-    *,
-    no_thread_cancel_exception: bool = False,
-    dynamic: bool = False,
 ) -> Callable[[Callable[P, T]], Callable[P, T]]:
     def decorator(activity_fn: Callable[P, T]) -> Callable[P, T]:
         activity_fn = positional_args_only(activity_fn)
         activity_fn = with_retriables(retriables)(activity_fn)
-        activity_fn = with_progress(progress_weight)(activity_fn)
-        activity_fn = activity.defn(
-            activity_fn,
-            name=name,
-            no_thread_cancel_exception=no_thread_cancel_exception,
-            dynamic=dynamic,
-        )
+        if supports_progress(activity_fn):
+            activity_fn = with_progress(progress_weight)(activity_fn)
+        activity_fn = activity.defn(activity_fn, name=name)
         return activity_fn
 
     return decorator
@@ -370,3 +429,71 @@ class LogWithWorkerIDMixin:
                 handler.addFilter(worker_id_filter)
             handler.setLevel(log_level)
         return handlers
+
+
+def safe_dir(filename: str) -> Path:
+    filename = filename.split(".", maxsplit=1)[0]
+    parts = (p for p in (filename[:2], filename[2:4]) if p)
+    return Path(*parts)
+
+
+def artifacts_dir(project: str, *, filename: str) -> Path:
+    return Path(project, safe_dir(filename), filename)
+
+
+def metadata_path(filename: str, *, project: str) -> Path:
+    metadata_path = artifacts_dir(project, filename=filename) / METADATA_JSON
+    return metadata_path
+
+
+def _read_artifact_metadata(root: Path, project: str, *, filename: str) -> dict:
+    m_path = root / metadata_path(filename, project=project)
+    return json.loads(m_path.read_text())
+
+
+def write_artifact(
+    artifact: bytes,
+    root: Path,
+    *,
+    project: str,
+    filename: str,
+    metadata_key: str,
+    metadata_value: str,
+) -> Path:
+    artif_dir = root / artifacts_dir(project, filename=filename)
+    artif_dir.mkdir(exist_ok=True, parents=True)
+    # TODO: if transcriptions are too large we could also serialize them
+    #  as jsonl
+    transcription_path = artif_dir / metadata_value
+    transcription_path.write_bytes(artifact)
+    try:
+        meta = _read_artifact_metadata(root, project, filename=filename)
+    except FileNotFoundError:
+        meta = dict()
+    meta[metadata_key] = metadata_value
+    meta_path = root / artifacts_dir(project, filename=filename) / METADATA_JSON
+    meta_path.write_text(json.dumps(meta))
+    return transcription_path.relative_to(artif_dir)
+
+
+def debuggable_name(
+    path: str, component_size_limit: int = 10, *, deterministic: bool = False
+) -> str:
+    path = Path(path)
+    displayable_file_name = [c[:component_size_limit] for c in path.parts]
+    uuid = sha256(str(path).encode()).hexdigest() if deterministic else uuid4().hex
+    uuid = uuid[:20]
+    return f"{uuid}-{'__'.join(displayable_file_name)}"
+
+
+def activity_contextual_id(*, wf_context: bool = False) -> str:
+    act_info = activity.info()
+    act_id = act_info.activity_id
+    act_run_id = act_info.activity_id
+    act_type = act_info.activity_type
+    contextual_id = f"{act_type}-{act_id}-{act_run_id}"
+    if wf_context:
+        wf_id = act_info.workflow_id
+        wf_run_id = act_info.workflow_run_id
+        contextual_id += f"-{wf_id}-{wf_run_id}"
+    return contextual_id

datashare_python/worker.py

@@ -1,13 +1,28 @@
+import inspect
 import logging
+import os
+import socket
+import sys
+import threading
+from asyncio import AbstractEventLoop
+from collections.abc import AsyncGenerator, Callable
 from concurrent.futures import ThreadPoolExecutor
+from contextlib import asynccontextmanager
+from typing import Any
 
 from temporalio.worker import PollerBehaviorSimpleMaximum, Worker
 
+from .config import WorkerConfig
+from .dependencies import with_dependencies
 from .discovery import Activity
-from .types_ import TemporalClient
+from .types_ import ContextManagerFactory, TemporalClient
 
 logger = logging.getLogger(__name__)
 
+_TEMPORAL_CLIENT = "temporal_client"
+_EVENT_LOOP = "event_loop"
+_EXPECTED_INIT_ARGS = {"self", _TEMPORAL_CLIENT, _EVENT_LOOP, "args", "kwargs"}
+
 _SEPARATE_IO_AND_CPU_WORKERS = """The worker will run sync (CPU-bound) activities as \
 well as IO-bound workflows.
 To avoid deadlocks due to the GIL, we advise to run all CPU-bound activities inside a \
@@ -30,20 +45,21 @@ _ACTIVITY_THREAD_NAME_PREFIX = "datashare-activity-worker-"
 
 def datashare_worker(
     client: TemporalClient,
+    worker_id: str,
     *,
     workflows: list[type] | None = None,
     activities: list[Activity] | None = None,
     task_queue: str,
     # Scale horizontally be default for activities, each worker processes one activity
     # at a time
-    max_concurrent_activities: int = 1,
+    max_concurrent_io_activities: int = 10,
 ) -> Worker:
     if workflows is None:
         workflows = []
     if activities is None:
         activities = []
     are_async = [a.__temporal_activity_definition.is_async for a in activities]
-    if all(not a for a in are_async):
+    if are_async and all(not a for a in are_async):
         activity_executor = ThreadPoolExecutor(
             thread_name_prefix=_ACTIVITY_THREAD_NAME_PREFIX
         )
@@ -55,11 +71,15 @@ def datashare_worker(
         )
         logger.warning(_SEPARATE_IO_AND_CPU_ACTIVITIES)
 
-    if isinstance(activity_executor, ThreadPoolExecutor) and workflows:
-        logger.warning(_SEPARATE_IO_AND_CPU_WORKERS)
+    max_concurrent_activities = max_concurrent_io_activities
+    if isinstance(activity_executor, ThreadPoolExecutor):
+        max_concurrent_activities = 1
+        if workflows:
+            logger.warning(_SEPARATE_IO_AND_CPU_WORKERS)
 
     return Worker(
         client,
+        identity=worker_id,
         workflows=workflows,
         activities=activities,
         task_queue=task_queue,
@@ -72,3 +92,84 @@ def datashare_worker(
         # several of them
         workflow_task_poller_behavior=PollerBehaviorSimpleMaximum(5),
     )
+
+
+def create_worker_id(prefix: str) -> str:
+    pid = os.getpid()
+    threadid = threading.get_ident()
+    hostname = socket.gethostname()
+    # TODO: this might not be unique when using asyncio
+    return f"{prefix}-{hostname}-{pid}-{threadid}"
+
+
+def init_activity(
+    activity: Callable, client: TemporalClient, event_loop: AbstractEventLoop
+) -> Callable:
+    is_object_method = "." not in activity.__qualname__
+    if is_object_method:
+        return activity
+    cls = _get_class_from_method(activity)
+    init_args = inspect.signature(cls.__init__).parameters
+    invalid = [p for p in init_args if p not in _EXPECTED_INIT_ARGS]
+    if invalid:
+        msg = f"invalid activity arguments: {invalid}"
+        raise ValueError(msg)
+    kwargs = {_TEMPORAL_CLIENT: client, _EVENT_LOOP: event_loop}
+    kwargs = {k: v for k, v in kwargs.items() if k in _EXPECTED_INIT_ARGS}
+    if not kwargs:
+        return activity
+    act_instance = cls(**kwargs)
+    act_method = getattr(act_instance, activity.__name__)
+    return act_method
+
+
+@asynccontextmanager
+async def bootstrap_worker(
+    worker_id: str,
+    *,
+    activities: list[Callable[..., Any] | None] | None = None,
+    workflows: list[type] | None = None,
+    bootstrap_config: WorkerConfig,
+    client: TemporalClient,
+    event_loop: AbstractEventLoop,
+    task_queue: str,
+    dependencies: list[ContextManagerFactory] | None = None,
+) -> AsyncGenerator[Worker, None]:
+    deps_cm = (
+        with_dependencies(
+            dependencies,
+            worker_config=bootstrap_config,
+            worker_id=worker_id,
+            event_loop=event_loop,
+        )
+        if dependencies
+        else _do_nothing_cm()
+    )
+    async with deps_cm:
+        if activities is not None:
+            acts = [
+                init_activity(a, client=client, event_loop=event_loop)
+                for a in activities
+            ]
+        else:
+            acts = None
+        worker = datashare_worker(
+            client,
+            worker_id,
+            workflows=workflows,
+            activities=acts,
+            task_queue=task_queue,
+            max_concurrent_io_activities=bootstrap_config.max_concurrent_io_activities,
+        )
+        yield worker
+
+
+@asynccontextmanager
+async def _do_nothing_cm() -> AsyncGenerator[None, None]:
+    yield
+
+
+def _get_class_from_method(method: Callable) -> type:
+    class_name = method.__qualname__.rsplit(".", 1)[0]
+    module = sys.modules[method.__module__]
+    return getattr(module, class_name)

{datashare_python-0.2.23.dist-info → datashare_python-0.2.25.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datashare-python
-Version: 0.2.23
+Version: 0.2.25
 Summary: Manage Pythoœn tasks and local resources in Datashare
 Project-URL: Homepage, https://icij.github.io/datashare-python/
 Project-URL: Documentation, https://icij.github.io/datashare-python/
@@ -9,13 +9,13 @@ Project-URL: Issues, https://github.com/ICIJ/datashare-python/issues
 Author-email: Clément Doumouro <cdoumouro@icij.org>, Clément Doumouro <clement.doumouro@gmail.com>, Lion Summerbell <lsummerbell@icij.org>
 Requires-Python: <4,>=3.11
 Requires-Dist: aiohttp~=3.11.9
-Requires-Dist: aiostream~=0.6.4
 Requires-Dist: alive-progress~=3.2.0
-Requires-Dist: datashare-worker-template~=0.1
+Requires-Dist: datashare-worker-template[ml]~=0.1
 Requires-Dist: hatchling~=1.27.0
-Requires-Dist: icij-common[elasticsearch]~=0.7.3
+Requires-Dist: icij-common[elasticsearch]~=0.8.2
 Requires-Dist: nest-asyncio~=1.6.0
 Requires-Dist: python-json-logger~=4.0.0
+Requires-Dist: pyyaml~=6.0
 Requires-Dist: temporalio~=1.23.0
 Requires-Dist: tomlkit~=0.14.0
 Requires-Dist: typer~=0.15.4

datashare_python-0.2.25.dist-info/RECORD

@@ -0,0 +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=u6iyOeSXzIO30Yja8Vj9LjM-cq8ESRBy3Kse6UadAMg,3763
+datashare_python/conftest.py,sha256=BdRLjy9eJtxAKLDCcon1Nyhzn54CIw2z4s3ZOupNYGo,8256
+datashare_python/constants.py,sha256=e6Px11OUee9GSHwTgsgFMszGCMwpW-OznHSMgINvepc,338
+datashare_python/dependencies.py,sha256=Diu7alKGaFWyC_ajp0fKU-xp8u5f_8x1axAHVBlppD0,3707
+datashare_python/discovery.py,sha256=R0wws-_QgiK3GyTSf3p2UV41Ok-iuKeZauh6EJAdzGc,6401
+datashare_python/exceptions.py,sha256=bVHEAXxDPKfxeeMC0hJXEsrJkgsKO2ESAhxWU96GA4M,496
+datashare_python/local_client.py,sha256=GP9MTcHVQ1mcb2eO6TiQ7mzQdx199lZRhK8DRuJqJVQ,2359
+datashare_python/objects.py,sha256=GMi2hlKuWFbWWoC2r8ITQGQcMsobHHChGvm4ZfIjMl0,4537
+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=DQt-rBwC3Ok72u8VyerG3rqwUTx3ftLfPdMQ5cnRrgs,16801
+datashare_python/worker-template.tar.gz,sha256=CTokaLmqZZ6tJPGT4N7YG9aDsP66yS1HvJN-IzWhh0c,142348
+datashare_python/worker.py,sha256=A4SnmDB4y0ck6Wp_UZWdsSOyTvW54Z2Bq76gxtp-_PE,6070
+datashare_python/cli/__init__.py,sha256=5MGSE_0SwlOiwbyPwsP8RIXlTBB2_GGP0zDg4l6UAIY,1479
+datashare_python/cli/local.py,sha256=S-7qMpSqzi0oMvu01TCFEb8tayEvpw4pXMdCszKEYtU,986
+datashare_python/cli/project.py,sha256=w32Gy9AOL5B00uDT4in7YUCt2g68FnNbvwg2M3a8G6o,946
+datashare_python/cli/task.py,sha256=9If5OC7loG4C4gWWl4iOeqPJ4GOLlCWXQfuNLUHORrQ,5860
+datashare_python/cli/utils.py,sha256=p69CQb0zfixuyBkiZprhdMCc_NuYwXyAn6vC9H1UzAw,911
+datashare_python/cli/worker.py,sha256=tJ2xj_TCyjZVh1Jlb_AknHEg8xn9Js90Vb39slew8t4,5160
+datashare_python-0.2.25.dist-info/METADATA,sha256=yIqCI2amsvCqNeKqnJ0RJtHUV1OnrUn7Lue5DgbYWk0,958
+datashare_python-0.2.25.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
+datashare_python-0.2.25.dist-info/entry_points.txt,sha256=ILE7auxabHWiu3GC-AunWnzjhOI_SbZp7D4GqZHlLw4,68
+datashare_python-0.2.25.dist-info/RECORD,,

datashare_python-0.2.23.dist-info/RECORD

@@ -1,27 +0,0 @@
-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=tTdLeBxefsLbXBy0sZS1n1wYuLPGMITm8mklo7CDraM,2954
-datashare_python/conftest.py,sha256=TVa57J4FQcwVIYgjCiNShvu4F1dUi94qJ7Q-SGlML6k,7853
-datashare_python/constants.py,sha256=QrtQEf-qHAkdzs1xJuyeN6XbCalTRVcgVvEZNdT-soU,305
-datashare_python/dependencies.py,sha256=3KDZ71x3KGnik-QPfXgSPRRLYfbj0oSQ0jCG-3x1c3M,3017
-datashare_python/discovery.py,sha256=zCTbTxFmAaY4aRwCyUAGh-HQVoZAnxVOygp5cl2Oha4,2957
-datashare_python/exceptions.py,sha256=bVHEAXxDPKfxeeMC0hJXEsrJkgsKO2ESAhxWU96GA4M,496
-datashare_python/local_client.py,sha256=GP9MTcHVQ1mcb2eO6TiQ7mzQdx199lZRhK8DRuJqJVQ,2359
-datashare_python/objects.py,sha256=UDc9PDaSXADLgmnfvqGBrDTmPVL-rRpyNHvY277-1x8,4147
-datashare_python/task_client.py,sha256=oTmP8bvZW0UyhLNMi1AV3XIAx7hrdbxNRss2Mw2azEc,8435
-datashare_python/template.py,sha256=EaFX7A6sw-HCb3RWrRvMg61cnKzfZlIB6Liom4mJXys,3725
-datashare_python/types_.py,sha256=9Hk1XqpdXbM1TnEzwvJ5G9ABbaCZW9KgBTtiPBVn_7k,649
-datashare_python/utils.py,sha256=mFegsYFoIRzRtPhx12748qrahe-c0rKNx02-r8tTG8Q,12777
-datashare_python/worker-template.tar.gz,sha256=THxySB554TAzBl2K5IIIj7yKulv0lhodM-zVN_Zk1Sg,142868
-datashare_python/worker.py,sha256=E1cECV0q2XcGI234gZZ5eBvqeULlTK0FUb0VSWkRQUI,2924
-datashare_python/cli/__init__.py,sha256=5MGSE_0SwlOiwbyPwsP8RIXlTBB2_GGP0zDg4l6UAIY,1479
-datashare_python/cli/local.py,sha256=S-7qMpSqzi0oMvu01TCFEb8tayEvpw4pXMdCszKEYtU,986
-datashare_python/cli/project.py,sha256=w32Gy9AOL5B00uDT4in7YUCt2g68FnNbvwg2M3a8G6o,946
-datashare_python/cli/task.py,sha256=9If5OC7loG4C4gWWl4iOeqPJ4GOLlCWXQfuNLUHORrQ,5860
-datashare_python/cli/utils.py,sha256=p69CQb0zfixuyBkiZprhdMCc_NuYwXyAn6vC9H1UzAw,911
-datashare_python/cli/worker.py,sha256=lmxeZDhLbGb-98XplCERN3R3Yo-sU0UkhO7KIo9Zbo4,4401
-datashare_python-0.2.23.dist-info/METADATA,sha256=yzNfCxfVgTNJLV2RigdLh8ltF64_h5JNpm6FBhru7Sg,959
-datashare_python-0.2.23.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
-datashare_python-0.2.23.dist-info/entry_points.txt,sha256=ILE7auxabHWiu3GC-AunWnzjhOI_SbZp7D4GqZHlLw4,68
-datashare_python-0.2.23.dist-info/RECORD,,