Add exponential backoff middleware for ArtifactBackend

python_tests/artifact/stubs.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+import copy
+import io
+import shutil
+import threading
+from typing import TYPE_CHECKING
+
+from optuna_dashboard.artifact.exceptions import ArtifactNotFound
+
+
+if TYPE_CHECKING:
+    from typing import BinaryIO
+
+
+class FailBackend:
+    def open(self, artifact_id: str) -> BinaryIO:
+        raise Exception("something error raised")
+
+    def write(self, artifact_id: str, content_body: BinaryIO) -> None:
+        raise Exception("something error raised")
+
+    def remove(self, artifact_id: str) -> None:
+        raise Exception("something error raised")
+
+
+class InMemoryBackend:
+    def __init__(self) -> None:
+        self._data: dict[str, io.BytesIO] = {}
+        self._lock = threading.Lock()
+
+    def open(self, artifact_id: str) -> BinaryIO:
+        with self._lock:
+            data = self._data.get(artifact_id)
+            if data is None:
+                raise ArtifactNotFound("not found")
+            return copy.deepcopy(data)
+
+    def write(self, artifact_id: str, content_body: BinaryIO) -> None:
+        buf = io.BytesIO()
+        shutil.copyfileobj(content_body, buf)
+        buf.seek(0)
+        with self._lock:
+            self._data[artifact_id] = buf
+
+    def remove(self, artifact_id: str) -> None:
+        with self._lock:
+            if artifact_id not in self._data:
+                raise ArtifactNotFound("not found")
+            del self._data[artifact_id]
+
+
+if TYPE_CHECKING:
+    # A mypy-runtime assertion to ensure that SCSBackend
+    # implements all abstract methods in ArtifactBackendProtocol.
+    from optuna_dashboard.artifact.protocol import ArtifactBackend
+
+    _fail: ArtifactBackend = FailBackend()
+    _inmemory: ArtifactBackend = InMemoryBackend()

python_tests/artifact/test_backoff.py

@@ -0,0 +1,35 @@
+import io
+import uuid
+
+from optuna_dashboard.artifact.backoff import Backoff
+
+from .stubs import FailBackend
+from .stubs import InMemoryBackend
+
+
+def test_backoff_time() -> None:
+    backend = Backoff(
+        backend=FailBackend(),
+        min_delay=0.1,
+        multiplier=10,
+        max_delay=10,
+    )
+    assert backend._get_sleep_secs(0) == 0.1
+    assert backend._get_sleep_secs(1) == 1
+    assert backend._get_sleep_secs(2) == 10
+
+
+def test_read_and_write() -> None:
+    artifact_id = f"test-{uuid.uuid4()}"
+    dummy_content = b"Hello World"
+
+    backend = Backoff(
+        backend=InMemoryBackend(),
+        min_delay=0.1,
+        multiplier=10,
+        max_delay=10,
+    )
+    backend.write(artifact_id, io.BytesIO(dummy_content))
+    with backend.open(artifact_id) as f:
+        actual = f.read()
+    assert actual == dummy_content

python_tests/artifact/test_boto3.py

@@ -0,0 +1,54 @@
+import io
+from unittest import TestCase
+
+import boto3
+from moto import mock_s3
+from optuna_dashboard.artifact.boto3 import Boto3Backend
+from optuna_dashboard.artifact.exceptions import ArtifactNotFound
+
+
+@mock_s3
+class Boto3BackendTestCase(TestCase):
+    def setUp(self) -> None:
+        self.s3_client = boto3.client("s3")
+        self.bucket_name = "moto-bucket"
+        self.s3_client.create_bucket(Bucket=self.bucket_name)
+
+    def tearDown(self) -> None:
+        objects = self.s3_client.list_objects(Bucket=self.bucket_name).get("Contents", [])
+        if objects:
+            self.s3_client.delete_objects(
+                Bucket=self.bucket_name,
+                Delete={"Objects": [{"Key": obj["Key"] for obj in objects}], "Quiet": True},
+            )
+        self.s3_client.delete_bucket(Bucket=self.bucket_name)
+
+    def test_upload_download(self) -> None:
+        artifact_id = "dummy-uuid"
+        dummy_content = b"Hello World"
+
+        backend = Boto3Backend(self.bucket_name)
+        backend.write(artifact_id, io.BytesIO(dummy_content))
+        assert len(self.s3_client.list_objects(Bucket=self.bucket_name)["Contents"]) == 1
+        obj = self.s3_client.get_object(Bucket=self.bucket_name, Key=artifact_id)
+        assert obj["Body"].read() == dummy_content
+
+        with backend.open(artifact_id) as f:
+            actual = f.read()
+        self.assertEqual(actual, dummy_content)
+
+    def test_remove(self) -> None:
+        artifact_id = "dummy-uuid"
+        backend = Boto3Backend(self.bucket_name)
+        backend.write(artifact_id, io.BytesIO(b"Hello"))
+        objects = self.s3_client.list_objects(Bucket=self.bucket_name)["Contents"]
+        assert len([obj for obj in objects if obj["Key"] == artifact_id]) == 1
+
+        backend.remove(artifact_id)
+        objects = self.s3_client.list_objects(Bucket=self.bucket_name).get("Contents", [])
+        assert len([obj for obj in objects if obj["Key"] == artifact_id]) == 0
+
+    def test_file_not_found_exception(self) -> None:
+        backend = Boto3Backend(self.bucket_name)
+        with self.assertRaises(ArtifactNotFound):
+            backend.open("not-found-id")

python_tests/artifact/test_file_system.py

@@ -0,0 +1,30 @@
+import io
+import tempfile
+from unittest import TestCase
+
+from optuna_dashboard.artifact.exceptions import ArtifactNotFound
+from optuna_dashboard.artifact.file_system import FileSystemBackend
+
+
+class FileSystemBackendTestCase(TestCase):
+    def setUp(self) -> None:
+        self.dir = tempfile.TemporaryDirectory()
+
+    def tearDown(self) -> None:
+        self.dir.cleanup()
+
+    def test_upload_download(self) -> None:
+        artifact_id = "dummy-uuid"
+        dummy_content = b"Hello World"
+        backend = FileSystemBackend(self.dir.name)
+        backend.write(artifact_id, io.BytesIO(dummy_content))
+        with backend.open(artifact_id) as f:
+            actual = f.read()
+        self.assertEqual(actual, dummy_content)
+
+    def test_file_not_found(self) -> None:
+        backend = FileSystemBackend(self.dir.name)
+        with self.assertRaises(ArtifactNotFound):
+            backend.open("not-found-id")
+        with self.assertRaises(ArtifactNotFound):
+            backend.remove("not-found-id")