diff --git a/optuna_dashboard/_app.py b/optuna_dashboard/_app.py index b91ce5f0..ead27188 100644 --- a/optuna_dashboard/_app.py +++ b/optuna_dashboard/_app.py @@ -1,8 +1,12 @@ from __future__ import annotations +import csv import functools +import io +from itertools import chain import logging import os +import re import typing from typing import Any from typing import Optional @@ -449,6 +453,48 @@ def create_app( response.status = 204 # No content return {} + @app.get("/csv/") + def download_csv(study_id: int) -> BottleViewReturn: + # Create a CSV file + try: + study_name = storage.get_study_name_from_id(study_id) + study = optuna.load_study(storage=storage, study_name=study_name) + except KeyError: + response.status = 404 # Not found + return {"reason": f"study_id={study_id} is not found"} + trials = study.trials + param_names = sorted(set(chain.from_iterable([t.params.keys() for t in trials]))) + user_attr_names = sorted(set(chain.from_iterable([t.user_attrs.keys() for t in trials]))) + param_names_header = [f"Param {x}" for x in param_names] + user_attr_names_header = [f"UserAttribute {x}" for x in user_attr_names] + n_objs = len(study.directions) + if study.metric_names is not None: + value_header = study.metric_names + else: + value_header = ["Value"] if n_objs == 1 else [f"Objective {x}" for x in range(n_objs)] + column_names = ( + ["Number", "State"] + value_header + param_names_header + user_attr_names_header + ) + + buf = io.StringIO("") + writer = csv.writer(buf) + writer.writerow(column_names) + for frozen_trial in trials: + row = [frozen_trial.number, frozen_trial.state.name] + row.extend(frozen_trial.values if frozen_trial.values is not None else [None] * n_objs) + row.extend([frozen_trial.params.get(name, None) for name in param_names]) + row.extend([frozen_trial.user_attrs.get(name, None) for name in user_attr_names]) + writer.writerow(row) + + # Set response headers + output_name = "-".join(re.sub(r'[\\/:*?"<>|]+', "", study_name).split(" ")) + response.headers["Content-Type"] = "text/csv; chatset=cp932" + response.headers["Content-Disposition"] = f"attachment; filename={output_name}.csv" + + # Response body + buf.seek(0) + return buf.read() + @app.get("/favicon.ico") def favicon() -> BottleViewReturn: use_gzip = "gzip" in request.headers["Accept-Encoding"] diff --git a/optuna_dashboard/ts/components/StudyDetail.tsx b/optuna_dashboard/ts/components/StudyDetail.tsx index bfdeb577..79702f2e 100644 --- a/optuna_dashboard/ts/components/StudyDetail.tsx +++ b/optuna_dashboard/ts/components/StudyDetail.tsx @@ -12,6 +12,7 @@ import { import Grid2 from "@mui/material/Unstable_Grid2" import ChevronRightIcon from "@mui/icons-material/ChevronRight" import HomeIcon from "@mui/icons-material/Home" +import DownloadIcon from "@mui/icons-material/Download" import { StudyNote } from "./Note" import { actionCreator } from "../action" @@ -149,11 +150,39 @@ export const StudyDetail: FC<{ content = } else if (page === "trialTable") { content = ( - - - - - + + + + + + + Download CSV File + + + + + + + + + + ) } else if (page === "note" && studyDetail !== null) { content = ( diff --git a/python_tests/test_csv_download.py b/python_tests/test_csv_download.py new file mode 100644 index 00000000..019f18ba --- /dev/null +++ b/python_tests/test_csv_download.py @@ -0,0 +1,109 @@ +from __future__ import annotations + +from typing import Any + +import optuna +from optuna.trial import TrialState +from optuna_dashboard._app import create_app +import pytest + +from .wsgi_client import send_request + + +def _validate_output( + storage: optuna.storages.BaseStorage, + correct_status: int, + study_id: int, + expect_no_result: bool = False, + extra_col_names: list[str] | None = None, +) -> None: + app = create_app(storage) + status, _, body = send_request( + app, + f"/csv/{study_id}", + "GET", + content_type="application/json", + ) + assert status == correct_status + decoded_csv = str(body.decode("utf-8")) + if expect_no_result: + assert "is not found" in decoded_csv + else: + col_names = ["Number", "State"] + ([] if extra_col_names is None else extra_col_names) + assert all(col_name in decoded_csv for col_name in col_names) + + +def test_download_csv_no_trial() -> None: + def objective(trial: optuna.Trial) -> float: + x = trial.suggest_float("x", -100, 100) + y = trial.suggest_categorical("y", [-1, 0, 1]) + return x**2 + y + + storage = optuna.storages.InMemoryStorage() + study = optuna.create_study(storage=storage) + study.optimize(objective, n_trials=0) + _validate_output(storage, 200, 0) + + +def test_download_csv_all_waiting() -> None: + storage = optuna.storages.InMemoryStorage() + study = optuna.create_study(storage=storage) + study.add_trial(optuna.trial.create_trial(state=TrialState.WAITING)) + _validate_output(storage, 200, 0) + + +def test_download_csv_all_running() -> None: + storage = optuna.storages.InMemoryStorage() + study = optuna.create_study(storage=storage) + study.add_trial(optuna.trial.create_trial(state=TrialState.RUNNING)) + _validate_output(storage, 200, 0) + + +@pytest.mark.parametrize("study_id", [0, 1]) +def test_download_csv_fail(study_id: int) -> None: + def objective(trial: optuna.Trial) -> float: + x = trial.suggest_float("x", -100, 100) + y = trial.suggest_categorical("y", [-1, 0, 1]) + return x**2 + y + + storage = optuna.storages.InMemoryStorage() + study = optuna.create_study(storage=storage) + optuna.logging.set_verbosity(optuna.logging.ERROR) + study.optimize(objective, n_trials=10) + expect_no_result = study_id != 0 + cols = ["Param x", "Param y", "Value"] + _validate_output(storage, 404 if expect_no_result else 200, study_id, expect_no_result, cols) + + +@pytest.mark.parametrize("is_multi_obj", [True, False]) +def test_download_csv_multi_obj(is_multi_obj: bool) -> None: + def objective(trial: optuna.Trial) -> Any: + x = trial.suggest_float("x", -100, 100) + y = trial.suggest_categorical("y", [-1, 0, 1]) + if is_multi_obj: + return x**2, y + return x**2 + y + + storage = optuna.storages.InMemoryStorage() + directions = ["minimize", "minimize"] if is_multi_obj else ["minimize"] + study = optuna.create_study(storage=storage, directions=directions) + optuna.logging.set_verbosity(optuna.logging.ERROR) + study.optimize(objective, n_trials=10) + cols = ["Param x", "Param y"] + cols += ["Objective 0", "Objective 1"] if is_multi_obj else ["Value"] + _validate_output(storage, 200, 0, extra_col_names=cols) + + +def test_download_csv_user_attr() -> None: + def objective(trial: optuna.Trial) -> float: + x = trial.suggest_float("x", -100, 100) + y = trial.suggest_categorical("y", [-1, 0, 1]) + trial.set_user_attr("abs_y", abs(y)) + return x**2 + y + + storage = optuna.storages.InMemoryStorage() + study = optuna.create_study(storage=storage) + optuna.logging.set_verbosity(optuna.logging.ERROR) + study.optimize(objective, n_trials=10) + cols = ["Param x", "Param y", "Value", "UserAttribute abs_y"] + _validate_output(storage, 200, 0, extra_col_names=cols)