Fix artifact endpoint

This commit is contained in:
c-bata
2023-08-04 18:59:57 +09:00
parent 2c3fb7711d
commit 1f9b05e15d
3 changed files with 67 additions and 17 deletions
@@ -1,19 +1,22 @@
from __future__ import annotations
import os
import tempfile
import optuna
from optuna.version import __version__ as optuna_ver
from optuna_dashboard.artifact._backend import get_artifact_meta
from optuna_dashboard.artifact._backend import list_study_artifacts
from optuna_dashboard.artifact._backend import list_trial_artifacts
from packaging import version
import pytest
@pytest.mark.skipif(
version.parse(optuna_ver) < version.Version("3.3.0"),
version.parse(optuna_ver) < version.Version("3.3.0.dev"),
reason="Artifact is not implemented yet in Optuna",
)
def test_list_optuna_artifacts() -> None:
def test_list_optuna_trial_artifacts() -> None:
from optuna.artifacts import FileSystemArtifactStore
from optuna.artifacts import upload_artifact
@@ -40,3 +43,39 @@ def test_list_optuna_artifacts() -> None:
artifact_id = artifact_meta_list[0]["artifact_id"]
with artifact_store.open_reader(artifact_id) as reader:
assert reader.read() == dummy_content
artifact_meta = get_artifact_meta(
storage=storage,
study_id=study._study_id,
trial_id=trial._trial_id,
artifact_id=artifact_id,
)
assert artifact_meta is not None
@pytest.mark.skipif(
version.parse(optuna_ver) < version.Version("3.3.0.dev"),
reason="Artifact is not implemented yet in Optuna",
)
def test_list_optuna_study_artifacts() -> None:
from optuna.artifacts import FileSystemArtifactStore
from optuna.artifacts import upload_artifact
storage = optuna.storages.InMemoryStorage()
study = optuna.create_study(storage=storage)
with tempfile.TemporaryDirectory() as tmpdir:
dummy_file_path = os.path.join(tmpdir, "dummy.txt")
with open(dummy_file_path, "wb") as f:
f.write(b"dummy content")
f.flush()
artifact_store = FileSystemArtifactStore(tmpdir)
def objective(trial: optuna.Trial) -> float:
upload_artifact(trial, dummy_file_path, artifact_store=artifact_store)
return 0.0
study.optimize(objective, n_trials=10)
artifact_meta_list = list_study_artifacts(storage, study_id=study._study_id)
assert len(artifact_meta_list) == 10