datashare-python 0.1.0__tar.gz

datashare_python-0.1.0/PKG-INFO

@@ -0,0 +1,80 @@
+Metadata-Version: 2.1
+Name: datashare-python
+Version: 0.1.0
+Summary: Implement Datashare task in Python
+Author-Email: =?utf-8?q?Cl=C3=A9ment_Doumouro?= <cdoumouro@icij.org>, =?utf-8?q?Cl=C3=A9ment_Doumouro?= <clement.doumouro@gmail.com>
+Requires-Python: ~=3.11
+Requires-Dist: aiostream~=0.6.4
+Requires-Dist: aiohttp~=3.11.9
+Requires-Dist: icij-common[elasticsearch]~=0.5.5
+Requires-Dist: icij-worker[amqp]~=0.12
+Requires-Dist: torch==2.6.0.dev20241101; sys_platform != "darwin"
+Requires-Dist: torch!=2.6.0.dev20241101+cpu,<=2.6.0.dev20241101; sys_platform == "darwin"
+Requires-Dist: transformers~=4.46.3
+Requires-Dist: pycountry>=24.6.1
+Requires-Dist: sentencepiece>=0.2.0
+Requires-Dist: typer>=0.13.1
+Requires-Dist: alive-progress>=3.2.0
+Description-Content-Type: text/markdown
+
+<div style="background-image: linear-gradient(45deg, #193d87, #fa4070);">
+  <br/>
+  <p align="center">
+    <a href="https://datashare.icij.org/">
+      <img align="center" src="docs/assets/datashare-logo.svg" alt="Datashare" style="max-width: 60%">
+    </a>
+  </p>
+  <p align="center">
+    <em>Better analyze information, in all its forms</em>  
+  </p>
+  <br/>
+</div>
+<br/>
+
+---
+
+**Documentation**: <a href="https://icij.github.io/datashare-python" target="_blank">https://icij.github.io/datashare-python</a>
+
+---
+
+# Implement **your own Datashare tasks**, written in Python
+
+Most AI, Machine Learning, Data Engineering happens in Python.
+[Datashare](https://icij.gitbook.io/datashare) now lets you extend its backend with your own tasks implemented in Python.
+
+Turning your own ML pipelines into Datashare tasks is **very simple**, learn about it inside [documentation](https://icij.github.io/datashare-python).
+
+Turning your own ML pipelines into Datashare tasks is **very simple**.
+
+Actually, it's *almost* as simple as cloning our [template repo](https://github.com/ICIJ/datashare-python):
+
+```
+$ git clone git@github.com:ICIJ/datashare-python.git
+```
+
+replacing existing [app](https://github.com/ICIJ/datashare-python/blob/main/datashare_python/app.py) tasks with your own:   
+```python
+from icij_worker import AsyncApp
+
+app = AsyncApp("app")
+
+
+@app.task
+def hello_world() -> str:
+    return "Hello world"
+```
+
+installing [`uv`](https://docs.astral.sh/uv/) to set up dependencies and running your async Datashare worker:
+```console
+$ cd datashare-python
+$ curl -LsSf https://astral.sh/uv/install.sh | sh
+$ uv run ./scripts/worker_entrypoint.sh
+[INFO][icij_worker.backend.backend]: Loading worker configuration from env...
+...
+}
+[INFO][icij_worker.backend.mp]: starting 1 worker for app datashare_python.app.app
+...
+```
+you'll then be able to execute task by starting using our [HTTP client]() (and soon using Datashare's UI).
+
+## Learn more reading our [documentation](https://icij.github.io/datashare-python) !

datashare_python-0.1.0/README.md

@@ -0,0 +1,61 @@
+<div style="background-image: linear-gradient(45deg, #193d87, #fa4070);">
+  <br/>
+  <p align="center">
+    <a href="https://datashare.icij.org/">
+      <img align="center" src="docs/assets/datashare-logo.svg" alt="Datashare" style="max-width: 60%">
+    </a>
+  </p>
+  <p align="center">
+    <em>Better analyze information, in all its forms</em>  
+  </p>
+  <br/>
+</div>
+<br/>
+
+---
+
+**Documentation**: <a href="https://icij.github.io/datashare-python" target="_blank">https://icij.github.io/datashare-python</a>
+
+---
+
+# Implement **your own Datashare tasks**, written in Python
+
+Most AI, Machine Learning, Data Engineering happens in Python.
+[Datashare](https://icij.gitbook.io/datashare) now lets you extend its backend with your own tasks implemented in Python.
+
+Turning your own ML pipelines into Datashare tasks is **very simple**, learn about it inside [documentation](https://icij.github.io/datashare-python).
+
+Turning your own ML pipelines into Datashare tasks is **very simple**.
+
+Actually, it's *almost* as simple as cloning our [template repo](https://github.com/ICIJ/datashare-python):
+
+```
+$ git clone git@github.com:ICIJ/datashare-python.git
+```
+
+replacing existing [app](https://github.com/ICIJ/datashare-python/blob/main/datashare_python/app.py) tasks with your own:   
+```python
+from icij_worker import AsyncApp
+
+app = AsyncApp("app")
+
+
+@app.task
+def hello_world() -> str:
+    return "Hello world"
+```
+
+installing [`uv`](https://docs.astral.sh/uv/) to set up dependencies and running your async Datashare worker:
+```console
+$ cd datashare-python
+$ curl -LsSf https://astral.sh/uv/install.sh | sh
+$ uv run ./scripts/worker_entrypoint.sh
+[INFO][icij_worker.backend.backend]: Loading worker configuration from env...
+...
+}
+[INFO][icij_worker.backend.mp]: starting 1 worker for app datashare_python.app.app
+...
+```
+you'll then be able to execute task by starting using our [HTTP client]() (and soon using Datashare's UI).
+
+## Learn more reading our [documentation](https://icij.github.io/datashare-python) !

datashare_python-0.1.0/datashare_python/__main__.py

@@ -0,0 +1,4 @@
+from datashare_python.cli import cli_app
+
+if __name__ == "__main__":
+    cli_app()

datashare_python-0.1.0/datashare_python/app.py

@@ -0,0 +1,85 @@
+from typing import Optional
+
+from icij_worker import AsyncApp
+from icij_worker.typing_ import PercentProgress
+from pydantic import parse_obj_as
+
+from datashare_python.constants import PYTHON_TASK_GROUP
+from datashare_python.objects import ClassificationConfig, TranslationConfig
+from datashare_python.tasks import (
+    classify_docs as classify_docs_,
+    create_classification_tasks as create_classification_tasks_,
+    create_translation_tasks as create_translation_tasks_,
+    translate_docs as translate_docs_,
+)
+from datashare_python.tasks.dependencies import APP_LIFESPAN_DEPS
+
+app = AsyncApp("ml", dependencies=APP_LIFESPAN_DEPS)
+
+
+@app.task(group=PYTHON_TASK_GROUP)
+async def create_translation_tasks(
+    project: str,
+    target_language: str,
+    config: dict | None = None,
+    user: dict | None = None,  # pylint: disable=unused-argument
+) -> list[str]:
+    # Parse the incoming config
+    config = parse_obj_as(Optional[TranslationConfig], config)
+    return await create_translation_tasks_(
+        project=project, target_language=target_language, config=config
+    )
+
+
+@app.task(group=PYTHON_TASK_GROUP)
+async def translate_docs(
+    docs: list[str],
+    project: str,
+    target_language: str,
+    progress: PercentProgress,
+    config: dict | None = None,
+    user: dict | None = None,  # pylint: disable=unused-argument
+) -> int:
+    config = parse_obj_as(Optional[TranslationConfig], config)
+    return await translate_docs_(
+        docs, target_language, project=project, config=config, progress=progress
+    )
+
+
+@app.task(group=PYTHON_TASK_GROUP)
+async def create_classification_tasks(
+    project: str,
+    language: str,
+    n_workers: int,
+    progress: PercentProgress,
+    config: dict | None = None,
+    user: dict | None = None,  # pylint: disable=unused-argument
+) -> list[str]:
+    config = parse_obj_as(Optional[ClassificationConfig], config)
+    return await create_classification_tasks_(
+        project=project,
+        language=language,
+        n_workers=n_workers,
+        config=config,
+        progress=progress,
+    )
+
+
+@app.task(group=PYTHON_TASK_GROUP)
+async def classify_docs(
+    docs: list[str],
+    language: str,
+    project: str,
+    progress: PercentProgress,
+    config: dict | None = None,
+    user: dict | None = None,  # pylint: disable=unused-argument
+) -> int:
+    config = parse_obj_as(Optional[ClassificationConfig], config)
+    return await classify_docs_(
+        docs, language=language, project=project, config=config, progress=progress
+    )
+
+
+@app.task(group=PYTHON_TASK_GROUP)
+def ping() -> str:
+    return "pong"

datashare_python-0.1.0/datashare_python/cli/__init__.py

@@ -0,0 +1,30 @@
+import importlib.metadata
+from typing import Annotated, Optional
+
+import typer
+
+import datashare_python
+from datashare_python.cli.tasks import task_app
+from datashare_python.cli.utils import AsyncTyper
+
+cli_app = AsyncTyper(context_settings={"help_option_names": ["-h", "--help"]})
+cli_app.add_typer(task_app)
+
+
+def version_callback(value: bool):
+    if value:
+        package_version = importlib.metadata.version(datashare_python.__name__)
+        print(package_version)
+        raise typer.Exit()
+
+
+@cli_app.callback(name="datashare-python")
+def main(
+    version: Annotated[  # pylint: disable=unused-argument
+        Optional[bool],
+        typer.Option(  # pylint: disable=unused-argument
+            "--version", callback=version_callback, is_eager=True
+        ),
+    ] = None
+):
+    """Datashare Python CLI"""

datashare_python-0.1.0/datashare_python/cli/tasks.py

@@ -0,0 +1,182 @@
+import asyncio
+import json
+import logging
+import sys
+from pathlib import Path
+from traceback import FrameSummary, StackSummary
+from typing import Annotated, Any, Optional
+
+import typer
+from alive_progress import alive_bar
+from icij_worker import TaskState
+from icij_worker.objects import READY_STATES, Task, TaskError
+
+from datashare_python.cli.utils import AsyncTyper, eprint
+from datashare_python.constants import PYTHON_TASK_GROUP
+from datashare_python.task_client import DatashareTaskClient
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_DS_ADDRESS = "http://localhost:8080"
+
+_ARGS_HELP = "task argument as a JSON string or file path"
+_GROUP_HELP = "task group"
+_DS_API_KEY_HELP = "datashare API key"
+_DS_URL_HELP = "datashare address"
+_POLLING_INTERVAL_S_HELP = "task state polling interval in seconds"
+_NAME_HELP = "registered task name"
+_RESULT_HELP = "get a task result"
+_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
+
+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,
+    group: Annotated[
+        Optional[str],
+        typer.Option("--group", "-g", help=_GROUP_HELP),
+    ] = PYTHON_TASK_GROUP.name,
+    ds_address: Annotated[
+        str, typer.Option("--ds-address", "-a", help=_DS_URL_HELP)
+    ] = DEFAULT_DS_ADDRESS,
+    ds_api_key: Annotated[
+        Optional[str], typer.Option("--ds-api-key", "-k", help=_DS_API_KEY_HELP)
+    ] = None,
+):
+    match args:
+        case str():
+            as_path = Path(name)
+            if as_path.exists():
+                args = json.loads(as_path.read_text())
+            else:
+                args = json.loads(args)
+        case None:
+            args = dict()
+        case _:
+            raise TypeError(f"Invalid args {args}")
+    client = DatashareTaskClient(ds_address, api_key=ds_api_key)
+    async with client:
+        task_id = await client.create_task(name, args, group=group)
+    eprint(f"Task({task_id}) started !")
+    eprint(f"Task({task_id}) 🛫")
+    print(task_id)
+
+
+@task_app.async_command(help=_WATCH_HELP)
+async def watch(
+    task_id: Annotated[str, typer.Argument(help=_TASK_ID_HELP)],
+    ds_address: Annotated[
+        str, typer.Option("--ds-address", "-a", help=_DS_URL_HELP)
+    ] = DEFAULT_DS_ADDRESS,
+    ds_api_key: Annotated[
+        Optional[str], typer.Option("--ds-api-key", "-k", help=_DS_API_KEY_HELP)
+    ] = None,
+    polling_interval_s: Annotated[
+        float, typer.Option("--polling-interval-s", "-p", help=_POLLING_INTERVAL_S_HELP)
+    ] = 1.0,
+):
+    client = DatashareTaskClient(ds_address, api_key=ds_api_key)
+    async with client:
+        task = await client.get_task(task_id)
+        if task.state is READY_STATES:
+            await _handle_ready(task, client, already_done=True)
+        await _handle_alive(task, client, polling_interval_s)
+    print(task_id)
+
+
+@task_app.async_command(help=_RESULT_HELP)
+async def result(
+    task_id: Annotated[str, typer.Argument(help=_TASK_ID_HELP)],
+    ds_address: Annotated[
+        str, typer.Option("--ds-address", "-a", help=_DS_URL_HELP)
+    ] = DEFAULT_DS_ADDRESS,
+    ds_api_key: Annotated[
+        Optional[str], typer.Option("--ds-api-key", "-k", help=_DS_API_KEY_HELP)
+    ] = None,
+) -> Any:
+    client = DatashareTaskClient(ds_address, api_key=ds_api_key)
+    async with client:
+        res = await client.get_task_result(task_id)
+        if isinstance(res, (dict, list)):
+            res = json.dumps(res, indent=2)
+        print(res)
+
+
+async def _handle_ready(
+    task: Task, client: DatashareTaskClient, already_done: bool = False
+) -> None:
+    match task.state:
+        case TaskState.ERROR:
+            await _handle_error(task, client)
+        case TaskState.CANCELLED:
+            await _handle_cancelled(task)
+        case TaskState.DONE:
+            if already_done:
+                await _handle_already_done(task)
+            else:
+                await _handle_done(task)
+        case _:
+            raise ValueError(f"Unexpected task state {task.state}")
+
+
+async def _handle_error(task, client: DatashareTaskClient):
+    error = await client.get_task_error(task.id)
+    eprint(
+        f"Task({task.id}) failed with the following"
+        f" error:\n\n{_format_error(error)}"
+    )
+    eprint(f"Task({task.id}) ❌")
+    raise typer.Exit(code=1)
+
+
+async def _handle_cancelled(task):
+    eprint(f"Task({task.id}) was cancelled !")
+    eprint(f"Task({task.id}) 🛑")
+    raise typer.Exit(code=1)
+
+
+async def _handle_already_done(task):
+    eprint(f"Task({task.id}) ✅ is already completed !")
+
+
+async def _handle_done(task):
+    eprint(f"Task({task.id}) 🛬")
+    eprint(f"Task({task.id}) ✅")
+
+
+async def _handle_alive(
+    task: Task, client: DatashareTaskClient, polling_interval_s: float
+) -> None:
+    title = f"Task({task.id}) 🛫"
+    stats = "(ETA: {eta})"
+    monitor = "{percent}"
+    progress_bar = alive_bar(
+        title=title, manual=True, stats=stats, monitor=monitor, file=sys.stderr
+    )
+    with progress_bar as bar:
+        task_state = task.state
+        while task_state not in READY_STATES:
+            task = await client.get_task(task.id)
+            task_state = task.state
+            progress = task.progress or 0.0
+            bar(progress)  # pylint: disable=not-callable
+            await asyncio.sleep(polling_interval_s)
+    if task_state in READY_STATES:
+        await _handle_ready(task, client)
+
+
+def _format_error(error: TaskError) -> str:
+    stack = StackSummary.from_list(
+        [FrameSummary(f.name, f.lineno, f.name) for f in error.stacktrace]
+    )
+    msg = f"{error.name}:\n{stack}\n{error.message}"
+    if error.cause:
+        msg += "\n cause by {error.cause}"
+    return msg

datashare_python-0.1.0/datashare_python/cli/utils.py

@@ -0,0 +1,33 @@
+import asyncio
+import concurrent.futures
+import sys
+from functools import wraps
+
+import typer
+
+
+class AsyncTyper(typer.Typer):
+    def async_command(self, *args, **kwargs):
+        def decorator(async_func):
+            @wraps(async_func)
+            def sync_func(*_args, **_kwargs):
+                res = asyncio.run(async_func(*_args, **_kwargs))
+                return res
+
+            self.command(*args, **kwargs)(sync_func)
+            return async_func
+
+        return decorator
+
+
+def eprint(*args, **kwargs):
+    print(*args, file=sys.stderr, **kwargs)
+
+
+def _to_concurrent(
+    fut: asyncio.Future, loop: asyncio.AbstractEventLoop
+) -> concurrent.futures.Future:
+    async def wait():
+        await fut
+
+    return asyncio.run_coroutine_threadsafe(wait(), loop)

datashare_python-0.1.0/datashare_python/config.py

@@ -0,0 +1,60 @@
+from typing import ClassVar
+
+from icij_common.pydantic_utils import ICIJSettings, NoEnumModel
+from icij_worker.utils.logging_ import LogWithWorkerIDMixin
+from pydantic import Field
+
+import datashare_python
+
+_ALL_LOGGERS = [datashare_python.__name__]
+
+
+class AppConfig(ICIJSettings, LogWithWorkerIDMixin, NoEnumModel):
+    class Config:
+        env_prefix = "DS_DOCKER_ML_"
+
+    loggers: ClassVar[list[str]] = Field(_ALL_LOGGERS, const=True)
+
+    log_level: str = Field(default="INFO")
+
+    batch_size: int = 1024
+    pipeline_batch_size: int = 1024
+    ne_buffer_size: int = 1000
+
+    # DS
+    ds_api_key: str | None = None
+    ds_url: str = "http://datashare:8080"
+    # ES
+    es_address: str = "http://localhost:9200"
+    es_default_page_size: int = 1000
+    es_keep_alive: str = "10m"
+    es_max_concurrency: int = 5
+    es_max_retries: int = 0
+    es_max_retry_wait_s: int | float = 60
+    es_timeout_s: int | float = 60 * 5
+
+    def to_es_client(self, address: str | None = None) -> "ESClient":
+        from icij_common.es import ESClient
+
+        if address is None:
+            address = self.es_address
+
+        client = ESClient(
+            hosts=[address],
+            pagination=self.es_default_page_size,
+            max_concurrency=self.es_max_concurrency,
+            keep_alive=self.es_keep_alive,
+            timeout=self.es_timeout_s,
+            max_retries=self.es_max_retries,
+            max_retry_wait_s=self.es_max_retry_wait_s,
+            api_key=self.ds_api_key,
+        )
+        client.transport._verified_elasticsearch = (  # pylint: disable=protected-access
+            True
+        )
+        return client
+
+    def to_task_client(self) -> "DatashareTaskClient":
+        from datashare_python.task_client import DatashareTaskClient
+
+        return DatashareTaskClient(self.ds_url)

datashare_python-0.1.0/datashare_python/constants.py

@@ -0,0 +1,6 @@
+from pathlib import Path
+
+from icij_worker.app import TaskGroup
+
+DATA_DIR = Path(__file__).parent.joinpath(".data")
+PYTHON_TASK_GROUP = TaskGroup(name="PYTHON")

datashare_python-0.1.0/datashare_python/objects.py

@@ -0,0 +1,49 @@
+from typing import Self
+
+import pycountry
+from icij_common.es import DOC_CONTENT, DOC_LANGUAGE, DOC_ROOT_ID, ID_, SOURCE
+from icij_common.pydantic_utils import ICIJModel, LowerCamelCaseModel
+from pydantic import Field
+
+
+class Document(LowerCamelCaseModel):
+    id: str
+    root_document: str
+    content: str
+    language: str
+    tags: list[str] = Field(default_factory=list)
+    content_translated: dict[str, str] = Field(
+        default_factory=dict, alias="content_translated"
+    )
+
+    @classmethod
+    def from_es(cls, es_doc: dict) -> Self:
+        sources = es_doc[SOURCE]
+        return cls(
+            id=es_doc[ID_],
+            content=sources[DOC_CONTENT],
+            content_translated=sources.get("content_translated", dict()),
+            language=sources[DOC_LANGUAGE],
+            root_document=sources[DOC_ROOT_ID],
+            tags=sources.get("tags", []),
+        )
+
+
+class ClassificationConfig(ICIJModel):
+    task: str = Field(const=True, default="text-classification")
+    model: str = "distilbert/distilbert-base-uncased-finetuned-sst-2-english"
+    batch_size: int = 16
+
+
+class TranslationConfig(ICIJModel):
+    task: str = Field(const=True, default="translation")
+    model: str = "Helsinki-NLP/opus-mt"
+    batch_size: int = 16
+
+    def to_pipeline_args(self, source_language: str, *, target_language: str) -> dict:
+        as_dict = self.dict()
+        source_alpha2 = pycountry.languages.get(name=source_language).alpha_2
+        target_alpha2 = pycountry.languages.get(name=target_language).alpha_2
+        as_dict["task"] = f"translation_{source_alpha2}_to_{target_alpha2}"
+        as_dict["model"] = f"{self.model}-{source_alpha2}-{target_alpha2}"
+        return as_dict