Extend Storage tests to multiple backends

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/base_storage.py

@@ -25,10 +25,12 @@
 from __future__ import annotations
 
 import logging
+import os
 from abc import ABCMeta, abstractmethod
 from collections.abc import Iterator, Mapping
 from contextlib import AbstractContextManager as ContextManager
 from datetime import datetime
+from subprocess import CalledProcessError
 from types import TracebackType
 from typing import Any, Literal
 
@@ -38,7 +40,7 @@
 from mlos_bench.services.base_service import Service
 from mlos_bench.storage.base_experiment_data import ExperimentData
 from mlos_bench.tunables.tunable_groups import TunableGroups
-from mlos_bench.util import get_git_info
+from mlos_bench.util import get_git_info, get_git_root, path_join
 
 _LOG = logging.getLogger(__name__)
 
@@ -187,16 +189,61 @@ def __init__(  # pylint: disable=too-many-arguments
             tunables: TunableGroups,
             experiment_id: str,
             trial_id: int,
-            root_env_config: str,
+            root_env_config: str | None,
             description: str,
             opt_targets: dict[str, Literal["min", "max"]],
+            git_repo: str | None = None,
+            git_commit: str | None = None,
+            rel_root_env_config: str | None = None,
         ):
             self._tunables = tunables.copy()
             self._trial_id = trial_id
             self._experiment_id = experiment_id
-            (self._git_repo, self._git_commit, self._root_env_config) = get_git_info(
-                root_env_config
-            )
+            if root_env_config is None:
+                # Restoring from DB.
+                if not (git_repo and git_commit and rel_root_env_config):
+                    raise ValueError(
+                        "Missing required args: git_repo, git_commit, rel_root_env_config"
+                    )
+                self._git_repo = git_repo
+                self._git_commit = git_commit
+                self._rel_root_env_config = rel_root_env_config
+
+                # Currently we only store the relative path of the root env config
+                # from the git repo it came from.
+                git_root = git_repo
+                if not os.path.exists(git_root):
+                    try:
+                        git_root = get_git_root(os.curdir)
+                    except CalledProcessError:
+                        _LOG.warning(
+                            "Failed to find a git repo in the current working directory: %s",
+                            os.curdir,
+                        )
+                        git_root = get_git_root(__file__)
+
+                self._abs_root_env_config = path_join(
+                    git_root,
+                    self._rel_root_env_config,
+                    abs_path=True,
+                )
+                _LOG.info(
+                    "Resolved relative root_config %s for experiment %s to %s",
+                    self._rel_root_env_config,
+                    self._experiment_id,
+                    self._abs_root_env_config,
+                )
+            else:
+                if git_repo or git_commit or rel_root_env_config:
+                    raise ValueError("Unexpected args: git_repo, git_commit, rel_root_env_config")
+                (
+                    self._git_repo,
+                    self._git_commit,
+                    self._rel_root_env_config,
+                    self._abs_root_env_config,
+                ) = get_git_info(
+                    root_env_config,
+                )
             self._description = description
             self._opt_targets = opt_targets
             self._in_context = False
@@ -278,9 +325,21 @@ def description(self) -> str:
             return self._description
 
         @property
-        def root_env_config(self) -> str:
-            """Get the Experiment's root Environment config file path."""
-            return self._root_env_config
+        def rel_root_env_config(self) -> str:
+            """Get the Experiment's root Environment config's relative file path to the
+            git repo root.
+            """
+            return self._rel_root_env_config
+
+        @property
+        def abs_root_env_config(self) -> str:
+            """
+            Get the Experiment's root Environment config file path.
+
+            This returns the current absolute path to the root config for this process
+            instead of the path relative to the git repo root.
+            """
+            return self._abs_root_env_config
 
         @property
         def tunables(self) -> TunableGroups:

mlos_bench/mlos_bench/storage/sql/alembic.ini

@@ -63,7 +63,10 @@ version_path_separator = os  # Use os.pathsep. Default configuration used for ne
 # output_encoding = utf-8
 
 # See README.md for details.
+# Uncomment one of these:
 sqlalchemy.url = sqlite:///mlos_bench.sqlite
+#sqlalchemy.url = mysql+mysqlconnector://root:password@localhost:3306/mlos_bench
+#sqlalchemy.url = postgresql+psycopg2://root:password@localhost:5432/mlos_bench
 
 
 [post_write_hooks]
@@ -72,10 +75,10 @@ sqlalchemy.url = sqlite:///mlos_bench.sqlite
 # detail and examples
 
 # format using "black" - use the console_scripts runner, against the "black" entrypoint
-# hooks = black
-# black.type = console_scripts
-# black.entrypoint = black
-# black.options = -l 79 REVISION_SCRIPT_FILENAME
+hooks = black
+black.type = console_scripts
+black.entrypoint = black
+black.options = REVISION_SCRIPT_FILENAME
 
 # lint with attempts to fix using "ruff" - use the exec runner, execute a binary
 # hooks = ruff

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

@@ -4,49 +4,144 @@ This document contains some notes on how to use [`alembic`](https://alembic.sqla
 
 ## Overview
 
-1. Create a blank `mlos_bench.sqlite` database file in the [`mlos_bench/storage/sql`](../) directory with the current schema using the following command:
+1. Create a blank database instance in the [`mlos_bench/storage/sql`](../) directory with the current schema using the following command:
 
-   ```sh
-   cd mlos_bench/storage/sql
-   rm mlos_bench.sqlite
-   mlos_bench --storage storage/sqlite.jsonc --create-update-storage-schema-only
-   ```
+   This allows `alembic` to automatically generate a migration script from the current schema.
+
+   > NOTE: If your schema changes target a particular backend engine (e.g., using `with_variant`) you will need to use an engine with that config for this step.
+   > \
+   > In the remainder of this document we should some examples for different DB types.
+   > Pick the one you're targeting and stick with it thru the example.
+   > You may need to repeat the process several times to test all of them.
+   >
+   > - [ ] TODO: Add scripts to automatically do this for several different backend engines all at once.
+
+   For instance:
+
+   1. Start a temporary server either as a local file or in a docker instance
+
+      ```sh
+      # sqlite
+      cd mlos_bench/storage/sql
+      rm -f mlos_bench.sqlite
+      ```
+
+      ```sh
+      # mysql
+      docker run -it --rm --name mysql-alembic --env MYSQL_ROOT_PASSWORD=password --env MYSQL_DATABASE=mlos_bench -p 3306:3306 mysql:latest
+      ```
+
+      ```sh
+      # postgres
+      docker run -it --rm --name postgres-alembic --env POSTGRES_PASSWORD=password --env POSTGRES_DB=mlos_bench -p 5432:5432 postgres:latest
+      ```
+
+   1. Adjust the `sqlalchemy.url` in the [`alembic.ini`](../alembic.ini) file.
 
-   > This allows `alembic` to automatically generate a migration script from the current schema.
+      ```ini
+      # Uncomment one of these.
+      # See README.md for details.
 
-1. Adjust the [`mlos_bench/storage/sql/schema.py`](../schema.py) file to reflect the new desired schema.
+      #sqlalchemy.url = sqlite:///mlos_bench.sqlite
+      sqlalchemy.url = mysql+pymysql://root:password@localhost:3306/mlos_bench
+      #sqlalchemy.url = postgresql+psycopg2://root:password@localhost:5432/mlos_bench
+      ```
 
+   1. Prime the DB schema
+
+      > Note: you may want to `git checkout main` first to make sure you're using the current schema here.
+
+      ```sh
+      # sqlite
+      mlos_bench --storage storage/sqlite.jsonc --create-update-storage-schema-only --password=password
+      ```
+
+      ```sh
+      # mysql
+      mlos_bench --storage storage/mysql.jsonc --create-update-storage-schema-only --password=password
+      ```
+
+      ```sh
+      # postgres
+      mlos_bench --storage storage/postgresql.jsonc --create-update-storage-schema-only --password=password
+      ```
+
+1. Now, adjust the [`mlos_bench/storage/sql/schema.py`](../schema.py) file to reflect the new desired schema.
+
+   > Don't forget to do this on a new branch.
+   > \
    > Keep each change small and atomic.
+   > \
    > For example, if you want to add a new column, do that in one change.
    > If you want to rename a column, do that in another change.
 
 1. Generate a new migration script with the following command:
 
    ```sh
-   alembic revision --autogenerate -m "Descriptive text about the change."
+   alembic revision --autogenerate -m "CHANGEME: Descriptive text about the change."
    ```
 
 1. Review the generated migration script in the [`mlos_bench/storage/sql/alembic/versions`](./versions/) directory.
 
 1. Verify that the migration script works by running the following command:
 
    ```sh
+   # sqlite
    mlos_bench --storage storage/sqlite.jsonc --create-update-storage-schema-only
    ```
 
+   ```sh
+   # mysql:
+   mlos_bench --storage storage/mysql.jsonc --create-update-storage-schema-only --password=password
+   ```
+
+   ```sh
+   # postgres:
+   mlos_bench --storage storage/postgresql.jsonc --create-update-storage-schema-only --password=password
+   ```
+
    > Normally this would be done with `alembic upgrade head`, but this command is convenient to ensure if will work with the `mlos_bench` command line interface as well.
 
    Examine the results using something like:
 
    ```sh
+   # For sqlite:
    sqlite3 mlos_bench.sqlite .schema
    sqlite3 mlos_bench.sqlite "SELECT * FROM alembic_version;"
    ```
 
+   ```sh
+   # For mysql:
+   mysql --user root --password=password --host localhost --protocol tcp --database mlos_bench -e "SHOW TABLES; SELECT * FROM alembic_version;"
+   ```
+
+   ```sh
+   # For postgres:
+   PGPASSWORD=password psql -h localhost -p 5432 -U postgres mlos_bench -c "SELECT table_name FROM information_schema.tables WHERE table_schema='public' and table_catalog='mlos_bench'; SELECT * FROM alembic_version;"
+   ```
+
+   > Use different CLI clients for targeting other engines.
+
 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.
 
+1. Cleanup any server instances you started.
+
+   For instance:
+
+   ```sh
+   rm mlos_bench/storage/sql/mlos_bench.sqlite
+   ```
+
+   ```sh
+   docker kill mysql-alembic
+   ```
+
+   ```sh
+   docker kill postgres-alembic
+   ```
+
 1. Merge that to the `main` branch.
 
 1. Might be good to cut a new `mlos_bench` release at this point as well.