diff --git a/tslib/storage/src/journal.ts b/tslib/storage/src/journal.ts index 29e59964..2bd9540d 100644 --- a/tslib/storage/src/journal.ts +++ b/tslib/storage/src/journal.ts @@ -29,6 +29,13 @@ interface JournalOpDeleteStudy extends JournalOpBase { study_id: number } +interface JournalOpSetStudySystemAttr extends JournalOpBase { + study_id: number + system_attr: { + "study:metric_names": string[] + } +} + interface JournalOpCreateTrial extends JournalOpBase { study_id: number datetime_start?: string @@ -179,6 +186,14 @@ class JournalStorage { this.studies = this.studies.filter((item) => item.id !== log.study_id) } + public applyStudySystemAttr(log: JournalOpSetStudySystemAttr): void { + const thisStudy = this.studies.find((item) => item.id === log.study_id) + if (thisStudy === undefined) { + return + } + thisStudy.metric_names = log.system_attr["study:metric_names"] + } + public applyCreateTrial(log: JournalOpCreateTrial): void { const thisStudy = this.studies.find((item) => item.id === log.study_id) if (thisStudy === undefined) { @@ -393,7 +408,9 @@ const loadJournalStorage = ( // Unsupported break case JournalOperation.SET_STUDY_SYSTEM_ATTR: - // Unsupported + journalStorage.applyStudySystemAttr( + parsedLog as JournalOpSetStudySystemAttr + ) break case JournalOperation.CREATE_TRIAL: journalStorage.applyCreateTrial(parsedLog as JournalOpCreateTrial) diff --git a/tslib/storage/src/sqlite.ts b/tslib/storage/src/sqlite.ts index 4756060f..3cd2cdc1 100644 --- a/tslib/storage/src/sqlite.ts +++ b/tslib/storage/src/sqlite.ts @@ -144,6 +144,11 @@ const getStudy = ( trials: [], } + const studySystemAttrs = getStudySystemAttributes(db, summary.id) + if (studySystemAttrs !== undefined) { + study.metric_names = studySystemAttrs.metric_names + } + let intersection_search_space: Set = new Set() study.trials = getTrials(db, summary.id, schemaVersion) for (const trial of study.trials) { @@ -367,6 +372,20 @@ const parseDistributionJSON = (t: string): Optuna.Distribution => { } } +const getStudySystemAttributes = (db: SQLite3DB, studyId: number) => { + let attrs: { metric_names: string[] } | undefined + db.exec({ + sql: `SELECT key, value_json FROM study_system_attributes WHERE study_id = ${studyId} AND key = 'dashboard:objective_names'`, + // biome-ignore lint/suspicious/noExplicitAny: + callback: (vals: any[]) => { + attrs = { + metric_names: JSON.parse(vals[1]), + } + }, + }) + return attrs +} + const getTrialUserAttributes = ( db: SQLite3DB, trialId: number diff --git a/tslib/storage/test/generate_assets.py b/tslib/storage/test/generate_assets.py index bd135720..c7cfaca4 100644 --- a/tslib/storage/test/generate_assets.py +++ b/tslib/storage/test/generate_assets.py @@ -78,6 +78,24 @@ def create_optuna_storage(storage: BaseStorage) -> None: study.optimize(objective_single_nan_report, n_trials=100) + # Multi-objective study with metric names + study = optuna.create_study( + study_name="multi-objective-metric-names", + storage=storage, + directions=["minimize", "minimize"], + ) + print(f"Generating {study.study_name} for {type(storage).__name__}...") + study.set_metric_names(["value1", "value2"]) + + def objective_multi(trial: optuna.Trial) -> tuple[float, float]: + x = trial.suggest_float("x", 0, 5) + y = trial.suggest_float("y", 0, 3) + v0 = 4 * x**2 + 4 * y**2 + v1 = (x - 5) ** 2 + (y - 5) ** 2 + return v0, v1 + + study.optimize(objective_multi, n_trials=50) + if __name__ == "__main__": remove_assets() diff --git a/tslib/storage/test/journal.test.mjs b/tslib/storage/test/journal.test.mjs index e1539247..5df85e7f 100644 --- a/tslib/storage/test/journal.test.mjs +++ b/tslib/storage/test/journal.test.mjs @@ -52,8 +52,13 @@ describe("Test Journal File Storage", async () => { ) }) + it("Check metric_names function", () => { + const study = studies.find((s) => s.name === "multi-objective-metric-names") + assert.deepStrictEqual(study.metric_names, ["value1", "value2"]) + }) + it("Check the number of studies", () => { - const N_STUDIES = 4 + const N_STUDIES = 5 assert.strictEqual(studies.length, N_STUDIES) }) }) diff --git a/tslib/types/src/index.ts b/tslib/types/src/index.ts index 4272a4d5..fab91d7d 100644 --- a/tslib/types/src/index.ts +++ b/tslib/types/src/index.ts @@ -60,6 +60,7 @@ export type Study = { union_user_attrs: AttributeSpec[] datetime_start?: Date trials: Trial[] + metric_names?: string[] } export type Trial = {