from __future__ import annotations
import os
from pathlib import Path
import shutil
from typing import TYPE_CHECKING
from optuna.artifacts.exceptions import ArtifactNotFound
if TYPE_CHECKING:
from typing import BinaryIO
[docs]
class FileSystemArtifactStore:
"""An artifact store for file systems.
Args:
base_path:
The base path to a directory to store artifacts.
Example:
.. code-block:: python
import os
import optuna
from optuna.artifacts import FileSystemArtifactStore
from optuna.artifacts import upload_artifact
base_path = "./artifacts"
os.makedirs(base_path, exist_ok=True)
artifact_store = FileSystemArtifactStore(base_path=base_path)
def objective(trial: optuna.Trial) -> float:
... = trial.suggest_float("x", -10, 10)
file_path = generate_example(...)
upload_artifact(
artifact_store=artifact_store,
file_path=file_path,
study_or_trial=trial,
)
return ...
"""
def __init__(self, base_path: str | Path) -> None:
if isinstance(base_path, str):
base_path = Path(base_path)
# TODO(Shinichi): Check if the base_path is valid directory.
self._base_path = base_path
def open_reader(self, artifact_id: str) -> BinaryIO:
filepath = os.path.join(self._base_path, artifact_id)
try:
f = open(filepath, "rb")
except FileNotFoundError as e:
raise ArtifactNotFound("not found") from e
return f
def write(self, artifact_id: str, content_body: BinaryIO) -> None:
filepath = os.path.join(self._base_path, artifact_id)
with open(filepath, "wb") as f:
shutil.copyfileobj(content_body, f)
def remove(self, artifact_id: str) -> None:
filepath = os.path.join(self._base_path, artifact_id)
try:
os.remove(filepath)
except FileNotFoundError as e:
raise ArtifactNotFound("not found") from e
if TYPE_CHECKING:
# A mypy-runtime assertion to ensure that LocalArtifactBackend
# implements all abstract methods in ArtifactBackendProtocol.
from optuna.artifacts._protocol import ArtifactStore
_: ArtifactStore = FileSystemArtifactStore("")