arkindex-base-worker 0.3.7rc2__py3-none-any.whl → 0.3.7rc4__py3-none-any.whl

{arkindex_base_worker-0.3.7rc2.dist-info → arkindex_base_worker-0.3.7rc4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: arkindex-base-worker
-Version: 0.3.7rc2
+Version: 0.3.7rc4
 Summary: Base Worker to easily build Arkindex ML workflows
 Author-email: Teklia <contact@teklia.com>
 Maintainer-email: Teklia <contact@teklia.com>

arkindex_base_worker-0.3.7rc4.dist-info/RECORD

@@ -0,0 +1,41 @@
+arkindex_worker/__init__.py,sha256=OlgCtTC9MaWeejviY0a3iQpALcRQGMVArFVVYwTF6I8,162
+arkindex_worker/cache.py,sha256=FTlB0coXofn5zTNRTcVIvh709mcw4a1bPGqkwWjKs3w,11248
+arkindex_worker/image.py,sha256=9-k_Wojk-sLbgvBSi7tWiiDc9YApWauJpHGKRay_nmo,14166
+arkindex_worker/models.py,sha256=HdKFw3qk4WIWC-DrHDkhsw0mHP3OILuCLFf7aTjruZU,9526
+arkindex_worker/utils.py,sha256=VSO8c21nsSaUCkyJaFX8wOwDQ0tztLOBFtiGvqlT0zU,6900
+arkindex_worker/worker/__init__.py,sha256=I8QmdAs659SalxNjtCu2K2ItdyUlXYm3mK_WhZdjgBs,19498
+arkindex_worker/worker/base.py,sha256=7ii3rZai6IB0-eB0TJ6pg-IhxMmW4izoJAKJKczbyZ4,19934
+arkindex_worker/worker/classification.py,sha256=0OiwxV9lb97Zs3kODm3hzyk0V7IxBTiW5SL6AYgRH1M,10351
+arkindex_worker/worker/dataset.py,sha256=qzjaXJtfeNCP2acsHbqp5tjQk-KpLHwVzjDAExeAmVg,3228
+arkindex_worker/worker/element.py,sha256=AWK3YJSHWy3j4ajntJloi_2X4zxsgXZ6c6dzphgq3OI,33848
+arkindex_worker/worker/entity.py,sha256=YT2Ttdn-L5TRoDdhOI3Z4GE1vtkWl7tKZqbYrtxZ2Ug,14630
+arkindex_worker/worker/metadata.py,sha256=SC6apVaOjFrmYw5b-njhqIlH-_r0ExbNpZeQZzlUjBE,6669
+arkindex_worker/worker/task.py,sha256=cz3wJNPgogZv1lm_3lm7WScitQtYQtL6H6I7Xokq208,1475
+arkindex_worker/worker/training.py,sha256=SOs3YKGikTr3rdWYp9H-jbtgRnZxQAoqtwB26ztx9j8,10235
+arkindex_worker/worker/transcription.py,sha256=6R7ofcGnNqX4rjT0kRKIE-G9FHq2TJ1tfztNM5sTqYE,20464
+arkindex_worker/worker/version.py,sha256=cs2pdlDxpKRO2Oldvcu54w-D_DQhf1cdeEt4tKX_QYs,1927
+tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tests/conftest.py,sha256=wzKXRnS7OKQCNHrlDaQhMC8EXlsQTY_S4L9U_hXbjpM,22004
+tests/test_base_worker.py,sha256=Uq6_MpLW23gmKFXkU-SyDUaA_4dlViLBGG4e3gpBBz0,24512
+tests/test_cache.py,sha256=ii0gyr0DrG7ChEs7pmT8hMdSguAOAcCze4bRMiFQxuk,10640
+tests/test_dataset_worker.py,sha256=Q-3gVu2FNa6mJVkUW-PUVgyUAvRkxSzLCJrPhwoJlxQ,28273
+tests/test_element.py,sha256=2G9M15TLxQRmvrWM9Kw2ucnElh4kSv_oF_5FYwwAxTY,13181
+tests/test_image.py,sha256=FZv8njLxh45sVgmY71UFHt0lv1cHr0cK4rrtPhQleX8,16262
+tests/test_merge.py,sha256=Q4zCbtZbe0wBfqE56gvAD06c6pDuhqnjKaioFqIgAQw,8331
+tests/test_utils.py,sha256=pFXegcBvIuy1tJDDSgQtCbC_tRaoLjd2055R5lu3hS0,1236
+tests/test_elements_worker/__init__.py,sha256=Fh4nkbbyJSMv_VtjQxnWrOqTnxXaaWI8S9WU0VrzCHs,179
+tests/test_elements_worker/test_classifications.py,sha256=PE88fsdra8QsWcKjSyao-pTHlaIWNxlbfF0CrLe9LBA,26517
+tests/test_elements_worker/test_cli.py,sha256=BsFTswLti63WAZ2pf6ipiZKWJJyCQuSfuKnSlESuK8g,2878
+tests/test_elements_worker/test_dataset.py,sha256=-kVll1NcMPWkIx8D7r-Z5neEGkFiZ9YQfC4eTMIfjg0,13475
+tests/test_elements_worker/test_elements.py,sha256=6XKtgXSVQJnTSgTHWwEVsAtIwLBapjYjUYPUdjxcHsY,84971
+tests/test_elements_worker/test_entities.py,sha256=yi1mXzvKvNwUNMzo0UZ56YOIJstYHcLyeepPJ8f10MQ,34557
+tests/test_elements_worker/test_metadata.py,sha256=b9CNv4W31TRJqYauvX_pRIN2SvnybaLqF-FWoFwa2Vc,18672
+tests/test_elements_worker/test_task.py,sha256=FCpxE9UpouKXgjGvWgNHEai_Hiy2d1YmqRG-_v2s27s,6312
+tests/test_elements_worker/test_training.py,sha256=WeG-cDuJ-YhPgfKH47TtXBxyargtLuk7c8tsik2WnL8,8414
+tests/test_elements_worker/test_transcriptions.py,sha256=WVJG26sZyY66fu-Eka9A1_WWIeNI2scogjypzURnp8A,73468
+tests/test_elements_worker/test_worker.py,sha256=7-jGJVT3yMGpIyN96Uafz5eIUrO4ieNLgw0k1D8BhGc,17163
+arkindex_base_worker-0.3.7rc4.dist-info/LICENSE,sha256=NVshRi1efwVezMfW7xXYLrdDr2Li1AfwfGOd5WuH1kQ,1063
+arkindex_base_worker-0.3.7rc4.dist-info/METADATA,sha256=ilh4IdFYSXepgr0imEMH3ZbewlFJlbg97VKnvhKXMVQ,3411
+arkindex_base_worker-0.3.7rc4.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
+arkindex_base_worker-0.3.7rc4.dist-info/top_level.txt,sha256=TtagLI8LSv7GE7nG8MQqDFAJ5bNDPJn7Z5vizOgrWkA,22
+arkindex_base_worker-0.3.7rc4.dist-info/RECORD,,

arkindex_worker/image.py

@@ -1,6 +1,7 @@
 """
 Helper methods to download and open IIIF images, and manage polygons.
 """
+
 import re
 from collections import namedtuple
 from io import BytesIO
@@ -114,32 +115,38 @@ def download_image(url: str) -> Image:
             )
         else:
             raise e
-    except requests.exceptions.SSLError:
-        logger.warning(
-            "An SSLError occurred during image download, retrying with a weaker and unsafe SSL configuration"
-        )
-
-        # Saving current ciphers
-        previous_ciphers = requests.packages.urllib3.util.ssl_.DEFAULT_CIPHERS
-
-        # Downgrading ciphers to download the image
-        requests.packages.urllib3.util.ssl_.DEFAULT_CIPHERS = "ALL:@SECLEVEL=1"
-        resp = _retried_request(url)
-
-        # Restoring previous ciphers
-        requests.packages.urllib3.util.ssl_.DEFAULT_CIPHERS = previous_ciphers
 
     # Preprocess the image and prepare it for classification
     image = Image.open(BytesIO(resp.content))
     logger.info(
-        "Downloaded image {} - size={}x{} in {}".format(
-            url, image.size[0], image.size[1], resp.elapsed
-        )
+        f"Downloaded image {url} - size={image.size[0]}x{image.size[1]} in {resp.elapsed}"
     )
 
     return image
 
 
+def upload_image(image: Image, url: str) -> requests.Response:
+    """
+    Upload a Pillow image to a URL.
+
+    :param image: Pillow image to upload.
+    :param url: Destination URL.
+    :returns: The upload response.
+    """
+    assert url.startswith("http"), "Destination URL for the image must be HTTP(S)"
+
+    # Retrieve a binarized version of the image
+    image_bytes = BytesIO()
+    image.save(image_bytes, format="jpeg")
+    image_bytes.seek(0)
+
+    # Upload the image
+    resp = _retried_request(url, method=requests.put, data=image_bytes)
+    logger.info(f"Uploaded image to {url} in {resp.elapsed}")
+
+    return resp
+
+
 def polygon_bounding_box(polygon: list[list[int | float]]) -> BoundingBox:
     """
     Compute the rectangle bounding box of a polygon.
@@ -167,8 +174,8 @@ def _retry_log(retry_state, *args, **kwargs):
     before_sleep=_retry_log,
     reraise=True,
 )
-def _retried_request(url):
-    resp = requests.get(url, timeout=DOWNLOAD_TIMEOUT)
+def _retried_request(url, *args, method=requests.get, **kwargs):
+    resp = method(url, *args, timeout=DOWNLOAD_TIMEOUT, **kwargs)
     resp.raise_for_status()
     return resp
 

arkindex_worker/models.py

@@ -75,10 +75,10 @@ class Element(MagicDict):
 
     def image_url(self, size: str = "full") -> str | None:
         """
-        Build an URL to access the image.
+        Build a URL to access the image.
         When possible, will return the S3 URL for images, so an ML worker can bypass IIIF servers.
         :param size: Subresolution of the image, following the syntax of the IIIF resize parameter.
-        :returns: An URL to the image, or None if the element does not have an image.
+        :returns: A URL to the image, or None if the element does not have an image.
         """
         if not self.get("zone"):
             return

arkindex_worker/utils.py

@@ -31,9 +31,10 @@ def decompress_zst_archive(compressed_archive: Path) -> tuple[int, Path]:
 
     logger.debug(f"Uncompressing file to {archive_path}")
     try:
-        with compressed_archive.open("rb") as compressed, archive_path.open(
-            "wb"
-        ) as decompressed:
+        with (
+            compressed_archive.open("rb") as compressed,
+            archive_path.open("wb") as decompressed,
+        ):
             dctx.copy_stream(compressed, decompressed)
         logger.debug(f"Successfully uncompressed archive {compressed_archive}")
     except zstandard.ZstdError as e:

arkindex_worker/worker/__init__.py

@@ -1,6 +1,7 @@
 """
 Base classes to implement Arkindex workers.
 """
+
 import contextlib
 import json
 import os
@@ -229,12 +230,13 @@ 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'
         if failed:
-            logger.error(
-                f"Ran on {count} elements: {count - failed} completed, {failed} failed"
-            )
+            logger.error(message)
             if failed >= count:  # Everything failed!
                 sys.exit(1)
+        else:
+            logger.info(message)
 
     def process_element(self, element: Element | CachedElement):
         """
@@ -504,9 +506,10 @@ class DatasetWorker(BaseWorker, DatasetMixin, TaskMixin):
                 if dataset_artifact:
                     dataset_artifact.unlink(missing_ok=True)
 
+        message = f'Ran on {count} dataset{"s"[:count>1]}: {count - failed} completed, {failed} failed'
         if failed:
-            logger.error(
-                f"Ran on {count} datasets: {count - failed} completed, {failed} failed"
-            )
+            logger.error(message)
             if failed >= count:  # Everything failed!
                 sys.exit(1)
+        else:
+            logger.info(message)

arkindex_worker/worker/base.py

@@ -1,6 +1,7 @@
 """
 The base class for all Arkindex workers.
 """
+
 import argparse
 import json
 import logging
@@ -148,6 +149,13 @@ class BaseWorker:
         # there is at least one available sqlite database either given or in the parent tasks
         self.use_cache = False
 
+        # model_version_id will be updated in configure() using the worker_run's model version
+        # or in configure_for_developers() from the environment
+        self.model_version_id = None
+        # model_details will be updated in configure() using the worker_run's model version
+        # or in configure_for_developers() from the environment
+        self.model_details = {}
+
         # task_parents will be updated in configure_cache() if the cache is supported,
         # if the task ID is set and if no database is passed as argument
         self.task_parents = []
@@ -257,15 +265,15 @@ class BaseWorker:
 
         # Load model version configuration when available
         model_version = worker_run.get("model_version")
-        if model_version and model_version.get("configuration"):
+        if model_version:
             logger.info("Loaded model version configuration from WorkerRun")
-            self.model_configuration.update(model_version.get("configuration"))
+            self.model_configuration.update(model_version["configuration"])
 
             # Set model_version ID as worker attribute
-            self.model_version_id = model_version.get("id")
+            self.model_version_id = model_version["id"]
 
             # Set model details as worker attribute
-            self.model_details = model_version.get("model")
+            self.model_details = model_version["model"]
 
         # Retrieve initial configuration from API
         self.config = worker_version["configuration"].get("configuration", {})

arkindex_worker/worker/dataset.py

@@ -51,7 +51,7 @@ class DatasetMixin:
 
         return map(
             lambda result: Dataset(**result["dataset"], selected_sets=result["sets"]),
-            list(results),
+            results,
         )
 
     def list_dataset_elements(self, dataset: Dataset) -> Iterator[tuple[str, Element]]:
@@ -65,14 +65,20 @@ class DatasetMixin:
             dataset, Dataset
         ), "dataset shouldn't be null and should be a Dataset"
 
-        results = self.api_client.paginate("ListDatasetElements", id=dataset.id)
+        if dataset.sets == dataset.selected_sets:
+            results = self.api_client.paginate("ListDatasetElements", id=dataset.id)
+        else:
+            results = iter(
+                element
+                for selected_set in dataset.selected_sets
+                for element in self.api_client.paginate(
+                    "ListDatasetElements", id=dataset.id, set=selected_set
+                )
+            )
 
-        def format_result(result):
-            if result["set"] not in dataset.selected_sets:
-                return
-            return (result["set"], Element(**result["element"]))
-
-        return filter(None, map(format_result, list(results)))
+        return map(
+            lambda result: (result["set"], Element(**result["element"])), results
+        )
 
     @unsupported_cache
     def update_dataset_state(self, dataset: Dataset, state: DatasetState) -> Dataset:

arkindex_worker/worker/element.py

@@ -1,6 +1,7 @@
 """
 ElementsWorker methods for elements and element types.
 """
+
 from collections.abc import Iterable
 from typing import NamedTuple
 from uuid import UUID

arkindex_worker/worker/metadata.py

@@ -50,7 +50,7 @@ class MetaType(Enum):
 
     URL = "url"
     """
-    A metadata with a string value that should be interpreted as an URL.
+    A metadata with a string value that should be interpreted as a URL.
     Only the ``http`` and ``https`` schemes are allowed.
     """
 

arkindex_worker/worker/training.py

@@ -81,6 +81,10 @@ class TrainingMixin:
 
     model_version = None
 
+    @property
+    def is_finetuning(self) -> bool:
+        return bool(self.model_version_id)
+
     @skip_if_read_only
     def publish_model_version(
         self,

arkindex_worker/worker/version.py

@@ -1,6 +1,7 @@
 """
 ElementsWorker methods for worker versions.
 """
+
 import functools
 from warnings import warn
 

tests/test_dataset_worker.py

@@ -195,7 +195,7 @@ def test_list_dataset_elements_per_split_api_error(
 ):
     responses.add(
         responses.GET,
-        f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/",
+        f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?set=set_1&with_count=true",
         status=500,
     )
 
@@ -211,23 +211,23 @@ def test_list_dataset_elements_per_split_api_error(
         # The API call is retried 5 times
         (
             "GET",
-            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?with_count=true",
+            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?set=set_1&with_count=true",
         ),
         (
             "GET",
-            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?with_count=true",
+            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?set=set_1&with_count=true",
         ),
         (
             "GET",
-            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?with_count=true",
+            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?set=set_1&with_count=true",
         ),
         (
             "GET",
-            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?with_count=true",
+            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?set=set_1&with_count=true",
         ),
         (
             "GET",
-            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?with_count=true",
+            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?set=set_1&with_count=true",
         ),
     ]
 
@@ -235,110 +235,60 @@ def test_list_dataset_elements_per_split_api_error(
 def test_list_dataset_elements_per_split(
     responses, mock_dataset_worker, default_dataset
 ):
-    expected_results = [
-        {
-            "set": "set_1",
-            "element": {
-                "id": "0000",
-                "type": "page",
-                "name": "Test",
-                "corpus": {},
-                "thumbnail_url": None,
-                "zone": {},
-                "best_classes": None,
-                "has_children": None,
-                "worker_version_id": None,
-                "worker_run_id": None,
-            },
-        },
-        {
-            "set": "set_1",
-            "element": {
-                "id": "1111",
-                "type": "page",
-                "name": "Test 2",
-                "corpus": {},
-                "thumbnail_url": None,
-                "zone": {},
-                "best_classes": None,
-                "has_children": None,
-                "worker_version_id": None,
-                "worker_run_id": None,
-            },
-        },
-        {
-            "set": "set_2",
-            "element": {
-                "id": "2222",
-                "type": "page",
-                "name": "Test 3",
-                "corpus": {},
-                "thumbnail_url": None,
-                "zone": {},
-                "best_classes": None,
-                "has_children": None,
-                "worker_version_id": None,
-                "worker_run_id": None,
-            },
-        },
-        {
-            "set": "set_3",
-            "element": {
-                "id": "3333",
-                "type": "page",
-                "name": "Test 4",
-                "corpus": {},
-                "thumbnail_url": None,
-                "zone": {},
-                "best_classes": None,
-                "has_children": None,
-                "worker_version_id": None,
-                "worker_run_id": None,
-            },
-        },
-        # `set_4` is not in `default_dataset.selected_sets`
-        {
-            "set": "set_4",
-            "element": {
-                "id": "4444",
-                "type": "page",
-                "name": "Test 5",
-                "corpus": {},
-                "thumbnail_url": None,
-                "zone": {},
-                "best_classes": None,
-                "has_children": None,
-                "worker_version_id": None,
-                "worker_run_id": None,
+    expected_results = []
+    for selected_set in default_dataset.selected_sets:
+        index = selected_set[-1]
+        expected_results.append(
+            {
+                "set": selected_set,
+                "element": {
+                    "id": str(index) * 4,
+                    "type": "page",
+                    "name": f"Test {index}",
+                    "corpus": {},
+                    "thumbnail_url": None,
+                    "zone": {},
+                    "best_classes": None,
+                    "has_children": None,
+                    "worker_version_id": None,
+                    "worker_run_id": None,
+                },
+            }
+        )
+        responses.add(
+            responses.GET,
+            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?set={selected_set}&with_count=true",
+            status=200,
+            json={
+                "count": 1,
+                "next": None,
+                "results": [expected_results[-1]],
             },
-        },
-    ]
-    responses.add(
-        responses.GET,
-        f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/",
-        status=200,
-        json={
-            "count": 4,
-            "next": None,
-            "results": expected_results,
-        },
-    )
+        )
 
     assert list(
         mock_dataset_worker.list_dataset_elements_per_split(default_dataset)
     ) == [
-        ("set_1", [expected_results[0]["element"], expected_results[1]["element"]]),
-        ("set_2", [expected_results[2]["element"]]),
-        ("set_3", [expected_results[3]["element"]]),
+        ("set_1", [expected_results[0]["element"]]),
+        ("set_2", [expected_results[1]["element"]]),
+        ("set_3", [expected_results[2]["element"]]),
     ]
 
-    assert len(responses.calls) == len(BASE_API_CALLS) + 1
+    assert len(responses.calls) == len(BASE_API_CALLS) + 3
     assert [
         (call.request.method, call.request.url) for call in responses.calls
     ] == BASE_API_CALLS + [
         (
             "GET",
-            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?with_count=true",
+            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?set=set_1&with_count=true",
+        ),
+        (
+            "GET",
+            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?set=set_2&with_count=true",
+        ),
+        (
+            "GET",
+            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?set=set_3&with_count=true",
         ),
     ]
 
@@ -360,7 +310,7 @@ def test_list_datasets_api_error(responses, mock_dataset_worker):
     with pytest.raises(
         Exception, match="Stopping pagination as data will be incomplete"
     ):
-        mock_dataset_worker.list_datasets()
+        next(mock_dataset_worker.list_datasets())
 
     assert len(responses.calls) == len(BASE_API_CALLS) + 5
     assert [
@@ -512,7 +462,7 @@ def test_run_initial_dataset_state_error(
         if generator
         else []
     ) + [
-        (logging.ERROR, "Ran on 1 datasets: 0 completed, 1 failed"),
+        (logging.ERROR, "Ran on 1 dataset: 0 completed, 1 failed"),
     ]
 
 
@@ -577,7 +527,7 @@ def test_run_update_dataset_state_api_error(
         ],
         (
             logging.ERROR,
-            "Ran on 1 datasets: 0 completed, 1 failed",
+            "Ran on 1 dataset: 0 completed, 1 failed",
         ),
     ]
 
@@ -639,7 +589,7 @@ def test_run_download_dataset_artifact_api_error(
         ),
         (
             logging.ERROR,
-            "Ran on 1 datasets: 0 completed, 1 failed",
+            "Ran on 1 dataset: 0 completed, 1 failed",
         ),
     ]
 
@@ -690,7 +640,7 @@ def test_run_no_downloaded_artifact_error(
         ),
         (
             logging.ERROR,
-            "Ran on 1 datasets: 0 completed, 1 failed",
+            "Ran on 1 dataset: 0 completed, 1 failed",
         ),
     ]
 
@@ -792,7 +742,9 @@ def test_run(
     assert [(level, message) for _, level, message in caplog.record_tuples] == [
         (logging.INFO, "Loaded Worker Fake worker @ 123412 from API"),
         (logging.INFO, "Processing Dataset (dataset_id) (1/1)"),
-    ] + extra_logs
+        *extra_logs,
+        (logging.INFO, "Ran on 1 dataset: 1 completed, 0 failed"),
+    ]
 
 
 @pytest.mark.parametrize(
@@ -890,4 +842,6 @@ def test_run_read_only(
     assert [(level, message) for _, level, message in caplog.record_tuples] == [
         (logging.WARNING, "Running without any extra configuration"),
         (logging.INFO, "Processing Dataset (dataset_id) (1/1)"),
-    ] + extra_logs
+        *extra_logs,
+        (logging.INFO, "Ran on 1 dataset: 1 completed, 0 failed"),
+    ]