diff --git a/optuna_dashboard/artifact/_backend.py b/optuna_dashboard/artifact/_backend.py index eea1e988..b83a8a0c 100644 --- a/optuna_dashboard/artifact/_backend.py +++ b/optuna_dashboard/artifact/_backend.py @@ -40,8 +40,8 @@ if TYPE_CHECKING: }, ) - -ARTIFACTS_ATTR_PREFIX = "dashboard:artifacts:" +ARTIFACTS_ATTR_PREFIX = "artifacts:" +DASHBOARD_ARTIFACTS_ATTR_PREFIX = "dashboard:artifacts:" DEFAULT_MIME_TYPE = "application/octet-stream" BaseRequest.MEMFILE_MAX = int( os.environ.get("OPTUNA_DASHBOARD_MEMFILE_MAX", 1024 * 1024 * 128) @@ -81,6 +81,14 @@ def register_artifact_route( @app.post("/api/artifacts//") @json_api_view def upload_artifact_api(study_id: int, trial_id: int) -> dict[str, Any]: + trial = storage.get_trial(trial_id) + if trial is None: + response.status = 400 + return {"reason": "Invalid study_id or trial_id"} + elif trial.state.is_finished(): + response.status = 400 + return {"reason": "The trial is already finished."} + # TODO(c-bata): Use optuna.artifacts.upload_artifact() if artifact_store is None: response.status = 400 # Bad Request @@ -102,14 +110,10 @@ def register_artifact_route( "mimetype": mimetype or DEFAULT_MIME_TYPE, "encoding": encoding, } - attr_key = _artifact_prefix(trial_id=trial_id) + artifact_id - storage.set_study_system_attr(study_id, attr_key, json.dumps(artifact)) + attr_key = ARTIFACTS_ATTR_PREFIX + artifact_id + storage.set_trial_system_attr(trial_id, attr_key, json.dumps(artifact)) response.status = 201 - trial = storage.get_trial(trial_id) - if trial is None: - response.status = 400 - return {"reason": "Invalid study_id or trial_id"} return { "artifact_id": artifact_id, "artifacts": list_trial_artifacts(storage.get_study_system_attrs(study_id), trial), @@ -123,8 +127,14 @@ def register_artifact_route( return {"reason": "Cannot access to the artifacts."} artifact_store.remove(artifact_id) - attr_key = _artifact_prefix(trial_id) + artifact_id - storage.set_study_system_attr(study_id, attr_key, json.dumps(None)) + # The artifact's metadata is stored in one of the following two locations: + storage.set_study_system_attr( + study_id, _artifact_prefix(trial_id) + artifact_id, json.dumps(None) + ) + storage.set_trial_system_attr( + trial_id, ARTIFACTS_ATTR_PREFIX + artifact_id, json.dumps(None) + ) + response.status = 204 return {} @@ -169,7 +179,6 @@ def upload_artifact( filename = os.path.basename(file_path) storage = trial.storage trial_id = trial._trial_id - study_id = trial.study._study_id artifact_id = str(uuid.uuid4()) guess_mimetype, guess_encoding = mimetypes.guess_type(filename) artifact: ArtifactMeta = { @@ -178,8 +187,8 @@ def upload_artifact( "encoding": encoding or guess_encoding, "filename": filename, } - attr_key = _artifact_prefix(trial_id=trial_id) + artifact_id - storage.set_study_system_attr(study_id, attr_key, json.dumps(artifact)) + attr_key = ARTIFACTS_ATTR_PREFIX + artifact_id + storage.set_trial_system_attr(trial_id, attr_key, json.dumps(artifact)) with open(file_path, "rb") as f: backend.write(artifact_id, f) @@ -187,23 +196,27 @@ def upload_artifact( def _artifact_prefix(trial_id: int) -> str: - return ARTIFACTS_ATTR_PREFIX + f"{trial_id}:" + return DASHBOARD_ARTIFACTS_ATTR_PREFIX + f"{trial_id}:" def get_artifact_meta( storage: BaseStorage, study_id: int, trial_id: int, artifact_id: str ) -> Optional[ArtifactMeta]: - study_system_attr = storage.get_study_system_attrs(study_id) + # Search study_system_attrs due to backward compatibility. + study_system_attrs = storage.get_study_system_attrs(study_id) attr_key = _artifact_prefix(trial_id=trial_id) + artifact_id - artifact_meta = study_system_attr.get(attr_key) + artifact_meta = study_system_attrs.get(attr_key) if artifact_meta is not None: return json.loads(artifact_meta) + # Search trial_system_attrs. Note that artifacts uploaded via optuna.artifacts.upload_artifact + # have a different trial_system_attrs key prefix. # See https://github.com/optuna/optuna/blob/f827582a8/optuna/artifacts/_upload.py#L71 trial_system_attrs = storage.get_trial_system_attrs(trial_id) - value = trial_system_attrs.get("artifacts:" + artifact_id) + value = trial_system_attrs.get(ARTIFACTS_ATTR_PREFIX + artifact_id) if value is not None: return json.loads(value) + return None @@ -221,18 +234,20 @@ def delete_all_artifacts(backend: ArtifactStore, storage: BaseStorage, study_id: def list_trial_artifacts( study_system_attrs: dict[str, Any], trial: FrozenTrial ) -> list[ArtifactMeta]: + # Collect ArtifactMeta from study_system_attrs due to backward compatibility. dashboard_artifact_metas = [ json.loads(value) for key, value in study_system_attrs.items() if key.startswith(_artifact_prefix(trial._trial_id)) ] + # Collect ArtifactMeta from trial_system_attrs. Note that artifacts uploaded via + # optuna.artifacts.upload_artifacts have a different trial_system_attrs key prefix. # See https://github.com/optuna/optuna/blob/f827582a8/optuna/artifacts/_upload.py#L16 optuna_artifact_metas = [ json.loads(value) for key, value in trial.system_attrs.items() - if key.startswith("artifacts:") + if key.startswith(ARTIFACTS_ATTR_PREFIX) ] - artifact_metas = dashboard_artifact_metas + optuna_artifact_metas return [a for a in artifact_metas if a is not None] diff --git a/optuna_dashboard/ts/components/TrialList.tsx b/optuna_dashboard/ts/components/TrialList.tsx index eb700c13..4a54331d 100644 --- a/optuna_dashboard/ts/components/TrialList.tsx +++ b/optuna_dashboard/ts/components/TrialList.tsx @@ -710,55 +710,57 @@ const TrialArtifact: FC<{ trial: Trial }> = ({ trial }) => { ) } })} - - - - - - Upload a New File - - Drag your file here or click to browse. - - - - + + + Upload a New File + + Drag your file here or click to browse. + + + + + ) : null} {renderDeleteArtifactDialog()} diff --git a/python_tests/artifact/test_backend.py b/python_tests/artifact/test_backend.py new file mode 100644 index 00000000..e2d0f152 --- /dev/null +++ b/python_tests/artifact/test_backend.py @@ -0,0 +1,82 @@ +from unittest.mock import MagicMock + +from optuna.storages import BaseStorage +from optuna_dashboard.artifact import _backend +import pytest + + +def test_get_artifact_path() -> None: + study = MagicMock(_study_id=0) + trial = MagicMock(_trial_id=0, study=study) + assert _backend.get_artifact_path(trial=trial, artifact_id="id0") == "/artifacts/0/0/id0" + + +def test_artifact_prefix() -> None: + actual = _backend._artifact_prefix(trial_id=0) + assert actual == "dashboard:artifacts:0:" + + +@pytest.fixture() +def init_storage_with_artifact_meta() -> BaseStorage: + from optuna import create_study + from optuna.storages import InMemoryStorage + + storage = InMemoryStorage() + study = create_study(storage=storage) + + study_system_attrs = { + "dashboard:artifacts:0:id0": '{"artifact_id": "id0", "filename": "foo.txt"}', + "dashboard:artifacts:0:id1": '{"artifact_id": "id1", "filename": "bar.txt"}', + "baz": "baz", + } + for key, value in study_system_attrs.items(): + study.set_system_attr(key, value) + + trial_system_attrs = { + "artifacts:id2": '{"artifact_id": "id2", "filename": "baz.txt"}', + "artifacts:id3": '{"artifact_id": "id3", "filename": "qux.txt"}', + } + for key, value in trial_system_attrs.items(): + trial = study.ask() + trial.set_system_attr(key, value) + study.tell(trial, 0.0) + + return storage + + +def test_get_artifact_meta(init_storage_with_artifact_meta: MagicMock) -> None: + storage = init_storage_with_artifact_meta + + actual = _backend.get_artifact_meta(storage, study_id=0, trial_id=0, artifact_id="id0") + assert actual == {"artifact_id": "id0", "filename": "foo.txt"} + + actual = _backend.get_artifact_meta(storage, study_id=0, trial_id=1, artifact_id="id3") + assert actual == {"artifact_id": "id3", "filename": "qux.txt"} + + actual = _backend.get_artifact_meta(storage, study_id=0, trial_id=0, artifact_id="id4") + assert actual is None + + +def test_delete_all_artifacts(init_storage_with_artifact_meta: MagicMock) -> None: + backend = MagicMock() + storage = init_storage_with_artifact_meta + _backend.delete_all_artifacts(backend, storage, study_id=0) + + assert backend.remove.call_args_list == [ + (("id0",),), + (("id1",),), + (("id2",),), + (("id3",),), + ] + + +def test_list_trial_artifacts(init_storage_with_artifact_meta: MagicMock) -> None: + storage = init_storage_with_artifact_meta + trial = MagicMock(_trial_id=0, system_attrs=storage.get_trial_system_attrs(0)) + + actual = _backend.list_trial_artifacts(storage.get_study_system_attrs(0), trial) + assert actual == [ + {"artifact_id": "id0", "filename": "foo.txt"}, + {"artifact_id": "id1", "filename": "bar.txt"}, + {"artifact_id": "id2", "filename": "baz.txt"}, + ]