datashare-python 0.3.0__py3-none-any.whl → 0.4.0__py3-none-any.whl

datashare_python/cli/worker.py

@@ -7,9 +7,9 @@ import typer
 import yaml
 from icij_common.pydantic_utils import safe_copy
 
-from datashare_python.config import TemporalClientConfig, WorkerConfig
+from datashare_python.config import TemporalClientConfig
 from datashare_python.discovery import discover, discover_activities, discover_workflows
-from datashare_python.worker import bootstrap_worker, create_worker_id
+from datashare_python.worker import create_worker_id, worker_context
 
 from .utils import AsyncTyper
 
@@ -24,6 +24,7 @@ _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)"
 _START_WORKER_DEPS_HELP = "worker lifetime dependencies name in the registry"
+_START_WORKER_CONFIG_HELP = "worker config class key the worker configs registry"
 _START_WORKER_WORKER_ID_PREFIX_HELP = "worker ID prefix"
 _START_WORKER_CONFIG_PATH_HELP = (
     "path to a worker config YAML file,"
@@ -85,6 +86,9 @@ async def start(
     dependencies: Annotated[
         str | None, typer.Option(help=_START_WORKER_DEPS_HELP)
     ] = None,
+    worker_config_name: Annotated[
+        str | None, typer.Option(help=_START_WORKER_CONFIG_HELP)
+    ] = None,
     config_path: Annotated[
         Path | None,
         typer.Option(
@@ -102,19 +106,25 @@ async def start(
         typer.Option("--temporal-namespace", "-ns", help=_TEMPORAL_NAMESPACE_HELP),
     ] = None,
 ) -> None:
+    registered_wfs, registered_acts, registered_deps, worker_config_cls = discover(
+        workflows,
+        act_names=activities,
+        deps_name=dependencies,
+        worker_config_name=worker_config_name,
+    )
     if config_path is not None:
         with config_path.open() as f:
-            bootstrap_config = WorkerConfig.model_validate(
+            worker_config = worker_config_cls.model_validate(
                 yaml.load(f, Loader=yaml.Loader)
             )
     else:
-        bootstrap_config = WorkerConfig()
+        worker_config = worker_config_cls()
     worker_id = create_worker_id(worker_id_prefix or "worker")
     logger.info(
         "starting worker %s on queue %s, with config: %s",
         worker_id,
         queue,
-        bootstrap_config.model_dump_json(indent=2),
+        worker_config.model_dump_json(indent=2),
     )
     temporal_override = dict()
     if temporal_address is not None:
@@ -124,24 +134,18 @@ async def start(
     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
-    )
-    client = await bootstrap_config.to_temporal_client()
+        worker_config = safe_copy(worker_config, update=update)
+    client = await worker_config.to_temporal_client()
     event_loop = asyncio.get_event_loop()
-    async with bootstrap_worker(
+    worker_ctx = worker_context(
         worker_id,
         activities=registered_acts,
         workflows=registered_wfs,
         dependencies=registered_deps,
-        bootstrap_config=bootstrap_config,
+        worker_config=worker_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
+    )
+    async with worker_ctx as worker:
+        await worker.is_done()

datashare_python/config.py

@@ -74,10 +74,6 @@ class TemporalClientConfig(BaseModel):
             )
         return self._client
 
-    # For the lru_cache
-    def __hash__(self) -> int:
-        return id(self)
-
 
 class WorkerConfig(ICIJSettings, LogWithWorkerIDMixin, BaseModel):
     model_config = DS_WORKER_SETTINGS_CONFIG

datashare_python/conftest.py

@@ -17,7 +17,6 @@ from datashare_python.config import (
 from datashare_python.dependencies import (
     lifespan_es_client,
     lifespan_task_client,
-    lifespan_temporal_client,
     set_es_client,
     set_event_loop,
     set_loggers,
@@ -170,10 +169,11 @@ async def test_task_client(
 
 
 @pytest.fixture(scope="session")
-def test_temporal_client_session(
-    worker_lifetime_deps,  # noqa: ANN001, ARG001
-) -> TemporalClient:
-    return lifespan_temporal_client()
+async def test_temporal_client_session(
+    test_worker_config: WorkerConfig,
+    event_loop: AbstractEventLoop,  # noqa: ARG001
+) -> TemporalClient:  # noqa: ANN001
+    return await test_worker_config.to_temporal_client()
 
 
 @pytest.fixture

datashare_python/dependencies.py

@@ -21,6 +21,7 @@ EVENT_LOOP: ContextVar[AbstractEventLoop] = ContextVar("event_loop")
 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")
 
 
 def set_event_loop(event_loop: AbstractEventLoop) -> None:
@@ -40,6 +41,17 @@ def set_loggers(worker_config: WorkerConfig, worker_id: str) -> None:
     logger.info("app config: %s", worker_config.model_dump_json(indent=2))
 
 
+def set_worker_config(worker_config: WorkerConfig) -> None:
+    WORKER_CONFIG.set(worker_config)
+
+
+def lifespan_worker_config() -> WorkerConfig:
+    try:
+        return WORKER_CONFIG.get()
+    except LookupError as e:
+        raise DependencyInjectionError("worker config") from e
+
+
 async def set_es_client(worker_config: WorkerConfig) -> ESClient:
     client = worker_config.to_es_client()
     ES_CLIENT.set(client)

datashare_python/discovery.py

@@ -3,6 +3,8 @@ import re
 from collections.abc import Callable, Iterable
 from importlib.metadata import entry_points
 
+from .config import WorkerConfig
+from .dependencies import set_worker_config
 from .types_ import ContextManagerFactory
 from .utils import ActivityWithProgress
 
@@ -11,9 +13,11 @@ logger = logging.getLogger(__name__)
 Activity = ActivityWithProgress | Callable | type
 
 _DEPENDENCIES = "dependencies"
-_WORKFLOW_GROUPS = "datashare.workflows"
-_ACTIVITIES_GROUPS = "datashare.activities"
-_DEPENDENCIES_GROUPS = "datashare.dependencies"
+_WORKER_CONFIGS = "worker_configs"
+_WORKFLOW_GROUP = "datashare.workflows"
+_ACTIVITIES_GROUP = "datashare.activities"
+_DEPENDENCIES_GROUP = "datashare.dependencies"
+_WORKER_CONFIGS_GROUP = "datashare.worker_configs"
 
 _RegisteredWorkflow = tuple[str, type]
 _RegisteredActivity = tuple[str, Activity]
@@ -22,11 +26,16 @@ _Discovery = tuple[
     Iterable[_RegisteredWorkflow] | None,
     Iterable[_RegisteredActivity] | None,
     _Dependencies | None,
+    type[WorkerConfig],
 ]
 
 
 def discover(
-    wf_names: list[str] | None, *, act_names: list[str] | None, deps_name: str | None
+    wf_names: list[str] | None,
+    *,
+    act_names: list[str] | None,
+    deps_name: str | None,
+    worker_config_name: str | None,
 ) -> _Discovery:
     discovered = ""
     wfs = None
@@ -58,11 +67,13 @@ def discover(
         if wf_names:
             msg += "workflow patterns " + ", ".join(wf_names) + " "
         if act_names:
-            msg = "activity patterns " + ", ".join(act_names)
+            msg += "activity patterns " + ", ".join(act_names)
         raise ValueError(msg)
-    deps = None
+    deps = []
     if deps_name is not None:
         deps = discover_dependencies(deps_name)
+    if set_worker_config not in deps:
+        deps.append(set_worker_config)
     if deps:
         n_deps = len(deps)
         discovered += "\n"
@@ -71,13 +82,18 @@ def discover(
             f"- {n_deps} dependenc{'ies' if n_deps > 1 else 'y'}:"
             f" {', '.join(deps_names)}"
         )
+    if worker_config_name is not None:
+        worker_config_cls = discover_worker_configs(worker_config_name)
+        discovered += f"- worker config class: {worker_config_cls}"
+    else:
+        worker_config_cls = WorkerConfig
     logger.info("discovered:\n%s", discovered)
-    return wfs, acts, deps
+    return wfs, acts, deps, worker_config_cls
 
 
 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)
+    impls = entry_points(group=_WORKFLOW_GROUP)
     registered = []
     for wf_impls in impls:
         wf_impls = wf_impls.load()  # noqa: PLW2901
@@ -93,7 +109,7 @@ def discover_workflows(names: list[str]) -> list[_RegisteredWorkflow]:
 
 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)
+    impls = entry_points(group=_ACTIVITIES_GROUP)
     registered = []
     for act_impls in impls:
         act_impls = act_impls.load()  # noqa: PLW2901
@@ -108,9 +124,9 @@ def discover_activities(names: list[str]) -> list[_RegisteredActivity]:
 
 
 def discover_dependencies(name: str) -> _Dependencies:
-    impls = entry_points(name=_DEPENDENCIES, group=_DEPENDENCIES_GROUPS)
+    impls = entry_points(name=_DEPENDENCIES, group=_DEPENDENCIES_GROUP)
     if not impls:
-        available_impls = entry_points(group=_DEPENDENCIES_GROUPS)
+        available_impls = entry_points(group=_DEPENDENCIES_GROUP)
         msg = (
             f'failed to find dependency: "{name}", '
             f"available dependencies: {available_impls}"
@@ -131,6 +147,30 @@ def discover_dependencies(name: str) -> _Dependencies:
         raise LookupError(msg) from e
 
 
+def discover_worker_configs(name: str) -> type[WorkerConfig]:
+    impls = entry_points(name=_WORKER_CONFIGS, group=_WORKER_CONFIGS_GROUP)
+    if not impls:
+        available_impls = entry_points(group=_WORKER_CONFIGS_GROUP)
+        msg = (
+            f'failed to find worker config: "{name}", '
+            f"available dependencies: {available_impls}"
+        )
+        raise LookupError(msg)
+    if len(impls) > 1:
+        msg = f'found multiple worker configs for name "{name}": {impls}'
+        raise ValueError(msg)
+    deps_registry = impls[_WORKER_CONFIGS].load()
+    try:
+        return deps_registry[name]
+    except KeyError as e:
+        available = list(deps_registry)
+        msg = (
+            f'failed to find worker config for name "{name}", available worker '
+            f"configs: {available}"
+        )
+        raise LookupError(msg) from e
+
+
 def _parse_wf_name(wf_type: type) -> str:
     if not isinstance(wf_type, type):
         msg = (

datashare_python/worker.py

@@ -43,6 +43,13 @@ documentation for more details:
 _ACTIVITY_THREAD_NAME_PREFIX = "datashare-activity-worker-"
 
 
+class DatashareWorker(Worker):
+    async def is_done(self) -> None:
+        if self._async_context_run_task is None:
+            raise ValueError("worker is not running")
+        await self._async_context_run_task
+
+
 def datashare_worker(
     client: TemporalClient,
     worker_id: str,
@@ -53,7 +60,7 @@ def datashare_worker(
     # Scale horizontally be default for activities, each worker processes one activity
     # at a time
     max_concurrent_io_activities: int = 10,
-) -> Worker:
+) -> DatashareWorker:
     if workflows is None:
         workflows = []
     if activities is None:
@@ -77,7 +84,7 @@ def datashare_worker(
         if workflows:
             logger.warning(_SEPARATE_IO_AND_CPU_WORKERS)
 
-    return Worker(
+    return DatashareWorker(
         client,
         identity=worker_id,
         workflows=workflows,
@@ -124,21 +131,21 @@ def init_activity(
 
 
 @asynccontextmanager
-async def bootstrap_worker(
+async def worker_context(
     worker_id: str,
     *,
     activities: list[Callable[..., Any] | None] | None = None,
     workflows: list[type] | None = None,
-    bootstrap_config: WorkerConfig,
+    worker_config: WorkerConfig,
     client: TemporalClient,
     event_loop: AbstractEventLoop,
     task_queue: str,
     dependencies: list[ContextManagerFactory] | None = None,
-) -> AsyncGenerator[Worker, None]:
+) -> AsyncGenerator[DatashareWorker, None]:
     deps_cm = (
         with_dependencies(
             dependencies,
-            worker_config=bootstrap_config,
+            worker_config=worker_config,
             worker_id=worker_id,
             event_loop=event_loop,
         )
@@ -159,9 +166,10 @@ async def bootstrap_worker(
             workflows=workflows,
             activities=acts,
             task_queue=task_queue,
-            max_concurrent_io_activities=bootstrap_config.max_concurrent_io_activities,
+            max_concurrent_io_activities=worker_config.max_concurrent_io_activities,
         )
-        yield worker
+        async with worker:
+            yield worker
 
 
 @asynccontextmanager

{datashare_python-0.3.0.dist-info → datashare_python-0.4.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datashare-python
-Version: 0.3.0
+Version: 0.4.0
 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/

{datashare_python-0.3.0.dist-info → datashare_python-0.4.0.dist-info}/RECORD

@@ -1,11 +1,11 @@
 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/config.py,sha256=p-uBTle30kbUdqrj8rXcYv2gHiNwtqcYMnGi4Kctumk,3683
+datashare_python/conftest.py,sha256=rQneF4ms-Gsx4NeRrtrU1bH-rjz3eJ6_LoLB9w-eRG8,8306
 datashare_python/constants.py,sha256=e6Px11OUee9GSHwTgsgFMszGCMwpW-OznHSMgINvepc,338
-datashare_python/dependencies.py,sha256=Diu7alKGaFWyC_ajp0fKU-xp8u5f_8x1axAHVBlppD0,3707
-datashare_python/discovery.py,sha256=khWTm11NlMVkRyxPV1lJimcTTicuVdxWHOC4H6PfFwE,6128
+datashare_python/dependencies.py,sha256=4UsVFKRjd2Q0ghg_fUU24P26tFYhg_SnAENj2mKErrY,4060
+datashare_python/discovery.py,sha256=EdWBIz5v0uDRG3zJO6h618qW8fB8tzg6NfK0pDZdV08,7546
 datashare_python/exceptions.py,sha256=bVHEAXxDPKfxeeMC0hJXEsrJkgsKO2ESAhxWU96GA4M,496
 datashare_python/local_client.py,sha256=GP9MTcHVQ1mcb2eO6TiQ7mzQdx199lZRhK8DRuJqJVQ,2359
 datashare_python/objects.py,sha256=MbTdBHbNNYeWbc7vrWlu2DcJT0uCSvtrDIiPpfgI_s4,4356
@@ -13,15 +13,15 @@ datashare_python/task_client.py,sha256=oTmP8bvZW0UyhLNMi1AV3XIAx7hrdbxNRss2Mw2az
 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=uarhX-BHryewAFwBZ98_PdA_QUKhApwlXfRYVIULbm0,142242
-datashare_python/worker.py,sha256=A4SnmDB4y0ck6Wp_UZWdsSOyTvW54Z2Bq76gxtp-_PE,6070
+datashare_python/worker-template.tar.gz,sha256=NaCH0r1Cxe82mZ4q41YYsv4p3qQL6_siJHdpAd9YNK0,274867
+datashare_python/worker.py,sha256=UdSCWZw8qhkhzok89oU7J35VTDZwxRAqef-5Z8yt95A,6333
 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=_LZFtv6O57Ez7XdRKmmwQ0b_5_8uFajPySv_jbzXbjI,5384
-datashare_python-0.3.0.dist-info/METADATA,sha256=z7AVC4VL38lRzJZlcLHvrj6DVj1H4lWx35HBlZH0btQ,907
-datashare_python-0.3.0.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
-datashare_python-0.3.0.dist-info/entry_points.txt,sha256=ILE7auxabHWiu3GC-AunWnzjhOI_SbZp7D4GqZHlLw4,68
-datashare_python-0.3.0.dist-info/RECORD,,
+datashare_python/cli/worker.py,sha256=ms69-vSgB7HOspH9D5h3-bLhg1mJNFvdIiBCtLtq3wc,5543
+datashare_python-0.4.0.dist-info/METADATA,sha256=mNCY9TciwKu-E9toJPb_B6VBrN8W0jREW_fKl034CHU,907
+datashare_python-0.4.0.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
+datashare_python-0.4.0.dist-info/entry_points.txt,sha256=ILE7auxabHWiu3GC-AunWnzjhOI_SbZp7D4GqZHlLw4,68
+datashare_python-0.4.0.dist-info/RECORD,,