optuna

1
from __future__ import annotations
2

3
import copy
4
import io
5
import shutil
6
import threading
7
from typing import TYPE_CHECKING
8

9
from optuna.artifacts.exceptions import ArtifactNotFound
10

11

12
if TYPE_CHECKING:
13
    from typing import BinaryIO
14

15

16
class FailArtifactStore:
17
    def open_reader(self, artifact_id: str) -> BinaryIO:
18
        raise Exception("something error raised")
19

20
    def write(self, artifact_id: str, content_body: BinaryIO) -> None:
21
        raise Exception("something error raised")
22

23
    def remove(self, artifact_id: str) -> None:
24
        raise Exception("something error raised")
25

26

27
class InMemoryArtifactStore:
28
    def __init__(self) -> None:
29
        self._data: dict[str, io.BytesIO] = {}
30
        self._lock = threading.Lock()
31

32
    def open_reader(self, artifact_id: str) -> BinaryIO:
33
        with self._lock:
34
            data = self._data.get(artifact_id)
35
            if data is None:
36
                raise ArtifactNotFound("not found")
37
            return copy.deepcopy(data)
38

39
    def write(self, artifact_id: str, content_body: BinaryIO) -> None:
40
        buf = io.BytesIO()
41
        shutil.copyfileobj(content_body, buf)
42
        buf.seek(0)
43
        with self._lock:
44
            self._data[artifact_id] = buf
45

46
    def remove(self, artifact_id: str) -> None:
47
        with self._lock:
48
            if artifact_id not in self._data:
49
                raise ArtifactNotFound("not found")
50
            del self._data[artifact_id]
51

52

53
if TYPE_CHECKING:
54
    from optuna.artifacts._protocol import ArtifactStore
55

56
    _fail: ArtifactStore = FailArtifactStore()
57
    _inmemory: ArtifactStore = InMemoryArtifactStore()
58