Add an Artifact column

optuna_dashboard/artifact/_backend.py

@@ -47,12 +47,12 @@ BaseRequest.MEMFILE_MAX = int(
 def register_artifact_route(
     app: Bottle, storage: BaseStorage, artifact_backend: Optional[ArtifactBackend]
 ) -> None:
-    @app.get("/artifacts/<trial_id:int>/<artifact_id:re:[0-9a-fA-F-]+>")
-    def proxy_artifact(trial_id: int, artifact_id: str) -> bytes:
+    @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) -> bytes:
         if artifact_backend is None:
             response.status = 400  # Bad Request
             return b"Cannot access to the artifacts."
-        artifact_dict = _get_artifact_meta(storage, trial_id, artifact_id)
+        artifact_dict = _get_artifact_meta(storage, study_id, trial_id, artifact_id)
         response.set_header("Content-Type", artifact_dict["mimetype"])
         if artifact_dict.get("encoding"):
             response.set_header("Content-Encodings", artifact_dict.get("encoding"))
@@ -150,14 +150,13 @@ def _artifact_prefix(trial_id: int) -> str:
     return ARTIFACTS_ATTR_PREFIX + f"{trial_id}:"
 
 
-def _get_artifact_meta(storage: BaseStorage, trial_id: int, artifact_id: str) -> ArtifactMeta:
-    artifact_key = ARTIFACTS_ATTR_PREFIX + artifact_id
-    storage.get_trial_system_attrs(trial_id)
-
-    for key, value in storage.get_trial_system_attrs(trial_id).items():
-        if key == artifact_key:
-            return json.loads(value)
-    raise ValueError("Artifact not found")
+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 None:
+        return None
+    return json.loads(artifact_meta)
 
 
 def delete_all_artifacts(backend: ArtifactBackend, study_system_attrs: dict[str, Any]) -> None:

optuna_dashboard/ts/action.ts

@@ -42,17 +42,40 @@ export const actionCreator = () => {
     setStudyDetails(newVal)
   }
 
+  const setTrial = (studyId: number, trialIndex: number, trial: Trial) => {
+    const newTrials: Trial[] = [...studyDetails[studyId].trials]
+    newTrials[trialIndex] = trial
+    const newStudy: StudyDetail = Object.assign({}, studyDetails[studyId])
+    newStudy.trials = newTrials
+    setStudyDetailState(studyId, newStudy)
+  }
+
   const setTrialNote = (studyId: number, index: number, note: Note) => {
     const newTrial: Trial = Object.assign(
       {},
       studyDetails[studyId].trials[index]
     )
     newTrial.note = note
-    const newTrials: Trial[] = [...studyDetails[studyId].trials]
-    newTrials[index] = newTrial
-    const newStudy: StudyDetail = Object.assign({}, studyDetails[studyId])
-    newStudy.trials = newTrials
-    setStudyDetailState(studyId, newStudy)
+    setTrial(studyId, index, newTrial)
+  }
+
+  const appendTrialArtifacts = (
+    studyId: number,
+    trialId: number,
+    artifact: Artifact
+  ) => {
+    const index = studyDetails[studyId].trials.findIndex(
+      (t) => t.trial_id === trialId
+    )
+    if (index === -1) {
+      return
+    }
+    const newTrial: Trial = Object.assign(
+      {},
+      studyDetails[studyId].trials[index]
+    )
+    newTrial.artifacts = [...newTrial.artifacts, artifact]
+    setTrial(studyId, index, newTrial)
   }
 
   const setStudyParamImportanceState = (
@@ -289,8 +312,7 @@ export const actionCreator = () => {
       uploadArtifactAPI(studyId, trialId, file.name, upload.target.result)
         .then((artifact) => {
           setUploading(false)
-          // TODO: update global state
-          console.dir(artifact)
+          appendTrialArtifacts(studyId, trialId, artifact)
         })
         .catch((err) => {
           setUploading(false)

optuna_dashboard/ts/apiClient.ts

@@ -29,6 +29,7 @@ interface TrialResponse {
   user_attrs: Attribute[]
   system_attrs: Attribute[]
   note: Note
+  artifacts: Artifact[]
 }
 
 const convertTrialResponse = (res: TrialResponse): Trial => {
@@ -50,6 +51,7 @@ const convertTrialResponse = (res: TrialResponse): Trial => {
     user_attrs: res.user_attrs,
     system_attrs: res.system_attrs,
     note: res.note,
+    artifacts: res.artifacts,
   }
 }
 

optuna_dashboard/ts/components/Note.tsx

@@ -10,6 +10,9 @@ import {
   DialogContentText,
   DialogTitle,
   IconButton,
+  ImageList,
+  ImageListItem,
+  ImageListItemBar,
   SxProps,
   TextField,
   Typography,
@@ -38,12 +41,13 @@ import {
 } from "react-markdown/lib/ast-to-react"
 import HtmlIcon from "@mui/icons-material/Html"
 import ModeEditIcon from "@mui/icons-material/ModeEdit"
+import UploadFileIcon from "@mui/icons-material/UploadFile"
 import { Prism as SyntaxHighlighter } from "react-syntax-highlighter"
 import { darcula } from "react-syntax-highlighter/dist/esm/styles/prism"
 
 import { actionCreator } from "../action"
 import { useRecoilValue } from "recoil"
-import { isFileUploading } from "../state"
+import { artifactIsAvailable, isFileUploading, useArtifacts } from "../state"
 
 const placeholder = `## What is this feature for?
 
@@ -179,13 +183,12 @@ const MarkdownEditorModal: FC<{
       window.onbeforeunload = null
     })
 
-  const [dragOver, setDragOver] = useState<boolean>(false)
   const [saving, setSaving] = useState(false)
   const [edited, setEdited] = useState(false)
   const [curNote, setCurNote] = useState({ version: 0, body: "" })
-  const uploading = useRecoilValue<boolean>(isFileUploading)
   const textAreaRef = createRef<HTMLTextAreaElement>()
   const notLatest = latestNote.version > curNote.version
+  const artifactEnabled = useRecoilValue<boolean>(artifactIsAvailable)
 
   const [previewMarkdown, setPreviewMarkdown] = useState<string>("")
   const [preview, setPreview] = useState<boolean>(false)
@@ -235,37 +238,6 @@ const MarkdownEditorModal: FC<{
     window.onbeforeunload = null
   }
 
-  const handleDrop: DragEventHandler = (e) => {
-    if (trialId === undefined) {
-      return
-    }
-    e.stopPropagation()
-    e.preventDefault()
-    const file = e.dataTransfer.files[0]
-    setDragOver(false)
-    action.uploadArtifact(studyId, trialId, file)
-  }
-
-  const handleDragOver: DragEventHandler = (e) => {
-    if (trialId === undefined) {
-      return
-    }
-    e.stopPropagation()
-    e.preventDefault()
-    e.dataTransfer.dropEffect = "copy"
-    setDragOver(true)
-  }
-
-  const handleDragLeave: DragEventHandler = (e) => {
-    if (trialId === undefined) {
-      return
-    }
-    e.stopPropagation()
-    e.preventDefault()
-    e.dataTransfer.dropEffect = "copy"
-    setDragOver(false)
-  }
-
   // See https://github.com/iamhosseindhv/notistack/issues/231#issuecomment-825924840
   const zIndex = theme.zIndex.snackbar - 2
 
@@ -314,36 +286,42 @@ const MarkdownEditorModal: FC<{
       >
         <MarkdownRenderer body={previewMarkdown} />
       </Box>
-      <Typography>{dragOver ? "DragOver=true" : "DragOver=false"}</Typography>
-      <Typography>{uploading ? "uploading=true" : "not uploading"}</Typography>
-      <TextField
-        disabled={saving}
-        multiline={true}
-        placeholder={placeholder}
-        onDragOver={handleDragOver}
-        onDragLeave={handleDragLeave}
-        onDrop={handleDrop}
+      <Box
         sx={{
-          position: "relative",
-          resize: "none",
           width: "100%",
           height: "100%",
+          display: preview ? "none" : "flex",
+          flexDirection: "row",
           margin: theme.spacing(1, 0),
-          display: preview ? "none" : "default",
-          "& .MuiInputBase-root": { height: "100%" },
         }}
-        inputProps={{
-          style: { resize: "none", overflow: "scroll", height: "100%" },
-        }}
-        inputRef={textAreaRef}
-        defaultValue={latestNote.body}
-        onChange={() => {
-          const cur = textAreaRef.current ? textAreaRef.current.value : ""
-          if (edited !== (cur !== curNote.body)) {
-            setEdited(cur !== curNote.body)
-          }
-        }}
-      />
+      >
+        <TextField
+          disabled={saving}
+          multiline={true}
+          placeholder={placeholder}
+          sx={{
+            position: "relative",
+            resize: "none",
+            width: "100%",
+            height: "100%",
+            "& .MuiInputBase-root": { height: "100%" },
+          }}
+          inputProps={{
+            style: { resize: "none", overflow: "scroll", height: "100%" },
+          }}
+          inputRef={textAreaRef}
+          defaultValue={latestNote.body}
+          onChange={() => {
+            const cur = textAreaRef.current ? textAreaRef.current.value : ""
+            if (edited !== (cur !== curNote.body)) {
+              setEdited(cur !== curNote.body)
+            }
+          }}
+        />
+        {artifactEnabled && trialId !== undefined && (
+          <ArtifactUploader studyId={studyId} trialId={trialId} />
+        )}
+      </Box>
       <Box sx={{ display: "flex", flexDirection: "row", alignItems: "center" }}>
         {notLatest && !saving && (
           <>
@@ -399,6 +377,93 @@ const MarkdownEditorModal: FC<{
   )
 }
 
+const ArtifactUploader: FC<{
+  studyId: number
+  trialId: number
+}> = ({ studyId, trialId }) => {
+  const theme = useTheme()
+  const action = actionCreator()
+
+  const uploading = useRecoilValue<boolean>(isFileUploading)
+  const artifacts = useArtifacts(studyId, trialId)
+  const [dragOver, setDragOver] = useState<boolean>(false)
+
+  const handleDrop: DragEventHandler = (e) => {
+    e.stopPropagation()
+    e.preventDefault()
+    const file = e.dataTransfer.files[0]
+    setDragOver(false)
+    action.uploadArtifact(studyId, trialId, file)
+  }
+
+  const handleDragOver: DragEventHandler = (e) => {
+    e.stopPropagation()
+    e.preventDefault()
+    e.dataTransfer.dropEffect = "copy"
+    setDragOver(true)
+  }
+
+  const handleDragLeave: DragEventHandler = (e) => {
+    e.stopPropagation()
+    e.preventDefault()
+    e.dataTransfer.dropEffect = "copy"
+    setDragOver(false)
+  }
+
+  return (
+    <Box
+      sx={{
+        width: "300px",
+        padding: theme.spacing(0, 1),
+        display: "flex",
+        flexDirection: "column",
+      }}
+      onDragOver={handleDragOver}
+      onDragLeave={handleDragLeave}
+      onDrop={handleDrop}
+    >
+      <Typography
+        sx={{
+          fontWeight: theme.typography.fontWeightBold,
+          margin: theme.spacing(1, 0),
+        }}
+      >
+        Artifact
+      </Typography>
+      <LoadingButton
+        loading={uploading}
+        loadingPosition="start"
+        startIcon={<UploadFileIcon />}
+        variant="outlined"
+      >
+        Upload
+      </LoadingButton>
+      <Box
+        sx={{
+          border: dragOver ? `2px dashed #ffffff` : `1px solid #fff`,
+          height: "100%",
+          margin: theme.spacing(1, 0),
+          borderRadius: "4px",
+        }}
+      >
+        <ImageList cols={1}>
+          {artifacts
+            .filter((a) => a.mimetype.startsWith("image"))
+            .map((a) => (
+              <ImageListItem key={a.artifact_id}>
+                <img
+                  src={`/artifacts/${studyId}/${trialId}/${a.artifact_id}`}
+                />
+                <ImageListItemBar title={a.filename} />
+              </ImageListItem>
+            ))}
+        </ImageList>
+      </Box>
+      <Button variant="outlined">Insert an image</Button>
+    </Box>
+  )
+}
+
 const NoteBase: FC<{
   studyId: number
   trialId?: number

optuna_dashboard/ts/state.ts

@@ -80,3 +80,12 @@ export const useStudyName = (studyId: number): string | null => {
   const studySummary = useStudySummaryValue(studyId)
   return studyDetail?.name || studySummary?.study_name || null
 }
+
+export const useArtifacts = (studyId: number, trialId: number): Artifact[] => {
+  const study = useStudyDetailValue(studyId)
+  const trial = study?.trials.find((t) => t.trial_id === trialId)
+  if (trial === undefined) {
+    return []
+  }
+  return trial.artifacts
+}

optuna_dashboard/ts/types/index.d.ts

@@ -112,6 +112,7 @@ type Trial = {
   user_attrs: Attribute[]
   system_attrs: Attribute[]
   note: Note
+  artifacts: Artifact[]
 }
 
 type StudySummary = {