Add exponential backoff middleware for ArtifactBackend

This commit is contained in:
c-bata
2023-06-09 11:37:14 +09:00
parent 917f980d89
commit 5a4240c004
7 changed files with 205 additions and 0 deletions
View File
+59
View File
@@ -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()
+35
View File
@@ -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