Refactor storage tests to support testing multiple backends on schema changes (#983)

# Pull Request

## Title

Refactor storage tests to support testing multiple backends on schema
changes

______________________________________________________________________

## Description

- adds the ability to run mysql and postgres servers in a container as a
part of pytest
- checks that they succeed with schema setup
- fixes the ones that were broken with mysql

______________________________________________________________________

## Type of Change

- 🛠️ Bug fix
- 🔄 Refactor
- 🧪 Tests

______________________________________________________________________

## Testing

- CI tests

______________________________________________________________________

## Additional Notes (optional)

Part of a series of changes to help improve testing for schedulers.

Next will be schema changes to allow storing fractional seconds in
MySQL.

Merge after #987 
______________________________________________________________________

---------

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

Author: bpkroth

mlos_bench/mlos_bench/config/storage/mysql.jsonc

@@ -6,7 +6,7 @@
         "log_sql": false,  // Write all SQL statements to the log.
         // Parameters below must match kwargs of `sqlalchemy.URL.create()`:
         "drivername": "mysql+mysqlconnector",
-        "database": "osat",
+        "database": "mlos_bench",
         "username": "root",
         "password": "PLACERHOLDER PASSWORD",  // Comes from global config
         "host": "localhost",

mlos_bench/mlos_bench/config/storage/postgresql.jsonc

@@ -8,7 +8,7 @@
         "log_sql": false,  // Write all SQL statements to the log.
         // Parameters below must match kwargs of `sqlalchemy.URL.create()`:
         "drivername": "postgresql+psycopg2",
-        "database": "osat",
+        "database": "mlos_bench",
         "username": "postgres",
         "password": "PLACERHOLDER PASSWORD",  // Comes from global config
         "host": "localhost",

mlos_bench/mlos_bench/storage/sql/alembic/README.md

@@ -45,7 +45,7 @@ This document contains some notes on how to use [`alembic`](https://alembic.sqla
 
 1. If the migration script works, commit the changes to the [`mlos_bench/storage/sql/schema.py`](../schema.py) and [`mlos_bench/storage/sql/alembic/versions`](./versions/) files.
 
-   > Be sure to update the latest version in the [`test_storage_schemas.py`](../../../tests/storage/test_storage_schemas.py) file as well.
+   > Be sure to update the latest version in the [`test_storage_schemas.py`](../../../tests/storage/sql/test_storage_schemas.py) file as well.
 
 1. Merge that to the `main` branch.
 

mlos_bench/mlos_bench/storage/sql/schema.py

@@ -72,12 +72,6 @@ class DbSchema:
     # for all DB tables, so it's ok to disable the warnings.
     # pylint: disable=too-many-instance-attributes
 
-    # Common string column sizes.
-    _ID_LEN = 512
-    _PARAM_VALUE_LEN = 1024
-    _METRIC_VALUE_LEN = 255
-    _STATUS_LEN = 16
-
     def __init__(self, engine: Engine | None):
         """
         Declare the SQLAlchemy schema for the database.
@@ -95,10 +89,24 @@ def __init__(self, engine: Engine | None):
         self._engine = engine
         self._meta = MetaData()
 
+        # Common string column sizes.
+        self._exp_id_len = 512
+        self._param_id_len = 512
+        self._param_value_len = 1024
+        self._metric_id_len = 512
+        self._metric_value_len = 255
+        self._status_len = 16
+
+        # Some overrides for certain DB engines:
+        if engine and engine.dialect.name in {"mysql", "mariadb"}:
+            self._exp_id_len = 255
+            self._param_id_len = 255
+            self._metric_id_len = 255
+
         self.experiment = Table(
             "experiment",
             self._meta,
-            Column("exp_id", String(self._ID_LEN), nullable=False),
+            Column("exp_id", String(self._exp_id_len), nullable=False),
             Column("description", String(1024)),
             Column("root_env_config", String(1024), nullable=False),
             Column("git_repo", String(1024), nullable=False),
@@ -108,7 +116,7 @@ def __init__(self, engine: Engine | None):
             Column("ts_end", DateTime),
             # Should match the text IDs of `mlos_bench.environments.Status` enum:
             # For backwards compatibility, we allow NULL for status.
-            Column("status", String(self._STATUS_LEN)),
+            Column("status", String(self._status_len)),
             # There may be more than one mlos_benchd_service running on different hosts.
             # This column stores the host/container name of the driver that
             # picked up the experiment.
@@ -126,7 +134,7 @@ def __init__(self, engine: Engine | None):
             "objectives",
             self._meta,
             Column("exp_id"),
-            Column("optimization_target", String(self._ID_LEN), nullable=False),
+            Column("optimization_target", String(self._metric_id_len), nullable=False),
             Column("optimization_direction", String(4), nullable=False),
             # TODO: Note: weight is not fully supported yet as currently
             # multi-objective is expected to explore each objective equally.
@@ -175,14 +183,14 @@ def __init__(self, engine: Engine | None):
         self.trial = Table(
             "trial",
             self._meta,
-            Column("exp_id", String(self._ID_LEN), nullable=False),
+            Column("exp_id", String(self._exp_id_len), nullable=False),
             Column("trial_id", Integer, nullable=False),
             Column("config_id", Integer, nullable=False),
             Column("trial_runner_id", Integer, nullable=True, default=None),
             Column("ts_start", DateTime, nullable=False),
             Column("ts_end", DateTime),
             # Should match the text IDs of `mlos_bench.environments.Status` enum:
-            Column("status", String(self._STATUS_LEN), nullable=False),
+            Column("status", String(self._status_len), nullable=False),
             PrimaryKeyConstraint("exp_id", "trial_id"),
             ForeignKeyConstraint(["exp_id"], [self.experiment.c.exp_id]),
             ForeignKeyConstraint(["config_id"], [self.config.c.config_id]),
@@ -197,8 +205,8 @@ def __init__(self, engine: Engine | None):
             "config_param",
             self._meta,
             Column("config_id", Integer, nullable=False),
-            Column("param_id", String(self._ID_LEN), nullable=False),
-            Column("param_value", String(self._PARAM_VALUE_LEN)),
+            Column("param_id", String(self._param_id_len), nullable=False),
+            Column("param_value", String(self._param_value_len)),
             PrimaryKeyConstraint("config_id", "param_id"),
             ForeignKeyConstraint(["config_id"], [self.config.c.config_id]),
         )
@@ -212,10 +220,10 @@ def __init__(self, engine: Engine | None):
         self.trial_param = Table(
             "trial_param",
             self._meta,
-            Column("exp_id", String(self._ID_LEN), nullable=False),
+            Column("exp_id", String(self._exp_id_len), nullable=False),
             Column("trial_id", Integer, nullable=False),
-            Column("param_id", String(self._ID_LEN), nullable=False),
-            Column("param_value", String(self._PARAM_VALUE_LEN)),
+            Column("param_id", String(self._param_id_len), nullable=False),
+            Column("param_value", String(self._param_value_len)),
             PrimaryKeyConstraint("exp_id", "trial_id", "param_id"),
             ForeignKeyConstraint(
                 ["exp_id", "trial_id"],
@@ -230,10 +238,10 @@ def __init__(self, engine: Engine | None):
         self.trial_status = Table(
             "trial_status",
             self._meta,
-            Column("exp_id", String(self._ID_LEN), nullable=False),
+            Column("exp_id", String(self._exp_id_len), nullable=False),
             Column("trial_id", Integer, nullable=False),
             Column("ts", DateTime(timezone=True), nullable=False, default="now"),
-            Column("status", String(self._STATUS_LEN), nullable=False),
+            Column("status", String(self._status_len), nullable=False),
             UniqueConstraint("exp_id", "trial_id", "ts"),
             ForeignKeyConstraint(
                 ["exp_id", "trial_id"],
@@ -247,10 +255,10 @@ def __init__(self, engine: Engine | None):
         self.trial_result = Table(
             "trial_result",
             self._meta,
-            Column("exp_id", String(self._ID_LEN), nullable=False),
+            Column("exp_id", String(self._exp_id_len), nullable=False),
             Column("trial_id", Integer, nullable=False),
-            Column("metric_id", String(self._ID_LEN), nullable=False),
-            Column("metric_value", String(self._METRIC_VALUE_LEN)),
+            Column("metric_id", String(self._metric_id_len), nullable=False),
+            Column("metric_value", String(self._metric_value_len)),
             PrimaryKeyConstraint("exp_id", "trial_id", "metric_id"),
             ForeignKeyConstraint(
                 ["exp_id", "trial_id"],
@@ -265,11 +273,11 @@ def __init__(self, engine: Engine | None):
         self.trial_telemetry = Table(
             "trial_telemetry",
             self._meta,
-            Column("exp_id", String(self._ID_LEN), nullable=False),
+            Column("exp_id", String(self._exp_id_len), nullable=False),
             Column("trial_id", Integer, nullable=False),
             Column("ts", DateTime(timezone=True), nullable=False, default="now"),
-            Column("metric_id", String(self._ID_LEN), nullable=False),
-            Column("metric_value", String(self._METRIC_VALUE_LEN)),
+            Column("metric_id", String(self._metric_id_len), nullable=False),
+            Column("metric_value", String(self._metric_value_len)),
             UniqueConstraint("exp_id", "trial_id", "ts", "metric_id"),
             ForeignKeyConstraint(
                 ["exp_id", "trial_id"],
@@ -296,6 +304,31 @@ def _get_alembic_cfg(conn: Connection) -> config.Config:
         alembic_cfg.attributes["connection"] = conn
         return alembic_cfg
 
+    def drop_all_tables(self, *, force: bool = False) -> None:
+        """
+        Helper method used in testing to reset the DB schema.
+
+        Notes
+        -----
+        This method is not intended for production use, as it will drop all tables
+        in the database. Use with caution.
+
+        Parameters
+        ----------
+        force : bool
+            If True, drop all tables in the target database.
+            If False, this method will not drop any tables and will log a warning.
+        """
+        assert self._engine
+        self.meta.reflect(bind=self._engine)
+        if force:
+            self.meta.drop_all(bind=self._engine)
+        else:
+            _LOG.warning(
+                "Resetting the schema without force is not implemented. "
+                "Use force=True to drop all tables."
+            )
+
     def create(self) -> "DbSchema":
         """Create the DB schema."""
         _LOG.info("Create the DB schema")

mlos_bench/mlos_bench/storage/sql/storage.py

@@ -5,6 +5,7 @@
 """Saving and restoring the benchmark data in SQL database."""
 
 import logging
+from types import TracebackType
 from typing import Literal
 
 from sqlalchemy import URL, Engine, create_engine
@@ -72,6 +73,22 @@ def __setstate__(self, state: dict) -> None:
         # Recreate the engine and schema.
         self._init_engine()
 
+    def dispose(self) -> None:
+        """Closes the database connection pool."""
+        if self._engine:
+            self._engine.dispose()
+            _LOG.info("Closed the database connection: %s", self)
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,  # pylint: disable=unused-argument
+        exc_val: BaseException | None,  # pylint: disable=unused-argument
+        exc_tb: TracebackType | None,  # pylint: disable=unused-argument
+    ) -> Literal[False]:
+        """Close the engine connection when exiting the context."""
+        self.dispose()
+        return False
+
     @property
     def _schema(self) -> DbSchema:
         """Lazily create schema upon first access."""
@@ -82,6 +99,33 @@ def _schema(self) -> DbSchema:
                 _LOG.debug("DDL statements:\n%s", self._db_schema)
         return self._db_schema
 
+    def _reset_schema(self, *, force: bool = False) -> None:
+        """
+        Helper method used in testing to reset the DB schema.
+
+        Notes
+        -----
+        This method is not intended for production use, as it will drop all tables
+        in the database. Use with caution.
+
+        Parameters
+        ----------
+        force : bool
+            If True, drop all tables in the target database.
+            If False, this method will not drop any tables and will log a warning.
+        """
+        assert self._engine
+        if force:
+            self._schema.drop_all_tables(force=force)
+            self._db_schema = DbSchema(self._engine)
+            self._schema_created = False
+            self._schema_updated = False
+        else:
+            _LOG.warning(
+                "Resetting the schema without force is not implemented. "
+                "Use force=True to drop all tables."
+            )
+
     def update_schema(self) -> None:
         """Update the database schema."""
         if not self._schema_updated:

mlos_bench/mlos_bench/tests/__init__.py

@@ -8,6 +8,7 @@
 Used to make mypy happy about multiple conftest.py modules.
 """
 import filecmp
+import json
 import os
 import shutil
 import socket
@@ -17,6 +18,7 @@
 
 import pytest
 import pytz
+from pytest_docker.plugin import Services as DockerServices
 
 from mlos_bench.util import get_class_from_name, nullable
 
@@ -87,6 +89,48 @@ def check_class_name(obj: object, expected_class_name: str) -> bool:
     return full_class_name == try_resolve_class_name(expected_class_name)
 
 
+def is_docker_service_healthy(
+    compose_project_name: str,
+    service_name: str,
+) -> bool:
+    """Check if a docker service is healthy."""
+    docker_ps_out = run(
+        f"docker compose -p {compose_project_name} " f"ps --format json {service_name}",
+        shell=True,
+        check=True,
+        capture_output=True,
+    )
+    docker_ps_json = json.loads(docker_ps_out.stdout.decode().strip())
+    state = docker_ps_json["State"]
+    assert isinstance(state, str)
+    health = docker_ps_json["Health"]
+    assert isinstance(health, str)
+    return state == "running" and health == "healthy"
+
+
+def wait_docker_service_healthy(
+    docker_services: DockerServices,
+    project_name: str,
+    service_name: str,
+    timeout: float = 30.0,
+) -> None:
+    """Wait until a docker service is healthy."""
+    docker_services.wait_until_responsive(
+        check=lambda: is_docker_service_healthy(project_name, service_name),
+        timeout=timeout,
+        pause=0.5,
+    )
+
+
+def wait_docker_service_socket(docker_services: DockerServices, hostname: str, port: int) -> None:
+    """Wait until a docker service is ready."""
+    docker_services.wait_until_responsive(
+        check=lambda: check_socket(hostname, port),
+        timeout=30.0,
+        pause=0.5,
+    )
+
+
 def check_socket(host: str, port: int, timeout: float = 1.0) -> bool:
     """
     Test to see if a socket is open.

mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py

@@ -14,7 +14,7 @@
 from mlos_bench.schedulers.base_scheduler import Scheduler
 from mlos_bench.schedulers.trial_runner import TrialRunner
 from mlos_bench.services.config_persistence import ConfigPersistenceService
-from mlos_bench.storage.sql.storage import SqlStorage
+from mlos_bench.storage.base_storage import Storage
 from mlos_bench.tests.config import BUILTIN_TEST_CONFIG_PATH, locate_config_examples
 from mlos_bench.util import get_class_from_name
 
@@ -58,7 +58,7 @@ def test_load_scheduler_config_examples(
     config_path: str,
     mock_env_config_path: str,
     trial_runners: list[TrialRunner],
-    storage: SqlStorage,
+    storage: Storage,
     mock_opt: MockOptimizer,
 ) -> None:
     """Tests loading a config example."""

mlos_bench/mlos_bench/tests/conftest.py

@@ -15,7 +15,7 @@
 from pytest_docker.plugin import get_docker_services
 
 from mlos_bench.environments.mock_env import MockEnv
-from mlos_bench.tests import SEED, tunable_groups_fixtures
+from mlos_bench.tests import SEED, resolve_host_name, tunable_groups_fixtures
 from mlos_bench.tunables.tunable_groups import TunableGroups
 
 # pylint: disable=redefined-outer-name
@@ -29,6 +29,22 @@
 covariant_group = tunable_groups_fixtures.covariant_group
 
 
+HOST_DOCKER_NAME = "host.docker.internal"
+
+
+@pytest.fixture(scope="session")
+def docker_hostname() -> str:
+    """Returns the local hostname to use to connect to the test ssh server."""
+    if sys.platform != "win32" and resolve_host_name(HOST_DOCKER_NAME):
+        # On Linux, if we're running in a docker container, we can use the
+        # --add-host (extra_hosts in docker-compose.yml) to refer to the host IP.
+        return HOST_DOCKER_NAME
+    # Docker (Desktop) for Windows (WSL2) uses a special networking magic
+    # to refer to the host machine as `localhost` when exposing ports.
+    # In all other cases, assume we're executing directly inside conda on the host.
+    return "127.0.0.1"  # "localhost"
+
+
 @pytest.fixture
 def mock_env(tunable_groups: TunableGroups) -> MockEnv:
     """Test fixture for MockEnv."""
@@ -90,6 +106,7 @@ def docker_compose_file(pytestconfig: pytest.Config) -> list[str]:
     _ = pytestconfig  # unused
     return [
         os.path.join(os.path.dirname(__file__), "services", "remote", "ssh", "docker-compose.yml"),
+        os.path.join(os.path.dirname(__file__), "storage", "sql", "docker-compose.yml"),
         # Add additional configs as necessary here.
     ]
 

mlos_bench/mlos_bench/tests/environments/remote/conftest.py

@@ -8,4 +8,3 @@
 
 # Expose some of those as local names so they can be picked up as fixtures by pytest.
 ssh_test_server = ssh_fixtures.ssh_test_server
-ssh_test_server_hostname = ssh_fixtures.ssh_test_server_hostname