mirror of
https://github.com/wassname/optuna-dashboard.git
synced 2026-08-21 11:18:45 +08:00
89 lines
3.0 KiB
Python
89 lines
3.0 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import tempfile
|
|
|
|
import optuna
|
|
from optuna.version import __version__ as optuna_ver
|
|
from optuna_dashboard.artifact._backend import delete_all_artifacts
|
|
from optuna_dashboard.artifact._backend import get_trial_artifact_meta
|
|
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"),
|
|
reason="Artifact is not implemented yet in Optuna",
|
|
)
|
|
def test_list_optuna_trial_artifacts() -> None:
|
|
from optuna.artifacts import FileSystemArtifactStore
|
|
from optuna.artifacts import upload_artifact
|
|
|
|
storage = optuna.storages.InMemoryStorage()
|
|
study = optuna.create_study(storage=storage)
|
|
dummy_content = b"dummy content"
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
artifact_store = FileSystemArtifactStore(tmpdir)
|
|
trial = study.ask()
|
|
|
|
with tempfile.NamedTemporaryFile() as f:
|
|
f.write(dummy_content)
|
|
f.flush()
|
|
upload_artifact(study_or_trial=trial, file_path=f.name, artifact_store=artifact_store)
|
|
|
|
study.tell(trial, 0.0)
|
|
|
|
study_system_attrs = storage.get_study_system_attrs(study._study_id)
|
|
trial_system_attrs = storage.get_trial_system_attrs(trial._trial_id)
|
|
frozen_trial = storage.get_trial(trial._trial_id)
|
|
artifact_meta_list = list_trial_artifacts(
|
|
study_system_attrs, trial_system_attrs, frozen_trial
|
|
)
|
|
assert len(artifact_meta_list) == 1
|
|
|
|
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_trial_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"),
|
|
reason="Artifact is not implemented yet in Optuna",
|
|
)
|
|
def test_delete_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(
|
|
study_or_trial=trial, file_path=dummy_file_path, artifact_store=artifact_store
|
|
)
|
|
return 0.0
|
|
|
|
study.optimize(objective, n_trials=10)
|
|
assert len(os.listdir(tmpdir)) == 11 # 10 artifacts + dummy.txt
|
|
|
|
delete_all_artifacts(artifact_store, storage, study._study_id)
|
|
assert len(os.listdir(tmpdir)) == 1 # dummy.txt only
|