wassname/optuna-dashboard
1
0
Fork 0
mirror of https://github.com/wassname/optuna-dashboard.git synced 2026-09-12 12:40:33 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #564 from Alnusjaponica/deprecate-artifact

Browse Source 
Store artifact meta in `trial_system_attr`
This commit is contained in:
c-bata
2023-08-25 15:07:10 +09:00
committed by GitHub
parent 5c5270b4bd af7abc913d
commit 273d4fe60b
3 changed files with 159 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files
optuna_dashboard/artifact/_backend.py
+34 -19
View File
@@ -40,8 +40,8 @@ if TYPE_CHECKING:
},
)
ARTIFACTS_ATTR_PREFIX = "dashboard:artifacts:"
ARTIFACTS_ATTR_PREFIX = "artifacts:"
DASHBOARD_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)
@@ -81,6 +81,14 @@ def register_artifact_route(
@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]:
trial = storage.get_trial(trial_id)
if trial is None:
response.status = 400
return {"reason": "Invalid study_id or trial_id"}
elif trial.state.is_finished():
response.status = 400
return {"reason": "The trial is already finished."}
# TODO(c-bata): Use optuna.artifacts.upload_artifact()
if artifact_store is None:
response.status = 400 # Bad Request
@@ -102,14 +110,10 @@ def register_artifact_route(
"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))
attr_key = ARTIFACTS_ATTR_PREFIX + artifact_id
storage.set_trial_system_attr(trial_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),
@@ -123,8 +127,14 @@ def register_artifact_route(
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))
# The artifact's metadata is stored in one of the following two locations:
storage.set_study_system_attr(
study_id, _artifact_prefix(trial_id) + artifact_id, json.dumps(None)
)
storage.set_trial_system_attr(
trial_id, ARTIFACTS_ATTR_PREFIX + artifact_id, json.dumps(None)
)
response.status = 204
return {}
@@ -169,7 +179,6 @@ def upload_artifact(
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 = {
@@ -178,8 +187,8 @@ def upload_artifact(
"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))
attr_key = ARTIFACTS_ATTR_PREFIX + artifact_id
storage.set_trial_system_attr(trial_id, attr_key, json.dumps(artifact))
with open(file_path, "rb") as f:
backend.write(artifact_id, f)
@@ -187,23 +196,27 @@ def upload_artifact(
def _artifact_prefix(trial_id: int) -> str:
return ARTIFACTS_ATTR_PREFIX + f"{trial_id}:"
return DASHBOARD_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)
# Search study_system_attrs due to backward compatibility.
study_system_attrs = 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)
artifact_meta = study_system_attrs.get(attr_key)
if artifact_meta is not None:
return json.loads(artifact_meta)
# Search trial_system_attrs. Note that artifacts uploaded via optuna.artifacts.upload_artifact
# have a different trial_system_attrs key prefix.
# 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)
value = trial_system_attrs.get(ARTIFACTS_ATTR_PREFIX + artifact_id)
if value is not None:
return json.loads(value)
return None
@@ -221,18 +234,20 @@ def delete_all_artifacts(backend: ArtifactStore, storage: BaseStorage, study_id:
def list_trial_artifacts(
study_system_attrs: dict[str, Any], trial: FrozenTrial
) -> list[ArtifactMeta]:
# Collect ArtifactMeta from study_system_attrs due to backward compatibility.
dashboard_artifact_metas = [
json.loads(value)
for key, value in study_system_attrs.items()
if key.startswith(_artifact_prefix(trial._trial_id))
]
# Collect ArtifactMeta from trial_system_attrs. Note that artifacts uploaded via
# optuna.artifacts.upload_artifacts have a different trial_system_attrs key prefix.
# 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:")
if key.startswith(ARTIFACTS_ATTR_PREFIX)
]
artifact_metas = dashboard_artifact_metas + optuna_artifact_metas
return [a for a in artifact_metas if a is not None]
optuna_dashboard/ts/components/TrialList.tsx
+43 -41
View File
@@ -710,55 +710,57 @@ const TrialArtifact: FC<{ trial: Trial }> = ({ trial }) => {
)
}
})}
<Card
sx={{
marginBottom: theme.spacing(2),
width: width,
minHeight: height,
margin: theme.spacing(0, 1, 1, 0),
border: dragOver
? `3px dashed ${
theme.palette.mode === "dark" ? "white" : "black"
}`
: `1px solid ${theme.palette.divider}`,
}}
onDragOver={handleDragOver}
onDragLeave={handleDragLeave}
onDrop={handleDrop}
>
<CardActionArea
onClick={handleClick}
{trial.state === "Running" || trial.state === "Waiting" ? (
<Card
sx={{
height: "100%",
marginBottom: theme.spacing(2),
width: width,
minHeight: height,
margin: theme.spacing(0, 1, 1, 0),
border: dragOver
? `3px dashed ${
theme.palette.mode === "dark" ? "white" : "black"
}`
: `1px solid ${theme.palette.divider}`,
}}
onDragOver={handleDragOver}
onDragLeave={handleDragLeave}
onDrop={handleDrop}
>
<CardContent
<CardActionArea
onClick={handleClick}
sx={{
display: "flex",
height: "100%",
flexDirection: "column",
justifyContent: "center",
alignItems: "center",
}}
>
<UploadFileIcon
sx={{ fontSize: 80, marginBottom: theme.spacing(2) }}
/>
<input
type="file"
ref={inputRef}
onChange={handleOnChange}
style={{ display: "none" }}
/>
<Typography>Upload a New File</Typography>
<Typography
sx={{ textAlign: "center", color: theme.palette.grey.A400 }}
<CardContent
sx={{
display: "flex",
height: "100%",
flexDirection: "column",
justifyContent: "center",
alignItems: "center",
}}
>
Drag your file here or click to browse.
</Typography>
</CardContent>
</CardActionArea>
</Card>
<UploadFileIcon
sx={{ fontSize: 80, marginBottom: theme.spacing(2) }}
/>
<input
type="file"
ref={inputRef}
onChange={handleOnChange}
style={{ display: "none" }}
/>
<Typography>Upload a New File</Typography>
<Typography
sx={{ textAlign: "center", color: theme.palette.grey.A400 }}
>
Drag your file here or click to browse.
</Typography>
</CardContent>
</CardActionArea>
</Card>
) : null}
</Box>
{renderDeleteArtifactDialog()}
</>
python_tests/artifact/test_backend.py
+82
View File
@@ -0,0 +1,82 @@
from unittest.mock import MagicMock
from optuna.storages import BaseStorage
from optuna_dashboard.artifact import _backend
import pytest
def test_get_artifact_path() -> None:
study = MagicMock(_study_id=0)
trial = MagicMock(_trial_id=0, study=study)
assert _backend.get_artifact_path(trial=trial, artifact_id="id0") == "/artifacts/0/0/id0"
def test_artifact_prefix() -> None:
actual = _backend._artifact_prefix(trial_id=0)
assert actual == "dashboard:artifacts:0:"
@pytest.fixture()
def init_storage_with_artifact_meta() -> BaseStorage:
from optuna import create_study
from optuna.storages import InMemoryStorage
storage = InMemoryStorage()
study = create_study(storage=storage)
study_system_attrs = {
"dashboard:artifacts:0:id0": '{"artifact_id": "id0", "filename": "foo.txt"}',
"dashboard:artifacts:0:id1": '{"artifact_id": "id1", "filename": "bar.txt"}',
"baz": "baz",
}
for key, value in study_system_attrs.items():
study.set_system_attr(key, value)
trial_system_attrs = {
"artifacts:id2": '{"artifact_id": "id2", "filename": "baz.txt"}',
"artifacts:id3": '{"artifact_id": "id3", "filename": "qux.txt"}',
}
for key, value in trial_system_attrs.items():
trial = study.ask()
trial.set_system_attr(key, value)
study.tell(trial, 0.0)
return storage
def test_get_artifact_meta(init_storage_with_artifact_meta: MagicMock) -> None:
storage = init_storage_with_artifact_meta
actual = _backend.get_artifact_meta(storage, study_id=0, trial_id=0, artifact_id="id0")
assert actual == {"artifact_id": "id0", "filename": "foo.txt"}
actual = _backend.get_artifact_meta(storage, study_id=0, trial_id=1, artifact_id="id3")
assert actual == {"artifact_id": "id3", "filename": "qux.txt"}
actual = _backend.get_artifact_meta(storage, study_id=0, trial_id=0, artifact_id="id4")
assert actual is None
def test_delete_all_artifacts(init_storage_with_artifact_meta: MagicMock) -> None:
backend = MagicMock()
storage = init_storage_with_artifact_meta
_backend.delete_all_artifacts(backend, storage, study_id=0)
assert backend.remove.call_args_list == [
(("id0",),),
(("id1",),),
(("id2",),),
(("id3",),),
]
def test_list_trial_artifacts(init_storage_with_artifact_meta: MagicMock) -> None:
storage = init_storage_with_artifact_meta
trial = MagicMock(_trial_id=0, system_attrs=storage.get_trial_system_attrs(0))
actual = _backend.list_trial_artifacts(storage.get_study_system_attrs(0), trial)
assert actual == [
{"artifact_id": "id0", "filename": "foo.txt"},
{"artifact_id": "id1", "filename": "bar.txt"},
{"artifact_id": "id2", "filename": "baz.txt"},
]