diff --git a/.git-blame-ignore-revs b/.git-blame-ignore-revs index ceac94dfa3..5f9f4b63f5 100644 --- a/.git-blame-ignore-revs +++ b/.git-blame-ignore-revs @@ -1,2 +1,4 @@ # Reformat with black, isort, docformatter, etc. e40ac28317c61ea90345d3499986957b0e1c9134 +# Reformat with pyupgrade +795d1b03c6d7a9272018413564a8f02eef6fdec6 diff --git a/mlos_bench/mlos_bench/storage/sql/alembic.ini b/mlos_bench/mlos_bench/storage/sql/alembic.ini index 642970df63..4d2a1120c5 100644 --- a/mlos_bench/mlos_bench/storage/sql/alembic.ini +++ b/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] diff --git a/mlos_bench/mlos_bench/storage/sql/alembic/README.md b/mlos_bench/mlos_bench/storage/sql/alembic/README.md index ac20a0a45d..f544b2b14f 100644 --- a/mlos_bench/mlos_bench/storage/sql/alembic/README.md +++ b/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. diff --git a/mlos_bench/mlos_bench/storage/sql/alembic/env.py b/mlos_bench/mlos_bench/storage/sql/alembic/env.py index 0b4066cee0..fc186b8cb1 100644 --- a/mlos_bench/mlos_bench/storage/sql/alembic/env.py +++ b/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 diff --git a/mlos_bench/mlos_bench/storage/sql/alembic/script.py.mako b/mlos_bench/mlos_bench/storage/sql/alembic/script.py.mako index c8e8aee1e2..b9ac9d59eb 100644 --- a/mlos_bench/mlos_bench/storage/sql/alembic/script.py.mako +++ b/mlos_bench/mlos_bench/storage/sql/alembic/script.py.mako @@ -1,3 +1,7 @@ +# +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# """${message} Revision ID: ${up_revision} @@ -5,7 +9,9 @@ Revises: ${down_revision | comma,n} Create Date: ${create_date} """ -from typing import Sequence +# pylint: disable=no-member + +from collections.abc import Sequence from alembic import op import sqlalchemy as sa diff --git a/mlos_bench/mlos_bench/storage/sql/alembic/versions/f83fb8ae7fc4_add_trial_runner_id_column.py b/mlos_bench/mlos_bench/storage/sql/alembic/versions/f83fb8ae7fc4_add_trial_runner_id_column.py new file mode 100644 index 0000000000..b49eb2bee9 --- /dev/null +++ b/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 ### diff --git a/mlos_bench/mlos_bench/storage/sql/schema.py b/mlos_bench/mlos_bench/storage/sql/schema.py index 014e610a9f..f1c887713a 100644 --- a/mlos_bench/mlos_bench/storage/sql/schema.py +++ b/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 diff --git a/mlos_bench/mlos_bench/tests/storage/test_storage_schemas.py b/mlos_bench/mlos_bench/tests/storage/test_storage_schemas.py new file mode 100644 index 0000000000..3f3eb965eb --- /dev/null +++ b/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