mirror of
https://github.com/wassname/optuna-dashboard.git
synced 2026-09-23 13:30:25 +08:00
239 lines
8.0 KiB
Python
239 lines
8.0 KiB
Python
from __future__ import annotations
|
|
|
|
import io
|
|
import json
|
|
import mimetypes
|
|
import os.path
|
|
from typing import TYPE_CHECKING
|
|
import uuid
|
|
import warnings
|
|
|
|
from bottle import BaseRequest
|
|
from bottle import Bottle
|
|
from bottle import HTTPResponse
|
|
from bottle import request
|
|
from bottle import response
|
|
import optuna
|
|
from optuna.trial import FrozenTrial
|
|
|
|
from .._bottle_util import json_api_view
|
|
from .._bottle_util import parse_data_uri
|
|
|
|
|
|
if TYPE_CHECKING:
|
|
from typing import Any
|
|
from typing import Optional
|
|
from typing import TypedDict
|
|
|
|
from optuna.artifacts._protocol import ArtifactStore
|
|
from optuna.storages import BaseStorage
|
|
|
|
from .protocol import ArtifactBackend
|
|
|
|
ArtifactMeta = TypedDict(
|
|
"ArtifactMeta",
|
|
{
|
|
"artifact_id": str,
|
|
"filename": str,
|
|
"mimetype": str,
|
|
"encoding": Optional[str],
|
|
},
|
|
)
|
|
|
|
|
|
ARTIFACTS_ATTR_PREFIX = "dashboard:artifacts:"
|
|
DEFAULT_MIME_TYPE = "application/octet-stream"
|
|
BaseRequest.MEMFILE_MAX = int(
|
|
os.environ.get("OPTUNA_DASHBOARD_MEMFILE_MAX", 1024 * 1024 * 128)
|
|
) # 128MB
|
|
|
|
|
|
def get_artifact_path(
|
|
trial: optuna.Trial,
|
|
artifact_id: str,
|
|
) -> str:
|
|
"""Get the URL path for a given artifact ID."""
|
|
study_id = trial.study._study_id
|
|
trial_id = trial._trial_id
|
|
return f"/artifacts/{study_id}/{trial_id}/{artifact_id}"
|
|
|
|
|
|
def register_artifact_route(
|
|
app: Bottle, storage: BaseStorage, artifact_store: Optional[ArtifactStore]
|
|
) -> None:
|
|
@app.get("/artifacts/<study_id:int>/<trial_id:int>/<artifact_id:re:[0-9a-fA-F-]+>")
|
|
def proxy_artifact(study_id: int, trial_id: int, artifact_id: str) -> HTTPResponse | bytes:
|
|
if artifact_store is None:
|
|
response.status = 400 # Bad Request
|
|
return b"Cannot access to the artifacts."
|
|
artifact_dict = get_artifact_meta(storage, study_id, trial_id, artifact_id)
|
|
if artifact_dict is None:
|
|
response.status = 404
|
|
return b"Not Found"
|
|
headers = {"Content-Type": artifact_dict["mimetype"]}
|
|
encoding = artifact_dict.get("encoding")
|
|
if encoding:
|
|
headers["Content-Encodings"] = encoding
|
|
|
|
fp = artifact_store.open_reader(artifact_id)
|
|
return HTTPResponse(fp, headers=headers)
|
|
|
|
@app.post("/api/artifacts/<study_id:int>/<trial_id:int>")
|
|
@json_api_view
|
|
def upload_artifact_api(study_id: int, trial_id: int) -> dict[str, Any]:
|
|
# TODO(c-bata): Use optuna.artifacts.upload_artifact()
|
|
if artifact_store is None:
|
|
response.status = 400 # Bad Request
|
|
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())
|
|
artifact_store.write(artifact_id, io.BytesIO(data))
|
|
|
|
mimetype, encoding = mimetypes.guess_type(filename)
|
|
artifact = {
|
|
"artifact_id": artifact_id,
|
|
"filename": filename,
|
|
"mimetype": mimetype or DEFAULT_MIME_TYPE,
|
|
"encoding": encoding,
|
|
}
|
|
attr_key = _artifact_prefix(trial_id=trial_id) + artifact_id
|
|
storage.set_study_system_attr(study_id, attr_key, json.dumps(artifact))
|
|
response.status = 201
|
|
|
|
trial = storage.get_trial(trial_id)
|
|
if trial is None:
|
|
response.status = 400
|
|
return {"reason": "Invalid study_id or trial_id"}
|
|
return {
|
|
"artifact_id": artifact_id,
|
|
"artifacts": list_trial_artifacts(storage.get_study_system_attrs(study_id), trial),
|
|
}
|
|
|
|
@app.delete("/api/artifacts/<study_id:int>/<trial_id:int>/<artifact_id:re:[0-9a-fA-F-]+>")
|
|
@json_api_view
|
|
def delete_artifact(study_id: int, trial_id: int, artifact_id: str) -> dict[str, Any]:
|
|
if artifact_store is None:
|
|
response.status = 400 # Bad Request
|
|
return {"reason": "Cannot access to the artifacts."}
|
|
artifact_store.remove(artifact_id)
|
|
|
|
attr_key = _artifact_prefix(trial_id) + artifact_id
|
|
storage.set_study_system_attr(study_id, attr_key, json.dumps(None))
|
|
response.status = 204
|
|
return {}
|
|
|
|
|
|
def upload_artifact(
|
|
backend: ArtifactBackend,
|
|
trial: optuna.Trial,
|
|
file_path: str,
|
|
*,
|
|
mimetype: Optional[str] = None,
|
|
encoding: Optional[str] = None,
|
|
) -> str:
|
|
"""Upload an artifact (files), which is associated with the trial.
|
|
|
|
.. warning::
|
|
|
|
This function is deprecated. Please use `optuna.artifacts.upload_artifact
|
|
<https://optuna.readthedocs.io/en/latest/reference/generated/optuna.artifacts.
|
|
upload_artifact.html>`_ instead.
|
|
|
|
Example:
|
|
.. code-block:: python
|
|
|
|
import optuna
|
|
from optuna_dashboard.artifact import upload_artifact
|
|
from optuna_dashboard.artifact.file_system import FileSystemBackend
|
|
|
|
artifact_backend = FileSystemBackend("./tmp/")
|
|
|
|
def objective(trial: optuna.Trial) -> float:
|
|
... = trial.suggest_float("x", -10, 10)
|
|
file_path = generate_example_png(...)
|
|
upload_artifact(artifact_backend, trial, file_path)
|
|
return ...
|
|
"""
|
|
warnings.warn(
|
|
"upload_artifact() is deprecated. Please use optuna.artifacts.upload_artifact() instead.\n"
|
|
"See https://optuna-dashboard.readthedocs.io/en/latest/errors.html for details",
|
|
DeprecationWarning,
|
|
)
|
|
|
|
filename = os.path.basename(file_path)
|
|
storage = trial.storage
|
|
trial_id = trial._trial_id
|
|
study_id = trial.study._study_id
|
|
artifact_id = str(uuid.uuid4())
|
|
guess_mimetype, guess_encoding = mimetypes.guess_type(filename)
|
|
artifact: ArtifactMeta = {
|
|
"artifact_id": artifact_id,
|
|
"mimetype": mimetype or guess_mimetype or DEFAULT_MIME_TYPE,
|
|
"encoding": encoding or guess_encoding,
|
|
"filename": filename,
|
|
}
|
|
attr_key = _artifact_prefix(trial_id=trial_id) + artifact_id
|
|
storage.set_study_system_attr(study_id, attr_key, json.dumps(artifact))
|
|
|
|
with open(file_path, "rb") as f:
|
|
backend.write(artifact_id, f)
|
|
return artifact_id
|
|
|
|
|
|
def _artifact_prefix(trial_id: int) -> str:
|
|
return ARTIFACTS_ATTR_PREFIX + f"{trial_id}:"
|
|
|
|
|
|
def get_artifact_meta(
|
|
storage: BaseStorage, study_id: int, trial_id: int, artifact_id: str
|
|
) -> Optional[ArtifactMeta]:
|
|
study_system_attr = storage.get_study_system_attrs(study_id)
|
|
attr_key = _artifact_prefix(trial_id=trial_id) + artifact_id
|
|
artifact_meta = study_system_attr.get(attr_key)
|
|
if artifact_meta is not None:
|
|
return json.loads(artifact_meta)
|
|
|
|
# See https://github.com/optuna/optuna/blob/f827582a8/optuna/artifacts/_upload.py#L71
|
|
trial_system_attrs = storage.get_trial_system_attrs(trial_id)
|
|
value = trial_system_attrs.get("artifacts:" + artifact_id)
|
|
if value is not None:
|
|
return json.loads(value)
|
|
return None
|
|
|
|
|
|
def delete_all_artifacts(backend: ArtifactStore, storage: BaseStorage, study_id: int) -> None:
|
|
artifact_metas = []
|
|
study_system_attrs = storage.get_study_system_attrs(study_id)
|
|
for trial in storage.get_all_trials(study_id):
|
|
trial_artifacts = list_trial_artifacts(study_system_attrs, trial)
|
|
artifact_metas.extend(trial_artifacts)
|
|
|
|
for meta in artifact_metas:
|
|
backend.remove(meta["artifact_id"])
|
|
|
|
|
|
def list_trial_artifacts(
|
|
study_system_attrs: dict[str, Any], trial: FrozenTrial
|
|
) -> list[ArtifactMeta]:
|
|
dashboard_artifact_metas = [
|
|
json.loads(value)
|
|
for key, value in study_system_attrs.items()
|
|
if key.startswith(_artifact_prefix(trial._trial_id))
|
|
]
|
|
|
|
# See https://github.com/optuna/optuna/blob/f827582a8/optuna/artifacts/_upload.py#L16
|
|
optuna_artifact_metas = [
|
|
json.loads(value)
|
|
for key, value in trial.system_attrs.items()
|
|
if key.startswith("artifacts:")
|
|
]
|
|
|
|
artifact_metas = dashboard_artifact_metas + optuna_artifact_metas
|
|
return [a for a in artifact_metas if a is not None]
|