An Optuna storage that uses multiprocessing module for inter process communication

multi_process_storage.py

import copy
from datetime import datetime
import multiprocessing

from optuna import distributions  # NOQA
from optuna.storages import base
from optuna.storages.base import DEFAULT_STUDY_NAME_PREFIX
from optuna import structs
from optuna import type_checking

if type_checking.TYPE_CHECKING:
    from typing import Any  # NOQA
    from typing import Dict  # NOQA
    from typing import List  # NOQA
    from typing import Optional  # NOQA

IN_MEMORY_STORAGE_STUDY_ID = 0
IN_MEMORY_STORAGE_STUDY_UUID = '00000000-0000-0000-0000-000000000000'


class MultiProcessStorage(base.BaseStorage):
    def __init__(self):
        # type: () -> None

        self._manager = multiprocessing.Manager()
        self._next_trial_id = self._manager.Value('i', 0)
        self._lock = self._manager.Lock()

        self._cached_trials = {}

        self.trials = self._manager.list()
        self.param_distribution = self._manager.dict()
        self.direction = structs.StudyDirection.NOT_SET
        self.study_user_attrs = self._manager.dict()
        self.study_system_attrs = self._manager.dict()
        self.study_name = DEFAULT_STUDY_NAME_PREFIX + IN_MEMORY_STORAGE_STUDY_UUID  # type: str

    def create_new_study_id(self, study_name=None):
        # type: (Optional[str]) -> int

        if study_name is not None:
            self.study_name = study_name

        return IN_MEMORY_STORAGE_STUDY_ID  # TODO(akiba)

    def set_study_direction(self, study_id, direction):
        # type: (int, structs.StudyDirection) -> None

        if self.direction != structs.StudyDirection.NOT_SET and self.direction != direction:
            raise ValueError('Cannot overwrite study direction from {} to {}.'.format(
                self.direction, direction))
        self.direction = direction

    def set_study_user_attr(self, study_id, key, value):
        # type: (int, str, Any) -> None

        self.study_user_attrs[key] = value

    def set_study_system_attr(self, study_id, key, value):
        # type: (int, str, Any) -> None

        self.study_system_attrs[key] = value

    def get_study_id_from_name(self, study_name):
        # type: (str) -> int

        if study_name != self.study_name:
            raise ValueError("No such study {}.".format(study_name))

        return IN_MEMORY_STORAGE_STUDY_ID

    def get_study_id_from_trial_id(self, trial_id):
        # type: (int) -> int

        return IN_MEMORY_STORAGE_STUDY_ID

    def get_study_name_from_id(self, study_id):
        # type: (int) -> str

        self._check_study_id(study_id)
        return self.study_name

    def get_study_direction(self, study_id):
        # type: (int) -> structs.StudyDirection

        return self.direction

    def get_study_user_attrs(self, study_id):
        # type: (int) -> Dict[str, Any]

        return copy.copy(self.study_user_attrs)

    def get_study_system_attrs(self, study_id):
        # type: (int) -> Dict[str, Any]

        return copy.copy(self.study_system_attrs)

    def get_all_study_summaries(self):
        # type: () -> List[structs.StudySummary]

        best_trial = None
        n_complete_trials = len([t for t in self.trials if t.state == structs.TrialState.COMPLETE])
        if n_complete_trials > 0:
            best_trial = self.get_best_trial(IN_MEMORY_STORAGE_STUDY_ID)

        datetime_start = None
        if len(self.trials) > 0:
            datetime_start = min([t.datetime_start for t in self.trials])

        return [
            structs.StudySummary(
                study_id=IN_MEMORY_STORAGE_STUDY_ID,
                study_name=self.study_name,
                direction=self.direction,
                best_trial=best_trial,
                user_attrs=copy.copy(self.study_user_attrs),
                system_attrs=copy.copy(self.study_system_attrs),
                n_trials=len(self.trials),
                datetime_start=datetime_start)
        ]

    def create_new_trial_id(self, study_id):
        # type: (int) -> int

        self._check_study_id(study_id)

        with self._lock:
            trial_id = self._next_trial_id.value

            self.trials.append(
                structs.FrozenTrial(
                    number=trial_id,
                    state=structs.TrialState.RUNNING,
                    params={},
                    distributions={},
                    user_attrs={},
                    system_attrs={'_number': trial_id},
                    value=None,
                    intermediate_values={},
                    datetime_start=datetime.now(),
                    datetime_complete=None,
                    trial_id=trial_id))

            self._next_trial_id.value += 1

        return trial_id

    def set_trial_state(self, trial_id, state):
        # type: (int, structs.TrialState) -> None

        self.check_trial_is_updatable(trial_id, self.trials[trial_id].state)

        self.trials[trial_id] = self.trials[trial_id]._replace(state=state)
        if state.is_finished():
            self.trials[trial_id] = self.trials[trial_id]._replace(datetime_complete=datetime.now())

    def set_trial_param(self, trial_id, param_name, param_value_internal, distribution):
        # type: (int, str, float, distributions.BaseDistribution) -> bool

        self.check_trial_is_updatable(trial_id, self.trials[trial_id].state)

        # Check param distribution compatibility with previous trial(s).
        if param_name in self.param_distribution:
            distributions.check_distribution_compatibility(self.param_distribution[param_name],
                                                           distribution)

        # Check param has not been set; otherwise, return False.
        if param_name in self.trials[trial_id].params:
            return False

        # Set param distribution.
        self.param_distribution[param_name] = distribution

        # Set param.
        trial = copy.copy(self.trials[trial_id])
        trial.params[param_name] = distribution.to_external_repr(param_value_internal)
        trial.distributions[param_name] = distribution
        self.trials[trial_id] = trial

        return True

    def get_trial_number_from_id(self, trial_id):
        # type: (int) -> int

        return trial_id

    def get_trial_param(self, trial_id, param_name):
        # type: (int, str) -> float

        distribution = self.trials[trial_id].distributions[param_name]
        return distribution.to_internal_repr(self.trials[trial_id].params[param_name])

    def set_trial_value(self, trial_id, value):
        # type: (int, float) -> None

        self.check_trial_is_updatable(trial_id, self.trials[trial_id].state)

        self.trials[trial_id] = self.trials[trial_id]._replace(value=value)

    def set_trial_intermediate_value(self, trial_id, step, intermediate_value):
        # type: (int, int, float) -> bool

        self.check_trial_is_updatable(trial_id, self.trials[trial_id].state)

        if step in self.trials[trial_id].intermediate_values:
            return False

        trial = copy.copy(self.trials[trial_id])
        trial.intermediate_values[step] = intermediate_value
        self.trials[trial_id] = trial

        return True

    def set_trial_user_attr(self, trial_id, key, value):
        # type: (int, str, Any) -> None

        self.check_trial_is_updatable(trial_id, self.trials[trial_id].state)

        trial = copy.copy(self.trials[trial_id])
        trial.user_attrs[key] = value
        self.trials[trial_id] = trial

    def set_trial_system_attr(self, trial_id, key, value):
        # type: (int, str, Any) -> None

        self.check_trial_is_updatable(trial_id, self.trials[trial_id].state)

        trial = copy.copy(self.trials[trial_id])
        trial.system_attrs[key] = value
        self.trials[trial_id] = trial

    def get_trial(self, trial_id):
        # type: (int) -> structs.FrozenTrial

        return copy.copy(self.trials[trial_id])

    def get_all_trials(self, study_id):
        # type: (int) -> List[structs.FrozenTrial]

        self._check_study_id(study_id)

        trials = []

        for i in range(self._next_trial_id.value):
            if i in self._cached_trials:
                trials.append(copy.deepcopy(self._cached_trials[i]))
            else:
                trial = copy.copy(self.trials[i])
                trials.append(trial)
                if trial.state.is_finished():
                    self._cached_trials[i] = copy.deepcopy(trial)

        return trials

    def get_n_trials(self, study_id, state=None):
        # type: (int, Optional[structs.TrialState]) -> int

        self._check_study_id(study_id)
        if state is None:
            return len(self.trials)

        return len([t for t in self.trials if t.state == state])

    def _check_study_id(self, study_id):
        # type: (int) -> None

        if study_id != IN_MEMORY_STORAGE_STUDY_ID:
            raise ValueError('study_id is supposed to be {} in {}.'.format(
                IN_MEMORY_STORAGE_STUDY_ID, self.__class__.__name__))

quadratic_simple.py

from multiprocessing import Process
from multi_process_storage import MultiProcessStorage
import optuna

# Define a simple 2-dimensional objective function whose minimum value is -1 when (x, y) = (0, -1).
def objective(trial):
    x = trial.suggest_uniform('x', -100, 100)
    y = trial.suggest_categorical('y', [-1, 0, 1])
    return x**2 + y


def optimize(storage, n_trials):
    study = optuna.create_study(storage=storage)
    study.optimize(objective, n_trials=n_trials)

if __name__ == '__main__':
    # Let us minimize the objective function above.
    print('Running 100 trials...')

    n_trials = 100
    concurrency = 10
    storage = MultiProcessStorage()

    workers = [Process(target=optimize, args=(storage, n_trials/concurrency)) for _ in range(concurrency)]
    for worker in workers:
        worker.start()

    for worker in workers:
        worker.join()

    study = optuna.create_study(storage=storage)
    print('Number of finished trials: ', len(study.trials))
    print('Best trial:')
    trial = study.best_trial

    print('  Value: ', trial.value)
    print('  Params: ')
    for key, value in trial.params.items():
        print('    {}: {}'.format(key, value))