Make Storage objects picklable and resumable by mp.Pool workers (microsoft#974)

# Pull Request

## Title

Make Storage objects picklable and resumable by mp.Pool workers.

______________________________________________________________________

## Description

To support a `ParallelTrialScheduler` with `Trials` running thru
`TrialRunners` in separate processes using `multiprocessing.Pool` we
need to be able to "send" all necessary state to "blank" python
processes using pickle since python doesn't actually use a `fork`
strategy for multiprocessing.

This change enables that by doing the following:

1. Allow pickling the core `Storage` class by hooking `__getstate__` and
`__setstate__` to disconnect the `Engine` prior to pickling and
reconnect on unpickling (restore).
2. Allow recreating an `Experiment` and `Trial` object by fetching it
from the `Storage` by `experiment_id` and `trial_id` respectively.

These two things allow gathering the details necessary for a
`TrialRunner` running in a separate child process to reconnect to the DB
independently, `run_trial`, and save the telemetry and results to the
DB.
After which, the "MainProcess" can `load` the results from the
Experiment from the Storage again to `bulk_register` the scores from
those completed `Trials` and start a new set.

______________________________________________________________________

## Type of Change

- 🛠️ Bug fix
- ✨ New feature
- 🔄 Refactor
- 🧪 Tests

______________________________________________________________________

## Testing

- [x] New CI tests

______________________________________________________________________

## Additional Notes (optional)

This change should have no impact on the existing `SyncScheduler`.

______________________________________________________________________

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: bpkroth

mlos_bench/mlos_bench/environments/status.py

@@ -5,6 +5,9 @@
 """Enum for the status of the benchmark/environment Trial or Experiment."""
 
 import enum
+import logging
+
+_LOG = logging.getLogger(__name__)
 
 
 class Status(enum.Enum):
@@ -19,6 +22,21 @@ class Status(enum.Enum):
     FAILED = 6
     TIMED_OUT = 7
 
+    @staticmethod
+    def from_str(status_str: str) -> "Status":
+        """Convert a string to a Status enum."""
+        if status_str.isdigit():
+            try:
+                return Status(int(status_str))
+            except ValueError:
+                _LOG.warning("Unknown status: %d", int(status_str))
+        try:
+            status_str = status_str.upper()
+            return Status[status_str]
+        except KeyError:
+            _LOG.warning("Unknown status: %s", status_str)
+        return Status.UNKNOWN
+
     def is_good(self) -> bool:
         """Check if the status of the benchmark/environment is good."""
         return self in {
@@ -73,4 +91,4 @@ def is_timed_out(self) -> bool:
         """Check if the status of the benchmark/environment Trial or Experiment is
         TIMED_OUT.
         """
-        return self == Status.FAILED
+        return self == Status.TIMED_OUT

mlos_bench/mlos_bench/storage/base_storage.py

@@ -22,6 +22,8 @@
     Base interface for accessing the stored benchmark trial data.
 """
 
+from __future__ import annotations
+
 import logging
 from abc import ABCMeta, abstractmethod
 from collections.abc import Iterator, Mapping
@@ -94,6 +96,42 @@ def experiments(self) -> dict[str, ExperimentData]:
             A dictionary of the experiments' data, keyed by experiment id.
         """
 
+    @abstractmethod
+    def get_experiment_by_id(
+        self,
+        experiment_id: str,
+        tunables: TunableGroups,
+        opt_targets: dict[str, Literal["min", "max"]],
+    ) -> Storage.Experiment | None:
+        """
+        Gets an Experiment by its ID.
+
+        Parameters
+        ----------
+        experiment_id : str
+            ID of the Experiment to retrieve.
+        tunables : TunableGroups
+            The tunables for the Experiment.
+        opt_targets : dict[str, Literal["min", "max"]]
+            The optimization targets for the Experiment's
+            :py:class:`~mlos_bench.optimizers.base_optimizer.Optimizer`.
+
+        Returns
+        -------
+        experiment : Storage.Experiment | None
+            The Experiment object, or None if it doesn't exist.
+
+        Notes
+        -----
+        Tunables are not stored in the database for the Experiment, only for the
+        Trials, so currently they can change if the user (incorrectly) adjusts
+        the configs on disk between resume runs.
+        Since this method is generally meant to load th Experiment from the
+        database for a child process to execute a Trial in the background we are
+        generally safe to simply pass these values from the parent process
+        rather than look them up in the database.
+        """
+
     @abstractmethod
     def experiment(  # pylint: disable=too-many-arguments
         self,
@@ -104,10 +142,12 @@ def experiment(  # pylint: disable=too-many-arguments
         description: str,
         tunables: TunableGroups,
         opt_targets: dict[str, Literal["min", "max"]],
-    ) -> "Storage.Experiment":
+    ) -> Storage.Experiment:
         """
-        Create a new experiment in the storage.
+        Create or reload an experiment in the Storage.
 
+        Notes
+        -----
         We need the `opt_target` parameter here to know what metric to retrieve
         when we load the data from previous trials. Later we will replace it with
         full metadata about the optimization direction, multiple objectives, etc.
@@ -161,7 +201,7 @@ def __init__(  # pylint: disable=too-many-arguments
             self._opt_targets = opt_targets
             self._in_context = False
 
-        def __enter__(self) -> "Storage.Experiment":
+        def __enter__(self) -> Storage.Experiment:
             """
             Enter the context of the experiment.
 
@@ -307,14 +347,33 @@ def load(
                 Trial ids, Tunable values, benchmark scores, and status of the trials.
             """
 
+        @abstractmethod
+        def get_trial_by_id(
+            self,
+            trial_id: int,
+        ) -> Storage.Trial | None:
+            """
+            Gets a Trial by its ID.
+
+            Parameters
+            ----------
+            trial_id : int
+                ID of the Trial to retrieve for this Experiment.
+
+            Returns
+            -------
+            trial : Storage.Trial | None
+                The Trial object, or None if it doesn't exist.
+            """
+
         @abstractmethod
         def pending_trials(
             self,
             timestamp: datetime,
             *,
             running: bool,
             trial_runner_assigned: bool | None = None,
-        ) -> Iterator["Storage.Trial"]:
+        ) -> Iterator[Storage.Trial]:
             """
             Return an iterator over :py:attr:`~.Status.PENDING`
             :py:class:`~.Storage.Trial` instances that have a scheduled start time to
@@ -345,7 +404,7 @@ def new_trial(
             tunables: TunableGroups,
             ts_start: datetime | None = None,
             config: dict[str, Any] | None = None,
-        ) -> "Storage.Trial":
+        ) -> Storage.Trial:
             """
             Create a new experiment run in the storage.
 
@@ -382,7 +441,7 @@ def _new_trial(
             tunables: TunableGroups,
             ts_start: datetime | None = None,
             config: dict[str, Any] | None = None,
-        ) -> "Storage.Trial":
+        ) -> Storage.Trial:
             """
             Create a new experiment run in the storage.
 
@@ -419,10 +478,11 @@ def __init__(  # pylint: disable=too-many-arguments
             tunable_config_id: int,
             trial_runner_id: int | None,
             opt_targets: dict[str, Literal["min", "max"]],
+            status: Status,
+            restoring: bool,
             config: dict[str, Any] | None = None,
-            status: Status = Status.UNKNOWN,
         ):
-            if status not in (Status.UNKNOWN, Status.PENDING):
+            if not restoring and status not in (Status.UNKNOWN, Status.PENDING):
                 raise ValueError(f"Invalid status for a new trial: {status}")
             self._tunables = tunables
             self._experiment_id = experiment_id
@@ -439,6 +499,11 @@ def __repr__(self) -> str:
                 f"{self._tunable_config_id}:{self.trial_runner_id}"
             )
 
+        @property
+        def experiment_id(self) -> str:
+            """Experiment ID of the Trial."""
+            return self._experiment_id
+
         @property
         def trial_id(self) -> int:
             """ID of the current trial."""

mlos_bench/mlos_bench/storage/sql/common.py

@@ -95,7 +95,7 @@ def get_trials(
                 config_id=trial.config_id,
                 ts_start=utcify_timestamp(trial.ts_start, origin="utc"),
                 ts_end=utcify_nullable_timestamp(trial.ts_end, origin="utc"),
-                status=Status[trial.status],
+                status=Status.from_str(trial.status),
                 trial_runner_id=trial.trial_runner_id,
             )
             for trial in trials.fetchall()

mlos_bench/mlos_bench/storage/sql/experiment.py

@@ -188,7 +188,7 @@ def load(
             status: list[Status] = []
 
             for trial in cur_trials.fetchall():
-                stat = Status[trial.status]
+                stat = Status.from_str(trial.status)
                 status.append(stat)
                 trial_ids.append(trial.trial_id)
                 configs.append(
@@ -235,6 +235,48 @@ def _get_key_val(conn: Connection, table: Table, field: str, **kwargs: Any) -> d
             row._tuple() for row in cur_result.fetchall()  # pylint: disable=protected-access
         )
 
+    def get_trial_by_id(
+        self,
+        trial_id: int,
+    ) -> Storage.Trial | None:
+        with self._engine.connect() as conn:
+            cur_trial = conn.execute(
+                self._schema.trial.select().where(
+                    self._schema.trial.c.exp_id == self._experiment_id,
+                    self._schema.trial.c.trial_id == trial_id,
+                )
+            )
+            trial = cur_trial.fetchone()
+            if trial is None:
+                return None
+            tunables = self._get_key_val(
+                conn,
+                self._schema.config_param,
+                "param",
+                config_id=trial.config_id,
+            )
+            config = self._get_key_val(
+                conn,
+                self._schema.trial_param,
+                "param",
+                exp_id=self._experiment_id,
+                trial_id=trial.trial_id,
+            )
+            return Trial(
+                engine=self._engine,
+                schema=self._schema,
+                # Reset .is_updated flag after the assignment:
+                tunables=self._tunables.copy().assign(tunables).reset(),
+                experiment_id=self._experiment_id,
+                trial_id=trial.trial_id,
+                config_id=trial.config_id,
+                trial_runner_id=trial.trial_runner_id,
+                opt_targets=self._opt_targets,
+                status=Status.from_str(trial.status),
+                restoring=True,
+                config=config,
+            )
+
     def pending_trials(
         self,
         timestamp: datetime,
@@ -288,6 +330,8 @@ def pending_trials(
                     config_id=trial.config_id,
                     trial_runner_id=trial.trial_runner_id,
                     opt_targets=self._opt_targets,
+                    status=Status.from_str(trial.status),
+                    restoring=True,
                     config=config,
                 )
 
@@ -363,8 +407,9 @@ def _new_trial(
                     config_id=config_id,
                     trial_runner_id=None,  # initially, Trials are not assigned to a TrialRunner
                     opt_targets=self._opt_targets,
-                    config=config,
                     status=new_trial_status,
+                    restoring=False,
+                    config=config,
                 )
                 self._trial_id += 1
                 return trial

mlos_bench/mlos_bench/storage/sql/storage.py

@@ -7,7 +7,7 @@
 import logging
 from typing import Literal
 
-from sqlalchemy import URL, create_engine
+from sqlalchemy import URL, Engine, create_engine
 
 from mlos_bench.services.base_service import Service
 from mlos_bench.storage.base_experiment_data import ExperimentData
@@ -25,28 +25,53 @@ class SqlStorage(Storage):
     backend.
     """
 
+    # pylint: disable=too-many-instance-attributes
+
     def __init__(
         self,
         config: dict,
         global_config: dict | None = None,
         service: Service | None = None,
     ):
         super().__init__(config, global_config, service)
-        lazy_schema_create = self._config.pop("lazy_schema_create", False)
+        self._lazy_schema_create = self._config.pop("lazy_schema_create", False)
         self._log_sql = self._config.pop("log_sql", False)
         self._url = URL.create(**self._config)
         self._repr = f"{self._url.get_backend_name()}:{self._url.database}"
+        self._engine: Engine
+        self._db_schema: DbSchema
+        self._schema_created = False
+        self._schema_updated = False
+        self._init_engine()
+
+    def _init_engine(self) -> None:
+        """Initialize the SQLAlchemy engine."""
+        # This is a no-op, as the engine is created in __init__.
         _LOG.info("Connect to the database: %s", self)
         self._engine = create_engine(self._url, echo=self._log_sql)
         self._db_schema = DbSchema(self._engine)
-        self._schema_created = False
-        self._schema_updated = False
-        if not lazy_schema_create:
+        if not self._lazy_schema_create:
             assert self._schema
             self.update_schema()
         else:
             _LOG.info("Using lazy schema create for database: %s", self)
 
+    # Make the object picklable.
+
+    def __getstate__(self) -> dict:
+        """Return the state of the object for pickling."""
+        state = self.__dict__.copy()
+        # Don't pickle the engine, as it cannot be pickled.
+        state.pop("_engine", None)
+        state.pop("_db_schema", None)
+        return state
+
+    def __setstate__(self, state: dict) -> None:
+        """Restore the state of the object from pickling."""
+        self.__dict__.update(state)
+        # Recreate the engine and schema.
+        self._init_engine()
+
     @property
     def _schema(self) -> DbSchema:
         """Lazily create schema upon first access."""
@@ -66,6 +91,32 @@ def update_schema(self) -> None:
     def __repr__(self) -> str:
         return self._repr
 
+    def get_experiment_by_id(
+        self,
+        experiment_id: str,
+        tunables: TunableGroups,
+        opt_targets: dict[str, Literal["min", "max"]],
+    ) -> Storage.Experiment | None:
+        with self._engine.connect() as conn:
+            cur_exp = conn.execute(
+                self._schema.experiment.select().where(
+                    self._schema.experiment.c.exp_id == experiment_id,
+                )
+            )
+            exp = cur_exp.fetchone()
+            if exp is None:
+                return None
+            return Experiment(
+                engine=self._engine,
+                schema=self._schema,
+                experiment_id=exp.exp_id,
+                trial_id=-1,  # will be loaded upon __enter__ which calls _setup()
+                description=exp.description,
+                root_env_config=exp.root_env_config,
+                tunables=tunables,
+                opt_targets=opt_targets,
+            )
+
     def experiment(  # pylint: disable=too-many-arguments
         self,
         *,