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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datashare-python
-Version: 0.9.8
+Version: 0.9.9
 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/
@@ -11,7 +11,7 @@ Requires-Python: <4,>=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

{datashare_python-0.9.8 → datashare_python-0.9.9}/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.9.9}/datashare_python/constants.py

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

{datashare_python-0.9.8 → datashare_python-0.9.9}/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:

{datashare_python-0.9.8 → datashare_python-0.9.9}/datashare_python/objects.py

@@ -1,5 +1,6 @@
 import logging
 import os
+from abc import ABC
 from asyncio import Lock
 from collections.abc import Awaitable, Callable
 from dataclasses import InitVar, dataclass, field
@@ -7,12 +8,15 @@ from datetime import UTC, datetime
 from enum import StrEnum, unique
 from io import BytesIO
 from pathlib import Path
-from typing import Annotated, Any, ClassVar, Literal, Self, TypeVar, cast
+from typing import Annotated, Any, ClassVar, Generic, Literal, Self, TypeVar, cast
 
 import langcodes
+from icij_common.registrable import Registrable
 from lru import LRU
 from pydantic_core import PydanticCustomError, ValidationError, core_schema
-from pydantic_core.core_schema import PlainValidatorFunctionSchema
+from pydantic_core.core_schema import (
+    PlainValidatorFunctionSchema,
+)
 from pydantic_extra_types.language_code import LanguageName
 from temporalio import workflow
 
@@ -34,12 +38,17 @@ with workflow.unsafe.imports_passed_through():
 from icij_common.pydantic_utils import (
     icij_config,
     lowercamel_case_config,
+    make_enum_discriminator,
     merge_configs,
     no_enum_values_config,
+    tagged_union,
 )
 from pydantic import (
     AfterValidator,
+    AliasChoices,
     BeforeValidator,
+    ConfigDict,
+    Discriminator,
     Field,
     GetCoreSchemaHandler,
     TypeAdapter,
@@ -256,13 +265,123 @@ def _is_absolute_path(v: bytes | BytesIO | Path) -> Any:
     return v
 
 
-@dataclass(frozen=True)
-class DocArtifact:
+class ArtifactType(StrEnum):
+    STRUCTURE = "structure"
+    ASR_TRANSCRIPTION = "transcription"
+
+
+class ManifestEntryStatus(StrEnum):
+    COMPLETE = "complete"
+
+
+class TaskArgs(DatashareModel, ABC):
+    def as_manifest_task_input(self) -> dict[str, Any]:
+        # This is a base implementation, if the input is too large to be dumped,
+        # override this and pop large keys
+        as_manifest = self.model_dump(by_alias=True)
+        return as_manifest
+
+
+A = TypeVar("A", bound=TaskArgs)
+
+
+class ManifestEntry(DatashareModel, Generic[A], ABC):
+    status: ManifestEntryStatus
+    label: str | None = None
+    input: Annotated[
+        dict[str, Any] | None,
+        Field(
+            validation_alias=AliasChoices("taskInput", "input"),
+            serialization_alias="taskInput",
+        ),
+    ]
+
+    @classmethod
+    def complete(cls, args: A, label: str | None = None, **kwargs) -> Self:
+        return cls(
+            input=args.as_manifest_task_input(),
+            label=label,
+            status=ManifestEntryStatus.COMPLETE,
+            **kwargs,
+        )
+
+
+class PaginationType(StrEnum):
+    FILESYSTEM = "filesystem"
+    BYTE_RANGES = "byteRanges"
+
+
+class BasePagination(DatashareModel, Registrable, ABC):
+    registry_key: ClassVar[str] = Field(frozen=True, default="type")
+
+    total: int
+    type: ClassVar[PaginationType] = Field(frozen=True)
+
+
+def _validate_pages_range(v: Any) -> None:
+    if not isinstance(v, list):
+        msg = f"expected a list, got {type(v)}"
+        raise TypeError(msg)
+    previous_end = None
+    for page_i, (start, end) in enumerate(v):
+        if not start <= end:
+            msg = "end of page must be >= start"
+            raise ValueError(msg)
+        if previous_end is not None and previous_end != start:
+            msg = (
+                f"start of page {page_i} doesn't match end of previous "
+                f"page {previous_end}"
+            )
+            raise ValueError(msg)
+    return v
+
+
+PagesRange = Annotated[list[tuple[int, int]], AfterValidator(_validate_pages_range)]
+
+
+@BasePagination.register(PaginationType.FILESYSTEM)
+class FilesystemPagination(BasePagination):
+    type: ClassVar[PaginationType] = Field(
+        default=PaginationType.FILESYSTEM, frozen=True
+    )
+
+
+@BasePagination.register(PaginationType.BYTE_RANGES)
+class ByteRangesPagination(BasePagination):
+    type: ClassVar[PaginationType] = Field(
+        default=PaginationType.BYTE_RANGES, frozen=True
+    )
+    byte_ranges: PagesRange
+
+    @model_validator(mode="after")
+    def byte_ranges_length_should_match_total(self) -> Self:
+        if len(self.byte_ranges) != self.total:
+            msg = (
+                f"byte_ranges must match total. Found {len(self.byte_ranges)} for"
+                f" byte_ranges and  {self.total} for total."
+            )
+            raise ValueError(msg)
+        return self
+
+
+pagination_discriminator = make_enum_discriminator("type", PaginationType)
+Pagination = Annotated[
+    tagged_union(BasePagination.__subclasses__(), lambda x: x.type.default),
+    Discriminator(pagination_discriminator),
+]
+
+
+class DocArtifact(BaseModel, ABC):
+    # This object is not used for serde, just as a container, it's OK to allow
+    # arbitrary types (to allow storing BytesIO)
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
     project: str
     doc_id: str
     artifact: Annotated[bytes | BytesIO | Path, AfterValidator(_is_absolute_path)]
-    filename: str
-    metadata_key: str
+    filename: ClassVar[str]  # Override this
+    type: ClassVar[ArtifactType]  # Override this
+    manifest_entry: ManifestEntry
 
 
 @unique

{datashare_python-0.9.8 → datashare_python-0.9.9}/datashare_python/utils.py

@@ -1,10 +1,13 @@
 import asyncio
 import contextlib
+import fcntl
 import inspect
 import json
 import os
 import shutil
-from collections.abc import Callable, Coroutine, Iterable, Sequence
+import threading
+import time
+from collections.abc import Callable, Coroutine, Generator, Iterable, Sequence
 from copy import deepcopy
 from dataclasses import dataclass
 from datetime import timedelta
@@ -36,7 +39,7 @@ from datashare_python.types_ import (
     SyncProgressRateHandler,
 )
 
-from .constants import METADATA_JSON
+from .constants import MANIFEST_JSON, METADATA_JSON
 from .objects import DocArtifact, DocumentLocation, FilesystemDocument
 from .types_ import RawAsyncProgressHandler
 
@@ -57,6 +60,11 @@ _NEVER_RETRIABLES = {
 }
 
 
+_ARTIFACT_LOCK = threading.Lock()
+# For test
+_LOCKED = threading.Event()
+
+
 @dataclass(frozen=True)
 class Progress:
     max_progress: float
@@ -338,34 +346,116 @@ def _metadata_path(doc_id: str, *, project: str) -> Path:
     return metadata_path
 
 
-def _read_artifact_metadata(root: Path, artifact: DocArtifact) -> dict:
-    m_path = root / _metadata_path(artifact.doc_id, project=artifact.project)
+def _manifest_path(doc_id: str, *, project: str) -> Path:
+    manifest_path = artifacts_dir(doc_id, project=project) / MANIFEST_JSON
+    return manifest_path
+
+
+def _read_artifact_manifest(root: Path, artifact: DocArtifact) -> dict:
+    m_path = root / _manifest_path(artifact.doc_id, project=artifact.project)
+    if not m_path.exists():
+        m_path = root / _metadata_path(artifact.doc_id, project=artifact.project)
+        if not m_path.exists():
+            msg = f"couldn't find manifest nor metadata for {artifact.doc_id}"
+            raise FileNotFoundError(msg)
     return json.loads(m_path.read_text())
 
 
-def write_artifact(root: Path, artifact: DocArtifact) -> Path:
+def write_artifact(
+    root: Path, artifact: DocArtifact, lock_timeout_ms: int = 30_000
+) -> Path:
+    # TODO: WARNING many writers could write at the time, to avoid inconsistent
+    #  states we should handle this somehow
     artif_dir = root / artifacts_dir(artifact.doc_id, project=artifact.project)
     artif_dir.mkdir(exist_ok=True, parents=True)
-    # TODO: if transcriptions are too large we could also serialize them
-    #  as jsonl
-    artifact_path: Path = artif_dir / artifact.filename
-    match artifact.artifact:
+    artifact_path = artif_dir / artifact.filename
+    # We're using POSIX locks which are not exclusive to filedescriptors
+    # but to processes, we hence have to ensure several threads from the same
+    # Python process acquire the lock at the same time. We're using a python mutex lock
+    with artifact_lock(artif_dir, lock_timeout_ms):
+        # Read the metadata first (things could go wrong here in case someone is reading
+        # at the same time). We read in a backward compatible wat and write to that same
+        # location. We don't take responsibility for migrating the data, the DS back
+        # will do it
+        manifest_path, manifest = _read_manifest_backward_compatible(root, artifact)
+        is_legacy = manifest_path.name == "metadata.json"
+        # Pop the status key from the manifest before writing
+        manifest_entry = manifest.get(artifact.type)
+        if manifest_entry is not None and not is_legacy:
+            manifest[artifact.type].pop("status", None)
+            manifest_path.write_text(json.dumps(manifest))
+        # Write the artifact
+        _write_artifact_bytes(artifact_path, artifact.artifact)
+        # Update the manifest entry with details and new states
+        if is_legacy:
+            manifest_entry = str(artifact_path.relative_to(artif_dir))
+        else:
+            manifest_entry = artifact.manifest_entry.model_dump(
+                mode="json", by_alias=True
+            )
+        manifest[artifact.type] = manifest_entry
+        manifest_path.write_text(json.dumps(manifest))
+        return artifact_path.relative_to(artif_dir)
+
+
+@contextlib.contextmanager
+def _set_event() -> Generator[None, None, None]:
+    try:
+        _LOCKED.set()
+        yield
+    finally:
+        _LOCKED.clear()
+
+
+@contextlib.contextmanager
+def artifact_lock(
+    artifacts_dir: Path, lock_timeout_ms: int
+) -> Generator[None, None, None]:
+    lock_path = artifacts_dir / f"{MANIFEST_JSON}.lock"
+    fd = os.open(lock_path, os.O_CREAT | os.O_WRONLY)
+    with _ARTIFACT_LOCK, _set_event():
+        start = time.time()
+        timeout_s = lock_timeout_ms / 1000
+        while time.time() - start < timeout_s:
+            try:
+                fcntl.lockf(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+            except OSError:
+                # TODO: we shouldn't block here, use asyncio instead
+                time.sleep(0.02)
+                continue
+            yield
+            return
+        msg = (
+            f"failed to acquire lock on {lock_path} in less than {lock_timeout_ms} ms."
+        )
+        raise TimeoutError(msg)
+
+
+def _read_manifest_backward_compatible(
+    root: Path, artifact: DocArtifact
+) -> tuple[Path, dict[str, Any]]:
+    manifest_path = root / _manifest_path(artifact.doc_id, project=artifact.project)
+    if manifest_path.exists():
+        return manifest_path, _read_artifact_manifest(root, artifact)
+    meta_path = root / _metadata_path(artifact.doc_id, project=artifact.project)
+    if meta_path.exists():
+        return meta_path, _read_artifact_manifest(root, artifact)
+    return manifest_path, dict()
+
+
+def _write_artifact_bytes(path: Path, artifact: bytes | BytesIO | Path) -> None:
+    match artifact:
         case bytes():
-            artifact_path.write_bytes(artifact.artifact)
+            path.write_bytes(artifact)
         case BytesIO():
-            with artifact_path.open("wb") as f:
-                f.write(artifact.artifact.read())
+            with path.open("wb") as f:
+                f.write(artifact.read())
         case Path():
-            artifact_path.unlink(missing_ok=True)
-            shutil.move(artifact.artifact, artifact_path)
+            path.unlink(missing_ok=True)
+            shutil.move(artifact, path)
         case _:
-            msg = f"unsupported artifact type: {artifact.artifact.__class__.__name__}"
+            msg = f"unsupported artifact type: {artifact.__class__.__name__}"
             raise ValueError(msg)
-    meta_path = root / _metadata_path(artifact.doc_id, project=artifact.project)
-    meta = _read_artifact_metadata(root, artifact) if meta_path.exists() else dict()
-    meta[artifact.metadata_key] = artifact.filename
-    meta_path.write_text(json.dumps(meta))
-    return artifact_path.relative_to(artif_dir)
 
 
 def debuggable_name(

{datashare_python-0.9.8 → datashare_python-0.9.9}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "datashare-python"
-version = "0.9.8"
+version = "0.9.9"
 description = "Manage Python tasks and local resources in Datashare"
 authors = [
   { name = "Clément Doumouro", email = "cdoumouro@icij.org" },
@@ -13,7 +13,7 @@ dependencies = [
   "alive-progress~=3.2",
   "aiohttp~=3.11",
   "hatchling~=1.27",
-  "icij-common[elasticsearch]~=0.8.2",
+  "icij-common[elasticsearch]~=0.8.3",
   "langcodes~=3.5",
   "python-json-logger~=4.0",
   "pyyaml~=6.0",