datashare-python 0.9.8__tar.gz → 0.10.0__tar.gz

{datashare_python-0.9.8 → datashare_python-0.10.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datashare-python
-Version: 0.9.8
+Version: 0.10.0
 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/
@@ -8,15 +8,16 @@ Project-URL: Repository, https://github.com/ICIJ/datashare-python
 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: aiofile~=3.11
 Requires-Dist: aiohttp~=3.11
 Requires-Dist: alive-progress~=3.2
 Requires-Dist: hatchling~=1.27
-Requires-Dist: icij-common[elasticsearch]~=0.8.2
+Requires-Dist: icij-common[elasticsearch]~=0.8.3
 Requires-Dist: langcodes~=3.5
 Requires-Dist: lru-dict~=1.4
 Requires-Dist: pydantic-extra-types[pycountry]>=2.11.1
 Requires-Dist: python-json-logger~=4.0
 Requires-Dist: pyyaml~=6.0
-Requires-Dist: temporalio~=1.23
+Requires-Dist: temporalio~=1.31
 Requires-Dist: tomlkit~=0.14.0
 Requires-Dist: typer<0.25.1,>=0.15.4

{datashare_python-0.9.8 → datashare_python-0.10.0}/datashare_python/cli/task.py

@@ -28,7 +28,7 @@ _START_HELP = "creates a new task and start it"
 _TASK_ID_HELP = "task ID"
 _WATCH_HELP = "watch a task until it's complete"
 
-TaskArgs = str
+StrTaskArgs = str
 
 task_app = AsyncTyper(name="task")
 
@@ -36,7 +36,7 @@ task_app = AsyncTyper(name="task")
 @task_app.async_command(help=_START_HELP)
 async def start(
     name: Annotated[str, typer.Argument(help=_NAME_HELP)],
-    args: Annotated[TaskArgs, typer.Argument(help=_ARGS_HELP)] = None,
+    args: Annotated[StrTaskArgs, typer.Argument(help=_ARGS_HELP)] = None,
     group: Annotated[
         str | None,
         typer.Option("--group", "-g", help=_GROUP_HELP),

{datashare_python-0.9.8 → datashare_python-0.10.0}/datashare_python/cli/worker.py

@@ -78,7 +78,7 @@ to the documentation to learn how to do so."""
 
 
 @worker_app.async_command(help=_START_WORKER_HELP)
-async def start(
+async def start(  # noqa: PLR0917
     queue: Annotated[str, typer.Option("--queue", "-q", help=_WORKER_QUEUE_HELP)],
     workflows: Annotated[
         list[str] | None,

{datashare_python-0.9.8 → datashare_python-0.10.0}/datashare_python/config.py

@@ -1,5 +1,4 @@
 from enum import StrEnum
-from pathlib import Path
 from typing import Literal
 
 from icij_common.es import ESClient
@@ -10,10 +9,10 @@ from temporalio.runtime import PrometheusConfig, Runtime, TelemetryConfig
 
 import datashare_python
 
-from .objects import BaseModel
+from .objects import BaseModel, WorkerPaths
 from .task_client import DatashareTaskClient
 from .types_ import TemporalClient
-from .utils import PYDANTIC_DATA_CONVERTER
+from .utils import PYDANTIC_DATA_CONVERTER, SharedResources, close_cm_callback
 
 _ALL_LOGGERS = [datashare_python.__name__]
 
@@ -94,6 +93,19 @@ class LoggingConfig(BaseModel):
     loggers: dict[str, LogLevel]
 
 
+class ResourceCacheConfig(BaseModel):
+    size: int = 1
+    exit_context_managers: bool = True
+
+    def to_resource_cache(self) -> SharedResources:
+        eviction_callback = None
+        if self.exit_context_managers:
+            eviction_callback = close_cm_callback
+        return SharedResources(
+            cache_size=self.size, eviction_callback=eviction_callback
+        )
+
+
 _DEFAULT_LOGGERS = {datashare_python.__name__: "INFO"}
 _DEFAULT_LOGGING_CONFIG = LoggingConfig(
     format=LogFormat.DEFAULT, loggers=_DEFAULT_LOGGERS
@@ -109,11 +121,9 @@ class WorkerConfig(ICIJSettings, BaseModel):
     elasticsearch: ESClientConfig = ESClientConfig()
     temporal: TemporalClientConfig = TemporalClientConfig()
 
-    max_concurrent_io_activities: int = 5
+    max_concurrent_activities: int = 5
 
-    docs_root: Path | None = None
-    artifacts_root: Path | None = None
-    workdir: Path | None = None
+    paths: WorkerPaths | None = None
 
     def to_es_client(self) -> ESClient:
         return self.elasticsearch.to_es_client(self.datashare.api_key)

{datashare_python-0.9.8 → datashare_python-0.10.0}/datashare_python/conftest.py

@@ -1,15 +1,22 @@
+import asyncio
+import faulthandler
+import logging
+import multiprocessing
 import shutil
+from asyncio import AbstractEventLoop
 from collections.abc import AsyncGenerator, Generator, Sequence
+from contextlib import AbstractAsyncContextManager, asynccontextmanager
+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
 from icij_common.test_utils import reset_env  # noqa: F401
 from pytest_asyncio import is_async_test
 from temporalio import workflow
+from temporalio.runtime import PrometheusConfig, Runtime, TelemetryConfig
 from temporalio.service import RPCError, RPCStatusCode
 
 from datashare_python.config import (
@@ -19,17 +26,14 @@ from datashare_python.config import (
     TemporalClientConfig,
     WorkerConfig,
 )
-from datashare_python.dependencies import (
-    lifespan_es_client,
-    lifespan_task_client,
-    set_es_client,
-    set_task_client,
-    with_dependencies,
-)
+from datashare_python.discovery import discover
 from datashare_python.objects import Document, TaskState
 from datashare_python.task_client import DatashareTaskClient
-from datashare_python.types_ import ContextManagerFactory, TemporalClient
-from datashare_python.utils import activity_defn
+from datashare_python.types_ import TemporalClient
+from datashare_python.utils import PYDANTIC_DATA_CONVERTER, activity_defn
+from datashare_python.worker import create_worker_id, worker_context
+
+logger = logging.getLogger(__name__)
 
 RABBITMQ_TEST_PORT = 5672
 RABBITMQ_TEST_HOST = "localhost"
@@ -83,11 +87,6 @@ class MockedWorkflow:
         return None
 
 
-@pytest.fixture(scope="session")
-def test_deps() -> list[ContextManagerFactory]:
-    return [set_es_client, set_task_client]
-
-
 @pytest.fixture(scope="session")
 def test_worker_config() -> WorkerConfig:
     logging_config = LoggingConfig(
@@ -113,25 +112,8 @@ def test_worker_config_path(test_worker_config: WorkerConfig, tmpdir: Path) -> P
 
 
 @pytest.fixture(scope="session")
-async def worker_lifetime_deps(
-    test_deps: list[ContextManagerFactory], test_worker_config: WorkerConfig
-) -> AsyncGenerator[None, None]:
-    worker_id = "test-worker-id"
-    ctx = "test application"
-    async with with_dependencies(
-        test_deps,
-        ctx=ctx,
-        worker_id=worker_id,
-        worker_config=test_worker_config,
-    ):
-        yield
-
-
-@pytest.fixture(scope="session")
-async def test_es_client_session(
-    worker_lifetime_deps,  # noqa: ANN001, ARG001
-) -> ESClient:
-    es = lifespan_es_client()
+async def test_es_client_session(test_worker_config: WorkerConfig) -> ESClient:
+    es = test_worker_config.to_es_client()
     await es.indices.delete(index="_all")
     await es.indices.create(index=TEST_PROJECT, body=_INDEX_BODY)
     return es
@@ -147,9 +129,9 @@ async def test_es_client(test_es_client_session: ESClient) -> ESClient:
 
 @pytest.fixture(scope="session")
 async def test_task_client_session(
-    worker_lifetime_deps,  # noqa: ANN001, ARG001
+    test_worker_config: WorkerConfig,
 ) -> AsyncGenerator[DatashareTaskClient, None]:
-    task_client = lifespan_task_client()
+    task_client = test_worker_config.to_task_client()
     async with task_client:
         user, key = await task_client.create_api_key()
         task_client.authenticate(user, key)
@@ -188,17 +170,24 @@ async def test_temporal_client(
 
 
 @pytest.fixture
-async def populate_es(
-    test_es_client: ESClient,
+async def indexed_docs(  # noqa: PLR0917
     doc_0: Document,
     doc_1: Document,
     doc_2: Document,
     doc_3: Document,
 ) -> list[Document]:
-    docs = [doc_0, doc_1, doc_2, doc_3]
-    async for _ in index_docs(test_es_client, docs=docs, index_name=TEST_PROJECT):
+    return [doc_0, doc_1, doc_2, doc_3]
+
+
+@pytest.fixture
+async def populate_es(
+    test_es_client: ESClient, indexed_docs: list[Document]
+) -> list[Document]:
+    async for _ in index_docs(
+        test_es_client, docs=indexed_docs, index_name=TEST_PROJECT
+    ):
         pass
-    return docs
+    return indexed_docs
 
 
 def index_docs_ops(
@@ -309,11 +298,155 @@ 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()
 
 
 def clear_dirs(config: WorkerConfig) -> None:
-    shutil.rmtree(str(config.artifacts_root))
-    config.artifacts_root.mkdir(parents=True, exist_ok=True)
-    shutil.rmtree(str(config.workdir))
-    config.workdir.mkdir(parents=True, exist_ok=True)
+    shutil.rmtree(str(config.paths.filesystem))
+    config.paths.filesystem.mkdir(parents=True, exist_ok=True)
+    shutil.rmtree(str(config.paths.artifacts))
+    config.paths.artifacts.mkdir(parents=True, exist_ok=True)
+    shutil.rmtree(str(config.paths.workdir))
+    config.paths.workdir.mkdir(parents=True, exist_ok=True)
+
+
+def dev_worker_context(
+    worker_id_prefix: str,
+    *,
+    is_async: bool,
+    client: TemporalClient,
+    activities: list[str] | None = None,
+    workflows: list[str] | None = None,
+    worker_config: WorkerConfig,
+    event_loop: AbstractEventLoop | None = None,
+    task_queue: str,
+    dependencies: str | None = None,
+    sandboxed: bool = True,
+) -> AbstractAsyncContextManager[None]:
+    worker_id = create_worker_id(worker_id_prefix)
+    if not is_async:
+        return _run_worker_in_separate_process(
+            worker_id=worker_id,
+            activities=activities,
+            workflows=workflows,
+            worker_config=worker_config.model_dump(),
+            task_queue=task_queue,
+            dependencies=dependencies,
+            sandboxed=sandboxed,
+        )
+    registered_wfs, registered_acts, registered_deps, worker_config_cls = discover(
+        workflows,
+        act_names=activities,
+        deps_name=dependencies,
+    )
+    cm = _erase_type(
+        worker_context(
+            worker_id,
+            activities=registered_acts,
+            workflows=registered_wfs,
+            dependencies=registered_deps,
+            worker_config=worker_config,
+            client=client,
+            event_loop=event_loop,
+            task_queue=task_queue,
+        )
+    )
+    return cm
+
+
+@asynccontextmanager
+async def _erase_type(cm: AbstractAsyncContextManager) -> AsyncGenerator[None, None]:
+    async with cm:
+        yield
+
+
+@asynccontextmanager
+async def _run_worker_in_separate_process(
+    worker_id: str,
+    *,
+    activities: list[str] | None = None,
+    workflows: list[str] | None = None,
+    worker_config: dict,
+    task_queue: str,
+    dependencies: str | None = None,
+    sandboxed: bool = True,
+) -> AsyncGenerator[None, None]:
+    ctx = multiprocessing.get_context("spawn")
+    shutdown_event = ctx.Event()
+    faulthandler.enable()
+    kwargs = {
+        "worker_id": worker_id,
+        "activities": activities,
+        "workflows": workflows,
+        "worker_config": worker_config,
+        "task_queue": task_queue,
+        "dependencies": dependencies,
+        "sandboxed": sandboxed,
+        "shutdown_event": shutdown_event,
+    }
+    p = ctx.Process(target=_run_worker_process, kwargs=kwargs)
+    p.start()
+    try:
+        yield
+    finally:
+        shutdown_event.set()
+        p.join()
+        if p.is_alive():
+            p.terminate()
+
+
+def _run_worker_process(
+    worker_id: str,
+    *,
+    activities: list[str] | None = None,
+    workflows: list[str] | None = None,
+    worker_config: dict,
+    event_loop: AbstractEventLoop | None = None,
+    task_queue: str,
+    dependencies: str | None = None,
+    sandboxed: bool = True,
+    shutdown_event: Event,
+) -> None:
+    async def _run() -> None:
+        registered_wfs, registered_acts, registered_deps, worker_config_cls = discover(
+            workflows,
+            act_names=activities,
+            deps_name=dependencies,
+        )
+        loaded_config = worker_config_cls.model_validate(worker_config)
+        temporal_config = loaded_config.temporal
+        runtime = Runtime(telemetry=TelemetryConfig())
+        # We don't use config.to_temporal_client since we get a problem with the
+        # temporal runtime not being properly recreated, we create it explicitely
+        if temporal_config.prometheus_host is not None:
+            telemetry_config = TelemetryConfig(
+                metrics=PrometheusConfig(bind_address="0.0.0.0:9000")
+            )
+            runtime = Runtime(telemetry=telemetry_config)
+        client = await TemporalClient.connect(
+            target_host=temporal_config.host,
+            namespace=temporal_config.namespace,
+            runtime=runtime,
+            data_converter=PYDANTIC_DATA_CONVERTER,
+        )
+        worker_ctx = worker_context(
+            worker_id,
+            activities=registered_acts,
+            workflows=registered_wfs,
+            worker_config=loaded_config,
+            client=client,
+            event_loop=event_loop,
+            task_queue=task_queue,
+            dependencies=registered_deps,
+            sandboxed=sandboxed,
+        )
+        async with worker_ctx as worker:
+            aws = [
+                asyncio.create_task(asyncio.to_thread(shutdown_event.wait)),
+                asyncio.create_task(worker.is_done()),
+            ]
+            await asyncio.wait(aws, return_when=asyncio.FIRST_COMPLETED)
+
+    asyncio.run(_run())

{datashare_python-0.9.8 → datashare_python-0.10.0}/datashare_python/constants.py

@@ -11,7 +11,6 @@ DEFAULT_DS_ADDRESS = "http://localhost:8080"
 DEFAULT_NAMESPACE = "datashare-default"
 
 METADATA_JSON = "metadata.json"
+MANIFEST_JSON = "manifest.json"
 
 TIKA_METADATA_RESOURCENAME = "tika_metadata_resourcename"
-
-DEFAULT_SHARED_RESOURCES_SIZE = 1

{datashare_python-0.9.8 → datashare_python-0.10.0}/datashare_python/dependencies.py

@@ -12,7 +12,6 @@ from icij_common.es import ESClient
 from .config import LogLevel, WorkerConfig
 from .exceptions import DependencyInjectionError
 from .logging_ import setup_worker_loggers
-from .objects import Shared
 from .task_client import DatashareTaskClient
 from .types_ import ContextManagerFactory, TemporalClient
 
@@ -24,7 +23,6 @@ ES_CLIENT: ContextVar[ESClient] = ContextVar("es_client")
 TASK_CLIENT: ContextVar[DatashareTaskClient] = ContextVar("task_client")
 TEMPORAL_CLIENT: ContextVar[TemporalClient] = ContextVar("temporal_client")
 WORKER_CONFIG: ContextVar[WorkerConfig] = ContextVar("worker_config")
-SHARED: ContextVar[Shared] = ContextVar("shared")
 
 
 def set_event_loop(event_loop: AbstractEventLoop) -> None:
@@ -101,20 +99,6 @@ def lifespan_temporal_client() -> TemporalClient:
         raise DependencyInjectionError("temporal client") from e
 
 
-# Setup shared resources
-async def set_shared_resources(shared: Shared) -> Shared:
-    SHARED.set(shared)
-    return shared
-
-
-# Return shared resources
-def lifespan_shared_resources() -> Shared:
-    try:
-        return SHARED.get()
-    except LookupError as e:
-        raise DependencyInjectionError("shared resources") from e
-
-
 @asynccontextmanager
 async def with_dependencies(
     dependencies: list[ContextManagerFactory], **kwargs

{datashare_python-0.9.8 → datashare_python-0.10.0}/datashare_python/interceptors.py

@@ -219,7 +219,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()

{datashare_python-0.9.8 → datashare_python-0.10.0}/datashare_python/logging_.py

@@ -132,7 +132,7 @@ def _encode_value(value: Any) -> str:
         return "true" if value else "false"
     if isinstance(value, numbers.Number):
         return str(value)
-    return json.dumps(value).decode()
+    return json.dumps(value)
 
 
 def _json_formatter(datefmt: str) -> BaseJsonFormatter: