Add trial_runner_id column

mlos_bench/mlos_bench/storage/sql/alembic.ini

@@ -94,9 +94,8 @@ keys = console
 keys = generic
 
 [logger_root]
-# Don't override the root logger's level, so that we can control it from mlos_bench configs.
-#level = WARNING
-handlers =
+level = WARNING
+handlers = console
 qualname =
 
 [logger_sqlalchemy]

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

@@ -38,6 +38,8 @@ This document contains some notes on how to use [`alembic`](https://alembic.sqla
 
 6. 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.
+
 7. Merge that to the `main` branch.
 
 8. Might be good to cut a new `mlos_bench` release at this point as well.

mlos_bench/mlos_bench/storage/sql/alembic/env.py

@@ -5,6 +5,7 @@
 """Alembic environment script."""
 # pylint: disable=no-member
 
+import sys
 from logging.config import fileConfig
 
 from alembic import context
@@ -18,7 +19,8 @@
 
 # Interpret the config file for Python logging.
 # This line sets up loggers basically.
-if config.config_file_name is not None:
+# Don't override the mlos_bench or pytest loggers though.
+if config.config_file_name is not None and "alembic" in sys.argv[0]:
     fileConfig(config.config_file_name)
 
 # add your model's MetaData object here

mlos_bench/mlos_bench/storage/sql/alembic/versions/f83fb8ae7fc4_add_trial_runner_id_column.py

@@ -0,0 +1,37 @@
+#
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+#
+"""
+Add trial_runner_id column.
+
+Revision ID: f83fb8ae7fc4
+Revises: d2a708351ba8
+Create Date: 2025-01-03 21:25:48.848196+00:00
+"""
+# pylint: disable=no-member
+
+from collections.abc import Sequence
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "f83fb8ae7fc4"
+down_revision: str | None = "d2a708351ba8"
+branch_labels: str | Sequence[str] | None = None
+depends_on: str | Sequence[str] | None = None
+
+
+def upgrade() -> None:
+    """The schema upgrade script for this revision."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column("trial", sa.Column("trial_runner_id", sa.Integer(), nullable=True, default=None))
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """The schema downgrade script for this revision."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column("trial", "trial_runner_id")
+    # ### end Alembic commands ###

mlos_bench/mlos_bench/storage/sql/schema.py

@@ -24,6 +24,7 @@
 from alembic import command, config
 from sqlalchemy import (
     Column,
+    Connection,
     DateTime,
     Dialect,
     Float,
@@ -36,6 +37,7 @@
     Table,
     UniqueConstraint,
     create_mock_engine,
+    inspect,
 )
 from sqlalchemy.engine import Engine
 
@@ -163,6 +165,7 @@ def __init__(self, engine: Engine | None):
             Column("exp_id", String(self._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:
@@ -272,11 +275,35 @@ def meta(self) -> MetaData:
         """Return the SQLAlchemy MetaData object."""
         return self._meta
 
+    def _get_alembic_cfg(self, conn: Connection) -> config.Config:  # pylint: disable=no-self-use
+        alembic_cfg = config.Config(
+            path_join(str(files("mlos_bench.storage.sql")), "alembic.ini", abs_path=True)
+        )
+        alembic_cfg.attributes["connection"] = conn
+        return alembic_cfg
+
     def create(self) -> "DbSchema":
         """Create the DB schema."""
         _LOG.info("Create the DB schema")
         assert self._engine
         self._meta.create_all(self._engine)
+        with self._engine.begin() as conn:
+            # If the trial table has the trial_runner_id column but no
+            # "alembic_version" table, then the schema is up to date as of initial
+            # create and we should mark it as such to avoid trying to run the
+            # (non-idempotent) upgrade scripts.
+            # Otherwise, either we already have an alembic_version table and can
+            # safely run the necessary upgrades or we are missing the
+            # trial_runner_id column (the first to introduce schema updates) and
+            # should run the upgrades.
+            if any(
+                column["name"] == "trial_runner_id"
+                for column in inspect(conn).get_columns(self.trial.name)
+            ) and not inspect(conn).has_table("alembic_version"):
+                # Mark the schema as up to date.
+                alembic_cfg = self._get_alembic_cfg(conn)
+                command.stamp(alembic_cfg, "heads")
+                # command.current(alembic_cfg)
         return self
 
     def update(self) -> "DbSchema":
@@ -289,11 +316,8 @@ def update(self) -> "DbSchema":
         for details on how to invoke only the schema creation/update routines.
         """
         assert self._engine
-        alembic_cfg = config.Config(
-            path_join(str(files("mlos_bench.storage.sql")), "alembic.ini", abs_path=True)
-        )
         with self._engine.connect() as conn:
-            alembic_cfg.attributes["connection"] = conn
+            alembic_cfg = self._get_alembic_cfg(conn)
             command.upgrade(alembic_cfg, "head")
         return self
 

mlos_bench/mlos_bench/tests/storage/test_storage_schemas.py

@@ -0,0 +1,54 @@
+#
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+#
+"""Test sql schemas for mlos_bench storage."""
+
+from alembic.migration import MigrationContext
+from sqlalchemy import inspect
+
+from mlos_bench.storage.base_experiment_data import ExperimentData
+from mlos_bench.storage.sql.storage import SqlStorage
+
+# NOTE: This value is hardcoded to the latest revision in the alembic versions directory.
+# It could also be obtained programmatically using the "alembic heads" command or heads() API.
+# See Also: schema.py for an example of programmatic alembic config access.
+CURRENT_ALEMBIC_HEAD = "f83fb8ae7fc4"
+
+
+def test_storage_schemas(storage: SqlStorage) -> None:
+    """Test storage schema creation."""
+    eng = storage._engine  # pylint: disable=protected-access
+    with eng.connect() as conn:  # pylint: disable=protected-access
+        inspector = inspect(conn)
+        # Make sure the "trial_runner_id" column exists.
+        # (i.e., the latest schema has been applied)
+        assert any(
+            column["name"] == "trial_runner_id" for column in inspect(conn).get_columns("trial")
+        )
+        # Make sure the "alembic_version" table exists and is appropriately stamped.
+        assert inspector.has_table("alembic_version")
+        context = MigrationContext.configure(conn)
+        current_rev = context.get_current_revision()
+        assert (
+            current_rev == CURRENT_ALEMBIC_HEAD
+        ), f"Expected {CURRENT_ALEMBIC_HEAD}, got {current_rev}"
+
+
+# Note: this is a temporary test.  It will be removed and replaced with a more
+# properly integrated test in #702.
+def test_trial_runner_id_default(storage: SqlStorage, exp_data: ExperimentData) -> None:
+    """Test that the new trial_runner_id column defaults to None."""
+    assert exp_data.trials
+    eng = storage._engine  # pylint: disable=protected-access
+    schema = storage._schema  # pylint: disable=protected-access
+    with eng.connect() as conn:
+        trials = conn.execute(
+            schema.trial_result.select().with_only_columns(
+                schema.trial.c.trial_runner_id,
+            )
+        )
+        # trial_runner_id is not currently fully implemented
+        trial_row = trials.fetchone()
+        assert trial_row
+        assert trial_row.trial_runner_id is None