Skip to content
Snippets Groups Projects

Compare revisions

Changes are shown as if the source revision was being merged into the target revision. Learn more about comparing revisions.

Source

Select target project
No results found

Target

Select target project
  • workers/base-worker
1 result
Show changes
Commits on Source (24)
Hide whitespace changes
Inline Side-by-side
Showing
with 1160 additions and 183 deletions
.gitlab-ci.yml
View file @ f67d28ff
......@@ -164,8 +164,57 @@ pages:
stage: release
only:
- master
- tags
docs-deploy:
image: node:16
stage: release
dependencies:
- docs-build
before_script:
- npm install -g surge
except:
- master
- tags
- schedules
environment:
name: ${CI_COMMIT_REF_SLUG}
url: https://${CI_COMMIT_REF_SLUG}-base-worker-arkindex.surge.sh
on_stop: docs-stop-surge
script:
- surge public ${CI_ENVIRONMENT_URL}
docs-stop-surge:
image: node:16
stage: release
when: manual
# Do not try to checkout the branch if it was deleted
variables:
GIT_STRATEGY: none
except:
- master
- tags
- schedules
environment:
name: ${CI_COMMIT_REF_SLUG}
url: https://${CI_COMMIT_REF_SLUG}-base-worker-arkindex.surge.sh
action: stop
before_script:
- npm install -g surge
script:
- surge teardown ${CI_ENVIRONMENT_URL}
release-notes:
stage: release
image: registry.gitlab.com/teklia/devops:latest
......
.isort.cfg
View file @ f67d28ff
......@@ -8,4 +8,4 @@ line_length = 88
default_section=FIRSTPARTY
known_first_party = arkindex,arkindex_common
known_third_party =PIL,apistar,gitlab,gnupg,peewee,playhouse,pytest,recommonmark,requests,setuptools,sh,tenacity,yaml
known_third_party =PIL,apistar,gitlab,gnupg,peewee,playhouse,pytest,recommonmark,requests,responses,setuptools,sh,shapely,tenacity,yaml,zstandard
VERSION
View file @ f67d28ff
0.2.3
0.3.0-rc1
arkindex_worker/cache.py
View file @ f67d28ff
# -*- coding: utf-8 -*-
"""
Database mappings and helper methods for the experimental worker caching feature.
On methods that support caching, the database will be used for all reads,
and writes will go both to the Arkindex API and the database,
reducing network usage.
"""
import json
import os
import sqlite3
......@@ -26,6 +34,10 @@ db = SqliteDatabase(None)
class JSONField(Field):
"""
A Peewee field that stores a JSON payload as a string and parses it automatically.
"""
field_type = "text"
def db_value(self, value):
......@@ -40,6 +52,11 @@ class JSONField(Field):
class Version(Model):
"""
Cache version table, used to warn about incompatible cache databases
when a worker uses an outdated version of ``base-worker``.
"""
version = IntegerField(primary_key=True)
class Meta:
......@@ -68,6 +85,7 @@ class CachedElement(Model):
mirrored = BooleanField(default=False)
initial = BooleanField(default=False)
worker_version_id = UUIDField(null=True)
confidence = FloatField(null=True)
class Meta:
database = db
......@@ -79,7 +97,12 @@ class CachedElement(Model):
This does not crop the image to the element's polygon.
IIIF servers with maxWidth, maxHeight or maxArea restrictions on image size are not supported.
:param \\*args: Positional arguments passed to :meth:`arkindex_worker.image.open_image`
:param max_size: Subresolution of the image.
:type max_size: int or None
:param \\**kwargs: Keyword arguments passed to :meth:`arkindex_worker.image.open_image`
:raises ValueError: When this element does not have an image ID or a polygon.
:returns PIL.Image: A Pillow image.
"""
if not self.image_id or not self.polygon:
raise ValueError(f"Element {self.id} has no image")
......@@ -134,7 +157,7 @@ class CachedTranscription(Model):
text = TextField()
confidence = FloatField()
orientation = CharField(max_length=50)
worker_version_id = UUIDField()
worker_version_id = UUIDField(null=True)
class Meta:
database = db
......@@ -147,7 +170,7 @@ class CachedClassification(Model):
class_name = TextField()
confidence = FloatField()
state = CharField(max_length=10)
worker_version_id = UUIDField()
worker_version_id = UUIDField(null=True)
class Meta:
database = db
......@@ -160,7 +183,7 @@ class CachedEntity(Model):
name = TextField()
validated = BooleanField(default=False)
metas = JSONField(null=True)
worker_version_id = UUIDField()
worker_version_id = UUIDField(null=True)
class Meta:
database = db
......@@ -174,7 +197,7 @@ class CachedTranscriptionEntity(Model):
entity = ForeignKeyField(CachedEntity, backref="transcription_entities")
offset = IntegerField(constraints=[Check("offset >= 0")])
length = IntegerField(constraints=[Check("length > 0")])
worker_version_id = UUIDField()
worker_version_id = UUIDField(null=True)
confidence = FloatField(null=True)
class Meta:
......
arkindex_worker/git.py
View file @ f67d28ff
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Helper classes for workers that interact with Git repositories and the GitLab API.
"""
import shutil
import time
from datetime import datetime
......@@ -205,7 +207,6 @@ class GitHelper:
in worker.configure() configure the git helper and start the cloning:
```
gitlab = GitlabHelper(...)
workflow_id = os.environ["ARKINDEX_PROCESS_ID"]
prepare_git_key(...)
self.git_helper = GitHelper(workflow_id=workflow_id, gitlab_helper=gitlab, ...)
self.git_helper.run_clone_in_background()
......
arkindex_worker/image.py
View file @ f67d28ff
# -*- coding: utf-8 -*-
"""
Helper methods to download and open IIIF images, and manage polygons.
"""
import os
from collections import namedtuple
from io import BytesIO
......@@ -6,6 +10,8 @@ from math import ceil
import requests
from PIL import Image
from shapely.affinity import rotate, scale, translate
from shapely.geometry import LinearRing
from tenacity import (
retry,
retry_if_exception_type,
......@@ -23,7 +29,20 @@ BoundingBox = namedtuple("BoundingBox", ["x", "y", "width", "height"])
def open_image(path, mode="RGB", rotation_angle=0, mirrored=False):
"""
Open an image from a path or a URL
Open an image from a path or a URL.
.. warning:: Prefer :meth:`Element.open_image` whenever possible.
:param path str: Path or URL to open the image from.
This parameter will be interpreted as a URL when it has a `http` or `https` scheme
and no file exist with this path locally.
:param mode str: Pillow mode for the image. See `the Pillow documentation
<https://pillow.readthedocs.io/en/stable/handbook/concepts.html#modes>`_.
:param int rotation_angle: Rotation angle to apply to the image, in degrees.
If it is not a multiple of 90°, then the rotation can cause empty pixels of
the mode's default color to be added for padding.
:param bool mirrored: Whether or not to mirror the image horizontally.
:returns PIL.Image: A Pillow image.
"""
if (
path.startswith("http://")
......@@ -51,7 +70,10 @@ def open_image(path, mode="RGB", rotation_angle=0, mirrored=False):
def download_image(url):
"""
Download an image and open it with Pillow
Download an image and open it with Pillow.
:param url str: URL of the image.
:returns PIL.Image: A Pillow image.
"""
assert url.startswith("http"), "Image URL must be HTTP(S)"
......@@ -87,6 +109,13 @@ def download_image(url):
def polygon_bounding_box(polygon):
"""
Compute the rectangle bounding box of a polygon.
:param polygon: Polygon to get the bounding box of.
:type polygon: list(list(int or float))
:returns BoundingBox: Bounding box of this polygon.
"""
x_coords, y_coords = zip(*polygon)
x, y = min(x_coords), min(y_coords)
width, height = max(x_coords) - x, max(y_coords) - y
......@@ -115,7 +144,10 @@ def _retried_request(url):
def download_tiles(url):
"""
Reconstruct a full IIIF image on servers that cannot serve the full-sized image using tiles.
Reconstruct a full IIIF image on servers that cannot serve the full-sized image, using tiles.
:param str url: URL of the image.
:returns PIL.Image: A Pillow image.
"""
if not url.endswith("/"):
url += "/"
......@@ -181,3 +213,120 @@ def download_tiles(url):
)
return full_image
def trim_polygon(polygon, image_width: int, image_height: int):
"""
Trim a polygon to an image's boundaries, with non-negative coordinates.
:param polygon: A polygon to trim.
:type: list(list(int or float) or tuple(int or float)) or tuple(tuple(int or float) or list(int or float))
:param image_width int: Width of the image.
:param image_height int: Height of the image.
:returns: A polygon trimmed to the image's bounds.
Some points may appear as missing, as the trimming can deduplicate points.
The first and last point are always equal, to reproduce the behavior
of the Arkindex backend.
:rtype: list(list(int or float))
:raises AssertionError: When argument types are invalid or when the trimmed polygon
is entirely outside of the image's bounds.
"""
assert isinstance(
polygon, (list, tuple)
), "Input polygon must be a valid list or tuple of points."
assert all(
isinstance(point, (list, tuple)) for point in polygon
), "Polygon points must be tuples or lists."
assert all(
len(point) == 2 for point in polygon
), "Polygon points must be tuples or lists of 2 elements."
assert all(
isinstance(point[0], int) and isinstance(point[1], int) for point in polygon
), "Polygon point coordinates must be integers."
assert any(
point[0] <= image_width and point[1] <= image_height for point in polygon
), "This polygon is entirely outside the image's bounds."
trimmed_polygon = [
[
min(image_width, max(0, x)),
min(image_height, max(0, y)),
]
for x, y in polygon
]
updated_polygon = []
for point in trimmed_polygon:
if point not in updated_polygon:
updated_polygon.append(point)
# Add back the matching last point, if it was present in the original polygon
if polygon[-1] == polygon[0]:
updated_polygon.append(updated_polygon[0])
return updated_polygon
def revert_orientation(element, polygon):
"""
Update the coordinates of the polygon of a child element based on the orientation of
its parent.
This method should be called before sending any polygon to Arkindex, to undo the possible
orientation applied by :meth:`Element.open_image`.
:param element: Parent element.
:type element: Element or CachedElement
:param polygon: Polygon corresponding to the child element.
:type polygon: list(list(int or float))
:return: A polygon with updated coordinates.
:rtype: list(list(int))
"""
from arkindex_worker.models import Element
from arkindex_worker.cache import CachedElement
assert element and isinstance(
element, (Element, CachedElement)
), "element shouldn't be null and should be an Element or CachedElement"
assert polygon and isinstance(
polygon, list
), "polygon shouldn't be null and should be a list"
# Rotating with Pillow can cause it to move the image around, as the image cannot have negative coordinates
# and must be a rectangle. This means the origin point of any coordinates from an image is invalid, and the
# center of the bounding box of the rotated image is different from the center of the element's bounding box.
# To properly undo the mirroring and rotation implicitly applied by open_image, we first need to find the center
# of the rotated bounding box.
if isinstance(element, Element):
assert (
element.zone and element.zone.polygon
), "element should have a zone and a polygon"
parent_ring = LinearRing(element.zone.polygon)
elif isinstance(element, CachedElement):
assert element.polygon, "cached element should have a polygon"
parent_ring = LinearRing(element.polygon)
rotated_ring = rotate(parent_ring, element.rotation_angle, origin="center")
# This rotated ring might have negative coordinates, so we get the vector that Pillow applies to offset the
# image to non-negative coordinates using the rotated bounding box.
offset_x, offset_y, _, _ = rotated_ring.bounds
# This uses the same calculation as what Shapely does for rotate/scale(origin='center').
# We will use this below to rotate around the center of the parent bounding box and not of each child polygon.
# https://github.com/Toblerity/Shapely/blob/462de3aa7a8bbd80408762a2d5aaf84b04476e4d/shapely/affinity.py#L98-L101
minx, miny, maxx, maxy = parent_ring.bounds
origin = ((maxx + minx) / 2.0, (maxy + miny) / 2.0)
ring = LinearRing(polygon)
# First undo the negative coordinates offset, since this is the last step of the original rotation
ring = translate(ring, xoff=offset_x, yoff=offset_y)
if element.rotation_angle:
ring = rotate(ring, -element.rotation_angle, origin=origin)
if element.mirrored:
ring = scale(ring, xfact=-1, origin=origin)
return [[int(x), int(y)] for x, y in ring.coords]
arkindex_worker/models.py
View file @ f67d28ff
# -*- coding: utf-8 -*-
"""
Wrappers around API results to provide more convenient attribute access and IIIF helpers.
"""
import tempfile
from contextlib import contextmanager
from typing import Optional
from requests import HTTPError
......@@ -55,7 +60,7 @@ class MagicDict(dict):
class Element(MagicDict):
"""
Describes any kind of element.
Describes an Arkindex element.
"""
def resize_zone_url(self, size="full"):
......@@ -66,10 +71,13 @@ class Element(MagicDict):
parts[-3] = size
return "/".join(parts)
def image_url(self, size="full"):
def image_url(self, size="full") -> Optional[str]:
"""
When possible, will return the S3 URL for images, so an ML worker can bypass IIIF servers
:param size: Subresolution of the image.
Build an URL to access the image.
When possible, will return the S3 URL for images, so an ML worker can bypass IIIF servers.
:param str 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.
:rtype: str or None
"""
if not self.get("zone"):
return
......@@ -82,7 +90,11 @@ class Element(MagicDict):
return "{}full/{}/0/default.jpg".format(url, size)
@property
def requires_tiles(self):
def requires_tiles(self) -> bool:
"""
:return bool: Whether or not downloading and combining IIIF tiles will be necessary
to retrieve this element's image. Will be False if the element has no image.
"""
if not self.get("zone") or self.zone.image.get("s3_url"):
return False
max_width = self.zone.image.server.max_width or float("inf")
......@@ -92,15 +104,39 @@ class Element(MagicDict):
def open_image(self, *args, max_size=None, use_full_image=False, **kwargs):
"""
Open this element's image as a Pillow image.
If use_full_image is False:
Use zone url, instead of the full image url to be able to get
single page from double page image.
Else:
This does not crop the image to the element's polygon.
:param max_size: Subresolution of the image.
Open this element's image using Pillow, rotating and mirroring it according
to the ``rotation_angle`` and ``mirrored`` attributes.
When tiling is not required to download the image, and no S3 URL is available
to bypass IIIF servers, the image will be cropped to the rectangle bounding box
of the ``zone.polygon`` attribute.
.. warning::
This method implicitly applies the element's orientation to the image.
If your process uses the returned image to find more polygons and send them
back to Arkindex, use the :meth:`arkindex_worker.image.revert_orientation`
helper to undo the orientation on all polygons before sending them, as the
Arkindex API expects unoriented polygons.
Although not recommended, you can bypass this behavior by passing
``rotation_angle=0, mirrored=False`` as keyword arguments.
:param max_size: The maximum size of the requested image.
:type max_size: int or None
:param bool use_full_image: Ignore the ``zone.polygon`` and always
retrieve the image without cropping.
:param \\*args: Positional arguments passed to :meth:`arkindex_worker.image.open_image`.
:param \\**kwargs: Keyword arguments passed to :meth:`arkindex_worker.image.open_image`.
:raises ValueError: When the element does not have an image.
:raises NotImplementedError: When the ``max_size`` parameter is set,
but the IIIF server's configuration requires downloading and combining tiles
to retrieve the image.
:raises NotImplementedError: When an S3 URL has been used to download the image,
but the URL has expired. Re-fetching the URL automatically is not supported.
:return PIL.Image: A Pillow image.
"""
if not self.get("zone"):
raise ValueError("Element {} has no zone".format(self.id))
......@@ -164,7 +200,15 @@ class Element(MagicDict):
def open_image_tempfile(self, format="jpeg", *args, **kwargs):
"""
Get the element's image as a temporary file stored on the disk.
To be used as a context manager: with element.open_image_tempfile() as f: ...
To be used as a context manager::
with element.open_image_tempfile() as f:
...
:param format str: File format to use the store the image on the disk.
Must be a format supported by Pillow.
:param \\*args: Positional arguments passed to :meth:`arkindex_worker.image.open_image`.
:param \\**kwargs: Keyword arguments passed to :meth:`arkindex_worker.image.open_image`.
"""
with tempfile.NamedTemporaryFile() as f:
self.open_image(*args, **kwargs).save(f, format=format)
......
arkindex_worker/reporting.py
View file @ f67d28ff
# -*- coding: utf-8 -*-
"""
Generator for the ``ml_report.json`` file, to report created worker results and exceptions.
"""
import json
import traceback
from collections import Counter
......@@ -10,6 +14,10 @@ from arkindex_worker import logger
class Reporter(object):
"""
Helper to generate an ``ml_report.json`` artifact.
"""
def __init__(
self, name="Unknown worker", slug="unknown-slug", version=None, **kwargs
):
......@@ -47,14 +55,21 @@ class Reporter(object):
def process(self, element_id):
"""
Report that a specific element ID is being processed.
:param element_id: ID of the element being processed.
:type element_id: str or uuid.UUID
"""
# Just call the element initializer
self._get_element(element_id)
def add_element(self, parent_id, type, type_count=1):
"""
Report creating a single element with a parent.
Multiple elements with the same type and parent can be declared with the type_count parameter.
Report creating an element as a child of another.
:param parent_id: ID of the parent element.
:type parent_id: str or uuid.UUID
:param type str: Slug of the type of the child element.
:param type_count int: How many elements of this type were created. Defaults to 1.
"""
elements = self._get_element(parent_id)["elements"]
elements.setdefault(type, 0)
......@@ -63,6 +78,10 @@ class Reporter(object):
def add_classification(self, element_id, class_name):
"""
Report creating a classification on an element.
:param element_id: ID of the element.
:type element_id: str or uuid.UUID
:param class_name str: Name of the ML class of the new classification.
"""
classifications = self._get_element(element_id)["classifications"]
classifications.setdefault(class_name, 0)
......@@ -70,7 +89,14 @@ class Reporter(object):
def add_classifications(self, element_id, classifications):
"""
Report one or more classifications at once.
Report creating one or more classifications at once on an element.
:param element_id: ID of the element.
:type element_id: str or uuid.UUID
:param classifications: List of classifications.
Each classification is represented as a ``dict`` with a ``class_name`` key
holding the name of the ML class being used.
:type classifications: List[Dict[str, str]]
"""
assert isinstance(
classifications, list
......@@ -87,31 +113,63 @@ class Reporter(object):
def add_transcription(self, element_id, count=1):
"""
Report creating a transcription on an element.
Multiple transcriptions with the same parent can be declared with the type_count parameter.
:param element_id: ID of the element.
:type element_id: str or uuid.UUID
:param count int: Number of transcriptions created at once, defaults to 1.
"""
self._get_element(element_id)["transcriptions"] += count
def add_entity(self, element_id, entity_id, type, name):
"""
Report creating an entity from an element.
Report creating an entity on an element.
:param element_id: ID of the element.
:type element_id: str or uuid.UUID
:param entity_id str: ID of the new entity.
:param type str: Type of the entity.
:param name str: Name of the entity.
"""
entities = self._get_element(element_id)["entities"]
entities.append({"id": entity_id, "type": type, "name": name})
def add_entity_link(self, *args, **kwargs):
"""
Report creating an entity link. Not currently supported.
:raises NotImplementedError:
"""
raise NotImplementedError
def add_entity_role(self, *args, **kwargs):
"""
Report creating an entity role. Not currently supported.
:raises NotImplementedError:
"""
raise NotImplementedError
def add_metadata(self, element_id, metadata_id, type, name):
"""
Report creating a metadata from an element.
:param element_id: ID of the element.
:type element_id: str or uuid.UUID
:param metadata_id str: ID of the new metadata.
:param type str: Type of the metadata.
:param name str: Name of the metadata.
"""
metadata = self._get_element(element_id)["metadata"]
metadata.append({"id": metadata_id, "type": type, "name": name})
def error(self, element_id, exception):
"""
Report that a Python exception occurred when processing an element.
:param element_id: ID of the element.
:type element_id: str or uuid.UUID
:param exception Exception: A Python exception.
"""
error_data = {
"class": exception.__class__.__name__,
"message": str(exception),
......@@ -129,6 +187,12 @@ class Reporter(object):
self._get_element(element_id)["errors"].append(error_data)
def save(self, path):
"""
Save the ML report to the specified path.
:param path: Path to save the ML report to.
:type path: str or pathlib.Path
"""
logger.info(f"Saving ML report to {path}")
with open(path, "w") as f:
json.dump(self.report_data, f)
arkindex_worker/utils.py
View file @ f67d28ff
# -*- coding: utf-8 -*-
"""
General utility functions and classes.
"""
import datetime
from timeit import default_timer
class Timer(object):
"""
A context manager to help measure execution times
A context manager to help measure execution times. Example usage::
with Timer() as t:
# do something interesting
print(t.delta) # X days, X:XX:XX
"""
def __init__(self):
......
arkindex_worker/worker/__init__.py
View file @ f67d28ff
# -*- coding: utf-8 -*-
"""
Base classes to implement Arkindex workers.
"""
import json
import os
import sys
......@@ -21,10 +25,29 @@ 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.
"""
class ElementsWorker(
......@@ -36,6 +59,13 @@ class ElementsWorker(
EntityMixin,
MetaDataMixin,
):
"""
Base class for ML workers that operate on Arkindex elements.
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.
"""
def __init__(self, description="Arkindex Elements Worker", support_cache=False):
super().__init__(description, support_cache)
......@@ -52,11 +82,20 @@ class ElementsWorker(
nargs="+",
help="One or more Arkindex element ID",
)
self.classes = {}
self._worker_version_cache = {}
def list_elements(self):
"""
List the elements to be processed, either from the CLI arguments or
the cache database when enabled.
:return: An iterable of :class:`CachedElement` when cache support is enabled,
and a list of strings representing element IDs otherwise.
:rtype: Iterable[CachedElement] or List[str]
"""
assert not (
self.args.elements_list and self.args.element
), "elements-list and element CLI args shouldn't be both set"
......@@ -83,6 +122,12 @@ class ElementsWorker(
@property
def store_activity(self):
"""
Whether or not WorkerActivity support has been enabled on the DataImport
used to run this worker.
:rtype: bool
"""
if self.args.dev:
return False
assert (
......@@ -91,7 +136,15 @@ class ElementsWorker(
return self.process_information.get("activity_state") == "ready"
def configure(self):
super().configure()
# 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()
# Add report concerning elements
self.report = Reporter(
**self.worker_details, version=getattr(self, "worker_version_id", None)
......@@ -99,7 +152,10 @@ class ElementsWorker(
def run(self):
"""
Process every elements from the provided list
Implements an Arkindex worker that goes through each element returned by
:meth:`list_elements`. It calls :meth:`process_element`, catching exceptions
and reporting them using the :class:`Reporter`, and handles saving the report
once the process is complete as well as WorkerActivity updates when enabled.
"""
self.configure()
......@@ -151,7 +207,9 @@ class ElementsWorker(
if isinstance(e, ErrorResponse):
message = f"An API error occurred while processing element {element_id}: {e.title} - {e.content}"
else:
message = f"Failed running worker on element {element_id}: {e}"
message = (
f"Failed running worker on element {element_id}: {repr(e)}"
)
logger.warning(
message,
......@@ -178,13 +236,25 @@ class ElementsWorker(
sys.exit(1)
def process_element(self, element):
"""Override this method to analyze an Arkindex element from the provided list"""
"""
Override this method to implement your worker and process a single Arkindex element at once.
def update_activity(self, element_id, state):
:param element: The element to process.
Will be a CachedElement instance if cache support is enabled,
and an Element instance otherwise.
:type element: Element or CachedElement
"""
Update worker activity for this element
Returns False when there is a conflict initializing the activity
Otherwise return True or the response payload
def update_activity(self, element_id, state) -> bool:
"""
Update the WorkerActivity for this element and worker.
:param element_id: ID of the element.
:type element_id: str or uuid.UUID
:param state ActivityState: New WorkerActivity state for this element.
:returns bool: True if the update has been successful or WorkerActivity support is disabled.
False if the update has failed due to a conflict; this worker might have already processed
this element.
"""
if not self.store_activity:
logger.debug(
......
arkindex_worker/worker/base.py
View file @ f67d28ff
# -*- coding: utf-8 -*-
"""
The base class for all Arkindex workers.
"""
import argparse
import json
import logging
......@@ -28,10 +32,11 @@ from arkindex_worker.cache import (
)
def _is_500_error(exc):
def _is_500_error(exc) -> bool:
"""
Check if an Arkindex API error is a 50x
This is used to retry most API calls implemented here
Check if an Arkindex API error has a HTTP 5xx error code.
Used to retry most API calls in :class:`BaseWorker`.
:rtype: bool
"""
if not isinstance(exc, ErrorResponse):
return False
......@@ -40,45 +45,20 @@ def _is_500_error(exc):
class BaseWorker(object):
def __init__(self, description="Arkindex Base Worker", support_cache=False):
self.parser = argparse.ArgumentParser(description=description)
# Setup workdir either in Ponos environment or on host's home
if os.environ.get("PONOS_DATA"):
self.work_dir = os.path.join(os.environ["PONOS_DATA"], "current")
else:
# We use the official XDG convention to store file for developers
# https://specifications.freedesktop.org/basedir-spec/basedir-spec-latest.html
xdg_data_home = os.environ.get(
"XDG_DATA_HOME", os.path.expanduser("~/.local/share")
)
self.work_dir = os.path.join(xdg_data_home, "arkindex")
os.makedirs(self.work_dir, exist_ok=True)
self.worker_version_id = os.environ.get("WORKER_VERSION_ID")
if not self.worker_version_id:
logger.warning(
"Missing WORKER_VERSION_ID environment variable, worker is in read-only mode"
)
logger.info(f"Worker will use {self.work_dir} as working directory")
self.process_information = None
self.user_configuration = {}
self.support_cache = support_cache
# use_cache will be updated in configure() if the cache is supported and if there
# is at least one available sqlite database either given or in the parent tasks
self.use_cache = False
@property
def is_read_only(self):
"""Worker cannot publish anything in dev mode or without a worker version ID"""
return self.args.dev or self.worker_version_id is None
"""
Base class for Arkindex workers.
"""
def configure(self):
def __init__(self, description="Arkindex Base Worker", support_cache=False):
"""
Configure worker using cli args and environment variables
Initialize the worker.
:param description str: Description shown in the ``worker-...`` command line tool.
:param support_cache bool: Whether or not this worker supports the cache database.
Override the constructor and set this parameter to start using the cache database.
"""
self.parser = argparse.ArgumentParser(description=description)
self.parser.add_argument(
"-c",
"--config",
......@@ -95,6 +75,7 @@ class BaseWorker(object):
self.parser.add_argument(
"-v",
"--verbose",
"--debug",
help="Display more information on events and errors",
action="store_true",
default=False,
......@@ -104,52 +85,77 @@ class BaseWorker(object):
help=(
"Run worker in developer mode. "
"Worker will be in read-only state even if a worker_version is supplied. "
"ARKINDEX_PROCESS_ID environment variable is not required with this flag."
),
action="store_true",
default=False,
)
# Call potential extra arguments
self.add_arguments()
# CLI args are stored on the instance so that implementations can access them
self.args = self.parser.parse_args()
# Setup workdir either in Ponos environment or on host's home
if os.environ.get("PONOS_DATA"):
self.work_dir = os.path.join(os.environ["PONOS_DATA"], "current")
else:
# We use the official XDG convention to store file for developers
# https://specifications.freedesktop.org/basedir-spec/basedir-spec-latest.html
xdg_data_home = os.environ.get(
"XDG_DATA_HOME", os.path.expanduser("~/.local/share")
)
self.work_dir = os.path.join(xdg_data_home, "arkindex")
os.makedirs(self.work_dir, exist_ok=True)
# Setup logging level
if self.args.verbose:
logger.setLevel(logging.DEBUG)
logger.debug("Debug output enabled")
self.worker_version_id = os.environ.get("WORKER_VERSION_ID")
if not self.worker_version_id:
logger.warning(
"Missing WORKER_VERSION_ID environment variable, worker is in read-only mode"
)
self.worker_run_id = os.environ.get("ARKINDEX_WORKER_RUN_ID")
if not self.worker_run_id:
logger.warning(
"Missing ARKINDEX_WORKER_RUN_ID environment variable, worker is in read-only mode"
)
logger.info(f"Worker will use {self.work_dir} as working directory")
self.process_information = None
self.user_configuration = {}
self.support_cache = support_cache
# use_cache will be updated in configure() if the cache is supported and if there
# is at least one available sqlite database either given or in the parent tasks
self.use_cache = False
# Define API Client
self.setup_api_client()
@property
def is_read_only(self) -> bool:
"""
Whether or not the worker can publish data.
:returns: False when dev mode is enabled with the ``--dev`` CLI argument,
or when no worker version ID is provided.
:rtype: bool
"""
return (
self.args.dev
or self.worker_version_id is None
or self.worker_run_id is None
)
def setup_api_client(self):
# Build Arkindex API client from environment variables
self.api_client = ArkindexClient(**options_from_env())
logger.debug(f"Setup Arkindex API client on {self.api_client.document.url}")
# Load features available on backend, and check authentication
user = self.request("RetrieveUser")
logger.debug(f"Connected as {user['display_name']} - {user['email']}")
self.features = user["features"]
# Load process information except in developer mode
if not self.args.dev:
assert os.environ.get(
"ARKINDEX_PROCESS_ID"
), "ARKINDEX_PROCESS_ID environment variable is not defined"
self.process_information = self.request(
"RetrieveDataImport", id=os.environ["ARKINDEX_PROCESS_ID"]
)
def configure_for_developers(self):
assert self.is_read_only
# Setup logging level if verbose or if ARKINDEX_DEBUG is set to true
if self.args.verbose or os.environ.get("ARKINDEX_DEBUG"):
logger.setLevel(logging.DEBUG)
logger.debug("Debug output enabled")
if self.worker_version_id:
# Retrieve initial configuration from API
worker_version = self.request(
"RetrieveWorkerVersion", id=self.worker_version_id
)
logger.info(
f"Loaded worker {worker_version['worker']['name']} revision {worker_version['revision']['hash'][0:7]} from API"
)
self.config = worker_version["configuration"]["configuration"]
self.worker_details = worker_version["worker"]
required_secrets = worker_version["configuration"].get("secrets", [])
elif self.args.config:
if self.args.config:
# Load config from YAML file
self.config = yaml.safe_load(self.args.config)
self.worker_details = {"name": "Local worker"}
......@@ -166,20 +172,54 @@ class BaseWorker(object):
# Load all required secrets
self.secrets = {name: self.load_secret(name) for name in required_secrets}
# Load worker run configuration when available and not in dev mode
if os.environ.get("ARKINDEX_WORKER_RUN_ID") and not self.args.dev:
worker_run = self.request(
"RetrieveWorkerRun", id=os.environ["ARKINDEX_WORKER_RUN_ID"]
)
configuration_id = worker_run.get("configuration_id")
if configuration_id:
worker_configuration = self.request(
"RetrieveWorkerConfiguration", id=configuration_id
)
self.user_configuration = worker_configuration.get("configuration")
if self.user_configuration:
logger.info("Loaded user configuration from WorkerRun")
def configure(self):
"""
Configure worker using CLI args and environment variables.
"""
assert not self.is_read_only
# Setup logging level if verbose or if ARKINDEX_DEBUG is set to true
if self.args.verbose or os.environ.get("ARKINDEX_DEBUG"):
logger.setLevel(logging.DEBUG)
logger.debug("Debug output enabled")
# Load worker run information
worker_run = self.request("RetrieveWorkerRun", id=self.worker_run_id)
# Load process information
self.process_information = worker_run["process"]
# Load worker version information
worker_version = worker_run["worker_version"]
self.worker_details = worker_version["worker"]
logger.info(
f"Loaded worker {self.worker_details['name']} revision {worker_version['revision']['hash'][0:7]} from API"
)
# Retrieve initial configuration from API
self.config = worker_version["configuration"]["configuration"]
if "user_configuration" in worker_version["configuration"]:
# Add default values (if set) to user_configuration
for key, value in worker_version["configuration"][
"user_configuration"
].items():
if "default" in value:
self.user_configuration[key] = value["default"]
# Load all required secrets
required_secrets = worker_version["configuration"].get("secrets", [])
self.secrets = {name: self.load_secret(name) for name in required_secrets}
# Load worker run configuration when available
worker_configuration = worker_run.get("configuration")
self.user_configuration = worker_configuration.get("configuration")
if self.user_configuration:
logger.info("Loaded user configuration from WorkerRun")
# if debug mode is set to true activate debug mode in logger
if self.user_configuration.get("debug"):
logger.setLevel(logging.DEBUG)
logger.debug("Debug output enabled")
def configure_cache(self):
task_id = os.environ.get("PONOS_TASK")
paths = None
if self.support_cache and self.args.database is not None:
......@@ -220,7 +260,12 @@ class BaseWorker(object):
logger.debug("Cache is disabled")
def load_secret(self, name):
"""Load all secrets described in the worker configuration"""
"""
Load a Ponos secret by name.
:param str name: Name of the Ponos secret.
:raises Exception: When the secret cannot be loaded from the API nor the local secrets directory.
"""
secret = None
# Load from the backend
......@@ -274,18 +319,19 @@ class BaseWorker(object):
)
def request(self, *args, **kwargs):
"""
Proxy all Arkindex API requests with a retry mechanism
in case of 50X errors
The same API call will be retried 5 times, with an exponential sleep time
going through 3, 4, 8 and 16 seconds of wait between call.
If the 5th call still gives a 50x, the exception is re-raised
and the caller should catch it
Log messages are displayed before sleeping (when at least one exception occurred)
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 argument to this worker"""
"""Override this method to add ``argparse`` arguments to this worker"""
def run(self):
"""Override this method to implement your own process"""
arkindex_worker/worker/classification.py
View file @ f67d28ff
# -*- coding: utf-8 -*-
"""
ElementsWorker methods for classifications and ML classes.
"""
import os
from apistar.exceptions import ErrorResponse
......@@ -10,9 +14,15 @@ from arkindex_worker.models import Element
class ClassificationMixin(object):
"""
Mixin for the :class:`ElementsWorker` to add ``MLClass`` and ``Classification`` helpers.
"""
def load_corpus_classes(self, corpus_id):
"""
Load ML classes for the given corpus ID
Load all ML classes for the given corpus ID and store them in the ``self.classes`` cache.
:param corpus_id str: ID of the corpus.
"""
corpus_classes = self.api_client.paginate(
"ListCorpusMLClasses",
......@@ -25,8 +35,14 @@ class ClassificationMixin(object):
def get_ml_class_id(self, corpus_id, ml_class):
"""
Return the ID corresponding to the given class name on a specific corpus
This method will automatically create missing classes
Return the MLClass ID corresponding to the given class name on a specific corpus.
If no MLClass exists for this class name, a new one is created.
:param corpus_id: ID of the corpus, or None to use the ``ARKINDEX_CORPUS_ID`` environment variable.
:type corpus_id: str or None
:param ml_class str: Name of the MLClass.
:returns str: ID of the retrieved or created MLClass.
"""
if corpus_id is None:
corpus_id = os.environ.get("ARKINDEX_CORPUS_ID")
......@@ -64,7 +80,14 @@ class ClassificationMixin(object):
self, element, ml_class, confidence, high_confidence=False
):
"""
Create a classification on the given element through API
Create a classification on the given element through the API.
:param element: The element to create a classification on.
:type element: Element or CachedElement
:param ml_class str: Name of the MLClass to use.
:param confidence float: Confidence score for the classification. Must be between 0 and 1.
:param high_confidence bool: Whether or not the classification is of high confidence.
:returns dict: The created classification, as returned by the ``CreateClassification`` API endpoint.
"""
assert element and isinstance(
element, (Element, CachedElement)
......@@ -136,7 +159,23 @@ class ClassificationMixin(object):
def create_classifications(self, element, classifications):
"""
Create multiple classifications at once on the given element through the API
Create multiple classifications at once on the given element through the API.
:param element: The element to create classifications on.
:type element: Element or CachedElement
:param classifications: The classifications to create, as a list of dicts with the following keys:
class_name (str)
Name of the MLClass for this classification.
confidence (float)
Confidence score, between 0 and 1.
high_confidence (bool)
High confidence state of the classification.
:type classifications: List[Dict[str, Union[str, float, bool]]]
:returns: List of created classifications, as returned in the ``classifications`` field by
the ``CreateClassifications`` API endpoint.
:rtype: List[Dict[str, Union[str, float, bool]]]
"""
assert element and isinstance(
element, (Element, CachedElement)
......
arkindex_worker/worker/element.py
View file @ f67d28ff
# -*- coding: utf-8 -*-
"""
ElementsWorker methods for elements and element types.
"""
import uuid
from typing import Dict, Iterable, List, Optional, Union
from peewee import IntegrityError
......@@ -16,10 +20,19 @@ class MissingTypeError(Exception):
class ElementMixin(object):
"""
Mixin for the :class:`ElementsWorker` to provide ``Element`` helpers.
"""
def check_required_types(self, corpus_id: str, *type_slugs: str) -> bool:
"""
Check that a corpus has a list of required element types,
and raise an exception if any of them are missing.
:param str corpus_id: ID of the corpus to check types on.
:param str \\*type_slugs: Type slugs to look for.
:returns bool: True if all of the specified type slugs have been found.
:raises MissingTypeError: If any of the specified type slugs were not found.
"""
assert isinstance(
corpus_id, (uuid.UUID, str)
......@@ -40,10 +53,25 @@ class ElementMixin(object):
return True
def create_sub_element(self, element, type, name, polygon):
def create_sub_element(
self,
element: Element,
type: str,
name: str,
polygon: List[List[Union[int, float]]],
confidence: Optional[float] = None,
) -> str:
"""
Create a child element on the given element through API
Return the ID of the created sub element
Create a child element on the given element through the API.
:param Element element: The parent element.
:param str type: Slug of the element type for this child element.
:param str name: Name of the child element.
:param polygon: Polygon of the child element.
:type polygon: list(list(int or float))
:param confidence: Optional confidence score, between 0.0 and 1.0.
:type confidence: float or None
:returns str: UUID of the created element.
"""
assert element and isinstance(
element, Element
......@@ -64,6 +92,10 @@ class ElementMixin(object):
assert all(
isinstance(coord, (int, float)) for point in polygon for coord in point
), "polygon points should be lists of two numbers"
assert confidence is None or (
isinstance(confidence, float) and 0 <= confidence <= 1
), "confidence should be None or a float in [0..1] range"
if self.is_read_only:
logger.warning("Cannot create element as this worker is in read-only mode")
return
......@@ -78,16 +110,43 @@ class ElementMixin(object):
"polygon": polygon,
"parent": element.id,
"worker_version": self.worker_version_id,
"confidence": confidence,
},
)
self.report.add_element(element.id, type)
return sub_element["id"]
def create_elements(self, parent, elements):
def create_elements(
self,
parent: Union[Element, CachedElement],
elements: List[
Dict[str, Union[str, List[List[Union[int, float]]], float, None]]
],
) -> List[Dict[str, str]]:
"""
Create children elements on the given element through API
Return the IDs of created elements
Create child elements on the given element in a single API request.
:param parent: Parent element for all the new child elements. The parent must have an image and a polygon.
:type parent: Element or CachedElement
:param elements: List of dicts, one per element. Each dict can have the following keys:
name (str)
Required. Name of the element.
type (str)
Required. Slug of the element type for this element.
polygon (list(list(int or float)))
Required. Polygon for this child element. Must have at least three points, with each point
having two non-negative coordinates and being inside of the parent element's image.
confidence (float or None)
Optional confidence score, between 0.0 and 1.0.
:type elements: list(dict(str, Any))
:return: List of dicts, with each dict having a single key, ``id``, holding the UUID of each created element.
:rtype: list(dict(str, str))
"""
if isinstance(parent, Element):
assert parent.get(
......@@ -135,6 +194,11 @@ class ElementMixin(object):
isinstance(coord, (int, float)) for point in polygon for coord in point
), f"Element at index {index} in elements: polygon points should be lists of two numbers"
confidence = element.get("confidence")
assert confidence is None or (
isinstance(confidence, float) and 0 <= confidence <= 1
), f"Element at index {index} in elements: confidence should be None or a float in [0..1] range"
if self.is_read_only:
logger.warning("Cannot create elements as this worker is in read-only mode")
return
......@@ -176,6 +240,7 @@ class ElementMixin(object):
"image_id": image_id,
"polygon": element["polygon"],
"worker_version_id": self.worker_version_id,
"confidence": element.get("confidence"),
}
for idx, element in enumerate(elements)
]
......@@ -187,19 +252,53 @@ class ElementMixin(object):
def list_element_children(
self,
element,
folder=None,
name=None,
recursive=None,
type=None,
with_classes=None,
with_corpus=None,
with_has_children=None,
with_zone=None,
worker_version=None,
):
element: Union[Element, CachedElement],
folder: Optional[bool] = None,
name: Optional[str] = None,
recursive: Optional[bool] = None,
type: Optional[str] = None,
with_classes: Optional[bool] = None,
with_corpus: Optional[bool] = None,
with_has_children: Optional[bool] = None,
with_zone: Optional[bool] = None,
worker_version: Optional[Union[str, bool]] = None,
) -> Union[Iterable[dict], Iterable[CachedElement]]:
"""
List children of an element
List children of an element.
:param element: Parent element to find children of.
:type element: Union[Element, CachedElement]
:param folder: Restrict to or exclude elements with folder types.
This parameter is not supported when caching is enabled.
:type folder: Optional[bool]
:param name: Restrict to elements whose name contain a substring (case-insensitive).
This parameter is not supported when caching is enabled.
:type name: Optional[str]
:param recursive: Look for elements recursively (grand-children, etc.)
This parameter is not supported when caching is enabled.
:type recursive: Optional[bool]
:param type: Restrict to elements with a specific type slug
This parameter is not supported when caching is enabled.
:type type: Optional[str]
:param with_classes: Include each element's classifications in the response.
This parameter is not supported when caching is enabled.
:type with_classes: Optional[bool]
:param with_corpus: Include each element's corpus in the response.
This parameter is not supported when caching is enabled.
:type with_corpus: Optional[bool]
:param with_has_children: Include the ``has_children`` attribute in the response,
indicating if this element has child elements of its own.
This parameter is not supported when caching is enabled.
:type with_has_children: Optional[bool]
:param with_zone: Include the ``zone`` attribute in the response,
holding the element's image and polygon.
This parameter is not supported when caching is enabled.
:type with_zone: Optional[bool]
:param worker_version: Restrict to elements created by a worker version with this UUID.
:type worker_version: Optional[Union[str, bool]]
:return: An iterable of dicts from the ``ListElementChildren`` API endpoint,
or an iterable of :class:`CachedElement` when caching is enabled.
:rtype: Union[Iterable[dict], Iterable[CachedElement]]
"""
assert element and isinstance(
element, (Element, CachedElement)
......@@ -231,10 +330,14 @@ class ElementMixin(object):
if with_zone is not None:
assert isinstance(with_zone, bool), "with_zone should be of type bool"
query_params["with_zone"] = with_zone
if worker_version:
if worker_version is not None:
assert isinstance(
worker_version, str
), "worker_version should be of type str"
worker_version, (str, bool)
), "worker_version should be of type str or bool"
if isinstance(worker_version, bool):
assert (
worker_version is False
), "if of type bool, worker_version can only be set to False"
query_params["worker_version"] = worker_version
if self.use_cache:
......@@ -247,8 +350,12 @@ class ElementMixin(object):
query = CachedElement.select().where(CachedElement.parent_id == element.id)
if type:
query = query.where(CachedElement.type == type)
if worker_version:
query = query.where(CachedElement.worker_version_id == worker_version)
if worker_version is not None:
# If worker_version=False, filter by manual worker_version e.g. None
worker_version_id = worker_version if worker_version else None
query = query.where(
CachedElement.worker_version_id == worker_version_id
)
return query
else:
......
arkindex_worker/worker/entity.py
View file @ f67d28ff
# -*- coding: utf-8 -*-
"""
ElementsWorker methods for entities.
"""
import os
from enum import Enum
......@@ -10,6 +14,10 @@ from arkindex_worker.models import Element
class EntityType(Enum):
"""
Type of an entity.
"""
Person = "person"
Location = "location"
Subject = "subject"
......@@ -20,12 +28,25 @@ class EntityType(Enum):
class EntityMixin(object):
"""
Mixin for the :class:`ElementsWorker` to add ``Entity`` helpers.
"""
def create_entity(
self, element, name, type, corpus=None, metas=dict(), validated=None
):
"""
Create an entity on the given corpus through API
Return the ID of the created entity
Create an entity on the given corpus.
If cache support is enabled, a :class:`CachedEntity` will also be created.
:param element: An element on which the entity will be reported with the :class:`Reporter`.
This does not have any effect on the entity itself.
:type element: Element or CachedElement
:param name str: Name of the entity.
:param type EntityType: Type of the entity.
:param corpus: UUID of the corpus to create an entity on, or None to use the
value of the ``ARKINDEX_CORPUS_ID`` environment variable.
:type corpus: str or None
"""
if corpus is None:
corpus = os.environ.get("ARKINDEX_CORPUS_ID")
......@@ -86,7 +107,19 @@ class EntityMixin(object):
self, transcription, entity, offset, length, confidence=None
):
"""
Create a link between an existing entity and an existing transcription through API
Create a link between an existing entity and an existing transcription.
If cache support is enabled, a :class:`CachedTranscriptionEntity` will also be created.
:param transcription str: UUID of the existing transcription.
:param entity str: UUID of the existing entity.
:param offset int: Starting position of the entity in the transcription's text,
as a 0-based index.
:param length int: Length of the entity in the transcription's text.
:param confidence: Optional confidence score between 0 or 1.
:type confidence: float or None
:returns: A dict as returned by the ``CreateTranscriptionEntity`` API endpoint,
or None if the worker is in read-only mode.
:rtype: dict(str, str or int) or None
"""
assert transcription and isinstance(
transcription, str
......
arkindex_worker/worker/metadata.py
View file @ f67d28ff
# -*- coding: utf-8 -*-
"""
ElementsWorker methods for metadata.
"""
from enum import Enum
from arkindex_worker import logger
......@@ -6,20 +10,67 @@ from arkindex_worker.models import Element
class MetaType(Enum):
"""
Type of a metadata.
"""
Text = "text"
"""
A regular string with no special interpretation.
"""
HTML = "html"
"""
A metadata with a string value that should be interpreted as HTML content.
The allowed HTML tags are restricted for security reasons.
"""
Date = "date"
"""
A metadata with a string value that should be interpreted as a date.
The date should be formatted as an ISO 8601 date (``YYYY-MM-DD``).
"""
Location = "location"
# Element's original structure reference (intended to be indexed)
"""
A metadata with a string value that should be interpreted as a location.
"""
Reference = "reference"
"""
A metadata with a string value that should be interpreted as an external identifier
to this element, for example to preserve a link to the original data before it was
imported into Arkindex.
"""
Numeric = "numeric"
"""
A metadata with a floating point value.
"""
URL = "url"
"""
A metadata with a string value that should be interpreted as an URL.
Only the ``http`` and ``https`` schemes are allowed.
"""
class MetaDataMixin(object):
"""
Mixin for the :class:`ElementsWorker` to add ``MetaData`` helpers.
"""
def create_metadata(self, element, type, name, value, entity=None):
"""
Create a metadata on the given element through API
Create a metadata on the given element through API.
:param element Element: The element to create a metadata on.
:param type MetaType: Type of the metadata.
:param name str: Name of the metadata.
:param value str: Value of the metadata.
:param entity: UUID of an entity this metadata is related to.
:type entity: str or None
:returns str: UUID of the created metadata.
"""
assert element and isinstance(
element, Element
......
arkindex_worker/worker/training.py 0 → 100644
View file @ f67d28ff
# -*- coding: utf-8 -*-
import hashlib
import os
import tarfile
import tempfile
from contextlib import contextmanager
from typing import NewType, Tuple
import requests
import zstandard as zstd
from apistar.exceptions import ErrorResponse
from arkindex_worker import logger
CHUNK_SIZE = 1024
DirPath = NewType("DirPath", str)
Hash = NewType("Hash", str)
FileSize = NewType("FileSize", int)
Archive = Tuple[DirPath, Hash, FileSize]
@contextmanager
def create_archive(path: DirPath) -> Archive:
"""First create a tar archive, then compress to a zst archive.
Finally, get its hash and size
"""
assert path.is_dir(), "create_archive needs a directory"
compressor = zstd.ZstdCompressor(level=3)
content_hasher = hashlib.md5()
archive_hasher = hashlib.md5()
# Remove extension from the model filename
_, path_to_tar_archive = tempfile.mkstemp(prefix="teklia-", suffix=".tar")
# Create an uncompressed tar archive with all the needed files
# Files hierarchy ifs kept in the archive.
file_list = []
with tarfile.open(path_to_tar_archive, "w") as tar:
for p in path.glob("**/*"):
x = p.relative_to(path)
tar.add(p, arcname=x, recursive=False)
file_list.append(p)
# Sort by path
file_list.sort()
# Compute hash of the files
for file_path in file_list:
with open(file_path, "rb") as file_data:
for chunk in iter(lambda: file_data.read(CHUNK_SIZE), b""):
content_hasher.update(chunk)
_, path_to_zst_archive = tempfile.mkstemp(prefix="teklia-", suffix=".tar.zst")
# Compress the archive
with open(path_to_zst_archive, "wb") as archive_file:
with open(path_to_tar_archive, "rb") as model_data:
for model_chunk in iter(lambda: model_data.read(CHUNK_SIZE), b""):
compressed_chunk = compressor.compress(model_chunk)
archive_hasher.update(compressed_chunk)
archive_file.write(compressed_chunk)
# Remove the tar archive
os.remove(path_to_tar_archive)
# Get content hash, archive size and hash
hash = content_hasher.hexdigest()
size = os.path.getsize(path_to_zst_archive)
archive_hash = archive_hasher.hexdigest()
yield path_to_zst_archive, hash, size, archive_hash
# Remove the zstd archive
os.remove(path_to_zst_archive)
class TrainingMixin(object):
"""
Mixin for the Training workers to add Model and ModelVersion helpers
"""
def publish_model_version(self, model_path: DirPath, model_id: str):
"""
This method creates a model archive and its associated hash,
to create a unique version that will be stored on a bucket and published on arkindex.
"""
# Create the zst archive, get its hash and size
with create_archive(path=model_path) as (
path_to_archive,
hash,
size,
archive_hash,
):
# Create a new model version with hash and size
model_version_details = self.create_model_version(
model_id=model_id,
hash=hash,
size=size,
archive_hash=archive_hash,
)
if model_version_details is None:
return
self.upload_to_s3(
archive_path=path_to_archive,
model_version_details=model_version_details,
)
# Update the model version with state, configuration parsed, tag, description (defaults to name of the worker)
self.update_model_version(
model_version_details=model_version_details,
)
def create_model_version(
self,
model_id: str,
hash: str,
size: int,
archive_hash: str,
) -> dict:
"""
Create a new version of the specified model with the given information (hashes and size).
If a version matching the information already exist, there are two cases:
- The version is in `Created` state: this version's details is used
- The version is in `Available` state: you cannot create twice the same version, an error is raised
"""
# Create a new model version with hash and size
try:
model_version_details = self.request(
"CreateModelVersion",
id=model_id,
body={"hash": hash, "size": size, "archive_hash": archive_hash},
)
except ErrorResponse as e:
if e.status_code >= 500:
logger.error(f"Failed to create model version: {e.content}")
model_version_details = e.content.get("hash")
# If the existing model is in Created state, this model is returned as a dict.
# Else an error in a list is returned.
if model_version_details and isinstance(
model_version_details, (list, tuple)
):
logger.error(model_version_details[0])
return
return model_version_details
def upload_to_s3(self, archive_path: str, model_version_details: dict) -> None:
"""
Upload the archive of the model's files to an Amazon s3 compatible storage
"""
s3_put_url = model_version_details.get("s3_put_url")
logger.info("Uploading to s3...")
# Upload the archive on s3
with open(archive_path, "rb") as archive:
r = requests.put(
url=s3_put_url,
data=archive,
headers={"Content-Type": "application/zstd"},
)
r.raise_for_status()
def update_model_version(
self,
model_version_details: dict,
description: str = None,
configuration: dict = None,
tag: str = None,
) -> None:
"""
Update the specified model version to the state `Available` and use the given information"
"""
logger.info("Updating the model version...")
try:
self.request(
"UpdateModelVersion",
id=model_version_details.get("id"),
body={
"state": "available",
"description": description,
"configuration": configuration,
"tag": tag,
},
)
except ErrorResponse as e:
logger.error(f"Failed to update model version: {e.content}")
arkindex_worker/worker/transcription.py
View file @ f67d28ff
# -*- coding: utf-8 -*-
"""
ElementsWorker methods for transcriptions.
"""
from enum import Enum
from typing import Iterable, Optional, Union
from peewee import IntegrityError
......@@ -10,13 +14,37 @@ from arkindex_worker.models import Element
class TextOrientation(Enum):
"""
Orientation of a transcription's text.
"""
HorizontalLeftToRight = "horizontal-lr"
"""
The text is read from top to bottom then left to right.
This is the default when no orientation is specified.
"""
HorizontalRightToLeft = "horizontal-rl"
"""
The text is read from top to bottom then right to left.
"""
VerticalRightToLeft = "vertical-rl"
"""
The text is read from right to left then top to bottom.
"""
VerticalLeftToRight = "vertical-lr"
"""
The text is read from left to right then top to bottom.
"""
class TranscriptionMixin(object):
"""
Mixin for the :class:`ElementsWorker` to provide ``Transcription`` helpers.
"""
def create_transcription(
self,
element,
......@@ -26,6 +54,15 @@ class TranscriptionMixin(object):
):
"""
Create a transcription on the given element through the API.
:param element: Element to create a transcription on.
:type element: Element or CachedElement
:param text str: Text of the transcription.
:param confidence float: Confidence score, between 0 and 1.
:param orientation TextOrientation: Orientation of the transcription's text.
:returns: A dict as returned by the ``CreateTranscription`` API endpoint,
or None if the worker is in read-only mode.
:rtype: Dict[str, Union[str, float]] or None
"""
assert element and isinstance(
element, (Element, CachedElement)
......@@ -82,7 +119,21 @@ class TranscriptionMixin(object):
def create_transcriptions(self, transcriptions):
"""
Create multiple transcriptions at once on existing elements through the API.
Create multiple transcriptions at once on existing elements through the API,
and creates :class:`CachedTranscription` instances if cache support is enabled.
:param transcriptions: A list of dicts representing a transcription each, with the following keys:
element_id (str)
Required. UUID of the element to create this transcription on.
text (str)
Required. Text of the transcription.
confidence (float)
Required. Confidence score between 0 and 1.
orientation (:class:`TextOrientation`)
Optional. Orientation of the transcription's text.
:returns: A list of dicts as returned in the ``transcriptions`` field by the ``CreateTranscriptions`` API endpoint.
"""
assert transcriptions and isinstance(
......@@ -154,7 +205,23 @@ class TranscriptionMixin(object):
def create_element_transcriptions(self, element, sub_element_type, transcriptions):
"""
Create multiple sub elements with their transcriptions on the given element through API
Create multiple elements and transcriptions at once on a single parent element through the API.
:param element: Element to create elements and transcriptions on.
:type element: Element or CachedElement
:param str sub_element_type: Slug of the element type to use for the new elements.
:param transcriptions: A list of dicts representing an element and transcription each, with the following keys:
polygon (list(list(int or float)))
Required. Polygon of the element.
text (str)
Required. Text of the transcription.
confidence (float)
Required. Confidence score between 0 and 1.
orientation (:class:`TextOrientation`)
Optional. Orientation of the transcription's text.
:returns: A list of dicts as returned by the ``CreateElementTranscriptions`` API endpoint.
"""
assert element and isinstance(
element, (Element, CachedElement)
......@@ -282,10 +349,26 @@ class TranscriptionMixin(object):
return annotations
def list_transcriptions(
self, element, element_type=None, recursive=None, worker_version=None
):
self,
element: Union[Element, CachedElement],
element_type: Optional[str] = None,
recursive: Optional[bool] = None,
worker_version: Optional[Union[str, bool]] = None,
) -> Union[Iterable[dict], Iterable[CachedTranscription]]:
"""
List transcriptions on an element
List transcriptions on an element.
:param element: The element to list transcriptions on.
:type element: Element or CachedElement
:param element_type: Restrict to transcriptions whose elements have an element type with this slug.
:type element_type: Optional[str]
:param recursive: Include transcriptions of any descendant of this element, recursively.
:type recursive: Optional[bool]
:param worker_version: Restrict to transcriptions created by a worker version with this UUID. Set to False to look for manually created transcriptions.
:type worker_version: Optional[Union[str, bool]]
:returns: An iterable of dicts representing each transcription,
or an iterable of CachedTranscription when cache support is enabled.
:rtype: Union[Iterable[dict], Iterable[CachedTranscription]]
"""
assert element and isinstance(
element, (Element, CachedElement)
......@@ -297,10 +380,14 @@ class TranscriptionMixin(object):
if recursive is not None:
assert isinstance(recursive, bool), "recursive should be of type bool"
query_params["recursive"] = recursive
if worker_version:
if worker_version is not None:
assert isinstance(
worker_version, str
), "worker_version should be of type str"
worker_version, (str, bool)
), "worker_version should be of type str or bool"
if isinstance(worker_version, bool):
assert (
worker_version is False
), "if of type bool, worker_version can only be set to False"
query_params["worker_version"] = worker_version
if self.use_cache:
......@@ -331,9 +418,11 @@ class TranscriptionMixin(object):
if element_type:
transcriptions = transcriptions.where(cte.c.type == element_type)
if worker_version:
if worker_version is not None:
# If worker_version=False, filter by manual worker_version e.g. None
worker_version_id = worker_version if worker_version else None
transcriptions = transcriptions.where(
CachedTranscription.worker_version_id == worker_version
CachedTranscription.worker_version_id == worker_version_id
)
else:
transcriptions = self.api_client.paginate(
......
arkindex_worker/worker/version.py
View file @ f67d28ff
# -*- coding: utf-8 -*-
"""
ElementsWorker methods for worker versions.
"""
class WorkerVersionMixin(object):
"""
Mixin for the :class:`ElementsWorker` to provide ``WorkerVersion`` helpers.
"""
def get_worker_version(self, worker_version_id: str) -> dict:
"""
Get worker version from cache if possible, otherwise make API request
Retrieve a worker version, using the :class:`ElementsWorker`'s internal cache when possible.
:param str worker_version_id: ID of the worker version to retrieve.
:returns dict: The requested worker version, as returned by the ``RetrieveWorkerVersion`` API endpoint.
"""
if worker_version_id is None:
raise ValueError("No worker version ID")
......@@ -19,11 +29,11 @@ class WorkerVersionMixin(object):
def get_worker_version_slug(self, worker_version_id: str) -> str:
"""
Helper function to get the worker slug from element, classification or transcription.
Gets the worker version slug from cache if possible, otherwise makes an API request.
Returns None if there is no associated worker version.
Retrieve the slug of the worker of a worker version, from a worker version UUID.
Uses a worker version from the internal cache if possible, otherwise makes an API request.
:type worker_version_id: A worker version UUID
:param worker_version_id str: ID of the worker version to find a slug for.
:returns str: Slug of the worker of this worker version.
"""
worker_version = self.get_worker_version(worker_version_id)
return worker_version["worker"]["slug"]
requirements.txt
View file @ f67d28ff
arkindex-client==1.0.8
arkindex-client==1.0.9
peewee==3.14.10
Pillow==9.1.0
Pillow>=9.0
python-gitlab==2.7.1
python-gnupg==0.4.8
sh==1.14.2
shapely==1.8.2
tenacity==8.0.1
zstandard==0.18.0
tests-requirements.txt
View file @ f67d28ff
pytest==7.1.1
pytest-mock==3.7.0
pytest-responses==0.5.0
requests==2.27.1