Apply suggestions

tests/conftest.py

@@ -5,7 +5,7 @@ import pytest
 
 from arkindex.mock import MockApiClient
 from arkindex_worker.worker.base import BaseWorker
-from worker_init_elements.worker import InitElementWorker
+from worker_init_elements.worker import InitElementsWorker
 
 
 @pytest.fixture()
@@ -94,7 +94,11 @@ def _mock_worker_run_api(mock_api_client: MockApiClient) -> None:
                 },
             },
             "configuration": None,
-            "process": {"id": "process_id", "corpus": os.getenv("ARKINDEX_CORPUS_ID")},
+            "process": {
+                "id": "process_id",
+                "corpus": os.getenv("ARKINDEX_CORPUS_ID"),
+                "activity_state": "disabled",
+            },
             "summary": os.getenv("ARKINDEX_WORKER_RUN_ID") + " @ version 1",
         },
     )
@@ -103,10 +107,10 @@ def _mock_worker_run_api(mock_api_client: MockApiClient) -> None:
 @pytest.fixture()
 def mock_worker(
     _mock_worker_run_api, tmp_path_factory, monkeypatch
-) -> InitElementWorker:
+) -> InitElementsWorker:
     monkeypatch.setattr(sys, "argv", ["worker-init-elements"])
 
-    worker = InitElementWorker()
+    worker = InitElementsWorker()
     worker.work_dir = tmp_path_factory.mktemp("data")
     worker.configure()
 

tests/test_worker.py

@@ -5,6 +5,7 @@ from pathlib import Path
 
 import pytest
 
+from arkindex_worker.cache import SQL_VERSION
 from worker_init_elements.worker import INIT_PAGE_SIZE
 
 
@@ -20,7 +21,7 @@ def check_db(db_path: Path, elements: list, images: list) -> None:
     db.row_factory = sqlite3.Row
 
     assert list(map(dict, db.execute("select * from version").fetchall())) == [
-        {"version": 3}
+        {"version": SQL_VERSION}
     ]
     assert (
         list(map(dict, db.execute("select * from elements order by id").fetchall()))
@@ -38,14 +39,6 @@ def check_db(db_path: Path, elements: list, images: list) -> None:
 def test_run_process(use_cache, mock_worker):
     mock_worker.use_cache = use_cache
 
-    mock_worker.api_client.add_response(
-        "RetrieveProcess",
-        id=mock_worker.process_information["id"],
-        response={
-            "activity_state": "ready",
-            "corpus": "corpusid",
-        },
-    )
     mock_worker.api_client.add_response(
         "RetrieveCorpus",
         id=mock_worker.process_information["corpus"],
@@ -72,6 +65,7 @@ def test_run_process(use_cache, mock_worker):
                 "polygon": [[0, 0], [0, 1], [1, 1], [1, 0], [0, 0]],
                 "rotation_angle": 42,
                 "mirrored": False,
+                "confidence": None,
             },
             {
                 "id": "22222222-2222-2222-2222-222222222222",
@@ -84,6 +78,7 @@ def test_run_process(use_cache, mock_worker):
                 "polygon": [[0, 0], [0, 2], [2, 2], [2, 0], [0, 0]],
                 "rotation_angle": 0,
                 "mirrored": True,
+                "confidence": None,
             },
             {
                 "id": "33333333-3333-3333-3333-333333333333",
@@ -92,6 +87,17 @@ def test_run_process(use_cache, mock_worker):
                 "rotation_angle": 17,
                 "mirrored": True,
                 "confidence": 0.42,
+                **(
+                    {
+                        "image_id": None,
+                        "image_width": None,
+                        "image_height": None,
+                        "image_url": None,
+                        "polygon": None,
+                    }
+                    if mock_worker.use_cache
+                    else {}
+                ),
             },
         ],
     )
@@ -108,7 +114,7 @@ def test_run_process(use_cache, mock_worker):
     )
 
     db_path = mock_worker.work_dir / "db.sqlite"
-    assert db_path.is_file() == use_cache
+    assert db_path.is_file() is use_cache
     if use_cache:
         check_db(
             db_path=db_path,
@@ -174,14 +180,6 @@ def test_run_distributed(mock_worker):
     mock_worker.use_cache = True
     mock_worker.chunks_number = 4
 
-    mock_worker.api_client.add_response(
-        "RetrieveProcess",
-        id=mock_worker.process_information["id"],
-        response={
-            "activity_state": "ready",
-            "corpus": "corpusid",
-        },
-    )
     mock_worker.api_client.add_response(
         "RetrieveCorpus",
         id=mock_worker.process_information["corpus"],
@@ -201,6 +199,14 @@ def test_run_distributed(mock_worker):
                 "id": "22222222-2222-2222-2222-222222222222",
                 "type_id": "A",
                 "name": "Class 2",
+                "image_id": None,
+                "image_width": None,
+                "image_height": None,
+                "image_url": None,
+                "polygon": None,
+                "rotation_angle": 0,
+                "mirrored": False,
+                "confidence": None,
             },
             {
                 "id": "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa",
@@ -211,6 +217,9 @@ def test_run_distributed(mock_worker):
                 "image_height": 1337,
                 "image_url": "http://cafe",
                 "polygon": [[0, 0], [0, 1], [1, 1], [1, 0], [0, 0]],
+                "rotation_angle": 0,
+                "mirrored": False,
+                "confidence": None,
             },
             {
                 "id": "bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb",
@@ -221,6 +230,9 @@ def test_run_distributed(mock_worker):
                 "image_height": 1337,
                 "image_url": "http://beef",
                 "polygon": [[0, 0], [0, 2], [2, 2], [2, 0], [0, 0]],
+                "rotation_angle": 0,
+                "mirrored": False,
+                "confidence": None,
             },
             {
                 "id": "cccccccc-cccc-cccc-cccc-cccccccccccc",
@@ -231,6 +243,9 @@ def test_run_distributed(mock_worker):
                 "image_height": 1337,
                 "image_url": "http://beef",
                 "polygon": [[0, 0], [0, 2], [2, 2], [2, 0], [0, 0]],
+                "rotation_angle": 0,
+                "mirrored": False,
+                "confidence": None,
             },
             {
                 "id": "dddddddd-dddd-dddd-dddd-dddddddddddd",
@@ -241,16 +256,35 @@ def test_run_distributed(mock_worker):
                 "image_height": 1337,
                 "image_url": "http://cafe",
                 "polygon": [[0, 0], [0, 1], [1, 1], [1, 0], [0, 0]],
+                "rotation_angle": 0,
+                "mirrored": False,
+                "confidence": None,
             },
             {
                 "id": "eeeeeeee-eeee-eeee-eeee-eeeeeeeeeeee",
                 "type_id": "B",
                 "name": "Student 5",
+                "image_id": None,
+                "image_width": None,
+                "image_height": None,
+                "image_url": None,
+                "polygon": None,
+                "rotation_angle": 0,
+                "mirrored": False,
+                "confidence": None,
             },
             {
                 "id": "ffffffff-ffff-ffff-ffff-ffffffffffff",
                 "type_id": "B",
                 "name": "Student 6",
+                "image_id": None,
+                "image_width": None,
+                "image_height": None,
+                "image_url": None,
+                "polygon": None,
+                "rotation_angle": 0,
+                "mirrored": False,
+                "confidence": None,
             },
         ],
     )
@@ -460,17 +494,11 @@ def test_not_enough_elements(mock_worker):
         mock_worker.process()
 
 
-def test_run_duplicates(mock_worker):
+def test_run_duplicates(mock_worker, caplog):
+    caplog.set_level(logging.WARNING)
+
     mock_worker.use_cache = True
 
-    mock_worker.api_client.add_response(
-        "RetrieveProcess",
-        id=mock_worker.process_information["id"],
-        response={
-            "activity_state": "ready",
-            "corpus": "corpusid",
-        },
-    )
     mock_worker.api_client.add_response(
         "RetrieveCorpus",
         id=mock_worker.process_information["corpus"],
@@ -495,6 +523,9 @@ def test_run_duplicates(mock_worker):
                 "image_height": 1337,
                 "image_url": "http://cafe",
                 "polygon": [[0, 0], [0, 1], [1, 1], [1, 0], [0, 0]],
+                "rotation_angle": 0,
+                "mirrored": False,
+                "confidence": None,
             },
             {
                 "id": "22222222-2222-2222-2222-222222222222",
@@ -505,6 +536,9 @@ def test_run_duplicates(mock_worker):
                 "image_height": 1337,
                 "image_url": "http://cafe",
                 "polygon": [[0, 0], [0, 1], [1, 1], [1, 0], [0, 0]],
+                "rotation_angle": 0,
+                "mirrored": False,
+                "confidence": None,
             },
             {
                 "id": "22222222-2222-2222-2222-222222222222",
@@ -515,11 +549,22 @@ def test_run_duplicates(mock_worker):
                 "image_height": 1337,
                 "image_url": "http://cafe",
                 "polygon": [[0, 0], [0, 1], [1, 1], [1, 0], [0, 0]],
+                "rotation_angle": 0,
+                "mirrored": False,
+                "confidence": None,
             },
             {
                 "id": "33333333-3333-3333-3333-333333333333",
                 "type_id": "A",
                 "name": "Class 3",
+                "image_id": None,
+                "image_width": None,
+                "image_height": None,
+                "image_url": None,
+                "polygon": None,
+                "rotation_angle": 0,
+                "mirrored": False,
+                "confidence": None,
             },
         ],
     )
@@ -588,6 +633,10 @@ def test_run_duplicates(mock_worker):
         ],
     )
 
+    assert [(record.levelname, record.message) for record in caplog.records] == [
+        ("WARNING", "1 duplicate elements have been ignored.")
+    ]
+
 
 def test_activity_state_awaiting(mock_worker, monkeypatch):
     """
@@ -649,6 +698,8 @@ def test_activity_state_timeout(mock_worker, caplog, monkeypatch):
     """
     caplog.set_level(logging.WARNING)
 
+    mock_worker.process_information["activity_state"] = "pending"
+
     sleep_args = []
     monkeypatch.setattr(
         "worker_init_elements.worker.sleep", lambda seconds: sleep_args.append(seconds)

worker_init_elements/worker.py

@@ -4,6 +4,7 @@ import sys
 import uuid
 from collections import OrderedDict
 from collections.abc import Iterator
+from enum import Enum
 from logging import Logger, getLogger
 from time import sleep
 
@@ -14,17 +15,13 @@ from arkindex_worker.cache import (
     create_version_table,
     init_cache_db,
 )
-from arkindex_worker.models import Element as BaseElement
 from arkindex_worker.worker.base import BaseWorker
 
 logger: Logger = getLogger(__name__)
 
+# Increases the number of elements returned per page by the API
 INIT_PAGE_SIZE = 500
 
-# Await worker activities initialization when the process activity state is tagged as pending
-PENDING_STATE = "pending"
-ERROR_STATE = "error"
-
 
 def split_chunks(items: list, n: int) -> Iterator[list]:
     """
@@ -35,22 +32,35 @@ def split_chunks(items: list, n: int) -> Iterator[list]:
         yield items[i::n]
 
 
-class Element(BaseElement):
+class ActivityState(Enum):
+    """
+    Store the state of the workers activity tracking for a process.
+    To support large elements set, the state is asynchronously set to `ready` after a process
+    has been started and worker activities have been initialized on its elements.
+    """
+
+    Disabled = "disabled"
+    """
+    Worker activities are disabled and will not be used
+    """
+
+    Pending = "pending"
     """
-    Some attributes in the `ListProcessElements` response conflict with the `BaseElement` property/function.
-    Override them to access the API response directly.
+    Worker activities are not yet initialized
     """
 
-    @property
-    def polygon(self) -> list[float]:
-        return self["polygon"]
+    Ready = "ready"
+    """
+    Worker activities are initialized and ready for use
+    """
 
-    @property
-    def image_url(self) -> list[float]:
-        return self["image_url"]
+    Error = "error"
+    """
+    An error occurred when initializing worker activities
+    """
 
 
-class InitElementWorker(BaseWorker):
+class InitElementsWorker(BaseWorker):
     def configure(self) -> None:
         # CLI args are stored on the instance so that implementations can access them
         self.args = self.parser.parse_args()
@@ -69,15 +79,21 @@ class InitElementWorker(BaseWorker):
         self.use_cache = self.config["use_cache"]
         self.api_client.sleep_duration = self.config["sleep"]
 
-    def dump_json(
-        self, elements: list[Element], filename: str = "elements.json"
-    ) -> None:
+    def dump_json(self, elements: list[dict], filename: str = "elements.json") -> None:
+        """
+        Store elements in a JSON file.
+        This file will become an artefact.
+        """
         path = self.work_dir / filename
         assert not path.exists(), f"JSON at {path} already exists"
 
         path.write_text(json.dumps(elements, indent=4))
 
-    def dump_sqlite(self, elements: list[Element], filename: str = "db.sqlite") -> None:
+    def dump_sqlite(self, elements: list[dict], filename: str = "db.sqlite") -> None:
+        """
+        Store elements in a SQLite database. Only images and elements will be added.
+        This file will become an artefact.
+        """
         if not self.use_cache:
             return
 
@@ -92,40 +108,40 @@ class InitElementWorker(BaseWorker):
 
         # Set of unique images found in the elements
         CachedImage.insert_many(
-            [
-                {
-                    "id": uuid.UUID(element.image_id).hex,
-                    "width": element.image_width,
-                    "height": element.image_height,
-                    "url": element.image_url,
-                }
-                for element in elements
-                if element.get("image_id") is not None
-            ]
+            {
+                "id": uuid.UUID(element["image_id"]).hex,
+                "width": element["image_width"],
+                "height": element["image_height"],
+                "url": element["image_url"],
+            }
+            for element in elements
+            if element["image_id"]
         ).on_conflict_ignore(ignore=True).execute()
 
         # Fastest way to INSERT multiple rows.
         CachedElement.insert_many(
-            [
-                {
-                    "id": uuid.UUID(element.id).hex,
-                    "type": element.type,
-                    "image_id": uuid.UUID(element.image_id).hex
-                    if element.get("image_id")
-                    else None,
-                    "polygon": (element.polygon if element.get("polygon") else None),
-                    "rotation_angle": element.get("rotation_angle") or 0,
-                    "mirrored": element.get("mirrored") or False,
-                    "confidence": element.get("confidence"),
-                    "initial": True,
-                }
-                for element in elements
-            ]
+            {
+                "id": uuid.UUID(element["id"]).hex,
+                "type": element["type"],
+                "image_id": (
+                    uuid.UUID(element["image_id"]).hex if element["image_id"] else None
+                ),
+                "polygon": element["polygon"],
+                "rotation_angle": element["rotation_angle"],
+                "mirrored": element["mirrored"],
+                "confidence": element["confidence"],
+                "initial": True,
+            }
+            for element in elements
         ).execute()
 
         db.close()
 
-    def dump_chunks(self, elements: list[Element]) -> None:
+    def dump_chunks(self, elements: list[dict]) -> None:
+        """
+        Store elements in a JSON file(s) and SQLite database(s).
+        If several chunks are requested, the files will be suffixed with the chunk index.
+        """
         assert (
             len(elements) >= self.chunks_number
         ), f"Too few elements have been retrieved to distribute workflow among {self.chunks_number} branches"
@@ -137,8 +153,8 @@ class InitElementWorker(BaseWorker):
             self.dump_json(
                 elements=[
                     {
-                        "id": element.id,
-                        "type": element.type,
+                        "id": element["id"],
+                        "type": element["type"],
                     }
                     for element in chunk_elts
                 ],
@@ -159,18 +175,21 @@ class InitElementWorker(BaseWorker):
             f"Added {len(elements)} element{'s'[:len(elements) > 1]} to workflow configuration"
         )
 
-    def list_process_elements(self) -> list[Element]:
+    def list_process_elements(self) -> list[dict]:
+        """
+        List all elements linked to this process and remove duplicates
+        """
         assert self.process_information.get(
             "corpus"
-        ), "init_elements only supports processes on corpora."
+        ), "This worker only supports processes on corpora."
 
         corpus = self.request("RetrieveCorpus", id=self.process_information["corpus"])
         type_slugs = {
             element_type["id"]: element_type["slug"] for element_type in corpus["types"]
         }
 
-        elements = list(
-            Element(**element, type=type_slugs[element["type_id"]])
+        elements = [
+            {**element, "type": type_slugs[element["type_id"]]}
             for element in self.api_client.paginate(
                 "ListProcessElements",
                 id=self.process_information["id"],
@@ -178,7 +197,7 @@ class InitElementWorker(BaseWorker):
                 allow_missing_data=True,
                 page_size=INIT_PAGE_SIZE,
             )
-        )
+        ]
         # Use a dict to make elements unique by ID, then turn them back into a elements.json-compatible list
         unique_elements = OrderedDict(
             [(element["id"], element) for element in elements]
@@ -199,17 +218,29 @@ class InitElementWorker(BaseWorker):
         return list(unique_elements.values())
 
     def check_worker_activity(self) -> bool:
-        # Check if workers activity associated to this process is in a pending state
+        """
+        Check if workers activity associated to this process is in a pending state
+        """
         process = self.request("RetrieveProcess", id=self.process_information["id"])
-        if process.get("activity_state") == ERROR_STATE:
+        if process.get("activity_state") == ActivityState.Error:
             logger.error(
-                "Workers activity could not be initialized. Please report this incident to an instance administrator."
+                "Worker activities could not be initialized. Please report this incident to an instance administrator."
             )
             sys.exit(1)
-        return process.get("activity_state") != PENDING_STATE
+        return process.get("activity_state") == ActivityState.Ready
 
     def await_worker_activity(self) -> None:
-        logger.info("Awaiting workers activity initialization")
+        """
+        Worker activities are initialized asynchronously after a process has been started.
+        This worker should be running until all activities have moved to `Ready`.
+        """
+        if (
+            ActivityState(self.process_information["activity_state"])
+            == ActivityState.Disabled
+        ):
+            return
+
+        logger.info("Awaiting worker activities initialization")
         # Await worker activities to be initialized for 0, 2, 4, 8 seconds up to an hour
         timer = 1
         while True:
@@ -236,7 +267,7 @@ class InitElementWorker(BaseWorker):
 
 
 def main() -> None:
-    InitElementWorker(
+    InitElementsWorker(
         description="Worker to initialize Arkindex elements to process"
     ).run()