arkindex-base-worker 0.3.7rc5__py3-none-any.whl → 0.5.0a1__py3-none-any.whl

arkindex_worker/worker/metadata.py

@@ -7,6 +7,7 @@ from enum import Enum
 from arkindex_worker import logger
 from arkindex_worker.cache import CachedElement, unsupported_cache
 from arkindex_worker.models import Element
+from arkindex_worker.utils import DEFAULT_BATCH_SIZE, batch_publication, make_batches
 
 
 class MetaType(Enum):
@@ -93,7 +94,7 @@ class MetaDataMixin:
             logger.warning("Cannot create metadata as this worker is in read-only mode")
             return
 
-        metadata = self.request(
+        metadata = self.api_client.request(
             "CreateMetaData",
             id=element.id,
             body={
@@ -108,33 +109,38 @@ class MetaDataMixin:
         return metadata["id"]
 
     @unsupported_cache
-    def create_metadatas(
+    @batch_publication
+    def create_metadata_bulk(
         self,
         element: Element | CachedElement,
-        metadatas: list[dict[str, MetaType | str | int | float | None]],
+        metadata_list: list[dict[str, MetaType | str | int | float | None]],
+        batch_size: int = DEFAULT_BATCH_SIZE,
     ) -> list[dict[str, str]]:
         """
         Create multiple metadata on an existing element.
         This method does not support cache.
 
         :param element: The element to create multiple metadata on.
-        :param metadatas: The list of dict whose keys are the following:
+        :param metadata_list: The list of dict whose keys are the following:
             - type: MetaType
             - name: str
             - value: str | int | float
             - entity_id: str | None
+        :param batch_size: The size of each batch, which will be used to split the publication to avoid API errors.
+
+        :returns: A list of dicts as returned in the ``metadata_list`` field by the ``CreateMetaDataBulk`` API endpoint.
         """
         assert element and isinstance(
             element, Element | CachedElement
         ), "element shouldn't be null and should be of type Element or CachedElement"
 
-        assert metadatas and isinstance(
-            metadatas, list
-        ), "type shouldn't be null and should be of type list of Dict"
+        assert metadata_list and isinstance(
+            metadata_list, list
+        ), "metadata_list shouldn't be null and should be of type list of dict"
 
         # Make a copy to avoid modifying the metadata_list argument
         metas = []
-        for index, metadata in enumerate(metadatas):
+        for index, metadata in enumerate(metadata_list):
             assert isinstance(
                 metadata, dict
             ), f"Element at index {index} in metadata_list: Should be of type dict"
@@ -168,14 +174,18 @@ class MetaDataMixin:
             logger.warning("Cannot create metadata as this worker is in read-only mode")
             return
 
-        created_metadata_list = self.request(
-            "CreateMetaDataBulk",
-            id=element.id,
-            body={
-                "worker_run_id": self.worker_run_id,
-                "metadata_list": metas,
-            },
-        )["metadata_list"]
+        created_metadata_list = [
+            created_metadata
+            for batch in make_batches(metas, "metadata", batch_size)
+            for created_metadata in self.api_client.request(
+                "CreateMetaDataBulk",
+                id=element.id,
+                body={
+                    "worker_run_id": self.worker_run_id,
+                    "metadata_list": batch,
+                },
+            )["metadata_list"]
+        ]
 
         return created_metadata_list
 

arkindex_worker/worker/process.py

@@ -0,0 +1,92 @@
+from collections.abc import Iterator
+from enum import Enum
+
+from arkindex_worker.cache import unsupported_cache
+
+# Increases the number of elements returned per page by the API
+PROCESS_ELEMENTS_PAGE_SIZE = 500
+
+
+class ActivityState(Enum):
+    """
+    Processing state of an element.
+    """
+
+    Queued = "queued"
+    """
+    The element has not yet been processed by a worker.
+    """
+
+    Started = "started"
+    """
+    The element is being processed by a worker.
+    """
+
+    Processed = "processed"
+    """
+    The element has been successfully processed by a worker.
+    """
+
+    Error = "error"
+    """
+    An error occurred while processing this element.
+    """
+
+
+class ProcessMode(Enum):
+    """
+    Mode of the process of the worker.
+    """
+
+    Files = "files"
+    """
+    Processes of files (images, PDFs, IIIF, ...) imports.
+    """
+
+    Workers = "workers"
+    """
+    Processes of worker executions.
+    """
+
+    Template = "template"
+    """
+    Process templates.
+    """
+
+    S3 = "s3"
+    """
+    Processes of imports from an S3-compatible storage.
+    """
+
+    Local = "local"
+    """
+    Local processes.
+    """
+
+    Dataset = "dataset"
+    """
+    Dataset processes.
+    """
+
+    Export = "export"
+    """
+    Export processes.
+    """
+
+
+class ProcessMixin:
+    @unsupported_cache
+    def list_process_elements(self, with_image: bool = False) -> Iterator[dict]:
+        """
+        List the elements of a process.
+
+        :param with_image: whether or not to include zone and image information in the elements response.
+        :returns: the process' elements.
+        """
+        return self.api_client.paginate(
+            "ListProcessElements",
+            id=self.process_information["id"],
+            with_image=with_image,
+            allow_missing_data=True,
+            page_size=PROCESS_ELEMENTS_PAGE_SIZE,
+        )

arkindex_worker/worker/task.py

@@ -5,8 +5,7 @@ BaseWorker methods for tasks.
 import uuid
 from collections.abc import Iterator
 
-from apistar.compat import DownloadedFile
-
+from arkindex.compat import DownloadedFile
 from arkindex_worker.models import Artifact
 
 
@@ -22,7 +21,7 @@ class TaskMixin:
             task_id, uuid.UUID
         ), "task_id shouldn't be null and should be an UUID"
 
-        results = self.request("ListArtifacts", id=task_id)
+        results = self.api_client.request("ListArtifacts", id=task_id)
 
         return map(Artifact, results)
 
@@ -43,4 +42,6 @@ class TaskMixin:
             artifact, Artifact
         ), "artifact shouldn't be null and should be an Artifact"
 
-        return self.request("DownloadArtifact", id=task_id, path=artifact.path)
+        return self.api_client.request(
+            "DownloadArtifact", id=task_id, path=artifact.path
+        )

arkindex_worker/worker/training.py

@@ -9,8 +9,8 @@ from typing import NewType
 from uuid import UUID
 
 import requests
-from apistar.exceptions import ErrorResponse
 
+from arkindex.exceptions import ErrorResponse
 from arkindex_worker import logger
 from arkindex_worker.utils import close_delete_file, create_tar_zst_archive
 
@@ -83,6 +83,9 @@ class TrainingMixin:
 
     @property
     def is_finetuning(self) -> bool:
+        """
+        Whether or not this worker is fine-tuning an existing model version.
+        """
         return bool(self.model_version_id)
 
     @skip_if_read_only
@@ -182,7 +185,7 @@ class TrainingMixin:
         assert not self.model_version, "A model version has already been created."
 
         configuration = configuration or {}
-        self.model_version = self.request(
+        self.model_version = self.api_client.request(
             "CreateModelVersion",
             id=model_id,
             body=build_clean_payload(
@@ -214,7 +217,7 @@ class TrainingMixin:
         :param parent: ID of the parent model version
         """
         assert self.model_version, "No model version has been created yet."
-        self.model_version = self.request(
+        self.model_version = self.api_client.request(
             "UpdateModelVersion",
             id=self.model_version["id"],
             body=build_clean_payload(
@@ -270,32 +273,44 @@ class TrainingMixin:
         """
         assert self.model_version, "You must create the model version and upload its archive before validating it."
         try:
-            self.model_version = self.request(
-                "ValidateModelVersion",
+            self.model_version = self.api_client.request(
+                "PartialUpdateModelVersion",
                 id=self.model_version["id"],
                 body={
+                    "state": "available",
                     "size": size,
                     "hash": hash,
                     "archive_hash": archive_hash,
                 },
             )
         except ErrorResponse as e:
-            if e.status_code != 409:
+            model_version = e.content
+            if not model_version or "id" not in model_version:
                 raise e
+
             logger.warning(
                 f"An available model version exists with hash {hash}, using it instead of the pending version."
             )
             pending_version_id = self.model_version["id"]
-            self.model_version = getattr(e, "content", None)
-            assert self.model_version is not None, "An unexpected error occurred."
-
             logger.warning("Removing the pending model version.")
             try:
-                self.request("DestroyModelVersion", id=pending_version_id)
+                self.api_client.request("DestroyModelVersion", id=pending_version_id)
             except ErrorResponse as e:
                 msg = getattr(e, "content", str(e))
                 logger.error(
                     f"An error occurred removing the pending version {pending_version_id}: {msg}."
                 )
 
+            logger.info("Retrieving the existing model version.")
+            existing_version_id = model_version["id"].pop()
+            try:
+                self.model_version = self.api_client.request(
+                    "RetrieveModelVersion", id=existing_version_id
+                )
+            except ErrorResponse as e:
+                logger.error(
+                    f"An error occurred retrieving the existing version {existing_version_id}: {e.status_code} - {e.content}."
+                )
+                raise
+
         logger.info(f"Model version {self.model_version['id']} is now available.")

arkindex_worker/worker/transcription.py

@@ -11,6 +11,7 @@ from peewee import IntegrityError
 from arkindex_worker import logger
 from arkindex_worker.cache import CachedElement, CachedTranscription
 from arkindex_worker.models import Element
+from arkindex_worker.utils import DEFAULT_BATCH_SIZE, batch_publication, make_batches
 
 
 class TextOrientation(Enum):
@@ -77,7 +78,7 @@ class TranscriptionMixin:
             )
             return
 
-        created = self.request(
+        created = self.api_client.request(
             "CreateTranscription",
             id=element.id,
             body={
@@ -109,9 +110,11 @@ class TranscriptionMixin:
 
         return created
 
+    @batch_publication
     def create_transcriptions(
         self,
         transcriptions: list[dict[str, str | float | TextOrientation | None]],
+        batch_size: int = DEFAULT_BATCH_SIZE,
     ) -> list[dict[str, str | float]]:
         """
         Create multiple transcriptions at once on existing elements through the API,
@@ -128,6 +131,8 @@ class TranscriptionMixin:
             orientation (TextOrientation)
                 Optional. Orientation of the transcription's text.
 
+        :param batch_size: The size of each batch, which will be used to split the publication to avoid API errors.
+
         :returns: A list of dicts as returned in the ``transcriptions`` field by the ``CreateTranscriptions`` API endpoint.
         """
 
@@ -171,13 +176,19 @@ class TranscriptionMixin:
             )
             return
 
-        created_trs = self.request(
-            "CreateTranscriptions",
-            body={
-                "worker_run_id": self.worker_run_id,
-                "transcriptions": transcriptions_payload,
-            },
-        )["transcriptions"]
+        created_trs = [
+            created_tr
+            for batch in make_batches(
+                transcriptions_payload, "transcription", batch_size
+            )
+            for created_tr in self.api_client.request(
+                "CreateTranscriptions",
+                body={
+                    "worker_run_id": self.worker_run_id,
+                    "transcriptions": batch,
+                },
+            )["transcriptions"]
+        ]
 
         if self.use_cache:
             # Store transcriptions in local cache
@@ -201,11 +212,13 @@ class TranscriptionMixin:
 
         return created_trs
 
+    @batch_publication
     def create_element_transcriptions(
         self,
         element: Element | CachedElement,
         sub_element_type: str,
         transcriptions: list[dict[str, str | float]],
+        batch_size: int = DEFAULT_BATCH_SIZE,
     ) -> dict[str, str | bool]:
         """
         Create multiple elements and transcriptions at once on a single parent element through the API.
@@ -225,6 +238,8 @@ class TranscriptionMixin:
             element_confidence (float)
                 Optional. Confidence score of the element between 0 and 1.
 
+        :param batch_size: The size of each batch, which will be used to split the publication to avoid API errors.
+
         :returns: A list of dicts as returned by the ``CreateElementTranscriptions`` API endpoint.
         """
         assert element and isinstance(
@@ -291,16 +306,22 @@ class TranscriptionMixin:
             )
             return
 
-        annotations = self.request(
-            "CreateElementTranscriptions",
-            id=element.id,
-            body={
-                "element_type": sub_element_type,
-                "worker_run_id": self.worker_run_id,
-                "transcriptions": transcriptions_payload,
-                "return_elements": True,
-            },
-        )
+        annotations = [
+            annotation
+            for batch in make_batches(
+                transcriptions_payload, "transcription", batch_size
+            )
+            for annotation in self.api_client.request(
+                "CreateElementTranscriptions",
+                id=element.id,
+                body={
+                    "element_type": sub_element_type,
+                    "worker_run_id": self.worker_run_id,
+                    "transcriptions": batch,
+                    "return_elements": True,
+                },
+            )
+        ]
 
         for annotation in annotations:
             if annotation["created"]:
@@ -420,60 +441,60 @@ class TranscriptionMixin:
                 ), "if of type bool, worker_run can only be set to False"
             query_params["worker_run"] = worker_run
 
-        if self.use_cache:
-            if not recursive:
-                # In this case we don't have to return anything, it's easier to use an
-                # impossible condition (False) rather than filtering by type for nothing
-                if element_type and element_type != element.type:
-                    return CachedTranscription.select().where(False)
-                transcriptions = CachedTranscription.select().where(
-                    CachedTranscription.element_id == element.id
+        if not self.use_cache:
+            return self.api_client.paginate(
+                "ListTranscriptions", id=element.id, **query_params
+            )
+
+        if not recursive:
+            # In this case we don't have to return anything, it's easier to use an
+            # impossible condition (False) rather than filtering by type for nothing
+            if element_type and element_type != element.type:
+                return CachedTranscription.select().where(False)
+            transcriptions = CachedTranscription.select().where(
+                CachedTranscription.element_id == element.id
+            )
+        else:
+            base_case = (
+                CachedElement.select()
+                .where(CachedElement.id == element.id)
+                .cte("base", recursive=True)
+            )
+            recursive = CachedElement.select().join(
+                base_case, on=(CachedElement.parent_id == base_case.c.id)
+            )
+            cte = base_case.union_all(recursive)
+            transcriptions = (
+                CachedTranscription.select()
+                .join(cte, on=(CachedTranscription.element_id == cte.c.id))
+                .with_cte(cte)
+            )
+
+            if element_type:
+                transcriptions = transcriptions.where(cte.c.type == element_type)
+
+        if worker_version is not None:
+            # If worker_version=False, filter by manual worker_version e.g. None
+            worker_version_id = worker_version or None
+            if worker_version_id:
+                transcriptions = transcriptions.where(
+                    CachedTranscription.worker_version_id == worker_version_id
                 )
             else:
-                base_case = (
-                    CachedElement.select()
-                    .where(CachedElement.id == element.id)
-                    .cte("base", recursive=True)
+                transcriptions = transcriptions.where(
+                    CachedTranscription.worker_version_id.is_null()
                 )
-                recursive = CachedElement.select().join(
-                    base_case, on=(CachedElement.parent_id == base_case.c.id)
+
+        if worker_run is not None:
+            # If worker_run=False, filter by manual worker_run e.g. None
+            worker_run_id = worker_run or None
+            if worker_run_id:
+                transcriptions = transcriptions.where(
+                    CachedTranscription.worker_run_id == worker_run_id
                 )
-                cte = base_case.union_all(recursive)
-                transcriptions = (
-                    CachedTranscription.select()
-                    .join(cte, on=(CachedTranscription.element_id == cte.c.id))
-                    .with_cte(cte)
+            else:
+                transcriptions = transcriptions.where(
+                    CachedTranscription.worker_run_id.is_null()
                 )
 
-                if element_type:
-                    transcriptions = transcriptions.where(cte.c.type == element_type)
-
-            if worker_version is not None:
-                # If worker_version=False, filter by manual worker_version e.g. None
-                worker_version_id = worker_version or None
-                if worker_version_id:
-                    transcriptions = transcriptions.where(
-                        CachedTranscription.worker_version_id == worker_version_id
-                    )
-                else:
-                    transcriptions = transcriptions.where(
-                        CachedTranscription.worker_version_id.is_null()
-                    )
-
-            if worker_run is not None:
-                # If worker_run=False, filter by manual worker_run e.g. None
-                worker_run_id = worker_run or None
-                if worker_run_id:
-                    transcriptions = transcriptions.where(
-                        CachedTranscription.worker_run_id == worker_run_id
-                    )
-                else:
-                    transcriptions = transcriptions.where(
-                        CachedTranscription.worker_run_id.is_null()
-                    )
-        else:
-            transcriptions = self.api_client.paginate(
-                "ListTranscriptions", id=element.id, **query_params
-            )
-
         return transcriptions

arkindex_worker/worker/version.py

@@ -34,7 +34,9 @@ class WorkerVersionMixin:
         if worker_version_id in self._worker_version_cache:
             return self._worker_version_cache[worker_version_id]
 
-        worker_version = self.request("RetrieveWorkerVersion", id=worker_version_id)
+        worker_version = self.api_client.request(
+            "RetrieveWorkerVersion", id=worker_version_id
+        )
         self._worker_version_cache[worker_version_id] = worker_version
 
         return worker_version

hooks/pre_gen_project.py

@@ -0,0 +1,3 @@
+# Normalize the slug to generate __package and __module private variables
+{{cookiecutter.update({"__package": cookiecutter.slug.lower().replace("_", "-")})}}  # noqa: F821
+{{cookiecutter.update({"__module": cookiecutter.slug.lower().replace("-", "_")})}}  # noqa: F821

tests/__init__.py

@@ -0,0 +1,8 @@
+from pathlib import Path
+
+BASE_DIR = Path(__file__).resolve().parent
+FIXTURES_DIR = BASE_DIR / "data"
+SAMPLES_DIR = BASE_DIR / "samples"
+
+CORPUS_ID = "11111111-1111-1111-1111-111111111111"
+PROCESS_ID = "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeffff"