Merge pull request #564 from Alnusjaponica/deprecate-artifact

Store artifact meta in `trial_system_attr`

optuna_dashboard/artifact/_backend.py

@@ -40,8 +40,8 @@
         },
     )
 
-
-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 proxy_artifact(study_id: int, trial_id: int, artifact_id: str) -> HTTPRespon
     @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 upload_artifact_api(study_id: int, trial_id: int) -> dict[str, Any]:
             "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 delete_artifact(study_id: int, trial_id: int, artifact_id: str) -> dict[str,
             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 objective(trial: optuna.Trial) -> float:
     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,32 +187,36 @@ def objective(trial: optuna.Trial) -> float:
         "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)
     return artifact_id
 
 
 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

@@ -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

@@ -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"},
+    ]