Store trial note in study_system_attrs

optuna_dashboard/_app.py

@@ -363,41 +363,38 @@ def create_app(storage: BaseStorage, debug: bool = False) -> Bottle:
             return {"reason": "Invalid request."}
 
         system_attrs = storage.get_study_system_attrs(study_id)
-        if not note.version_is_incremented(system_attrs, req_note_ver):
+        if not note.version_is_incremented(system_attrs, None, req_note_ver):
             response.status = 409  # Conflict
             return {
                 "reason": "The text you are editing has changed. "
                 "Please copy your edits and refresh the page.",
-                "note": note.get_note_from_system_attrs(system_attrs),
+                "note": note.get_note_from_system_attrs(system_attrs, None),
             }
 
-        note.save_note_in_study(storage, study_id, req_note_ver, req_note_body)
+        note.save_note(storage, study_id, None, req_note_ver, req_note_body)
         response.status = 204  # No content
         return {}
 
-    @app.put("/api/trials/<trial_id:int>/note")
+    @app.put("/api/studies/<study_id:int>/<trial_id:int>/note")
     @json_api_view
-    def save_trial_note(trial_id: int) -> BottleViewReturn:
-        trial = storage.get_trial(trial_id)
-        if trial.state.is_finished():
-            response.status = 400  # Bad request
-            return {"reason": "Cannot update the finished trials"}
-
+    def save_trial_note(study_id: int, trial_id: int) -> BottleViewReturn:
         req_note_ver = request.json.get("version", None)
         req_note_body = request.json.get("body", None)
         if req_note_ver is None or req_note_body is None:
             response.status = 400  # Bad request
             return {"reason": "Invalid request."}
 
-        if not note.version_is_incremented(trial.system_attrs, req_note_ver):
+        # Store note content in study system attrs since it's always updatable.
+        system_attrs = storage.get_study_system_attrs(study_id=study_id)
+        if not note.version_is_incremented(system_attrs, trial_id, req_note_ver):
             response.status = 409  # Conflict
             return {
                 "reason": "The text you are editing has changed. "
                 "Please copy your edits and refresh the page.",
-                "note": note.get_note_from_system_attrs(system_attrs),
+                "note": note.get_note_from_system_attrs(system_attrs, trial_id),
             }
 
-        note.save_note_in_trial(storage, trial_id, req_note_ver, req_note_body)
+        note.save_note(storage, study_id, trial_id, req_note_ver, req_note_body)
         response.status = 204  # No content
         return {}
 

optuna_dashboard/_note.py

@@ -9,6 +9,7 @@ from optuna.storages import BaseStorage
 
 if TYPE_CHECKING:
     from typing import TypedDict
+    from typing import Optional
 
     NoteType = TypedDict(
         "NoteType",
@@ -19,32 +20,50 @@ if TYPE_CHECKING:
     )
 
 SYSTEM_ATTR_MAX_LENGTH = 2045
-NOTE_VER_KEY = "dashboard:note_ver"
-NOTE_STR_KEY_PREFIX = "dashboard:note_str:"
 
 
-def get_note_from_system_attrs(system_attrs: dict[str, Any]) -> NoteType:
-    if NOTE_VER_KEY not in system_attrs:
+def note_ver_key(trial_id: Optional[int]) -> str:
+    prefix = "dashboard:note_ver"
+    if trial_id is None:
+        return prefix
+    return f"dashboard:{trial_id}:note_ver"
+
+
+def note_str_key_prefix(trial_id: int) -> str:
+    prefix = "dashboard:note_str:"
+    if trial_id is None:
+        return prefix
+    return f"dashboard:{trial_id}:note_str:"
+
+
+def get_note_from_system_attrs(system_attrs: dict[str, Any], trial_id: Optional[int]) -> NoteType:
+    if note_ver_key(trial_id) not in system_attrs:
         return {
             "version": 0,
             "body": "",
         }
-    note_ver = int(system_attrs[NOTE_VER_KEY])
+    note_ver = int(system_attrs[note_ver_key(trial_id)])
     note_attrs: dict[str, str] = {
-        key: value for key, value in system_attrs.items() if key.startswith(NOTE_STR_KEY_PREFIX)
+        key: value
+        for key, value in system_attrs.items()
+        if key.startswith(note_str_key_prefix(trial_id))
     }
-    return {"version": note_ver, "body": concat_body(note_attrs)}
+    return {"version": note_ver, "body": concat_body(note_attrs, trial_id)}
 
 
-def version_is_incremented(system_attrs: dict[str, Any], req_note_ver: int) -> bool:
-    db_note_ver = system_attrs.get(NOTE_VER_KEY, 0)
+def version_is_incremented(
+    system_attrs: dict[str, Any], trial_id: Optional[int], req_note_ver: int
+) -> bool:
+    db_note_ver = system_attrs.get(note_ver_key(trial_id), 0)
     return req_note_ver == db_note_ver + 1
 
 
-def save_note_in_study(storage: BaseStorage, study_id: int, ver: int, body: str) -> None:
-    storage.set_study_system_attr(study_id, NOTE_VER_KEY, ver)
+def save_note(
+    storage: BaseStorage, study_id: int, trial_id: Optional[int], ver: int, body: str
+) -> None:
+    storage.set_study_system_attr(study_id, note_ver_key(trial_id), ver)
 
-    attrs = split_body(body)
+    attrs = split_body(body, trial_id)
     for k, v in attrs.items():
         storage.set_study_system_attr(study_id, k, v)
 
@@ -52,40 +71,24 @@ def save_note_in_study(storage: BaseStorage, study_id: int, ver: int, body: str)
     all_note_attrs: dict[str, str] = {
         key: value
         for key, value in storage.get_study_system_attrs(study_id).items()
-        if key.startswith(NOTE_STR_KEY_PREFIX)
+        if key.startswith(note_str_key_prefix(trial_id))
     }
     if len(all_note_attrs) > len(attrs):
         for i in range(len(attrs), len(all_note_attrs)):
-            storage.set_study_system_attr(study_id, f"{NOTE_STR_KEY_PREFIX}{i}", "")
+            storage.set_study_system_attr(study_id, f"{note_str_key_prefix(trial_id)}{i}", "")
 
 
-def save_note_in_trial(storage: BaseStorage, trial_id: int, ver: int, body: str) -> None:
-    storage.set_trial_system_attr(trial_id, NOTE_VER_KEY, ver)
-
-    attrs = split_body(body)
-    for k, v in attrs.items():
-        storage.set_study_system_attr(trial_id, k, v)
-
-    # Clear previous messages
-    all_note_attrs: dict[str, str] = {
-        key: value
-        for key, value in storage.get_trial_system_attrs(trial_id).items()
-        if key.startswith(NOTE_STR_KEY_PREFIX)
-    }
-    if len(all_note_attrs) > len(attrs):
-        for i in range(len(attrs), len(all_note_attrs)):
-            storage.set_trial_system_attr(trial_id, f"{NOTE_STR_KEY_PREFIX}{i}", "")
-
-
-def split_body(note_str: str) -> dict[str, str]:
+def split_body(note_str: str, trial_id: Optional[int]) -> dict[str, str]:
     note_len = len(note_str)
     attrs = {}
     for i in range(math.ceil(note_len / SYSTEM_ATTR_MAX_LENGTH)):
         start = i * SYSTEM_ATTR_MAX_LENGTH
         end = min((i + 1) * SYSTEM_ATTR_MAX_LENGTH, note_len)
-        attrs[f"{NOTE_STR_KEY_PREFIX}{i}"] = note_str[start:end]
+        attrs[f"{note_str_key_prefix(trial_id)}{i}"] = note_str[start:end]
     return attrs
 
 
-def concat_body(note_attrs: dict[str, str]) -> str:
-    return "".join(note_attrs[f"{NOTE_STR_KEY_PREFIX}{i}"] for i in range(len(note_attrs)))
+def concat_body(note_attrs: dict[str, str], trial_id: Optional[int]) -> str:
+    return "".join(
+        note_attrs[f"{note_str_key_prefix(trial_id)}{i}"] for i in range(len(note_attrs))
+    )

optuna_dashboard/_serializer.py

@@ -84,22 +84,27 @@ def serialize_study_detail(
         "name": summary.study_name,
         "directions": [d.name.lower() for d in summary.directions],
     }
+    system_attrs = getattr(summary, "system_attrs", {})
     if summary.datetime_start is not None:
         serialized["datetime_start"] = summary.datetime_start.isoformat()
 
-    serialized["trials"] = [serialize_frozen_trial(summary._study_id, trial) for trial in trials]
+    serialized["trials"] = [
+        serialize_frozen_trial(summary._study_id, trial, system_attrs) for trial in trials
+    ]
     serialized["best_trials"] = [
-        serialize_frozen_trial(summary._study_id, trial) for trial in best_trials
+        serialize_frozen_trial(summary._study_id, trial, system_attrs) for trial in best_trials
     ]
     serialized["intersection_search_space"] = serialize_search_space(intersection)
     serialized["union_search_space"] = serialize_search_space(union)
     serialized["union_user_attrs"] = [{"key": a[0], "sortable": a[1]} for a in union_user_attrs]
     serialized["has_intermediate_values"] = has_intermediate_values
-    serialized["note"] = note.get_note_from_system_attrs(getattr(summary, "system_attrs", {}))
+    serialized["note"] = note.get_note_from_system_attrs(system_attrs, None)
     return serialized
 
 
-def serialize_frozen_trial(study_id: int, trial: FrozenTrial) -> dict[str, Any]:
+def serialize_frozen_trial(
+    study_id: int, trial: FrozenTrial, study_system_attrs: dict[str, Any]
+) -> dict[str, Any]:
     serialized = {
         "trial_id": trial._trial_id,
         "study_id": study_id,
@@ -108,7 +113,7 @@ def serialize_frozen_trial(study_id: int, trial: FrozenTrial) -> dict[str, Any]:
         "params": [{"name": name, "value": str(value)} for name, value in trial.params.items()],
         "user_attrs": serialize_attrs(trial.user_attrs),
         "system_attrs": serialize_attrs(getattr(trial, "_system_attrs", {})),
-        "note": note.get_note_from_system_attrs(getattr(trial, "_system_attrs", {}))
+        "note": note.get_note_from_system_attrs(study_system_attrs, trial._trial_id),
     }
 
     serialized_intermediate_values: list[IntermediateValue] = []

optuna_dashboard/ts/action.ts

@@ -35,6 +35,19 @@ export const actionCreator = () => {
     setStudyDetails(newVal)
   }
 
+  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)
+  }
+
   const setStudyParamImportanceState = (
     studyId: number,
     importance: ParamImportance[][]
@@ -189,27 +202,29 @@ export const actionCreator = () => {
     trialId: number,
     note: Note
   ): Promise<void> => {
-    return saveTrialNoteAPI(trialId, note)
+    return saveTrialNoteAPI(studyId, trialId, note)
       .then(() => {
-        const newStudy = Object.assign({}, studyDetails[studyId])
-        const trial = newStudy.trials.find((t) => t.trial_id === trialId)
-        if (trial === undefined) {
+        const index = studyDetails[studyId].trials.findIndex(
+          (t) => t.trial_id === trialId
+        )
+        if (index === -1) {
           enqueueSnackbar(`Unexpected error happens. Please reload the page.`, {
             variant: "error",
           })
           return
         }
-        trial.note = note
-        setStudyDetailState(studyId, newStudy)
+        setTrialNote(studyId, index, note)
         enqueueSnackbar(`Success to save the note`, {
           variant: "success",
         })
       })
       .catch((err) => {
+        console.dir(err)
         if (err.response.status === 409) {
-          const newStudy = Object.assign({}, studyDetails[studyId])
-          const trial = newStudy.trials.find((t) => t.trial_id === trialId)
-          if (trial === undefined) {
+          const index = studyDetails[studyId].trials.findIndex(
+            (t) => t.trial_id === trialId
+          )
+          if (index === -1) {
             enqueueSnackbar(
               `Unexpected error happens. Please reload the page.`,
               {
@@ -218,8 +233,7 @@ export const actionCreator = () => {
             )
             return
           }
-          trial.note = err.response.data.note
-          setStudyDetailState(studyId, newStudy)
+          setTrialNote(studyId, index, note)
         }
         const reason = err.response?.data.reason
         if (reason !== undefined) {

optuna_dashboard/ts/apiClient.ts

@@ -48,10 +48,7 @@ interface StudyDetailResponse {
   union_search_space: SearchSpace[]
   union_user_attrs: AttributeSpec[]
   has_intermediate_values: boolean
-  note: {
-    version: number
-    body: string
-  }
+  note: Note
 }
 
 export const getStudyDetailAPI = (
@@ -197,11 +194,12 @@ export const saveStudyNoteAPI = (
 }
 
 export const saveTrialNoteAPI = (
+  studyId: number,
   trialId: number,
   note: { version: number; body: string }
 ): Promise<void> => {
   return axiosInstance
-    .put<void>(`/api/trials/${trialId}/note`, note)
+    .put<void>(`/api/studies/${studyId}/${trialId}/note`, note)
     .then((res) => {
       return
     })

optuna_dashboard/ts/components/Note.tsx

@@ -55,15 +55,13 @@ export const TrialNote: FC<{
   studyId: number
   trialId: number
   latestNote: Note
-  editable: boolean
-}> = ({ studyId, trialId, latestNote, editable }) => {
+}> = ({ studyId, trialId, latestNote }) => {
   return (
     <NoteBase
       studyId={studyId}
       trialId={trialId}
       latestNote={latestNote}
       minRows={5}
-      editable={editable}
     />
   )
 }
@@ -80,19 +78,17 @@ export const StudyNote: FC<{
       latestNote={latestNote}
       minRows={minRows}
       cardSx={cardSx}
-      editable={true}
     />
   )
 }
 
-export const NoteBase: FC<{
+const NoteBase: FC<{
   studyId: number
   trialId?: number
   latestNote: Note
   minRows: number
-  editable: boolean
   cardSx?: SxProps<Theme>
-}> = ({ studyId, trialId, latestNote, minRows, editable, cardSx }) => {
+}> = ({ studyId, trialId, latestNote, minRows, cardSx }) => {
   const theme = useTheme()
   const [renderMarkdown, setRenderMarkdown] = useState(true)
   const [saving, setSaving] = useState(false)
@@ -147,16 +143,11 @@ export const NoteBase: FC<{
     setCurNote(latestNote)
     window.onbeforeunload = null
   }
-  let defaultBody: string
-  if (editable) {
-    defaultBody =
-      "*A markdown editor for taking a memo, related to the study. Click the 'Edit' button in the upper right corner to access the editor.*"
-  } else {
-    defaultBody = ""
-  }
 
   let content
   if (renderMarkdown) {
+    const defaultBody =
+      "*A markdown editor for taking a memo, related to the study. Click the 'Edit' button in the upper right corner to access the editor.*"
     content = (
       <ReactMarkdown
         children={latestNote.body || defaultBody}
@@ -239,10 +230,7 @@ export const NoteBase: FC<{
               <CloseIcon />
             </IconButton>
           ) : (
-            <IconButton
-              disabled={!editable}
-              onClick={() => setRenderMarkdown(false)}
-            >
+            <IconButton onClick={() => setRenderMarkdown(false)}>
               <EditIcon />
             </IconButton>
           )

optuna_dashboard/ts/components/TrialTable.tsx

@@ -230,20 +230,16 @@ export const TrialTable: FC<{
   ]
 
   const collapseBody = (index: number) => {
-    const editable =
-      trials[index].state === "Running" || trials[index].state === "Waiting"
+    console.dir(trials)
     return (
       <Grid container direction="row">
-        {trials[index].note.body !== "" || editable ? (
-          <Grid item xs={12}>
-            <TrialNote
-              studyId={trials[index].study_id}
-              trialId={trials[index].trial_id}
-              latestNote={trials[index].note}
-              editable={editable}
-            />
-          </Grid>
-        ) : null}
+        <Grid item xs={12}>
+          <TrialNote
+            studyId={trials[index].study_id}
+            trialId={trials[index].trial_id}
+            latestNote={trials[index].note}
+          />
+        </Grid>
         <Grid item xs={6}>
           <Box margin={1}>
             <Typography variant="h6" gutterBottom component="div">

python_tests/test_note.py

@@ -15,7 +15,7 @@ class NoteTestCase(TestCase):
             ("012345", 2),
         ]:
             with self.subTest(f"with_{dummy_body_str}_{attr_len}"):
-                attrs = note.split_body(dummy_body_str)
+                attrs = note.split_body(dummy_body_str, None)
                 assert len(attrs) == attr_len
-                actual = note.concat_body(attrs)
+                actual = note.concat_body(attrs, None)
                 assert dummy_body_str == actual