Schema changes example for mlos_benchd service (#931)

# Pull Request

## Title

Schema changes for mlos_benchd service.

______________________________________________________________________

## Description

Schema changes for mlos_benchd service.

Storage APIs to adjust these to come in a future PR.

- See #732

______________________________________________________________________

## Type of Change

- ✨ New feature

______________________________________________________________________

## Testing

Local, CI

______________________________________________________________________

## Additional Notes (optional)

@eujing have a look at the [commit history in the
PR](https://github.com/microsoft/MLOS/pull/931/commits) for a sense of
what's going on.
This can probably be merged as is. Happy to discuss further though.

______________________________________________________________________

Author: bpkroth

mlos_bench/mlos_bench/environments/status.py

@@ -2,13 +2,13 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 #
-"""Enum for the status of the benchmark/environment."""
+"""Enum for the status of the benchmark/environment Trial or Experiment."""
 
 import enum
 
 
 class Status(enum.Enum):
-    """Enum for the status of the benchmark/environment."""
+    """Enum for the status of the benchmark/environment Trial or Experiment."""
 
     UNKNOWN = 0
     PENDING = 1
@@ -29,8 +29,8 @@ def is_good(self) -> bool:
         }
 
     def is_completed(self) -> bool:
-        """Check if the status of the benchmark/environment is one of {SUCCEEDED,
-        CANCELED, FAILED, TIMED_OUT}.
+        """Check if the status of the benchmark/environment Trial or Experiment is one
+        of {SUCCEEDED, CANCELED, FAILED, TIMED_OUT}.
         """
         return self in {
             Status.SUCCEEDED,
@@ -40,25 +40,37 @@ def is_completed(self) -> bool:
         }
 
     def is_pending(self) -> bool:
-        """Check if the status of the benchmark/environment is PENDING."""
+        """Check if the status of the benchmark/environment Trial or Experiment is
+        PENDING.
+        """
         return self == Status.PENDING
 
     def is_ready(self) -> bool:
-        """Check if the status of the benchmark/environment is READY."""
+        """Check if the status of the benchmark/environment Trial or Experiment is
+        READY.
+        """
         return self == Status.READY
 
     def is_succeeded(self) -> bool:
-        """Check if the status of the benchmark/environment is SUCCEEDED."""
+        """Check if the status of the benchmark/environment Trial or Experiment is
+        SUCCEEDED.
+        """
         return self == Status.SUCCEEDED
 
     def is_failed(self) -> bool:
-        """Check if the status of the benchmark/environment is FAILED."""
+        """Check if the status of the benchmark/environment Trial or Experiment is
+        FAILED.
+        """
         return self == Status.FAILED
 
     def is_canceled(self) -> bool:
-        """Check if the status of the benchmark/environment is CANCELED."""
+        """Check if the status of the benchmark/environment Trial or Experiment is
+        CANCELED.
+        """
         return self == Status.CANCELED
 
     def is_timed_out(self) -> bool:
-        """Check if the status of the benchmark/environment is TIMED_OUT."""
+        """Check if the status of the benchmark/environment Trial or Experiment is
+        TIMED_OUT.
+        """
         return self == Status.FAILED

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

@@ -36,6 +36,13 @@ This document contains some notes on how to use [`alembic`](https://alembic.sqla
 
    > 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
+   sqlite3 mlos_bench.sqlite .schema
+   sqlite3 mlos_bench.sqlite "SELECT * FROM alembic_version;"
+   ```
+
 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.

mlos_bench/mlos_bench/storage/sql/alembic/versions/8928a401115b_adding_experiment_table_columns_to_.py

@@ -0,0 +1,56 @@
+#
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+#
+"""Adding Experiment table columns to support mlos_benchd service - See #732
+
+Revision ID: 8928a401115b
+Revises: f83fb8ae7fc4
+Create Date: 2025-01-14 17:06:36.181503+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 = "8928a401115b"
+down_revision: str | None = "f83fb8ae7fc4"
+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("experiment", sa.Column("ts_start", sa.DateTime(), nullable=True))
+    op.add_column("experiment", sa.Column("ts_end", sa.DateTime(), nullable=True))
+    op.add_column("experiment", sa.Column("status", sa.String(length=16), nullable=True))
+    op.add_column(
+        "experiment",
+        sa.Column(
+            "driver_name",
+            sa.String(length=40),
+            nullable=True,
+            comment="Driver Host/Container Name",
+        ),
+    )
+    op.add_column(
+        "experiment",
+        sa.Column("driver_pid", sa.Integer(), nullable=True, comment="Driver Process ID"),
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """The schema downgrade script for this revision."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column("experiment", "driver_pid")
+    op.drop_column("experiment", "driver_name")
+    op.drop_column("experiment", "status")
+    op.drop_column("experiment", "ts_end")
+    op.drop_column("experiment", "ts_start")
+    # ### end Alembic commands ###

mlos_bench/mlos_bench/storage/sql/schema.py

@@ -103,6 +103,19 @@ def __init__(self, engine: Engine | None):
             Column("root_env_config", String(1024), nullable=False),
             Column("git_repo", String(1024), nullable=False),
             Column("git_commit", String(40), nullable=False),
+            # For backwards compatibility, we allow NULL for ts_start.
+            Column("ts_start", DateTime),
+            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)),
+            # 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.
+            # They should use a transaction to update it to their own hostname when
+            # they start if and only if its NULL.
+            Column("driver_name", String(40), comment="Driver Host/Container Name"),
+            Column("driver_pid", Integer, comment="Driver Process ID"),
             PrimaryKeyConstraint("exp_id"),
         )
         """The Table storing

mlos_bench/mlos_bench/tests/storage/test_storage_schemas.py

@@ -12,7 +12,7 @@
 # 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"
+CURRENT_ALEMBIC_HEAD = "8928a401115b"
 
 
 def test_storage_schemas(storage: SqlStorage) -> None: