Repository: optuna/optuna

Issue:
Cannot store CMA object on PostgreSQL and MySQL.
<!-- Please write a clear and concise description of what the bug is. -->

## Problem details

This bug is reported by @maartenpants on gitter. Original report is here:

> Hi there, has anyone else experienced this issue with CmaEsSampler and the PostgreSQL storage backend:
> ```
> (psycopg2.errors.StringDataRightTruncation) value too long for type character varying(2048)
> [SQL: INSERT INTO trial_system_attributes (trial_id, key, value_json) VALUES (%(trial_id)s, %(key)s, %(value_json)s) RETURNING trial_system_attributes.trial_system_attribute_id]
> ```
> It seems the serialized version of the CMA class is too large to fit in the 2048 character column. This code is the culprit:
> ```
> optimizer_str = pickle.dumps(optimizer).hex()
> study._storage.set_trial_system_attr(trial._trial_id, "cma:optimizer", optimizer_str)
> ```

https://gitter.im/optuna/optuna?at=5f484dea36e6f709fd08a5fa

I checked the length of serialized CMA object.

```python
import pickle

import numpy as np
from cmaes import CMA


def get_cma_object_size(dim: int) -> int:
    optimizer = CMA(mean=np.zeros(dim), sigma=1.3)
    optimizer_str = pickle.dumps(optimizer).hex()
    return len(optimizer_str)


if __name__ == '__main__':
    for i in range(2, 32, 2):
        print(f"dim={i} size={get_cma_object_size(i)}")
```

```
$ python examples/cma_size.py 
dim=2 size=2516
dim=4 size=2916
dim=6 size=3364
dim=8 size=3850
dim=10 size=4442
dim=12 size=5018
dim=14 size=5738
dim=16 size=6378
dim=18 size=7194
dim=20 size=8074
dim=22 size=8906
dim=24 size=9916
dim=26 size=10956
dim=28 size=12092
dim=30 size=13148
```

This result means that the bug is raised even on low-dimensional search space.

## Additional context (optional)

I've ever tested on SQLite3 but never bumped into this bug. The reason why is:

> (9) What is the maximum size of a VARCHAR in SQLite?
> SQLite does not enforce the length of a VARCHAR. You can declare a VARCHAR(10) and SQLite will be happy to store a 500-million character string there. And it will keep all 500-million characters intact. Your content is never truncated. SQLite understands the column type of "VARCHAR(N)" to be the same as "TEXT", regardless of the value of N.
> https://www.sqlite.org/faq.html#:~:text=(9)%20What%20is%20the%20maximum,all%20500%2Dmillion%20characters%20intact.




Discussion:
I opened a pull request to fix this bug.
https://github.com/optuna/optuna/pull/1776

Resolving patch:
diff --git a/optuna/samplers/_cmaes.py b/optuna/samplers/_cmaes.py
--- a/optuna/samplers/_cmaes.py
+++ b/optuna/samplers/_cmaes.py
@@ -23,6 +23,8 @@
 _logger = logging.get_logger(__name__)
 
 _EPS = 1e-10
+# The value of system_attrs must be less than 2046 characters on RDBStorage.
+_SYSTEM_ATTR_MAX_LENGTH = 2045
 
 
 class CmaEsSampler(BaseSampler):
@@ -294,8 +296,11 @@ def sample_relative(
                     search_space, ordered_keys, population_size=popsize, randomize_start_point=True
                 )
 
+            # Store optimizer
             optimizer_str = pickle.dumps(optimizer).hex()
-            study._storage.set_trial_system_attr(trial._trial_id, "cma:optimizer", optimizer_str)
+            optimizer_attrs = _split_optimizer_str(optimizer_str)
+            for key in optimizer_attrs:
+                study._storage.set_trial_system_attr(trial._trial_id, key, optimizer_attrs[key])
 
         # Caution: optimizer should update its seed value
         seed = self._cma_rng.randint(1, 2 ** 16) + trial.number
@@ -317,13 +322,21 @@ def _restore_optimizer(
     ) -> Tuple[Optional[CMA], int]:
         # Restore a previous CMA object.
         for trial in reversed(completed_trials):
-            serialized_optimizer = trial.system_attrs.get(
-                "cma:optimizer", None
-            )  # type: Optional[str]
-            if serialized_optimizer is None:
+            optimizer_attrs = {
+                key: value
+                for key, value in trial.system_attrs.items()
+                if key.startswith("cma:optimizer")
+            }
+            if len(optimizer_attrs) == 0:
                 continue
+
+            # Check "cma:optimizer" key for backward compatibility.
+            optimizer_str = optimizer_attrs.get("cma:optimizer", None)
+            if optimizer_str is None:
+                optimizer_str = _concat_optimizer_attrs(optimizer_attrs)
+
             n_restarts = trial.system_attrs.get("cma:n_restarts", 0)  # type: int
-            return pickle.loads(bytes.fromhex(serialized_optimizer)), n_restarts
+            return pickle.loads(bytes.fromhex(optimizer_str)), n_restarts
         return None, 0
 
     def _init_optimizer(
@@ -413,6 +426,22 @@ def _get_trials(self, study: "optuna.Study") -> List[FrozenTrial]:
         return complete_trials
 
 
+def _split_optimizer_str(optimizer_str: str) -> Dict[str, str]:
+    optimizer_len = len(optimizer_str)
+    attrs = {}
+    for i in range(math.ceil(optimizer_len / _SYSTEM_ATTR_MAX_LENGTH)):
+        start = i * _SYSTEM_ATTR_MAX_LENGTH
+        end = min((i + 1) * _SYSTEM_ATTR_MAX_LENGTH, optimizer_len)
+        attrs["cma:optimizer:{}".format(i)] = optimizer_str[start:end]
+    return attrs
+
+
+def _concat_optimizer_attrs(optimizer_attrs: Dict[str, str]) -> str:
+    return "".join(
+        optimizer_attrs["cma:optimizer:{}".format(i)] for i in range(len(optimizer_attrs))
+    )
+
+
 def _to_cma_param(distribution: BaseDistribution, optuna_param: Any) -> float:
     if isinstance(distribution, optuna.distributions.LogUniformDistribution):
         return math.log(optuna_param)


Test patch:
diff --git a/tests/samplers_tests/test_cmaes.py b/tests/samplers_tests/test_cmaes.py
--- a/tests/samplers_tests/test_cmaes.py
+++ b/tests/samplers_tests/test_cmaes.py
@@ -1,3 +1,6 @@
+import pickle
+from typing import Any
+from typing import Dict
 from typing import List
 from unittest.mock import MagicMock
 from unittest.mock import Mock
@@ -8,12 +11,16 @@
 import pytest
 
 import optuna
+from optuna import create_trial
+from optuna.samplers._cmaes import _concat_optimizer_attrs
 from optuna.samplers._cmaes import _initialize_sigma0
 from optuna.samplers._cmaes import _initialize_x0
 from optuna.samplers._cmaes import _initialize_x0_randomly
+from optuna.samplers._cmaes import _split_optimizer_str
 from optuna.testing.distribution import UnsupportedDistribution
 from optuna.testing.sampler import DeterministicRelativeSampler
 from optuna.trial import FrozenTrial
+from optuna.trial import TrialState
 
 
 def test_consider_pruned_trials_experimental_warning() -> None:
@@ -213,3 +220,61 @@ def objective(trial: optuna.Trial) -> float:
 
         _, actual_kwargs = cma_class_mock.call_args
         assert actual_kwargs["population_size"] == inc_popsize * popsize
+
+
+def test_restore_optimizer_keeps_backward_compatibility() -> None:
+    sampler = optuna.samplers.CmaEsSampler()
+    optimizer = CMA(np.zeros(2), sigma=1.3)
+    optimizer_str = pickle.dumps(optimizer).hex()
+
+    completed_trials = [
+        create_trial(state=TrialState.COMPLETE, value=0.1),
+        create_trial(
+            state=TrialState.COMPLETE,
+            value=0.1,
+            system_attrs={"cma:optimizer": optimizer_str, "cma:n_restarts": 1},
+        ),
+        create_trial(state=TrialState.COMPLETE, value=0.1),
+    ]
+    optimizer, n_restarts = sampler._restore_optimizer(completed_trials)
+    assert isinstance(optimizer, CMA)
+    assert n_restarts == 1
+
+
+def test_restore_optimizer_from_substrings() -> None:
+    sampler = optuna.samplers.CmaEsSampler()
+    optimizer = CMA(np.zeros(10), sigma=1.3)
+    optimizer_str = pickle.dumps(optimizer).hex()
+
+    system_attrs = _split_optimizer_str(optimizer_str)  # type: Dict[str, Any]
+    assert len(system_attrs) > 1
+    system_attrs["cma:n_restarts"] = 1
+
+    completed_trials = [
+        create_trial(state=TrialState.COMPLETE, value=0.1),
+        create_trial(
+            state=TrialState.COMPLETE,
+            value=0.1,
+            system_attrs=system_attrs,
+        ),
+        create_trial(state=TrialState.COMPLETE, value=0.1),
+    ]
+    optimizer, n_restarts = sampler._restore_optimizer(completed_trials)
+    assert isinstance(optimizer, CMA)
+    assert n_restarts == 1
+
+
+@pytest.mark.parametrize(
+    "dummy_optimizer_str,attr_len",
+    [
+        ("012", 1),
+        ("01234", 1),
+        ("012345", 2),
+    ],
+)
+def test_split_and_concat_optimizer_string(dummy_optimizer_str: str, attr_len: int) -> None:
+    with patch("optuna.samplers._cmaes._SYSTEM_ATTR_MAX_LENGTH", 5):
+        attrs = _split_optimizer_str(dummy_optimizer_str)
+        assert len(attrs) == attr_len
+        actual = _concat_optimizer_attrs(attrs)
+        assert dummy_optimizer_str == actual