diff --git a/tslib/storage/src/journal.ts b/tslib/storage/src/journal.ts index 2bd9540d..f878352f 100644 --- a/tslib/storage/src/journal.ts +++ b/tslib/storage/src/journal.ts @@ -80,6 +80,13 @@ interface JournalOpSetTrialUserAttr extends JournalOpBase { user_attr: { [key: string]: any } // eslint-disable-line @typescript-eslint/no-explicit-any } +interface JournalOpSetTrialSystemAttr extends JournalOpBase { + trial_id: number + system_attr: { + constraints: number[] + } +} + const trialStateNumToTrialState = (state: number): Optuna.TrialState => { switch (state) { case 0: @@ -224,7 +231,7 @@ class JournalStorage { } }) - const userAtter = log.user_attrs + const userAttrs = log.user_attrs ? Object.entries(log.user_attrs).map(([key, value]) => { return { key: key, @@ -249,7 +256,8 @@ class JournalStorage { })(), params: params, intermediate_values: [], - user_attrs: userAtter, + user_attrs: userAttrs, + constraints: [], datetime_start: log.datetime_start ? new Date(log.datetime_start) : undefined, @@ -341,6 +349,14 @@ class JournalStorage { } } } + + public applySetTrialSystemAttr(log: JournalOpSetTrialSystemAttr) { + const [thisStudy, thisTrial] = this.getStudyAndTrial(log.trial_id) + if (thisStudy === undefined || thisTrial === undefined) { + return + } + thisTrial.constraints = log.system_attr.constraints + } } const loadJournalStorage = ( @@ -434,7 +450,9 @@ const loadJournalStorage = ( ) break case JournalOperation.SET_TRIAL_SYSTEM_ATTR: - // Unsupported + journalStorage.applySetTrialSystemAttr( + parsedLog as JournalOpSetTrialSystemAttr + ) break } } diff --git a/tslib/storage/src/sqlite.ts b/tslib/storage/src/sqlite.ts index 3cd2cdc1..cb4be1ff 100644 --- a/tslib/storage/src/sqlite.ts +++ b/tslib/storage/src/sqlite.ts @@ -159,6 +159,11 @@ const getStudy = ( } } + const systemAttrs = getTrialSystemAttributes(db, trial.trial_id) + if (systemAttrs !== undefined) { + trial.constraints = systemAttrs.constraints + } + const params = getTrialParams(db, trial.trial_id) const param_names = new Set() for (const param of params) { @@ -222,6 +227,7 @@ const getTrials = ( ), params: [], // Set this column later user_attrs: [], // Set this column later + constraints: [], datetime_start: vals[3], datetime_complete: vals[4], } @@ -404,6 +410,20 @@ const getTrialUserAttributes = ( return attrs } +const getTrialSystemAttributes = (db: SQLite3DB, trialId: number) => { + let attrs: { constraints: number[] } | undefined + db.exec({ + sql: `SELECT key, value_json FROM trial_system_attributes WHERE trial_id = ${trialId} AND key = 'constraints'`, + // biome-ignore lint/suspicious/noExplicitAny: + callback: (vals: any[]) => { + attrs = { + constraints: JSON.parse(vals[1]), + } + }, + }) + return attrs +} + const getTrialIntermediateValues = ( db: SQLite3DB, trialId: number,