arkindex-base-worker 0.3.7rc10__py3-none-any.whl → 0.4.0__py3-none-any.whl

arkindex_worker/worker/corpus.py

@@ -0,0 +1,86 @@
+"""
+BaseWorker methods for corpora.
+"""
+
+from enum import Enum
+from operator import itemgetter
+from tempfile import _TemporaryFileWrapper
+from uuid import UUID
+
+from arkindex_worker import logger
+
+
+class CorpusExportState(Enum):
+    """
+    State of a corpus export.
+    """
+
+    Created = "created"
+    """
+    The corpus export is created, awaiting its processing.
+    """
+
+    Running = "running"
+    """
+    The corpus export is being built.
+    """
+
+    Failed = "failed"
+    """
+    The corpus export failed.
+    """
+
+    Done = "done"
+    """
+    The corpus export ended in success.
+    """
+
+
+class CorpusMixin:
+    def download_export(self, export_id: str) -> _TemporaryFileWrapper:
+        """
+        Download an export.
+
+        :param export_id: UUID of the export to download
+        :returns: The downloaded export stored in a temporary file.
+        """
+        try:
+            UUID(export_id)
+        except ValueError as e:
+            raise ValueError("export_id is not a valid uuid.") from e
+
+        logger.info(f"Downloading export ({export_id})...")
+        export: _TemporaryFileWrapper = self.api_client.request(
+            "DownloadExport", id=export_id
+        )
+        logger.info(f"Downloaded export ({export_id}) @ `{export.name}`")
+        return export
+
+    def download_latest_export(self) -> _TemporaryFileWrapper:
+        """
+        Download the latest export in `done` state of the current corpus.
+
+        :returns: The downloaded export stored in a temporary file.
+        """
+        # List all exports on the corpus
+        exports = self.api_client.paginate("ListExports", id=self.corpus_id)
+
+        # Find the latest that is in "done" state
+        exports: list[dict] = sorted(
+            list(
+                filter(
+                    lambda export: export["state"] == CorpusExportState.Done.value,
+                    exports,
+                )
+            ),
+            key=itemgetter("updated"),
+            reverse=True,
+        )
+        assert (
+            len(exports) > 0
+        ), f'No available exports found for the corpus ({self.corpus_id}) with state "{CorpusExportState.Done.value.capitalize()}".'
+
+        # Download latest export
+        export_id: str = exports[0]["id"]
+
+        return self.download_export(export_id)

arkindex_worker/worker/dataset.py

@@ -2,6 +2,8 @@
 BaseWorker methods for datasets.
 """
 
+import uuid
+from argparse import ArgumentTypeError
 from collections.abc import Iterator
 from enum import Enum
 
@@ -36,7 +38,55 @@ class DatasetState(Enum):
     """
 
 
+class MissingDatasetArchive(Exception):
+    """
+    Exception raised when the compressed archive associated to
+    a dataset isn't found in its task artifacts.
+    """
+
+
+def check_dataset_set(value: str) -> tuple[uuid.UUID, str]:
+    """The `--set` argument should have the following format:
+    <dataset_id>:<set_name>
+
+    Args:
+        value (str): Provided argument.
+
+    Raises:
+        ArgumentTypeError: When the value is invalid.
+
+    Returns:
+        tuple[uuid.UUID, str]: The ID of the dataset parsed as UUID and the name of the set.
+    """
+    values = value.split(":")
+    if len(values) != 2:
+        raise ArgumentTypeError(
+            f"'{value}' is not in the correct format `<dataset_id>:<set_name>`"
+        )
+
+    dataset_id, set_name = values
+    try:
+        dataset_id = uuid.UUID(dataset_id)
+        return (dataset_id, set_name)
+    except (TypeError, ValueError) as e:
+        raise ArgumentTypeError(f"'{dataset_id}' should be a valid UUID") from e
+
+
 class DatasetMixin:
+    def add_arguments(self) -> None:
+        """Define specific ``argparse`` arguments for the worker using this mixin"""
+        self.parser.add_argument(
+            "--set",
+            type=check_dataset_set,
+            nargs="+",
+            help="""
+                One or more Arkindex dataset sets, format is <dataset_uuid>:<set_name>
+                (e.g.: "12341234-1234-1234-1234-123412341234:train")
+            """,
+            default=[],
+        )
+        super().add_arguments()
+
     def list_process_sets(self) -> Iterator[Set]:
         """
         List dataset sets associated to the worker's process. This helper is not available in developer mode.
@@ -73,6 +123,26 @@ class DatasetMixin:
 
         return map(lambda result: Element(**result["element"]), results)
 
+    def list_sets(self) -> Iterator[Set]:
+        """
+        List the sets to be processed, either from the CLI arguments or using the
+        [list_process_sets][arkindex_worker.worker.dataset.DatasetMixin.list_process_sets] method.
+
+        :returns: An iterator of ``Set`` objects.
+        """
+        if not self.is_read_only:
+            yield from self.list_process_sets()
+
+        datasets: dict[uuid.UUID, Dataset] = {}
+        for dataset_id, set_name in self.args.set:
+            # Retrieving dataset information if not already cached
+            if dataset_id not in datasets:
+                datasets[dataset_id] = Dataset(
+                    **self.api_client.request("RetrieveDataset", id=dataset_id)
+                )
+
+            yield Set(name=set_name, dataset=datasets[dataset_id])
+
     @unsupported_cache
     def update_dataset_state(self, dataset: Dataset, state: DatasetState) -> Dataset:
         """
@@ -93,7 +163,7 @@ class DatasetMixin:
             logger.warning("Cannot update dataset as this worker is in read-only mode")
             return
 
-        updated_dataset = self.request(
+        updated_dataset = self.api_client.request(
             "PartialUpdateDataset",
             id=dataset.id,
             body={"state": state.value},