Enhance delete artifacts APIs and add tests for them

This commit is contained in:
y0z
2024-01-15 18:30:45 +09:00
parent 7c4f67302d
commit 944b748ced
2 changed files with 60 additions and 2 deletions
+11 -2
View File
@@ -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)
+49
View File
@@ -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