diff --git a/optuna_dashboard/_app.py b/optuna_dashboard/_app.py index 414b4a52..b791655a 100644 --- a/optuna_dashboard/_app.py +++ b/optuna_dashboard/_app.py @@ -3,21 +3,13 @@ from __future__ import annotations from datetime import datetime from datetime import timedelta import functools -import json import logging import os import threading -import traceback import typing -from typing import Any -from typing import Callable -from typing import cast -from typing import Dict from typing import Optional -from typing import TypeVar from typing import Union -from bottle import BaseResponse from bottle import Bottle from bottle import redirect from bottle import request @@ -37,6 +29,8 @@ from packaging import version from . import _note as note from . import artifact +from ._bottleutil import BottleViewReturn +from ._bottleutil import json_api_view from ._cached_extra_study_property import get_cached_extra_study_property from ._importance import get_param_importance_from_trials_cache from ._pareto_front import get_pareto_front_trials @@ -52,8 +46,6 @@ if typing.TYPE_CHECKING: except ImportError: FrozenStudy = None # type: ignore -BottleViewReturn = Union[str, bytes, Dict[str, Any], BaseResponse] -BottleView = TypeVar("BottleView", bound=Callable[..., BottleViewReturn]) logger = logging.getLogger(__name__) @@ -127,23 +119,6 @@ def update_schema_compatibility_flags(storage: BaseStorage) -> None: rdb_schema_unsupported = current_version not in storage.get_all_versions() -def json_api_view(view: BottleView) -> BottleView: - @functools.wraps(view) - def decorated(*args: list[Any], **kwargs: dict[str, Any]) -> BottleViewReturn: - try: - response.content_type = "application/json" - response_body = view(*args, **kwargs) - return response_body - except Exception as e: - response.status = 500 - response.content_type = "application/json" - stacktrace = "\n".join(traceback.format_tb(e.__traceback__)) - logger.error(f"Exception: {e}\n{stacktrace}") - return json.dumps({"reason": "internal server error"}) - - return cast(BottleView, decorated) - - def get_study_summaries(storage: BaseStorage) -> list[StudySummary]: if version.parse(optuna_ver) >= version.Version("3.0.0rc0.dev"): frozen_studies = storage.get_all_studies() # type: ignore diff --git a/optuna_dashboard/_bottleutil.py b/optuna_dashboard/_bottleutil.py new file mode 100644 index 00000000..8d8423fd --- /dev/null +++ b/optuna_dashboard/_bottleutil.py @@ -0,0 +1,46 @@ +import base64 +import functools +import json +import logging +import traceback +from typing import Any +from typing import Callable +from typing import cast +from typing import Dict +from typing import TypeVar +from typing import Union + +from bottle import BaseResponse +from bottle import response + + +BottleViewReturn = Union[str, bytes, Dict[str, Any], BaseResponse] +BottleView = TypeVar("BottleView", bound=Callable[..., BottleViewReturn]) +logger = logging.getLogger(__name__) + + +def json_api_view(view: BottleView) -> BottleView: + @functools.wraps(view) + def decorated(*args: list[Any], **kwargs: dict[str, Any]) -> BottleViewReturn: + try: + response.content_type = "application/json" + response_body = view(*args, **kwargs) + return response_body + except Exception as e: + response.status = 500 + response.content_type = "application/json" + stacktrace = "\n".join(traceback.format_tb(e.__traceback__)) + logger.error(f"Exception: {e}\n{stacktrace}") + return json.dumps({"reason": "internal server error"}) + + return cast(BottleView, decorated) + + +def parse_data_uri(data_uri: str) -> tuple[str, bytes]: + prefix, a = data_uri.split(":", 1) + if prefix != "data": + raise ValueError("data url must start with 'data:' prefix") + mediatype_with_suffix, base64_data = a.split(",", 1) + mediatype = mediatype_with_suffix.split(";", 1)[0] + data = base64.standard_b64decode(base64_data) + return mediatype, data diff --git a/optuna_dashboard/artifact/__init__.py b/optuna_dashboard/artifact/__init__.py index 8fe81335..07518698 100644 --- a/optuna_dashboard/artifact/__init__.py +++ b/optuna_dashboard/artifact/__init__.py @@ -7,9 +7,14 @@ from typing import TYPE_CHECKING import uuid from bottle import Bottle +from bottle import request from bottle import response import optuna +from .._bottleutil import BottleViewReturn +from .._bottleutil import json_api_view +from .._bottleutil import parse_data_uri + if TYPE_CHECKING: from typing import Any @@ -24,9 +29,9 @@ if TYPE_CHECKING: "ArtifactMeta", { "artifact_id": str, - "mimetype": str, - "encoding": str, "filename": str, + "mimetype": Optional[str], + "encoding": Optional[str], }, ) @@ -43,17 +48,51 @@ def register_artifact_route( response.status = 400 # Bad Request return b"Cannot access to the artifacts." artifact_dict = _get_artifact_meta(storage, trial_id, artifact_id) - response.set_header("Content-Type", artifact_dict["mimetype"]) - response.set_header("Content-Encodings", artifact_dict["encoding"]) + mimetype, encoding = mimetypes.guess_type(artifact_dict["filename"]) + mimetype: str = artifact_dict.get("mimetype") or encoding or "application/octet-stream" + encoding: Optional[str] = artifact_dict.get("encoding") or encoding + response.set_header("Content-Type", mimetype) + if encoding: + response.set_header("Content-Encodings", encoding) + with artifact_backend.open(artifact_id) as f: body = f.read() return body - @app.delete("/artifacts/") - def delete_artifact(artifact_id: str) -> bytes: + @app.post("/api/artifacts//") + @json_api_view + def upload_artifact(trial_id: int) -> BottleViewReturn: if artifact_backend is None: response.status = 400 # Bad Request - return b"Cannot access to the artifacts." + return {"reason": "Cannot access to the artifacts."} + file = request.json.get("file") + if file is None: + response.status = 400 + return {"reason": "Please specify the 'file' key."} + + _, data = parse_data_uri(file) + filename = request.json.get("filename", "") + artifact_id = str(uuid.uuid4()) + with artifact_backend.open(artifact_id=artifact_id) as f: + f.write(data) + + artifact: ArtifactMeta = { + "artifact_id": artifact_id, + "mimetype": None, + "encoding": None, + "filename": filename, + } + attr_key = _artifact_prefix(trial_id=trial_id) + artifact_id + storage.set_study_system_attr(trial_id, attr_key, json.dumps(artifact)) + response.status = 201 + return artifact + + @app.delete("/api/artifacts/") + @json_api_view + def delete_artifact(artifact_id: str) -> BottleViewReturn: + if artifact_backend is None: + response.status = 400 # Bad Request + return {"reason": "Cannot access to the artifacts."} artifact_backend.remove(artifact_id) response.status = 204 return b"" @@ -85,13 +124,6 @@ def upload_artifact( return ... """ filename = os.path.basename(file_path) - - guess_mimetype, guess_encoding = mimetypes.guess_type(filename) - mimetype = mimetype or guess_mimetype - encoding = encoding or guess_encoding - if mimetype is None or encoding is None: - raise ValueError("Failed to guess mimetype and encoding. Please explicitly specify them.") - storage = trial.storage trial_id = trial._trial_id artifact_id = str(uuid.uuid4())