mirror of
https://github.com/wassname/optuna-dashboard.git
synced 2026-09-09 11:28:14 +08:00
Add exponential backoff middleware for ArtifactBackend
This commit is contained in:
@@ -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()
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user