We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Occasional test failure with `ChainerMNStudy` with Python 3.5.
## Expected behavior

Tests should pass.

## Error messages, stack traces, or logs

https://app.circleci.com/pipelines/github/optuna/optuna/1171/workflows/41e626df-8501-46ca-a724-9b1c27079253/jobs/45897

## Additional context (optional)

Seems to be specific to Python 3.5 and only happen randomly.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `optuna/integration/chainermn.py`
Content:
```
1 from typing import Optional
2 
3 from optuna._imports import try_import
4 from optuna.logging import get_logger
5 from optuna.storages import InMemoryStorage
6 from optuna.storages import RDBStorage
7 from optuna.trial import BaseTrial
8 from optuna import TrialPruned
9 from optuna import type_checking
10 
11 if type_checking.TYPE_CHECKING:
12     from datetime import datetime  # NOQA
13     from typing import Any  # NOQA
14     from typing import Callable  # NOQA
15     from typing import Dict  # NOQA
16     from typing import Sequence  # NOQA
17     from typing import Tuple  # NOQA
18     from typing import Type  # NOQA
19     from typing import Union  # NOQA
20 
21     from optuna.distributions import BaseDistribution  # NOQA
22     from optuna.distributions import CategoricalChoiceType  # NOQA
23     from optuna.study import Study  # NOQA
24     from optuna.trial import Trial  # NOQA
25 
26 
27 with try_import() as _imports:
28     from chainermn.communicators.communicator_base import CommunicatorBase  # NOQA
29 
30 
31 class _ChainerMNObjectiveFunc(object):
32     """A wrapper of an objective function to incorporate Optuna with ChainerMN.
33 
34     Note that this class is not supposed to be used by library users.
35 
36     Args:
37         func:
38             A callable that implements objective function.
39         comm:
40             A `ChainerMN communicator <https://docs.chainer.org/en/stable/chainermn/reference/
41             index.html#communicators>`_.
42     """
43 
44     def __init__(self, func, comm):
45         # type: (Callable[[ChainerMNTrial, CommunicatorBase], float], CommunicatorBase) -> None
46 
47         self.comm = comm
48         self.objective = func
49 
50     def __call__(self, trial):
51         # type: (Trial) -> float
52 
53         self.comm.mpi_comm.bcast(True)
54         return self.objective(ChainerMNTrial(trial, self.comm), self.comm)
55 
56 
57 class ChainerMNStudy(object):
58     """A wrapper of :class:`~optuna.study.Study` to incorporate Optuna with ChainerMN.
59 
60     .. seealso::
61         :class:`~optuna.integration.chainermn.ChainerMNStudy` provides the same interface as
62         :class:`~optuna.study.Study`. Please refer to :class:`optuna.study.Study` for further
63         details.
64 
65     See `the example <https://github.com/optuna/optuna/blob/master/
66     examples/pruning/chainermn_integration.py>`__
67     if you want to optimize an objective function that trains neural network
68     written with ChainerMN.
69 
70     Args:
71         study:
72             A :class:`~optuna.study.Study` object.
73         comm:
74             A `ChainerMN communicator <https://docs.chainer.org/en/stable/chainermn/reference/
75             index.html#communicators>`_.
76     """
77 
78     def __init__(
79         self,
80         study,  # type: Study
81         comm,  # type: CommunicatorBase
82     ):
83         # type: (...) -> None
84 
85         _imports.check()
86 
87         if isinstance(study._storage, InMemoryStorage):
88             raise ValueError("ChainerMN integration is not available with InMemoryStorage.")
89 
90         if isinstance(study._storage, RDBStorage):
91             if study._storage.engine.dialect.name == "sqlite":
92                 logger = get_logger(__name__)
93                 logger.warning(
94                     "SQLite may cause synchronization problems when used with "
95                     "ChainerMN integration. Please use other DBs like PostgreSQL."
96                 )
97 
98         study_names = comm.mpi_comm.allgather(study.study_name)
99         if len(set(study_names)) != 1:
100             raise ValueError("Please make sure an identical study name is shared among workers.")
101 
102         super(ChainerMNStudy, self).__setattr__("delegate", study)
103         super(ChainerMNStudy, self).__setattr__("comm", comm)
104 
105     def optimize(
106         self,
107         func,  # type: Callable[[ChainerMNTrial, CommunicatorBase], float]
108         n_trials=None,  # type: Optional[int]
109         timeout=None,  # type: Optional[float]
110         catch=(),  # type: Union[Tuple[()], Tuple[Type[Exception]]]
111     ):
112         # type: (...) -> None
113         """Optimize an objective function.
114 
115         This method provides the same interface as :func:`optuna.study.Study.optimize` except
116         the absence of ``n_jobs`` argument.
117         """
118 
119         if self.comm.rank == 0:
120             func_mn = _ChainerMNObjectiveFunc(func, self.comm)
121             self.delegate.optimize(func_mn, n_trials=n_trials, timeout=timeout, catch=catch)
122             self.comm.mpi_comm.bcast(False)
123         else:
124             while True:
125                 has_next_trial = self.comm.mpi_comm.bcast(None)
126                 if not has_next_trial:
127                     break
128                 try:
129                     func(ChainerMNTrial(None, self.comm), self.comm)
130 
131                     # We assume that if a node raises an exception,
132                     # all other nodes will do the same.
133                     #
134                     # The responsibility to handle acceptable exceptions (i.e., `TrialPruned` and
135                     # `catch`) is in the rank-0 node, so other nodes simply ignore them.
136                 except TrialPruned:
137                     pass
138                 except catch:
139                     pass
140 
141     def __getattr__(self, attr_name):
142         # type: (str) -> Any
143 
144         return getattr(self.delegate, attr_name)
145 
146     def __setattr__(self, attr_name, value):
147         # type: (str, Any) -> None
148 
149         setattr(self.delegate, attr_name, value)
150 
151 
152 class ChainerMNTrial(BaseTrial):
153     """A wrapper of :class:`~optuna.trial.Trial` to incorporate Optuna with ChainerMN.
154 
155     .. seealso::
156         :class:`~optuna.integration.chainermn.ChainerMNTrial` provides the same interface as
157         :class:`~optuna.trial.Trial`. Please refer to :class:`optuna.trial.Trial` for further
158         details.
159 
160     Args:
161         trial:
162             A :class:`~optuna.trial.Trial` object if the caller is rank0 worker,
163             :obj:`None` otherwise.
164         comm:
165             A `ChainerMN communicator <https://docs.chainer.org/en/stable/chainermn/reference/
166             index.html#communicators>`_.
167     """
168 
169     def __init__(self, trial, comm):
170         # type: (Optional[Trial], CommunicatorBase) -> None
171 
172         self.delegate = trial
173         self.comm = comm
174 
175     def suggest_float(
176         self,
177         name: str,
178         low: float,
179         high: float,
180         *,
181         step: Optional[float] = None,
182         log: bool = False
183     ) -> float:
184         def func() -> float:
185             assert self.delegate is not None
186             return self.delegate.suggest_float(name, low, high, log=log, step=step)
187 
188         return self._call_with_mpi(func)
189 
190     def suggest_uniform(self, name, low, high):
191         # type: (str, float, float) -> float
192 
193         def func():
194             # type: () -> float
195 
196             assert self.delegate is not None
197             return self.delegate.suggest_uniform(name, low, high)
198 
199         return self._call_with_mpi(func)
200 
201     def suggest_loguniform(self, name, low, high):
202         # type: (str, float, float) -> float
203 
204         def func():
205             # type: () -> float
206 
207             assert self.delegate is not None
208             return self.delegate.suggest_loguniform(name, low, high)
209 
210         return self._call_with_mpi(func)
211 
212     def suggest_discrete_uniform(self, name, low, high, q):
213         # type: (str, float, float, float) -> float
214 
215         def func():
216             # type: () -> float
217 
218             assert self.delegate is not None
219             return self.delegate.suggest_discrete_uniform(name, low, high, q)
220 
221         return self._call_with_mpi(func)
222 
223     def suggest_int(self, name, low, high, step=1, log=False):
224         # type: (str, int, int, int, bool) -> int
225 
226         def func():
227             # type: () -> int
228 
229             assert self.delegate is not None
230             return self.delegate.suggest_int(name, low, high, step=step, log=log)
231 
232         return self._call_with_mpi(func)
233 
234     def suggest_categorical(self, name, choices):
235         # type: (str, Sequence[CategoricalChoiceType]) -> Any
236 
237         def func():
238             # type: () -> CategoricalChoiceType
239 
240             assert self.delegate is not None
241             return self.delegate.suggest_categorical(name, choices)
242 
243         return self._call_with_mpi(func)
244 
245     def report(self, value, step):
246         # type: (float, int) -> None
247 
248         if self.comm.rank == 0:
249             assert self.delegate is not None
250             self.delegate.report(value, step)
251         self.comm.mpi_comm.barrier()
252 
253     def should_prune(self) -> bool:
254         def func():
255             # type: () -> bool
256 
257             assert self.delegate is not None
258             return self.delegate.should_prune()
259 
260         return self._call_with_mpi(func)
261 
262     def set_user_attr(self, key, value):
263         # type: (str, Any) -> None
264 
265         if self.comm.rank == 0:
266             assert self.delegate is not None
267             self.delegate.set_user_attr(key, value)
268         self.comm.mpi_comm.barrier()
269 
270     def set_system_attr(self, key, value):
271         # type: (str, Any) -> None
272 
273         if self.comm.rank == 0:
274             assert self.delegate is not None
275             self.delegate.set_system_attr(key, value)
276         self.comm.mpi_comm.barrier()
277 
278     @property
279     def number(self):
280         # type: () -> int
281 
282         def func():
283             # type: () -> int
284 
285             assert self.delegate is not None
286             return self.delegate.number
287 
288         return self._call_with_mpi(func)
289 
290     @property
291     def _trial_id(self):
292         # type: () -> int
293 
294         def func():
295             # type: () -> int
296 
297             assert self.delegate is not None
298             return self.delegate._trial_id
299 
300         return self._call_with_mpi(func)
301 
302     @property
303     def params(self):
304         # type: () -> Dict[str, Any]
305 
306         def func():
307             # type: () -> Dict[str, Any]
308 
309             assert self.delegate is not None
310             return self.delegate.params
311 
312         return self._call_with_mpi(func)
313 
314     @property
315     def distributions(self):
316         # type: () -> Dict[str, BaseDistribution]
317 
318         def func():
319             # type: () -> Dict[str, BaseDistribution]
320 
321             assert self.delegate is not None
322             return self.delegate.distributions
323 
324         return self._call_with_mpi(func)
325 
326     @property
327     def user_attrs(self):
328         # type: () -> Dict[str, Any]
329 
330         def func():
331             # type: () -> Dict[str, Any]
332 
333             assert self.delegate is not None
334             return self.delegate.user_attrs
335 
336         return self._call_with_mpi(func)
337 
338     @property
339     def system_attrs(self):
340         # type: () -> Dict[str, Any]
341 
342         def func():
343             # type: () -> Dict[str, Any]
344 
345             assert self.delegate is not None
346             return self.delegate.system_attrs
347 
348         return self._call_with_mpi(func)
349 
350     @property
351     def datetime_start(self):
352         # type: () -> Optional[datetime]
353 
354         def func():
355             # type: () -> Optional[datetime]
356 
357             assert self.delegate is not None
358             return self.delegate.datetime_start
359 
360         return self._call_with_mpi(func)
361 
362     def _call_with_mpi(self, func):
363         # type: (Callable) -> Any
364 
365         if self.comm.rank == 0:
366             try:
367                 result = func()
368                 self.comm.mpi_comm.bcast(result)
369                 return result
370             except Exception as e:
371                 self.comm.mpi_comm.bcast(e)
372                 raise
373         else:
374             result = self.comm.mpi_comm.bcast(None)
375             if isinstance(result, Exception):
376                 raise result
377             return result
378
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/optuna/integration/chainermn.py b/optuna/integration/chainermn.py
--- a/optuna/integration/chainermn.py
+++ b/optuna/integration/chainermn.py
@@ -118,11 +118,13 @@
 
         if self.comm.rank == 0:
             func_mn = _ChainerMNObjectiveFunc(func, self.comm)
-            self.delegate.optimize(func_mn, n_trials=n_trials, timeout=timeout, catch=catch)
-            self.comm.mpi_comm.bcast(False)
+            try:
+                self.delegate.optimize(func_mn, n_trials=n_trials, timeout=timeout, catch=catch)
+            finally:
+                self.comm.mpi_comm.bcast(False)
         else:
+            has_next_trial = self.comm.mpi_comm.bcast(None)
             while True:
-                has_next_trial = self.comm.mpi_comm.bcast(None)
                 if not has_next_trial:
                     break
                 try:
@@ -137,6 +139,8 @@
                     pass
                 except catch:
                     pass
+                finally:
+                    has_next_trial = self.comm.mpi_comm.bcast(None)
 
     def __getattr__(self, attr_name):
         # type: (str) -> Any

{"golden_diff": "diff --git a/optuna/integration/chainermn.py b/optuna/integration/chainermn.py\n--- a/optuna/integration/chainermn.py\n+++ b/optuna/integration/chainermn.py\n@@ -118,11 +118,13 @@\n \n         if self.comm.rank == 0:\n             func_mn = _ChainerMNObjectiveFunc(func, self.comm)\n-            self.delegate.optimize(func_mn, n_trials=n_trials, timeout=timeout, catch=catch)\n-            self.comm.mpi_comm.bcast(False)\n+            try:\n+                self.delegate.optimize(func_mn, n_trials=n_trials, timeout=timeout, catch=catch)\n+            finally:\n+                self.comm.mpi_comm.bcast(False)\n         else:\n+            has_next_trial = self.comm.mpi_comm.bcast(None)\n             while True:\n-                has_next_trial = self.comm.mpi_comm.bcast(None)\n                 if not has_next_trial:\n                     break\n                 try:\n@@ -137,6 +139,8 @@\n                     pass\n                 except catch:\n                     pass\n+                finally:\n+                    has_next_trial = self.comm.mpi_comm.bcast(None)\n \n     def __getattr__(self, attr_name):\n         # type: (str) -> Any\n", "issue": "Occasional test failure with `ChainerMNStudy` with Python 3.5.\n## Expected behavior\r\n\r\nTests should pass.\r\n\r\n## Error messages, stack traces, or logs\r\n\r\nhttps://app.circleci.com/pipelines/github/optuna/optuna/1171/workflows/41e626df-8501-46ca-a724-9b1c27079253/jobs/45897\r\n\r\n## Additional context (optional)\r\n\r\nSeems to be specific to Python 3.5 and only happen randomly.\n", "before_files": [{"content": "from typing import Optional\n\nfrom optuna._imports import try_import\nfrom optuna.logging import get_logger\nfrom optuna.storages import InMemoryStorage\nfrom optuna.storages import RDBStorage\nfrom optuna.trial import BaseTrial\nfrom optuna import TrialPruned\nfrom optuna import type_checking\n\nif type_checking.TYPE_CHECKING:\n    from datetime import datetime  # NOQA\n    from typing import Any  # NOQA\n    from typing import Callable  # NOQA\n    from typing import Dict  # NOQA\n    from typing import Sequence  # NOQA\n    from typing import Tuple  # NOQA\n    from typing import Type  # NOQA\n    from typing import Union  # NOQA\n\n    from optuna.distributions import BaseDistribution  # NOQA\n    from optuna.distributions import CategoricalChoiceType  # NOQA\n    from optuna.study import Study  # NOQA\n    from optuna.trial import Trial  # NOQA\n\n\nwith try_import() as _imports:\n    from chainermn.communicators.communicator_base import CommunicatorBase  # NOQA\n\n\nclass _ChainerMNObjectiveFunc(object):\n    \"\"\"A wrapper of an objective function to incorporate Optuna with ChainerMN.\n\n    Note that this class is not supposed to be used by library users.\n\n    Args:\n        func:\n            A callable that implements objective function.\n        comm:\n            A `ChainerMN communicator <https://docs.chainer.org/en/stable/chainermn/reference/\n            index.html#communicators>`_.\n    \"\"\"\n\n    def __init__(self, func, comm):\n        # type: (Callable[[ChainerMNTrial, CommunicatorBase], float], CommunicatorBase) -> None\n\n        self.comm = comm\n        self.objective = func\n\n    def __call__(self, trial):\n        # type: (Trial) -> float\n\n        self.comm.mpi_comm.bcast(True)\n        return self.objective(ChainerMNTrial(trial, self.comm), self.comm)\n\n\nclass ChainerMNStudy(object):\n    \"\"\"A wrapper of :class:`~optuna.study.Study` to incorporate Optuna with ChainerMN.\n\n    .. seealso::\n        :class:`~optuna.integration.chainermn.ChainerMNStudy` provides the same interface as\n        :class:`~optuna.study.Study`. Please refer to :class:`optuna.study.Study` for further\n        details.\n\n    See `the example <https://github.com/optuna/optuna/blob/master/\n    examples/pruning/chainermn_integration.py>`__\n    if you want to optimize an objective function that trains neural network\n    written with ChainerMN.\n\n    Args:\n        study:\n            A :class:`~optuna.study.Study` object.\n        comm:\n            A `ChainerMN communicator <https://docs.chainer.org/en/stable/chainermn/reference/\n            index.html#communicators>`_.\n    \"\"\"\n\n    def __init__(\n        self,\n        study,  # type: Study\n        comm,  # type: CommunicatorBase\n    ):\n        # type: (...) -> None\n\n        _imports.check()\n\n        if isinstance(study._storage, InMemoryStorage):\n            raise ValueError(\"ChainerMN integration is not available with InMemoryStorage.\")\n\n        if isinstance(study._storage, RDBStorage):\n            if study._storage.engine.dialect.name == \"sqlite\":\n                logger = get_logger(__name__)\n                logger.warning(\n                    \"SQLite may cause synchronization problems when used with \"\n                    \"ChainerMN integration. Please use other DBs like PostgreSQL.\"\n                )\n\n        study_names = comm.mpi_comm.allgather(study.study_name)\n        if len(set(study_names)) != 1:\n            raise ValueError(\"Please make sure an identical study name is shared among workers.\")\n\n        super(ChainerMNStudy, self).__setattr__(\"delegate\", study)\n        super(ChainerMNStudy, self).__setattr__(\"comm\", comm)\n\n    def optimize(\n        self,\n        func,  # type: Callable[[ChainerMNTrial, CommunicatorBase], float]\n        n_trials=None,  # type: Optional[int]\n        timeout=None,  # type: Optional[float]\n        catch=(),  # type: Union[Tuple[()], Tuple[Type[Exception]]]\n    ):\n        # type: (...) -> None\n        \"\"\"Optimize an objective function.\n\n        This method provides the same interface as :func:`optuna.study.Study.optimize` except\n        the absence of ``n_jobs`` argument.\n        \"\"\"\n\n        if self.comm.rank == 0:\n            func_mn = _ChainerMNObjectiveFunc(func, self.comm)\n            self.delegate.optimize(func_mn, n_trials=n_trials, timeout=timeout, catch=catch)\n            self.comm.mpi_comm.bcast(False)\n        else:\n            while True:\n                has_next_trial = self.comm.mpi_comm.bcast(None)\n                if not has_next_trial:\n                    break\n                try:\n                    func(ChainerMNTrial(None, self.comm), self.comm)\n\n                    # We assume that if a node raises an exception,\n                    # all other nodes will do the same.\n                    #\n                    # The responsibility to handle acceptable exceptions (i.e., `TrialPruned` and\n                    # `catch`) is in the rank-0 node, so other nodes simply ignore them.\n                except TrialPruned:\n                    pass\n                except catch:\n                    pass\n\n    def __getattr__(self, attr_name):\n        # type: (str) -> Any\n\n        return getattr(self.delegate, attr_name)\n\n    def __setattr__(self, attr_name, value):\n        # type: (str, Any) -> None\n\n        setattr(self.delegate, attr_name, value)\n\n\nclass ChainerMNTrial(BaseTrial):\n    \"\"\"A wrapper of :class:`~optuna.trial.Trial` to incorporate Optuna with ChainerMN.\n\n    .. seealso::\n        :class:`~optuna.integration.chainermn.ChainerMNTrial` provides the same interface as\n        :class:`~optuna.trial.Trial`. Please refer to :class:`optuna.trial.Trial` for further\n        details.\n\n    Args:\n        trial:\n            A :class:`~optuna.trial.Trial` object if the caller is rank0 worker,\n            :obj:`None` otherwise.\n        comm:\n            A `ChainerMN communicator <https://docs.chainer.org/en/stable/chainermn/reference/\n            index.html#communicators>`_.\n    \"\"\"\n\n    def __init__(self, trial, comm):\n        # type: (Optional[Trial], CommunicatorBase) -> None\n\n        self.delegate = trial\n        self.comm = comm\n\n    def suggest_float(\n        self,\n        name: str,\n        low: float,\n        high: float,\n        *,\n        step: Optional[float] = None,\n        log: bool = False\n    ) -> float:\n        def func() -> float:\n            assert self.delegate is not None\n            return self.delegate.suggest_float(name, low, high, log=log, step=step)\n\n        return self._call_with_mpi(func)\n\n    def suggest_uniform(self, name, low, high):\n        # type: (str, float, float) -> float\n\n        def func():\n            # type: () -> float\n\n            assert self.delegate is not None\n            return self.delegate.suggest_uniform(name, low, high)\n\n        return self._call_with_mpi(func)\n\n    def suggest_loguniform(self, name, low, high):\n        # type: (str, float, float) -> float\n\n        def func():\n            # type: () -> float\n\n            assert self.delegate is not None\n            return self.delegate.suggest_loguniform(name, low, high)\n\n        return self._call_with_mpi(func)\n\n    def suggest_discrete_uniform(self, name, low, high, q):\n        # type: (str, float, float, float) -> float\n\n        def func():\n            # type: () -> float\n\n            assert self.delegate is not None\n            return self.delegate.suggest_discrete_uniform(name, low, high, q)\n\n        return self._call_with_mpi(func)\n\n    def suggest_int(self, name, low, high, step=1, log=False):\n        # type: (str, int, int, int, bool) -> int\n\n        def func():\n            # type: () -> int\n\n            assert self.delegate is not None\n            return self.delegate.suggest_int(name, low, high, step=step, log=log)\n\n        return self._call_with_mpi(func)\n\n    def suggest_categorical(self, name, choices):\n        # type: (str, Sequence[CategoricalChoiceType]) -> Any\n\n        def func():\n            # type: () -> CategoricalChoiceType\n\n            assert self.delegate is not None\n            return self.delegate.suggest_categorical(name, choices)\n\n        return self._call_with_mpi(func)\n\n    def report(self, value, step):\n        # type: (float, int) -> None\n\n        if self.comm.rank == 0:\n            assert self.delegate is not None\n            self.delegate.report(value, step)\n        self.comm.mpi_comm.barrier()\n\n    def should_prune(self) -> bool:\n        def func():\n            # type: () -> bool\n\n            assert self.delegate is not None\n            return self.delegate.should_prune()\n\n        return self._call_with_mpi(func)\n\n    def set_user_attr(self, key, value):\n        # type: (str, Any) -> None\n\n        if self.comm.rank == 0:\n            assert self.delegate is not None\n            self.delegate.set_user_attr(key, value)\n        self.comm.mpi_comm.barrier()\n\n    def set_system_attr(self, key, value):\n        # type: (str, Any) -> None\n\n        if self.comm.rank == 0:\n            assert self.delegate is not None\n            self.delegate.set_system_attr(key, value)\n        self.comm.mpi_comm.barrier()\n\n    @property\n    def number(self):\n        # type: () -> int\n\n        def func():\n            # type: () -> int\n\n            assert self.delegate is not None\n            return self.delegate.number\n\n        return self._call_with_mpi(func)\n\n    @property\n    def _trial_id(self):\n        # type: () -> int\n\n        def func():\n            # type: () -> int\n\n            assert self.delegate is not None\n            return self.delegate._trial_id\n\n        return self._call_with_mpi(func)\n\n    @property\n    def params(self):\n        # type: () -> Dict[str, Any]\n\n        def func():\n            # type: () -> Dict[str, Any]\n\n            assert self.delegate is not None\n            return self.delegate.params\n\n        return self._call_with_mpi(func)\n\n    @property\n    def distributions(self):\n        # type: () -> Dict[str, BaseDistribution]\n\n        def func():\n            # type: () -> Dict[str, BaseDistribution]\n\n            assert self.delegate is not None\n            return self.delegate.distributions\n\n        return self._call_with_mpi(func)\n\n    @property\n    def user_attrs(self):\n        # type: () -> Dict[str, Any]\n\n        def func():\n            # type: () -> Dict[str, Any]\n\n            assert self.delegate is not None\n            return self.delegate.user_attrs\n\n        return self._call_with_mpi(func)\n\n    @property\n    def system_attrs(self):\n        # type: () -> Dict[str, Any]\n\n        def func():\n            # type: () -> Dict[str, Any]\n\n            assert self.delegate is not None\n            return self.delegate.system_attrs\n\n        return self._call_with_mpi(func)\n\n    @property\n    def datetime_start(self):\n        # type: () -> Optional[datetime]\n\n        def func():\n            # type: () -> Optional[datetime]\n\n            assert self.delegate is not None\n            return self.delegate.datetime_start\n\n        return self._call_with_mpi(func)\n\n    def _call_with_mpi(self, func):\n        # type: (Callable) -> Any\n\n        if self.comm.rank == 0:\n            try:\n                result = func()\n                self.comm.mpi_comm.bcast(result)\n                return result\n            except Exception as e:\n                self.comm.mpi_comm.bcast(e)\n                raise\n        else:\n            result = self.comm.mpi_comm.bcast(None)\n            if isinstance(result, Exception):\n                raise result\n            return result\n", "path": "optuna/integration/chainermn.py"}], "after_files": [{"content": "from typing import Optional\n\nfrom optuna._imports import try_import\nfrom optuna.logging import get_logger\nfrom optuna.storages import InMemoryStorage\nfrom optuna.storages import RDBStorage\nfrom optuna.trial import BaseTrial\nfrom optuna import TrialPruned\nfrom optuna import type_checking\n\nif type_checking.TYPE_CHECKING:\n    from datetime import datetime  # NOQA\n    from typing import Any  # NOQA\n    from typing import Callable  # NOQA\n    from typing import Dict  # NOQA\n    from typing import Sequence  # NOQA\n    from typing import Tuple  # NOQA\n    from typing import Type  # NOQA\n    from typing import Union  # NOQA\n\n    from optuna.distributions import BaseDistribution  # NOQA\n    from optuna.distributions import CategoricalChoiceType  # NOQA\n    from optuna.study import Study  # NOQA\n    from optuna.trial import Trial  # NOQA\n\n\nwith try_import() as _imports:\n    from chainermn.communicators.communicator_base import CommunicatorBase  # NOQA\n\n\nclass _ChainerMNObjectiveFunc(object):\n    \"\"\"A wrapper of an objective function to incorporate Optuna with ChainerMN.\n\n    Note that this class is not supposed to be used by library users.\n\n    Args:\n        func:\n            A callable that implements objective function.\n        comm:\n            A `ChainerMN communicator <https://docs.chainer.org/en/stable/chainermn/reference/\n            index.html#communicators>`_.\n    \"\"\"\n\n    def __init__(self, func, comm):\n        # type: (Callable[[ChainerMNTrial, CommunicatorBase], float], CommunicatorBase) -> None\n\n        self.comm = comm\n        self.objective = func\n\n    def __call__(self, trial):\n        # type: (Trial) -> float\n\n        self.comm.mpi_comm.bcast(True)\n        return self.objective(ChainerMNTrial(trial, self.comm), self.comm)\n\n\nclass ChainerMNStudy(object):\n    \"\"\"A wrapper of :class:`~optuna.study.Study` to incorporate Optuna with ChainerMN.\n\n    .. seealso::\n        :class:`~optuna.integration.chainermn.ChainerMNStudy` provides the same interface as\n        :class:`~optuna.study.Study`. Please refer to :class:`optuna.study.Study` for further\n        details.\n\n    See `the example <https://github.com/optuna/optuna/blob/master/\n    examples/pruning/chainermn_integration.py>`__\n    if you want to optimize an objective function that trains neural network\n    written with ChainerMN.\n\n    Args:\n        study:\n            A :class:`~optuna.study.Study` object.\n        comm:\n            A `ChainerMN communicator <https://docs.chainer.org/en/stable/chainermn/reference/\n            index.html#communicators>`_.\n    \"\"\"\n\n    def __init__(\n        self,\n        study,  # type: Study\n        comm,  # type: CommunicatorBase\n    ):\n        # type: (...) -> None\n\n        _imports.check()\n\n        if isinstance(study._storage, InMemoryStorage):\n            raise ValueError(\"ChainerMN integration is not available with InMemoryStorage.\")\n\n        if isinstance(study._storage, RDBStorage):\n            if study._storage.engine.dialect.name == \"sqlite\":\n                logger = get_logger(__name__)\n                logger.warning(\n                    \"SQLite may cause synchronization problems when used with \"\n                    \"ChainerMN integration. Please use other DBs like PostgreSQL.\"\n                )\n\n        study_names = comm.mpi_comm.allgather(study.study_name)\n        if len(set(study_names)) != 1:\n            raise ValueError(\"Please make sure an identical study name is shared among workers.\")\n\n        super(ChainerMNStudy, self).__setattr__(\"delegate\", study)\n        super(ChainerMNStudy, self).__setattr__(\"comm\", comm)\n\n    def optimize(\n        self,\n        func,  # type: Callable[[ChainerMNTrial, CommunicatorBase], float]\n        n_trials=None,  # type: Optional[int]\n        timeout=None,  # type: Optional[float]\n        catch=(),  # type: Union[Tuple[()], Tuple[Type[Exception]]]\n    ):\n        # type: (...) -> None\n        \"\"\"Optimize an objective function.\n\n        This method provides the same interface as :func:`optuna.study.Study.optimize` except\n        the absence of ``n_jobs`` argument.\n        \"\"\"\n\n        if self.comm.rank == 0:\n            func_mn = _ChainerMNObjectiveFunc(func, self.comm)\n            try:\n                self.delegate.optimize(func_mn, n_trials=n_trials, timeout=timeout, catch=catch)\n            finally:\n                self.comm.mpi_comm.bcast(False)\n        else:\n            has_next_trial = self.comm.mpi_comm.bcast(None)\n            while True:\n                if not has_next_trial:\n                    break\n                try:\n                    func(ChainerMNTrial(None, self.comm), self.comm)\n\n                    # We assume that if a node raises an exception,\n                    # all other nodes will do the same.\n                    #\n                    # The responsibility to handle acceptable exceptions (i.e., `TrialPruned` and\n                    # `catch`) is in the rank-0 node, so other nodes simply ignore them.\n                except TrialPruned:\n                    pass\n                except catch:\n                    pass\n                finally:\n                    has_next_trial = self.comm.mpi_comm.bcast(None)\n\n    def __getattr__(self, attr_name):\n        # type: (str) -> Any\n\n        return getattr(self.delegate, attr_name)\n\n    def __setattr__(self, attr_name, value):\n        # type: (str, Any) -> None\n\n        setattr(self.delegate, attr_name, value)\n\n\nclass ChainerMNTrial(BaseTrial):\n    \"\"\"A wrapper of :class:`~optuna.trial.Trial` to incorporate Optuna with ChainerMN.\n\n    .. seealso::\n        :class:`~optuna.integration.chainermn.ChainerMNTrial` provides the same interface as\n        :class:`~optuna.trial.Trial`. Please refer to :class:`optuna.trial.Trial` for further\n        details.\n\n    Args:\n        trial:\n            A :class:`~optuna.trial.Trial` object if the caller is rank0 worker,\n            :obj:`None` otherwise.\n        comm:\n            A `ChainerMN communicator <https://docs.chainer.org/en/stable/chainermn/reference/\n            index.html#communicators>`_.\n    \"\"\"\n\n    def __init__(self, trial, comm):\n        # type: (Optional[Trial], CommunicatorBase) -> None\n\n        self.delegate = trial\n        self.comm = comm\n\n    def suggest_float(\n        self,\n        name: str,\n        low: float,\n        high: float,\n        *,\n        step: Optional[float] = None,\n        log: bool = False\n    ) -> float:\n        def func() -> float:\n            assert self.delegate is not None\n            return self.delegate.suggest_float(name, low, high, log=log, step=step)\n\n        return self._call_with_mpi(func)\n\n    def suggest_uniform(self, name, low, high):\n        # type: (str, float, float) -> float\n\n        def func():\n            # type: () -> float\n\n            assert self.delegate is not None\n            return self.delegate.suggest_uniform(name, low, high)\n\n        return self._call_with_mpi(func)\n\n    def suggest_loguniform(self, name, low, high):\n        # type: (str, float, float) -> float\n\n        def func():\n            # type: () -> float\n\n            assert self.delegate is not None\n            return self.delegate.suggest_loguniform(name, low, high)\n\n        return self._call_with_mpi(func)\n\n    def suggest_discrete_uniform(self, name, low, high, q):\n        # type: (str, float, float, float) -> float\n\n        def func():\n            # type: () -> float\n\n            assert self.delegate is not None\n            return self.delegate.suggest_discrete_uniform(name, low, high, q)\n\n        return self._call_with_mpi(func)\n\n    def suggest_int(self, name, low, high, step=1, log=False):\n        # type: (str, int, int, int, bool) -> int\n\n        def func():\n            # type: () -> int\n\n            assert self.delegate is not None\n            return self.delegate.suggest_int(name, low, high, step=step, log=log)\n\n        return self._call_with_mpi(func)\n\n    def suggest_categorical(self, name, choices):\n        # type: (str, Sequence[CategoricalChoiceType]) -> Any\n\n        def func():\n            # type: () -> CategoricalChoiceType\n\n            assert self.delegate is not None\n            return self.delegate.suggest_categorical(name, choices)\n\n        return self._call_with_mpi(func)\n\n    def report(self, value, step):\n        # type: (float, int) -> None\n\n        if self.comm.rank == 0:\n            assert self.delegate is not None\n            self.delegate.report(value, step)\n        self.comm.mpi_comm.barrier()\n\n    def should_prune(self) -> bool:\n        def func():\n            # type: () -> bool\n\n            assert self.delegate is not None\n            return self.delegate.should_prune()\n\n        return self._call_with_mpi(func)\n\n    def set_user_attr(self, key, value):\n        # type: (str, Any) -> None\n\n        if self.comm.rank == 0:\n            assert self.delegate is not None\n            self.delegate.set_user_attr(key, value)\n        self.comm.mpi_comm.barrier()\n\n    def set_system_attr(self, key, value):\n        # type: (str, Any) -> None\n\n        if self.comm.rank == 0:\n            assert self.delegate is not None\n            self.delegate.set_system_attr(key, value)\n        self.comm.mpi_comm.barrier()\n\n    @property\n    def number(self):\n        # type: () -> int\n\n        def func():\n            # type: () -> int\n\n            assert self.delegate is not None\n            return self.delegate.number\n\n        return self._call_with_mpi(func)\n\n    @property\n    def _trial_id(self):\n        # type: () -> int\n\n        def func():\n            # type: () -> int\n\n            assert self.delegate is not None\n            return self.delegate._trial_id\n\n        return self._call_with_mpi(func)\n\n    @property\n    def params(self):\n        # type: () -> Dict[str, Any]\n\n        def func():\n            # type: () -> Dict[str, Any]\n\n            assert self.delegate is not None\n            return self.delegate.params\n\n        return self._call_with_mpi(func)\n\n    @property\n    def distributions(self):\n        # type: () -> Dict[str, BaseDistribution]\n\n        def func():\n            # type: () -> Dict[str, BaseDistribution]\n\n            assert self.delegate is not None\n            return self.delegate.distributions\n\n        return self._call_with_mpi(func)\n\n    @property\n    def user_attrs(self):\n        # type: () -> Dict[str, Any]\n\n        def func():\n            # type: () -> Dict[str, Any]\n\n            assert self.delegate is not None\n            return self.delegate.user_attrs\n\n        return self._call_with_mpi(func)\n\n    @property\n    def system_attrs(self):\n        # type: () -> Dict[str, Any]\n\n        def func():\n            # type: () -> Dict[str, Any]\n\n            assert self.delegate is not None\n            return self.delegate.system_attrs\n\n        return self._call_with_mpi(func)\n\n    @property\n    def datetime_start(self):\n        # type: () -> Optional[datetime]\n\n        def func():\n            # type: () -> Optional[datetime]\n\n            assert self.delegate is not None\n            return self.delegate.datetime_start\n\n        return self._call_with_mpi(func)\n\n    def _call_with_mpi(self, func):\n        # type: (Callable) -> Any\n\n        if self.comm.rank == 0:\n            try:\n                result = func()\n                self.comm.mpi_comm.bcast(result)\n                return result\n            except Exception as e:\n                self.comm.mpi_comm.bcast(e)\n                raise\n        else:\n            result = self.comm.mpi_comm.bcast(None)\n            if isinstance(result, Exception):\n                raise result\n            return result\n", "path": "optuna/integration/chainermn.py"}]}