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

arkindex_worker/worker/__init__.py

@@ -4,65 +4,47 @@ Base classes to implement Arkindex workers.
 
 import contextlib
 import json
-import os
 import sys
 import uuid
-from collections.abc import Iterable, Iterator
-from enum import Enum
-from itertools import groupby
-from operator import itemgetter
+from collections.abc import Iterable
+from itertools import chain
 from pathlib import Path
 
-from apistar.exceptions import ErrorResponse
-
+from arkindex.exceptions import ErrorResponse
 from arkindex_worker import logger
 from arkindex_worker.cache import CachedElement
-from arkindex_worker.models import Dataset, Element
+from arkindex_worker.models import Dataset, Element, Set
+from arkindex_worker.utils import pluralize
 from arkindex_worker.worker.base import BaseWorker
 from arkindex_worker.worker.classification import ClassificationMixin
-from arkindex_worker.worker.dataset import DatasetMixin, DatasetState
+from arkindex_worker.worker.corpus import CorpusMixin
+from arkindex_worker.worker.dataset import (
+    DatasetMixin,
+    DatasetState,
+    MissingDatasetArchive,
+)
 from arkindex_worker.worker.element import ElementMixin
-from arkindex_worker.worker.entity import EntityMixin  # noqa: F401
+from arkindex_worker.worker.entity import EntityMixin
+from arkindex_worker.worker.image import ImageMixin
 from arkindex_worker.worker.metadata import MetaDataMixin, MetaType  # noqa: F401
+from arkindex_worker.worker.process import ActivityState, ProcessMixin, ProcessMode
 from arkindex_worker.worker.task import TaskMixin
 from arkindex_worker.worker.transcription import TranscriptionMixin
-from arkindex_worker.worker.version import WorkerVersionMixin  # noqa: F401
-
-
-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.
-    """
+from arkindex_worker.worker.version import WorkerVersionMixin
 
 
 class ElementsWorker(
+    ElementMixin,
+    DatasetMixin,
     BaseWorker,
     ClassificationMixin,
-    ElementMixin,
+    CorpusMixin,
     TranscriptionMixin,
     WorkerVersionMixin,
     EntityMixin,
     MetaDataMixin,
+    ImageMixin,
+    ProcessMixin,
 ):
     """
     Base class for ML workers that operate on Arkindex elements.
@@ -80,39 +62,41 @@ class ElementsWorker(
         """
         super().__init__(description, support_cache)
 
-        # Add mandatory argument to process elements
-        self.parser.add_argument(
-            "--elements-list",
-            help="JSON elements list to use",
-            type=open,
-            default=os.environ.get("TASK_ELEMENTS"),
-        )
-        self.parser.add_argument(
-            "--element",
-            type=uuid.UUID,
-            nargs="+",
-            help="One or more Arkindex element ID",
-        )
-
         self.classes = {}
 
         self.entity_types = {}
         """Known and available entity types in processed corpus
         """
 
+        self.corpus_types = {}
+        """Known and available element types in processed corpus
+        """
+
         self._worker_version_cache = {}
 
-    def list_elements(self) -> Iterable[CachedElement] | list[str]:
+    def get_elements(self) -> Iterable[CachedElement] | list[str] | list[Element]:
         """
         List the elements to be processed, either from the CLI arguments or
         the cache database when enabled.
 
         :return: An iterable of [CachedElement][arkindex_worker.cache.CachedElement] when cache support is enabled,
-           and a list of strings representing element IDs otherwise.
+           or a list of strings representing element IDs otherwise.
         """
         assert not (
             self.args.elements_list and self.args.element
         ), "elements-list and element CLI args shouldn't be both set"
+
+        def invalid_element_id(value: str) -> bool:
+            """
+            Return whether the ID of an element is a valid UUID or not
+            """
+            try:
+                uuid.UUID(value)
+            except Exception:
+                return True
+
+            return False
+
         out = []
 
         # Load from the cache when available
@@ -122,15 +106,28 @@ class ElementsWorker(
         )
         if self.use_cache and cache_query.exists():
             return cache_query
-        # Process elements from JSON file
         elif self.args.elements_list:
+            # Process elements from JSON file
             data = json.load(self.args.elements_list)
             assert isinstance(data, list), "Elements list must be a list"
             assert len(data), "No elements in elements list"
             out += list(filter(None, [element.get("id") for element in data]))
-        # Add any extra element from CLI
         elif self.args.element:
+            # Add any extra element from CLI
             out += self.args.element
+        elif self.process_mode == ProcessMode.Dataset or self.args.set:
+            # Elements from datasets
+            return list(
+                chain.from_iterable(map(self.list_set_elements, self.list_sets()))
+            )
+        elif self.process_mode == ProcessMode.Export:
+            # For export mode processes, use list_process_elements and return element IDs
+            return {item["id"] for item in self.list_process_elements()}
+
+        invalid_element_ids = list(filter(invalid_element_id, out))
+        assert (
+            not invalid_element_ids
+        ), f"These element IDs are invalid: {', '.join(invalid_element_ids)}"
 
         return out
 
@@ -140,30 +137,22 @@ class ElementsWorker(
         Whether or not WorkerActivity support has been enabled on the DataImport
         used to run this worker.
         """
-        if self.is_read_only:
+        if self.is_read_only or self.process_mode in [
+            ProcessMode.Dataset,
+            ProcessMode.Export,
+        ]:
+            # Worker activities are also disabled when running an ElementsWorker in a Dataset process
+            # and when running export processes.
             return False
         assert (
             self.process_information
         ), "Worker must be configured to access its process activity state"
         return self.process_information.get("activity_state") == "ready"
 
-    def configure(self):
-        """
-        Setup the worker using CLI arguments and environment variables.
-        """
-        # CLI args are stored on the instance so that implementations can access them
-        self.args = self.parser.parse_args()
-
-        if self.is_read_only:
-            super().configure_for_developers()
-        else:
-            super().configure()
-            super().configure_cache()
-
     def run(self):
         """
         Implements an Arkindex worker that goes through each element returned by
-        [list_elements][arkindex_worker.worker.ElementsWorker.list_elements].
+        [get_elements][arkindex_worker.worker.ElementsWorker.get_elements].
         It calls [process_element][arkindex_worker.worker.ElementsWorker.process_element],
         catching exceptions, and handles saving WorkerActivity updates when enabled.
         """
@@ -171,7 +160,7 @@ class ElementsWorker(
 
         # List all elements either from JSON file
         # or direct list of elements on CLI
-        elements = self.list_elements()
+        elements = self.get_elements()
         if not elements:
             logger.warning("No elements to process, stopping.")
             sys.exit(1)
@@ -187,12 +176,14 @@ class ElementsWorker(
         for i, item in enumerate(elements, start=1):
             element = None
             try:
-                if self.use_cache:
-                    # Just use the result of list_elements as the element
+                if isinstance(item, CachedElement | Element):
+                    # Just use the result of get_elements as the element
                     element = item
                 else:
                     # Load element using the Arkindex API
-                    element = Element(**self.request("RetrieveElement", id=item))
+                    element = Element(
+                        **self.api_client.request("RetrieveElement", id=item)
+                    )
 
                 logger.info(f"Processing {element} ({i}/{count})")
 
@@ -230,7 +221,7 @@ class ElementsWorker(
                     with contextlib.suppress(Exception):
                         self.update_activity(element.id, ActivityState.Error)
 
-        message = f'Ran on {count} element{"s"[:count>1]}: {count - failed} completed, {failed} failed'
+        message = f'Ran on {count} {pluralize("element", count)}: {count - failed} completed, {failed} failed'
         if failed:
             logger.error(message)
             if failed >= count:  # Everything failed!
@@ -271,7 +262,7 @@ class ElementsWorker(
         assert isinstance(state, ActivityState), "state should be an ActivityState"
 
         try:
-            self.request(
+            self.api_client.request(
                 "UpdateWorkerActivity",
                 id=self.worker_run_id,
                 body={
@@ -301,16 +292,9 @@ class ElementsWorker(
         return True
 
 
-class MissingDatasetArchive(Exception):
-    """
-    Exception raised when the compressed archive associated to
-    a dataset isn't found in its task artifacts.
-    """
-
-
-class DatasetWorker(BaseWorker, DatasetMixin, TaskMixin):
+class DatasetWorker(DatasetMixin, BaseWorker, TaskMixin):
     """
-    Base class for ML workers that operate on Arkindex datasets.
+    Base class for ML workers that operate on Arkindex dataset sets.
 
     This class inherits from numerous mixin classes found in other modules of
     ``arkindex.worker``, which provide helpers to read and write to the Arkindex API.
@@ -320,193 +304,108 @@ class DatasetWorker(BaseWorker, DatasetMixin, TaskMixin):
         self,
         description: str = "Arkindex Dataset Worker",
         support_cache: bool = False,
-        generator: bool = False,
     ):
         """
         :param description: The worker's description.
         :param support_cache: Whether the worker supports cache.
-        :param generator: Whether the worker generates the dataset archive artifact.
         """
         super().__init__(description, support_cache)
 
-        self.parser.add_argument(
-            "--dataset",
-            type=uuid.UUID,
-            nargs="+",
-            help="One or more Arkindex dataset ID",
-        )
-
-        self.generator = generator
+        # Path to the dataset compressed archive (containing images and a SQLite database)
+        # Set as an instance variable as dataset workers might use it to easily extract its content
+        self.downloaded_dataset_artifact: Path | None = None
 
-    def configure(self):
+    def cleanup_downloaded_artifact(self) -> None:
         """
-        Setup the worker using CLI arguments and environment variables.
+        Cleanup the downloaded dataset artifact if any
         """
-        # CLI args are stored on the instance so that implementations can access them
-        self.args = self.parser.parse_args()
+        if not self.downloaded_dataset_artifact:
+            return
 
-        if self.is_read_only:
-            super().configure_for_developers()
-        else:
-            super().configure()
-            super().configure_cache()
+        self.downloaded_dataset_artifact.unlink(missing_ok=True)
 
-    def download_dataset_artifact(self, dataset: Dataset) -> Path:
+    def download_dataset_artifact(self, dataset: Dataset) -> None:
         """
         Find and download the compressed archive artifact describing a dataset using
         the [list_artifacts][arkindex_worker.worker.task.TaskMixin.list_artifacts] and
         [download_artifact][arkindex_worker.worker.task.TaskMixin.download_artifact] methods.
 
         :param dataset: The dataset to retrieve the compressed archive artifact for.
-        :returns: A path to the downloaded artifact.
         :raises MissingDatasetArchive: When the dataset artifact is not found.
         """
+        extra_dir = self.find_extras_directory()
+        archive = extra_dir / dataset.filepath
+        if archive.exists():
+            return
 
-        task_id = uuid.UUID(dataset.task_id)
+        # Cleanup the dataset artifact that was downloaded previously
+        self.cleanup_downloaded_artifact()
 
+        logger.info(f"Downloading artifact for {dataset}")
+        task_id = uuid.UUID(dataset.task_id)
         for artifact in self.list_artifacts(task_id):
             if artifact.path != dataset.filepath:
                 continue
 
-            extra_dir = self.find_extras_directory()
-            archive = extra_dir / dataset.filepath
             archive.write_bytes(self.download_artifact(task_id, artifact).read())
-            return archive
+            self.downloaded_dataset_artifact = archive
+            return
 
         raise MissingDatasetArchive(
             "The dataset compressed archive artifact was not found."
         )
 
-    def list_dataset_elements_per_split(
-        self, dataset: Dataset
-    ) -> Iterator[tuple[str, list[Element]]]:
-        """
-        List the elements in the dataset, grouped by split, using the
-        [list_dataset_elements][arkindex_worker.worker.dataset.DatasetMixin.list_dataset_elements] method.
-
-        :param dataset: The dataset to retrieve elements from.
-        :returns: An iterator of tuples containing the split name and the list of its elements.
-        """
-
-        def format_split(
-            split: tuple[str, Iterator[tuple[str, Element]]],
-        ) -> tuple[str, list[Element]]:
-            return (split[0], list(map(itemgetter(1), list(split[1]))))
-
-        return map(
-            format_split,
-            groupby(
-                sorted(self.list_dataset_elements(dataset), key=itemgetter(0)),
-                key=itemgetter(0),
-            ),
-        )
-
-    def process_dataset(self, dataset: Dataset):
-        """
-        Override this method to implement your worker and process a single Arkindex dataset at once.
-
-        :param dataset: The dataset to process.
+    def process_set(self, set: Set):
         """
+        Override this method to implement your worker and process a single Arkindex dataset set at once.
 
-    def list_datasets(self) -> Iterator[Dataset] | Iterator[str]:
+        :param set: The set to process.
         """
-        List the datasets to be processed, either from the CLI arguments or using the
-        [list_process_datasets][arkindex_worker.worker.dataset.DatasetMixin.list_process_datasets] method.
-
-        :returns: An iterator of strings if the worker is in read-only mode,
-        else an iterator of ``Dataset`` objects.
-        """
-        if self.is_read_only:
-            return map(str, self.args.dataset)
-
-        return self.list_process_datasets()
 
     def run(self):
         """
-        Implements an Arkindex worker that goes through each dataset returned by
-        [list_datasets][arkindex_worker.worker.DatasetWorker.list_datasets].
+        Implements an Arkindex worker that goes through each dataset set returned by
+        [list_sets][arkindex_worker.worker.dataset.DatasetMixin.list_sets].
 
-        It calls [process_dataset][arkindex_worker.worker.DatasetWorker.process_dataset],
-        catching exceptions, and handles updating the [DatasetState][arkindex_worker.worker.dataset.DatasetState]
-        when the worker is a generator.
+        It calls [process_set][arkindex_worker.worker.DatasetWorker.process_set],
+        catching exceptions.
         """
         self.configure()
 
-        datasets: list[Dataset] | list[str] = list(self.list_datasets())
-        if not datasets:
-            logger.warning("No datasets to process, stopping.")
+        dataset_sets: list[Set] = list(self.list_sets())
+        if not dataset_sets:
+            logger.warning("No sets to process, stopping.")
             sys.exit(1)
 
-        # Process every dataset
-        count = len(datasets)
+        # Process every set
+        count = len(dataset_sets)
         failed = 0
-        for i, item in enumerate(datasets, start=1):
-            dataset = None
-            dataset_artifact = None
-
+        for i, dataset_set in enumerate(dataset_sets, start=1):
             try:
-                if not self.is_read_only:
-                    # Just use the result of list_datasets as the dataset
-                    dataset = item
-                else:
-                    # Load dataset using the Arkindex API
-                    dataset = Dataset(**self.request("RetrieveDataset", id=item))
-
-                if self.generator:
-                    assert (
-                        dataset.state
-                        in [DatasetState.Open.value, DatasetState.Error.value]
-                    ), "When generating a new dataset, its state should be Open or Error."
-                else:
-                    assert (
-                        dataset.state == DatasetState.Complete.value
-                    ), "When processing an existing dataset, its state should be Complete."
-
-                logger.info(f"Processing {dataset} ({i}/{count})")
-
-                if self.generator:
-                    # Update the dataset state to Building
-                    logger.info(f"Building {dataset} ({i}/{count})")
-                    self.update_dataset_state(dataset, DatasetState.Building)
-                else:
-                    logger.info(f"Downloading data for {dataset} ({i}/{count})")
-                    dataset_artifact = self.download_dataset_artifact(dataset)
+                assert (
+                    dataset_set.dataset.state == DatasetState.Complete.value
+                ), "When processing a set, its dataset state should be Complete."
 
-                # Process the dataset
-                self.process_dataset(dataset)
+                logger.info(f"Retrieving data for {dataset_set} ({i}/{count})")
+                self.download_dataset_artifact(dataset_set.dataset)
 
-                if self.generator:
-                    # Update the dataset state to Complete
-                    logger.info(f"Completed {dataset} ({i}/{count})")
-                    self.update_dataset_state(dataset, DatasetState.Complete)
+                logger.info(f"Processing {dataset_set} ({i}/{count})")
+                self.process_set(dataset_set)
             except Exception as e:
-                # Handle errors occurring while retrieving, processing or patching the state for this dataset.
+                # Handle errors occurring while retrieving or processing this dataset set
                 failed += 1
 
-                # Handle the case where we failed retrieving the dataset
-                dataset_id = dataset.id if dataset else item
-
                 if isinstance(e, ErrorResponse):
-                    message = f"An API error occurred while processing dataset {dataset_id}: {e.title} - {e.content}"
+                    message = f"An API error occurred while processing {dataset_set}: {e.title} - {e.content}"
                 else:
-                    message = (
-                        f"Failed running worker on dataset {dataset_id}: {repr(e)}"
-                    )
+                    message = f"Failed running worker on {dataset_set}: {repr(e)}"
 
-                logger.warning(
-                    message,
-                    exc_info=e if self.args.verbose else None,
-                )
-                if dataset and self.generator:
-                    # Try to update the state to Error regardless of the response
-                    with contextlib.suppress(Exception):
-                        self.update_dataset_state(dataset, DatasetState.Error)
-            finally:
-                # Cleanup the dataset artifact if it was downloaded, no matter what
-                if dataset_artifact:
-                    dataset_artifact.unlink(missing_ok=True)
+                logger.warning(message, exc_info=e if self.args.verbose else None)
+
+        # Cleanup the latest downloaded dataset artifact
+        self.cleanup_downloaded_artifact()
 
-        message = f'Ran on {count} dataset{"s"[:count>1]}: {count - failed} completed, {failed} failed'
+        message = f'Ran on {count} {pluralize("set", count)}: {count - failed} completed, {failed} failed'
         if failed:
             logger.error(message)
             if failed >= count:  # Everything failed!

arkindex_worker/worker/base.py

@@ -12,16 +12,9 @@ from tempfile import mkdtemp
 
 import gnupg
 import yaml
-from apistar.exceptions import ErrorResponse
-from tenacity import (
-    before_sleep_log,
-    retry,
-    retry_if_exception,
-    stop_after_attempt,
-    wait_exponential,
-)
 
-from arkindex import ArkindexClient, options_from_env
+from arkindex import options_from_env
+from arkindex.exceptions import ErrorResponse
 from arkindex_worker import logger
 from arkindex_worker.cache import (
     check_version,
@@ -31,18 +24,8 @@ from arkindex_worker.cache import (
     merge_parents_cache,
 )
 from arkindex_worker.utils import close_delete_file, extract_tar_zst_archive
-
-
-def _is_500_error(exc: Exception) -> bool:
-    """
-    Check if an Arkindex API error has a HTTP 5xx error code.
-    Used to retry most API calls in [BaseWorker][arkindex_worker.worker.base.BaseWorker].
-    :param exc: Exception to check
-    """
-    if not isinstance(exc, ErrorResponse):
-        return False
-
-    return 500 <= exc.status_code < 600
+from arkindex_worker.worker.process import ProcessMode
+from teklia_toolbox.requests import get_arkindex_client
 
 
 class ExtrasDirNotFoundError(Exception):
@@ -174,6 +157,13 @@ class BaseWorker:
             raise Exception("Missing ARKINDEX_CORPUS_ID environment variable")
         return self._corpus_id
 
+    @property
+    def process_mode(self) -> ProcessMode | None:
+        """Mode of the process being run. Returns None when read-only."""
+        if self.is_read_only:
+            return
+        return ProcessMode(self.process_information["mode"])
+
     @property
     def is_read_only(self) -> bool:
         """
@@ -197,7 +187,7 @@ class BaseWorker:
         Create an ArkindexClient to make API requests towards Arkindex instances.
         """
         # Build Arkindex API client from environment variables
-        self.api_client = ArkindexClient(**options_from_env())
+        self.api_client = get_arkindex_client(**options_from_env())
         logger.debug(f"Setup Arkindex API client on {self.api_client.document.url}")
 
     def configure_for_developers(self):
@@ -237,7 +227,7 @@ class BaseWorker:
         # Load all required secrets
         self.secrets = {name: self.load_secret(Path(name)) for name in required_secrets}
 
-    def configure(self):
+    def configure_worker_run(self):
         """
         Setup the necessary configuration needed using CLI args and environment variables.
         This is the method called when running a worker on Arkindex.
@@ -249,7 +239,7 @@ class BaseWorker:
             logger.debug("Debug output enabled")
 
         # Load worker run information
-        worker_run = self.request("RetrieveWorkerRun", id=self.worker_run_id)
+        worker_run = self.api_client.request("RetrieveWorkerRun", id=self.worker_run_id)
 
         # Load process information
         self.process_information = worker_run["process"]
@@ -308,7 +298,7 @@ class BaseWorker:
         if self.support_cache and self.args.database is not None:
             self.use_cache = True
         elif self.support_cache and self.task_id:
-            task = self.request("RetrieveTaskFromAgent", id=self.task_id)
+            task = self.api_client.request("RetrieveTask", id=self.task_id)
             self.task_parents = task["parents"]
             paths = self.find_parents_file_paths(Path("db.sqlite"))
             self.use_cache = len(paths) > 0
@@ -338,6 +328,29 @@ class BaseWorker:
         else:
             logger.debug("Cache is disabled")
 
+    def configure(self):
+        """
+        Setup the worker using CLI arguments and environment variables.
+        """
+        # CLI args are stored on the instance so that implementations can access them
+        self.args = self.parser.parse_args()
+
+        if self.is_read_only:
+            self.configure_for_developers()
+        else:
+            self.configure_worker_run()
+            self.configure_cache()
+
+        # Retrieve the model configuration
+        if self.model_configuration:
+            self.config.update(self.model_configuration)
+            logger.info("Model version configuration retrieved")
+
+        # Retrieve the user configuration
+        if self.user_configuration:
+            self.config.update(self.user_configuration)
+            logger.info("User configuration retrieved")
+
     def load_secret(self, name: Path):
         """
         Load a Ponos secret by name.
@@ -349,7 +362,7 @@ class BaseWorker:
 
         # Load from the backend
         try:
-            resp = self.request("RetrieveSecret", name=str(name))
+            resp = self.api_client.request("RetrieveSecret", name=str(name))
             secret = resp["content"]
             logging.info(f"Loaded API secret {name}")
         except ErrorResponse as e:
@@ -489,26 +502,6 @@ class BaseWorker:
             # Clean up
             shutil.rmtree(base_extracted_path)
 
-    @retry(
-        retry=retry_if_exception(_is_500_error),
-        wait=wait_exponential(multiplier=2, min=3),
-        reraise=True,
-        stop=stop_after_attempt(5),
-        before_sleep=before_sleep_log(logger, logging.INFO),
-    )
-    def request(self, *args, **kwargs):
-        """
-        Wrapper around the ``ArkindexClient.request`` method.
-
-        The API call will be retried up to 5 times in case of HTTP 5xx errors,
-        with an exponential sleep time of 3, 4, 8 and 16 seconds between calls.
-        If the 5th call still causes an HTTP 5xx error, the exception is re-raised
-        and the caller should catch it.
-
-        Log messages are displayed when an HTTP 5xx error occurs, before waiting for the next call.
-        """
-        return self.api_client.request(*args, **kwargs)
-
     def add_arguments(self):
         """Override this method to add ``argparse`` arguments to this worker"""