Add remove_associated_artifacts query param into delete_study api

This commit is contained in:
porink0424
2024-05-31 15:08:05 +09:00
parent cffc86dac7
commit a6c7bfd437
2 changed files with 66 additions and 1 deletions
+4 -1
View File
@@ -180,7 +180,10 @@ def create_app(
@app.delete("/api/studies/<study_id:int>")
@json_api_view
def delete_study(study_id: int) -> dict[str, Any]:
if artifact_store is not None:
remove_associated_artifacts = (
True if request.params.get("remove_associated_artifacts") == "true" else False
)
if artifact_store is not None and remove_associated_artifacts:
delete_all_artifacts(artifact_store, storage, study_id)
try:
+62
View File
@@ -2,10 +2,13 @@ from __future__ import annotations
import json
import sys
import tempfile
from unittest import TestCase
import optuna
from optuna import get_all_study_summaries
from optuna.artifacts import upload_artifact
from optuna.artifacts.exceptions import ArtifactNotFound
from optuna.study import StudyDirection
from optuna_dashboard._app import create_app
from optuna_dashboard._app import create_new_study
@@ -523,6 +526,65 @@ class APITestCase(TestCase):
self.assertEqual(status, 204)
self.assertEqual(len(get_all_study_summaries(storage)), 1)
def test_delete_study_with_removing_artifacts(self) -> None:
storage = optuna.storages.InMemoryStorage()
study = optuna.create_study(storage=storage)
with tempfile.TemporaryDirectory() as tmpdir_name:
artifact_store = optuna.artifacts.FileSystemArtifactStore(base_path=tmpdir_name)
with tempfile.NamedTemporaryFile() as f:
f.write(b"dummy")
f.flush()
artifact_id = upload_artifact(study, f.name, artifact_store)
app = create_app(storage, artifact_store)
with artifact_store.open_reader(artifact_id) as reader:
self.assertEqual(reader.read(), b"dummy")
status, _, _ = send_request(
app,
f"/api/studies/{study._study_id}",
"DELETE",
queries={"remove_associated_artifacts": "true"},
content_type="application/json",
)
self.assertEqual(status, 204)
with self.assertRaises(ArtifactNotFound):
with artifact_store.open_reader(artifact_id) as reader:
reader.read()
self.assertEqual(len(get_all_study_summaries(storage)), 0)
def test_delete_study_without_removing_artifacts(self) -> None:
storage = optuna.storages.InMemoryStorage()
study = optuna.create_study(storage=storage)
with tempfile.TemporaryDirectory() as tmpdir_name:
artifact_store = optuna.artifacts.FileSystemArtifactStore(base_path=tmpdir_name)
with tempfile.NamedTemporaryFile() as f:
f.write(b"dummy")
f.flush()
artifact_id = upload_artifact(study, f.name, artifact_store)
app = create_app(storage, artifact_store)
with artifact_store.open_reader(artifact_id) as reader:
self.assertEqual(reader.read(), b"dummy")
status, _, _ = send_request(
app,
f"/api/studies/{study._study_id}",
"DELETE",
queries={"remove_associated_artifacts": "false"},
content_type="application/json",
)
self.assertEqual(status, 204)
with artifact_store.open_reader(artifact_id) as reader:
self.assertEqual(reader.read(), b"dummy")
self.assertEqual(len(get_all_study_summaries(storage)), 0)
def test_delete_study_not_found(self) -> None:
storage = optuna.storages.InMemoryStorage()
app = create_app(storage)