diff --git a/optuna_dashboard/artifact/_backend.py b/optuna_dashboard/artifact/_backend.py index 81cd766d..b8a44abd 100644 --- a/optuna_dashboard/artifact/_backend.py +++ b/optuna_dashboard/artifact/_backend.py @@ -14,6 +14,7 @@ from bottle import HTTPResponse from bottle import request from bottle import response import optuna +from optuna.artifacts.exceptions import ArtifactNotFound from optuna.trial import FrozenTrial from .._bottle_util import json_api_view @@ -184,7 +185,11 @@ def register_artifact_route( if artifact_store is None: response.status = 400 # Bad Request return {"reason": "Cannot access to the artifacts."} - artifact_store.remove(artifact_id) + try: + artifact_store.remove(artifact_id) + except ArtifactNotFound as e: + response.status = 404 + return {"reason": str(e)} # The artifact's metadata is stored in one of the following two locations: storage.set_study_system_attr( @@ -203,7 +208,11 @@ def register_artifact_route( if artifact_store is None: response.status = 400 # Bad Request return {"reason": "Cannot access to the artifacts."} - artifact_store.remove(artifact_id) + try: + artifact_store.remove(artifact_id) + except ArtifactNotFound as e: + response.status = 404 + return {"reason": str(e)} storage.set_study_system_attr( study_id, ARTIFACTS_ATTR_PREFIX + artifact_id, json.dumps(None) diff --git a/python_tests/artifact/test_backend.py b/python_tests/artifact/test_backend.py index 52fe6027..165c5029 100644 --- a/python_tests/artifact/test_backend.py +++ b/python_tests/artifact/test_backend.py @@ -249,3 +249,52 @@ def test_upload_artifact() -> None: with open(f"{tmpdir}/{res['artifact_id']}", "r") as f: data = f.read() assert data == "dummy_content" + + +def test_delete_study_artifact() -> None: + storage = optuna.storages.InMemoryStorage() + study = optuna.create_study(storage=storage) + with tempfile.TemporaryDirectory() as tmpdir: + artifact_store = FileSystemArtifactStore(tmpdir) + with tempfile.NamedTemporaryFile() as f: + f.write(b"dummy_content") + f.flush() + artifact_id = upload_artifact(study, f.name, artifact_store=artifact_store) + app = create_app(storage, artifact_store) + status, _, _ = send_request( + app, + f"/api/artifacts/{study._study_id}/{artifact_id}", + "DELETE", + ) + assert status == 204 + status, _, _ = send_request( + app, + f"/api/artifacts/{study._study_id}/{artifact_id}", + "DELETE", + ) + assert status == 404 + + +def test_delete_trial_artifact() -> None: + storage = optuna.storages.InMemoryStorage() + study = optuna.create_study(storage=storage) + trial = study.ask() + with tempfile.TemporaryDirectory() as tmpdir: + artifact_store = FileSystemArtifactStore(tmpdir) + with tempfile.NamedTemporaryFile() as f: + f.write(b"dummy_content") + f.flush() + artifact_id = upload_artifact(trial, f.name, artifact_store=artifact_store) + app = create_app(storage, artifact_store) + status, _, _ = send_request( + app, + f"/api/artifacts/{study._study_id}/{trial._trial_id}/{artifact_id}", + "DELETE", + ) + assert status == 204 + status, _, _ = send_request( + app, + f"/api/artifacts/{study._study_id}/{trial._trial_id}/{artifact_id}", + "DELETE", + ) + assert status == 404