Repository: microsoft/MLOS
Pull request #974: Make Storage objects picklable and resumable by mp.Pool workers
Author: login: bpkroth
Labels: name: ready for review
description: Ready for review
Created: 2025-05-09T19:04:01Z | Merged: 2025-05-13T00:44:54Z

# 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`.

______________________________________________________________________



Diff:
diff --git a/mlos_bench/mlos_bench/environments/status.py b/mlos_bench/mlos_bench/environments/status.py
index ca35b3473d..0341698d34 100644
--- a/mlos_bench/mlos_bench/environments/status.py
+++ b/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
diff --git a/mlos_bench/mlos_bench/storage/base_storage.py b/mlos_bench/mlos_bench/storage/base_storage.py
index 52ecf93b7d..75a84bf0b2 100644
--- a/mlos_bench/mlos_bench/storage/base_storage.py
+++ b/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,6 +347,25 @@ 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,
@@ -314,7 +373,7 @@ def pending_trials(
             *,
             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."""
diff --git a/mlos_bench/mlos_bench/storage/sql/common.py b/mlos_bench/mlos_bench/storage/sql/common.py
index 55d1fcac84..97eb270c9d 100644
--- a/mlos_bench/mlos_bench/storage/sql/common.py
+++ b/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()
diff --git a/mlos_bench/mlos_bench/storage/sql/experiment.py b/mlos_bench/mlos_bench/storage/sql/experiment.py
index 2d7db8e34e..acc2a497b4 100644
--- a/mlos_bench/mlos_bench/storage/sql/experiment.py
+++ b/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
diff --git a/mlos_bench/mlos_bench/storage/sql/storage.py b/mlos_bench/mlos_bench/storage/sql/storage.py
index b3bf63d0ed..6d98dc97fd 100644
--- a/mlos_bench/mlos_bench/storage/sql/storage.py
+++ b/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,6 +25,8 @@ class SqlStorage(Storage):
     backend.
     """
 
+    # pylint: disable=too-many-instance-attributes
+
     def __init__(
         self,
         config: dict,
@@ -32,21 +34,44 @@ def __init__(
         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,
         *,
diff --git a/mlos_bench/mlos_bench/storage/sql/trial.py b/mlos_bench/mlos_bench/storage/sql/trial.py
index a9960c466a..e1276c2881 100644
--- a/mlos_bench/mlos_bench/storage/sql/trial.py
+++ b/mlos_bench/mlos_bench/storage/sql/trial.py
@@ -40,8 +40,9 @@ def __init__(  # pylint: disable=too-many-arguments
         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,
     ):
         super().__init__(
             tunables=tunables,
@@ -50,8 +51,9 @@ def __init__(  # pylint: disable=too-many-arguments
             tunable_config_id=config_id,
             trial_runner_id=trial_runner_id,
             opt_targets=opt_targets,
-            config=config,
             status=status,
+            restoring=restoring,
+            config=config,
         )
         self._engine = engine
         self._schema = schema
diff --git a/mlos_bench/mlos_bench/tests/storage/test_storage_pickling.py b/mlos_bench/mlos_bench/tests/storage/test_storage_pickling.py
new file mode 100644
index 0000000000..3d5053837b
--- /dev/null
+++ b/mlos_bench/mlos_bench/tests/storage/test_storage_pickling.py
@@ -0,0 +1,97 @@
+#
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+#
+"""Test pickling and unpickling of Storage, and restoring Experiment and Trial by id."""
+import json
+import os
+import pickle
+import sys
+import tempfile
+from datetime import datetime
+from typing import Literal
+
+import pytest
+from pytz import UTC
+
+from mlos_bench.environments.status import Status
+from mlos_bench.storage.sql.storage import SqlStorage
+from mlos_bench.storage.storage_factory import from_config
+from mlos_bench.tunables.tunable_groups import TunableGroups
+
+
+# TODO: When we introduce ParallelTrialScheduler warn at config startup time
+# that it is incompatible with sqlite storage on Windows.
+@pytest.mark.skipif(
+    sys.platform == "win32",
+    reason="Windows doesn't support multiple processes accessing the same file.",
+)
+def test_storage_pickle_restore_experiment_and_trial(tunable_groups: TunableGroups) -> None:
+    """Check that we can pickle and unpickle the Storage object, and restore Experiment
+    and Trial by id.
+    """
+    # pylint: disable=too-many-locals
+    with tempfile.TemporaryDirectory() as tmpdir:
+        db_path = os.path.join(tmpdir, "mlos_bench.sqlite")
+        config_str = json.dumps(
+            {
+                "class": "mlos_bench.storage.sql.storage.SqlStorage",
+                "config": {
+                    "drivername": "sqlite",
+                    "database": db_path,
+                    "lazy_schema_create": False,
+                },
+            }
+        )
+
+        storage = from_config(config_str)
+        storage.update_schema()
+
+        # Create an Experiment and a Trial
+        opt_targets: dict[str, Literal["min", "max"]] = {"metric": "min"}
+        experiment = storage.experiment(
+            experiment_id="experiment_id",
+            trial_id=0,
+            root_env_config="dummy_env.json",
+            description="Pickle test experiment",
+            tunables=tunable_groups,
+            opt_targets=opt_targets,
+        )
+        with experiment:
+            trial = experiment.new_trial(tunable_groups)
+            trial_id_created = trial.trial_id
+            trial.set_trial_runner(1)
+            trial.update(Status.RUNNING, datetime.now(UTC))
+
+        # Pickle and unpickle the Storage object
+        pickled = pickle.dumps(storage)
+        restored_storage = pickle.loads(pickled)
+        assert isinstance(restored_storage, SqlStorage)
+
+        # Restore the Experiment from storage by id and check that it matches the original
+        restored_experiment = restored_storage.get_experiment_by_id(
+            experiment_id=experiment.experiment_id,
+            tunables=tunable_groups,
+            opt_targets=opt_targets,
+        )
+        assert restored_experiment is not None
+        assert restored_experiment is not experiment
+        assert restored_experiment.experiment_id == experiment.experiment_id
+        assert restored_experiment.description == experiment.description
+        assert restored_experiment.root_env_config == experiment.root_env_config
+        assert restored_experiment.tunables == experiment.tunables
+        assert restored_experiment.opt_targets == experiment.opt_targets
+        with restored_experiment:
+            # trial_id should have been restored during __enter__
+            assert restored_experiment.trial_id == experiment.trial_id
+
+        # Restore the Trial from storage by id and check that it matches the original
+        restored_trial = restored_experiment.get_trial_by_id(trial_id_created)
+        assert restored_trial is not None
+        assert restored_trial is not trial
+        assert restored_trial.trial_id == trial.trial_id
+        assert restored_trial.experiment_id == trial.experiment_id
+        assert restored_trial.tunables == trial.tunables
+        assert restored_trial.status == trial.status
+        assert restored_trial.config() == trial.config()
+        assert restored_trial.trial_runner_id == trial.trial_runner_id


Reviews:
user: login: copilot-pull-request-reviewer[bot]
body: ## Pull Request Overview

This PR enables Storage objects to be picklable and resumable by introducing custom __getstate__/__setstate__ methods and adding functionality to restore Experiment and Trial objects by their IDs for multiprocessing scenarios.  
- Added new pickling support in SqlStorage with engine reinitialization  
- Updated Trial and Experiment constructors to support restoration with new status and restoring parameters  
- Introduced test cases for verifying pickling and unpickling

### Reviewed Changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 2 comments.

<details>
<summary>Show a summary per file</summary>

| File                                                      | Description                                                            |
|-----------------------------------------------------------|------------------------------------------------------------------------|
| mlos_bench/mlos_bench/tests/storage/test_storage_pickling.py | New tests to ensure Storage objects can be pickled and restored         |
| mlos_bench/mlos_bench/storage/sql/trial.py                 | Modified Trial constructor to include new status and restoring params   |
| mlos_bench/mlos_bench/storage/sql/storage.py               | Added __getstate__/__setstate__ changes to support engine recreation     |
| mlos_bench/mlos_bench/storage/sql/experiment.py            | Added get_trial_by_id method and updated trial creation logic           |
| mlos_bench/mlos_bench/storage/base_storage.py              | Updated abstract methods for experiment/trial retrieval                 |
| mlos_bench/mlos_bench/environments/status.py                | Extended Status with from_str; existing is_timed_out implementation noted|
| .vscode/settings.json                                      | Updated code action settings by commenting out the unusedImports rule      |
</details>



<details>
<summary>Comments suppressed due to low confidence (1)</summary>

**mlos_bench/mlos_bench/environments/status.py:37**
* The is_timed_out method is checking for Status.FAILED instead of Status.TIMED_OUT, which appears to be a logic bug that could lead to incorrect status reporting.
```
return self == Status.FAILED
```
</details>


state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-09T23:16:42Z
user: login: bpkroth
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-10T00:20:13Z
user: login: bpkroth
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-10T14:54:51Z
user: login: bpkroth
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-10T14:55:19Z
user: login: bpkroth
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T00:08:28Z
user: login: bpkroth
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T00:08:40Z
user: login: motus
body: Looks great, thank you!
state: APPROVED
author_association: MEMBER
submitted_at: 2025-05-13T00:44:46Z
user: login: bpkroth
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T15:29:15Z
user: login: bpkroth
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T19:40:50Z

Review comments:
diff_hunk: @@ -147,7 +147,7 @@
     "[python]": {
         "editor.codeActionsOnSave": {
             "source.organizeImports": "explicit",
-            "source.unusedImports": "explicit"
+            //"source.unusedImports": "explicit"
path: .vscode/settings.json
user: login: Copilot
body: [nitpick] If disabling the unusedImports rule is intended, consider adding an inline comment explaining the rationale to improve future maintainability.
```suggestion
            //"source.unusedImports": "explicit" // Disabled to allow flexibility in handling unused imports during development.
```
created_at: 2025-05-09T23:16:41Z
author_association: CONTRIBUTOR
diff_hunk: @@ -40,8 +40,9 @@ def __init__(  # pylint: disable=too-many-arguments
         config_id: int,
         trial_runner_id: int | None,
         opt_targets: dict[str, Literal["min", "max"]],
+        status: Status,
path: mlos_bench/mlos_bench/storage/sql/trial.py
user: login: Copilot
body: [nitpick] Ensure that the new parameters 'status' and 'restoring' in the Trial constructor are clearly documented where this constructor is invoked to aid future developers.
created_at: 2025-05-09T23:16:42Z
author_association: CONTRIBUTOR
diff_hunk: @@ -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
path: mlos_bench/mlos_bench/environments/status.py
user: login: bpkroth
body: Unrelated bug that copilot found in its review.
created_at: 2025-05-10T00:20:13Z
author_association: CONTRIBUTOR
diff_hunk: @@ -19,6 +22,21 @@ class Status(enum.Enum):
     FAILED = 6
     TIMED_OUT = 7
 
+    @staticmethod
+    def from_str(status_str: str) -> "Status":
path: mlos_bench/mlos_bench/environments/status.py
user: login: bpkroth
body: Previously we did not restore `Status` on `Trials` restored from the DB.  This helps do that.
created_at: 2025-05-10T14:54:51Z
author_association: CONTRIBUTOR
diff_hunk: @@ -19,6 +22,21 @@ class Status(enum.Enum):
     FAILED = 6
     TIMED_OUT = 7
 
+    @staticmethod
+    def from_str(status_str: str) -> "Status":
path: mlos_bench/mlos_bench/environments/status.py
user: login: bpkroth
body: Should maybe check to see how we were doing that in the `TrialData` class.
created_at: 2025-05-10T14:55:19Z
author_association: CONTRIBUTOR
diff_hunk: @@ -19,6 +22,21 @@ class Status(enum.Enum):
     FAILED = 6
     TIMED_OUT = 7
 
+    @staticmethod
+    def from_str(status_str: str) -> "Status":
path: mlos_bench/mlos_bench/environments/status.py
user: login: bpkroth
body: Did that.
created_at: 2025-05-13T00:08:27Z
author_association: CONTRIBUTOR
diff_hunk: @@ -19,6 +22,21 @@ class Status(enum.Enum):
     FAILED = 6
     TIMED_OUT = 7
 
+    @staticmethod
+    def from_str(status_str: str) -> "Status":
path: mlos_bench/mlos_bench/environments/status.py
user: login: bpkroth
body: Hmm, probably need tests for this too.
created_at: 2025-05-13T00:08:40Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,97 @@
+#
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+#
+"""Test pickling and unpickling of Storage, and restoring Experiment and Trial by id."""
+import json
+import os
+import pickle
+import sys
+import tempfile
+from datetime import datetime
+from typing import Literal
+
+import pytest
+from pytz import UTC
+
+from mlos_bench.environments.status import Status
+from mlos_bench.storage.sql.storage import SqlStorage
+from mlos_bench.storage.storage_factory import from_config
+from mlos_bench.tunables.tunable_groups import TunableGroups
+
+
+# TODO: When we introduce ParallelTrialScheduler warn at config startup time
+# that it is incompatible with sqlite storage on Windows.
+@pytest.mark.skipif(
+    sys.platform == "win32",
+    reason="Windows doesn't support multiple processes accessing the same file.",
+)
+def test_storage_pickle_restore_experiment_and_trial(tunable_groups: TunableGroups) -> None:
+    """Check that we can pickle and unpickle the Storage object, and restore Experiment
+    and Trial by id.
+    """
+    # pylint: disable=too-many-locals
+    with tempfile.TemporaryDirectory() as tmpdir:
+        db_path = os.path.join(tmpdir, "mlos_bench.sqlite")
+        config_str = json.dumps(
+            {
+                "class": "mlos_bench.storage.sql.storage.SqlStorage",
+                "config": {
+                    "drivername": "sqlite",
+                    "database": db_path,
+                    "lazy_schema_create": False,
+                },
+            }
+        )
+
+        storage = from_config(config_str)
path: mlos_bench/mlos_bench/tests/storage/test_storage_pickling.py
user: login: motus
body: Maybe `from_config()` should be able to accept a `dict` so we don't have to convert configs back and forth
created_at: 2025-05-13T00:42:47Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,97 @@
+#
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+#
+"""Test pickling and unpickling of Storage, and restoring Experiment and Trial by id."""
+import json
+import os
+import pickle
+import sys
+import tempfile
+from datetime import datetime
+from typing import Literal
+
+import pytest
+from pytz import UTC
+
+from mlos_bench.environments.status import Status
+from mlos_bench.storage.sql.storage import SqlStorage
+from mlos_bench.storage.storage_factory import from_config
+from mlos_bench.tunables.tunable_groups import TunableGroups
+
+
+# TODO: When we introduce ParallelTrialScheduler warn at config startup time
+# that it is incompatible with sqlite storage on Windows.
+@pytest.mark.skipif(
+    sys.platform == "win32",
+    reason="Windows doesn't support multiple processes accessing the same file.",
+)
+def test_storage_pickle_restore_experiment_and_trial(tunable_groups: TunableGroups) -> None:
+    """Check that we can pickle and unpickle the Storage object, and restore Experiment
+    and Trial by id.
+    """
+    # pylint: disable=too-many-locals
+    with tempfile.TemporaryDirectory() as tmpdir:
+        db_path = os.path.join(tmpdir, "mlos_bench.sqlite")
+        config_str = json.dumps(
+            {
+                "class": "mlos_bench.storage.sql.storage.SqlStorage",
+                "config": {
+                    "drivername": "sqlite",
+                    "database": db_path,
+                    "lazy_schema_create": False,
+                },
+            }
+        )
+
+        storage = from_config(config_str)
path: mlos_bench/mlos_bench/tests/storage/test_storage_pickling.py
user: login: bpkroth
body: Sure, I can add that in a new PR.
created_at: 2025-05-13T15:29:15Z
author_association: CONTRIBUTOR
diff_hunk: @@ -19,6 +22,21 @@ class Status(enum.Enum):
     FAILED = 6
     TIMED_OUT = 7
 
+    @staticmethod
+    def from_str(status_str: str) -> "Status":
path: mlos_bench/mlos_bench/environments/status.py
user: login: bpkroth
body: Putting these in #976 
created_at: 2025-05-13T19:40:50Z
author_association: CONTRIBUTOR

PR comments:
user: login: bpkroth
created_at: 2025-05-09T23:57:14Z
body: Windows doesn't allow for multiple processes to open the same file. I think I will just exclude this on Windows with an error that it's not supported there.
author_association: CONTRIBUTOR
user: login: bpkroth
created_at: 2025-05-10T00:08:20Z
body:  
> <details>
> <summary>Comments suppressed due to low confidence (1)</summary>
> 
> **mlos_bench/mlos_bench/environments/status.py:37**
> * The is_timed_out method is checking for Status.FAILED instead of Status.TIMED_OUT, which appears to be a logic bug that could lead to incorrect status reporting.
> ```
> return self == Status.FAILED
> ```
> </details>
> 
> 

This last comment is correct. We should correct that too.
author_association: CONTRIBUTOR