Repository: optuna/optuna

Issue:
Race condition in set_study_param_distribution
Running a study from multiple processes with RDBStorage, it fails as the same param is set twice.

Error log:
```
(pfnopt) bash-3.2$ python experiment.py 6f42b83d-f4b5-4081-810b-e79d625fc7dc & python experiment.py 6f42b83d-f4b5-4081-810b-e79d625fc7dc &
[1] 39942
[2] 39943
(pfnopt) bash-3.2$
(pfnopt) bash-3.2$
(pfnopt) bash-3.2$
(pfnopt) bash-3.2$
(pfnopt) bash-3.2$ Traceback (most recent call last):
  File "/Users/sano/anaconda3/envs/pfnopt/lib/python3.6/site-packages/sqlalchemy/orm/query.py", line 2837, in one
    ret = self.one_or_none()
  File "/Users/sano/anaconda3/envs/pfnopt/lib/python3.6/site-packages/sqlalchemy/orm/query.py", line 2816, in one_or_none
    "Multiple rows were found for one_or_none()")
sqlalchemy.orm.exc.MultipleResultsFound: Multiple rows were found for one_or_none()

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "experiment.py", line 23, in <module>
    minimize(objective, study=study, storage=storage, n_trials=1000)
  File "/Users/sano/PycharmProjects/pfnopt/pfnopt/study.py", line 145, in minimize
    study.run(func, n_trials, timeout_seconds, n_jobs)
  File "/Users/sano/PycharmProjects/pfnopt/pfnopt/study.py", line 66, in run
    self._run_sequential(func, n_trials, timeout_seconds)
  File "/Users/sano/PycharmProjects/pfnopt/pfnopt/study.py", line 88, in _run_sequential
    result = func(client)
  File "experiment.py", line 9, in objective
    x = client.sample_uniform('x', -10000, 10000)
  File "/Users/sano/PycharmProjects/pfnopt/pfnopt/client.py", line 24, in sample_uniform
    return self._sample(name, distributions.UniformDistribution(low=low, high=high))
  File "/Users/sano/PycharmProjects/pfnopt/pfnopt/client.py", line 91, in _sample
    self.storage.set_trial_param(self.trial_id, name, param_value_in_internal_repr)
  File "/Users/sano/PycharmProjects/pfnopt/pfnopt/storage/rdb.py", line 170, in set_trial_param
    filter(StudyParam.param_name == param_name).one()
  File "/Users/sano/anaconda3/envs/pfnopt/lib/python3.6/site-packages/sqlalchemy/orm/query.py", line 2840, in one
    "Multiple rows were found for one()")
sqlalchemy.orm.exc.MultipleResultsFound: Multiple rows were found for one()
Traceback (most recent call last):
  File "/Users/sano/anaconda3/envs/pfnopt/lib/python3.6/site-packages/sqlalchemy/orm/query.py", line 2837, in one
    ret = self.one_or_none()
  File "/Users/sano/anaconda3/envs/pfnopt/lib/python3.6/site-packages/sqlalchemy/orm/query.py", line 2816, in one_or_none
    "Multiple rows were found for one_or_none()")
sqlalchemy.orm.exc.MultipleResultsFound: Multiple rows were found for one_or_none()

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "experiment.py", line 23, in <module>
    minimize(objective, study=study, storage=storage, n_trials=1000)
  File "/Users/sano/PycharmProjects/pfnopt/pfnopt/study.py", line 145, in minimize
    study.run(func, n_trials, timeout_seconds, n_jobs)
  File "/Users/sano/PycharmProjects/pfnopt/pfnopt/study.py", line 66, in run
    self._run_sequential(func, n_trials, timeout_seconds)
  File "/Users/sano/PycharmProjects/pfnopt/pfnopt/study.py", line 88, in _run_sequential
    result = func(client)
  File "experiment.py", line 9, in objective
    x = client.sample_uniform('x', -10000, 10000)
  File "/Users/sano/PycharmProjects/pfnopt/pfnopt/client.py", line 24, in sample_uniform
    return self._sample(name, distributions.UniformDistribution(low=low, high=high))
  File "/Users/sano/PycharmProjects/pfnopt/pfnopt/client.py", line 91, in _sample
    self.storage.set_trial_param(self.trial_id, name, param_value_in_internal_repr)
  File "/Users/sano/PycharmProjects/pfnopt/pfnopt/storage/rdb.py", line 170, in set_trial_param
    filter(StudyParam.param_name == param_name).one()
  File "/Users/sano/anaconda3/envs/pfnopt/lib/python3.6/site-packages/sqlalchemy/orm/query.py", line 2840, in one
    "Multiple rows were found for one()")
sqlalchemy.orm.exc.MultipleResultsFound: Multiple rows were found for one()

[1]-  Exit 1                  python experiment.py 6f42b83d-f4b5-4081-810b-e79d625fc7dc
[2]+  Exit 1                  python experiment.py 6f42b83d-f4b5-4081-810b-e79d625fc7dc
```

where experiment.py is as follows:

```python
import sys

from pfnopt import Study, minimize
from pfnopt.storage import RDBStorage


def objective(client):

    x = client.sample_uniform('x', -10000, 10000)
    y = client.sample_uniform('y', -10000, 10000)

    return (x - 2) ** 2 + y


if __name__ == '__main__':

    url = 'postgresql://user:password@dburl/dbname'  # modify here
    study_uuid = sys.argv[1]

    storage = RDBStorage(url)
    study = Study(study_uuid, storage)

    minimize(objective, study=study, storage=storage, n_trials=1000)
``` 


Discussion:

Resolving patch:
diff --git a/pfnopt/storage/rdb.py b/pfnopt/storage/rdb.py
--- a/pfnopt/storage/rdb.py
+++ b/pfnopt/storage/rdb.py
@@ -2,12 +2,14 @@
 from sqlalchemy import Column
 from sqlalchemy.engine import create_engine
 from sqlalchemy import Enum
+from sqlalchemy.exc import IntegrityError
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy import Float
 from sqlalchemy import ForeignKey
 from sqlalchemy import Integer
 from sqlalchemy import orm
 from sqlalchemy import String
+from sqlalchemy import UniqueConstraint
 from typing import Any  # NOQA
 from typing import Dict  # NOQA
 from typing import List  # NOQA
@@ -30,6 +32,7 @@ class Study(Base):
 
 class StudyParam(Base):
     __tablename__ = 'study_params'
+    __table_args__ = (UniqueConstraint('study_id', 'param_name'), {})  # type: Any
     study_param_id = Column(Integer, primary_key=True)
     study_id = Column(Integer, ForeignKey('studies.study_id'))
     param_name = Column(String(255))
@@ -51,6 +54,7 @@ class Trial(Base):
 
 class TrialParam(Base):
     __tablename__ = 'trial_params'
+    __table_args__ = (UniqueConstraint('trial_id', 'study_param_id'), {})  # type: Any
     trial_param_id = Column(Integer, primary_key=True)
     trial_id = Column(Integer, ForeignKey('trials.trial_id'))
     study_param_id = Column(Integer, ForeignKey('study_params.study_param_id'))
@@ -62,6 +66,7 @@ class TrialParam(Base):
 
 class TrialValue(Base):
     __tablename__ = 'trial_values'
+    __table_args__ = (UniqueConstraint('trial_id', 'step'), {})  # type: Any
     trial_value_id = Column(Integer, primary_key=True)
     trial_id = Column(Integer, ForeignKey('trials.trial_id'))
     step = Column(Integer)
@@ -133,7 +138,13 @@ def set_study_param_distribution(self, study_id, param_name, distribution):
         study_param.param_name = param_name
         study_param.distribution_json = distributions.distribution_to_json(distribution)
         self.session.add(study_param)
-        self.session.commit()
+
+        try:
+            self.session.commit()
+        except IntegrityError:
+            # This happens due to a known race condition.
+            # Another process/thread might have committed a record with the same unique key.
+            self.session.rollback()
 
     def create_new_trial_id(self, study_id):
         # type: (int) -> int
@@ -183,7 +194,12 @@ def set_trial_param(self, trial_id, param_name, param_value):
         trial_param.param_value = param_value
         self.session.add(trial_param)
 
-        self.session.commit()
+        try:
+            self.session.commit()
+        except IntegrityError:
+            # This happens due to a known race condition.
+            # Another process/thread might have committed a record with the same unique key.
+            self.session.rollback()
 
     def set_trial_value(self, trial_id, value):
         # type: (int, float) -> None
@@ -211,7 +227,13 @@ def set_trial_intermediate_value(self, trial_id, step, intermediate_value):
         trial_value.step = step
         trial_value.value = intermediate_value
         self.session.add(trial_value)
-        self.session.commit()
+
+        try:
+            self.session.commit()
+        except IntegrityError:
+            # This happens due to a known race condition.
+            # Another process/thread might have committed a record with the same unique key.
+            self.session.rollback()
 
     def set_trial_system_attrs(self, trial_id, system_attrs):
         # type: (int, trial_module.SystemAttributes) -> None


Test patch: