From cdee5c28a514bee4364dbc6f085abe75bcc48916 Mon Sep 17 00:00:00 2001
From: Masashi Shibata <c-bata@users.noreply.github.com>
Date: Thu, 27 Apr 2023 13:40:28 +0900
Subject: [PATCH] Merge pull request #428 from
 c-bata/support-optuna-metric-names

Support `metric_names` introduced in Optuna
---
 docs/errors.rst                       | 17 ++++++++++++++++-
 optuna_dashboard/_named_objectives.py | 18 ++++++++++++++++++
 python_tests/test_metric_names.py     | 24 ++++++++++++++++++++++++
 3 files changed, 58 insertions(+), 1 deletion(-)
 create mode 100644 python_tests/test_metric_names.py

diff --git a/docs/errors.rst b/docs/errors.rst
index ef936ef3..f1b6356a 100644
--- a/docs/errors.rst
+++ b/docs/errors.rst
@@ -7,7 +7,7 @@ Warning Messages
 ----------------
 
 Human-in-the-loop optimization will not work with ``_CachedStorage`` in Optuna prior to v3.2.
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 This warning occurs when the storage object associated with the Optuna Study is of the ``_CachedStorage`` class.
 
@@ -22,3 +22,18 @@ or use a following dirty hack to unwrap ``_CachedStorage`` class.
 
    if isinstance(study._storage, optuna.storages._CachedStorage):
        study._storage = study._storage._backend
+
+
+``set_objective_names()`` function is deprecated. Please use ``study.set_metric_names()`` instead.
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+:func:`~optuna_dashboard.set_objective_names` function has been ported to Optuna.
+Please use `study.set_metric_names() <https://optuna.readthedocs.io/en/latest/reference/generated/optuna.study.Study.html#optuna.study.Study>`_ function instead.
+
+.. list-table::
+
+   * - Deprecated APIs
+     - Corresponding Active APIs
+   * - ``optuna_dashboard.set_objective_names(study, ["objective 1", "objective 2"])``
+     - ``study.set_metric_names(["objective 1", "objective 2"])``
+
diff --git a/optuna_dashboard/_named_objectives.py b/optuna_dashboard/_named_objectives.py
index e61dcc82..ad37ed57 100644
--- a/optuna_dashboard/_named_objectives.py
+++ b/optuna_dashboard/_named_objectives.py
@@ -2,10 +2,15 @@ from __future__ import annotations
 
 from typing import Any
 from typing import Optional
+import warnings
 
 import optuna
 
 
+# Should be equivalent to `optuna.study.study._SYSTEM_ATTR_METRIC_NAMES`.
+# See https://github.com/optuna/optuna/pull/4383 for details.
+SYSTEM_ATTR_METRIC_NAMES = "study:metric_names"
+
 SYSTEM_ATTR_NAME = "dashboard:objective_names"
 
 
@@ -22,6 +27,17 @@ def set_objective_names(study: optuna.Study, names: list[str]) -> None:
           study = optuna.create_study(directions=["minimize", "minimize"])
           set_objective_names(study, ["val_loss", "flops"])
     """
+
+    if hasattr(study, "set_metric_names"):
+        warnings.warn(
+            "`set_objective_names()` function is deprecated."
+            " Please use `study.set_metric_names()` instead."
+            " See https://optuna-dashboard.readthedocs.io/en/latest/errors.html for details.",
+            category=FutureWarning,
+        )
+        study.set_metric_names(names)
+        return
+
     storage = study._storage
     study_id = study._study_id
 
@@ -32,4 +48,6 @@ def set_objective_names(study: optuna.Study, names: list[str]) -> None:
 
 
 def get_objective_names(system_attrs: dict[str, Any]) -> Optional[list[str]]:
+    if SYSTEM_ATTR_METRIC_NAMES in system_attrs:
+        return system_attrs[SYSTEM_ATTR_METRIC_NAMES]
     return system_attrs.get(SYSTEM_ATTR_NAME)
diff --git a/python_tests/test_metric_names.py b/python_tests/test_metric_names.py
new file mode 100644
index 00000000..ca373809
--- /dev/null
+++ b/python_tests/test_metric_names.py
@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+import unittest
+
+import optuna
+from optuna.version import __version__ as optuna_ver
+from optuna_dashboard._named_objectives import get_objective_names
+from packaging import version
+
+
+class MetricNamesTestCase(unittest.TestCase):
+    @unittest.skipIf(
+        version.parse(optuna_ver) < version.Version("3.2.0.dev"),
+        "study.set_metric_names() is not implemented yet",
+    )
+    def test_get_metric_names(self) -> None:
+        study = optuna.create_study(directions=["minimize", "minimize"])
+        # TODO(c-bata): Remove the following `type: ignore` after released Optuna v3.2.
+        study.set_metric_names(["val_loss", "flops"])  # type: ignore
+
+        study_system_attrs = study._storage.get_study_system_attrs(study._study_id)
+        metric_names = get_objective_names(study_system_attrs)
+
+        assert metric_names == ["val_loss", "flops"]