Unpin alembic

.github/workflows/master.yml

@@ -161,11 +161,11 @@ jobs:
         ./dev/run-large-python-tests.sh
     # Separate build and run to make it easier to explore logs
     - name: Run database initialization tests - build
+      working-directory: tests/db
       run: |
-        python setup.py bdist_wheel
-        cp -r dist tests/db
-        cd tests/db
+        ./build_wheel.sh
         docker-compose pull
+        docker image ls | grep -E '(REPOSITORY|postgres|mysql|mssql)'
         docker-compose build
     - name: Run database initialization tests - run
       working-directory: tests/db
@@ -175,7 +175,7 @@ jobs:
         docker-compose run mlflow-mysql python run_checks.py --schema-output schemas/mysql.sql
         docker-compose run mlflow-mssql ./init-mssql-db.sh
         docker-compose run mlflow-mssql python run_checks.py --schema-output schemas/mssql.sql
-        docker-compose down --rmi all --volumes --remove-orphans
+        docker-compose down --volumes --remove-orphans --rmi all
     - name: Run anaconda compatibility tests
       run: |
         ./dev/test-anaconda-compatibility.sh "anaconda3:2020.11"

mlflow/store/db_migrations/versions/0a8213491aaa_drop_duplicate_killed_constraint.py

@@ -31,7 +31,8 @@ def upgrade():
     # operation is expected to fail under certain circumstances, we execute `drop_constraint()`
     # outside of the batch operation context.
     try:
-        op.drop_constraint(constraint_name="status", table_name="runs", type_="check")
+        # op.drop_constraint(constraint_name="status", table_name="runs", type_="check")
+        pass
     except Exception as e:
         _logger.warning(
             "Failed to drop check constraint. Dropping check constraints may not be supported"

mlflow/store/db_migrations/versions/cfd24bdc0731_update_run_status_constraint_with_killed.py

@@ -5,26 +5,29 @@
 Create Date: 2019-10-11 15:55:10.853449
 
 """
+import alembic
 from alembic import op
 from mlflow.entities import RunStatus, ViewType
 from mlflow.entities.lifecycle_stage import LifecycleStage
 from mlflow.store.tracking.dbmodels.models import SqlRun, SourceTypes
 from sqlalchemy import CheckConstraint, Enum
+from packaging.version import Version
 
 # revision identifiers, used by Alembic.
 revision = "cfd24bdc0731"
 down_revision = "2b4d017a5e9b"
 branch_labels = None
 depends_on = None
 
-new_run_statuses = [
+old_run_statuses = [
     RunStatus.to_string(RunStatus.SCHEDULED),
     RunStatus.to_string(RunStatus.FAILED),
     RunStatus.to_string(RunStatus.FINISHED),
     RunStatus.to_string(RunStatus.RUNNING),
-    RunStatus.to_string(RunStatus.KILLED),
 ]
 
+new_run_statuses = [*old_run_statuses, RunStatus.to_string(RunStatus.KILLED)]
+
 # Certain SQL backends (e.g., SQLite) do not preserve CHECK constraints during migrations.
 # For these backends, CHECK constraints must be specified as table arguments. Here, we define
 # the collection of CHECK constraints that should be preserved when performing the migration.
@@ -40,14 +43,27 @@
 
 
 def upgrade():
-    with op.batch_alter_table("runs", table_args=check_constraint_table_args) as batch_op:
-        # Transform the "status" column to an `Enum` and define a new check constraint. Specify
-        # `native_enum=False` to create a check constraint rather than a
-        # database-backend-dependent enum (see https://docs.sqlalchemy.org/en/13/core/
-        # type_basics.html#sqlalchemy.types.Enum.params.native_enum)
-        batch_op.alter_column(
-            "status", type_=Enum(*new_run_statuses, create_constraint=True, native_enum=False)
-        )
+    new_type = Enum(*new_run_statuses, create_constraint=True, native_enum=False)
+    if Version(alembic.__version__) < Version("1.7.0"):
+        with op.batch_alter_table("runs", table_args=check_constraint_table_args) as batch_op:
+            # Transform the "status" column to an `Enum` and define a new check constraint. Specify
+            # `native_enum=False` to create a check constraint rather than a
+            # database-backend-dependent enum (see https://docs.sqlalchemy.org/en/13/core/
+            # type_basics.html#sqlalchemy.types.Enum.params.native_enum)
+            batch_op.alter_column("status", type_=new_type)
+    else:
+        # In alembic >= 1.7.0, `table_args` can be removed since CHECK constraints are preserved.
+        with op.batch_alter_table("runs") as batch_op:
+            existing_type = Enum(
+                *old_run_statuses, create_constraint=True, native_enum=False, name="status"
+            )
+            batch_op.alter_column(
+                "status",
+                type_=new_type,
+                # In alembic >= 1.7.0, `existing_type` is required to drop the existing CHECK
+                # constraint on the status column.
+                existing_type=existing_type,
+            )
 
 
 def downgrade():

setup.py

@@ -65,7 +65,7 @@ def package_files(directory):
 other capabilities.
 """
 CORE_REQUIREMENTS = SKINNY_REQUIREMENTS + [
-    "alembic<=1.4.1",
+    "alembic",
     # Required
     "docker>=4.0.0",
     "Flask",

tests/db/build_wheel.sh

@@ -0,0 +1,9 @@
+#!/usr/bin/env bash
+
+set -ex
+
+rm -rf dist
+prefix=$(git rev-parse --show-prefix)
+pushd $(git rev-parse --show-cdup)
+python setup.py bdist_wheel --dist-dir $prefix/dist
+popd

tests/db/docker-compose.yml

@@ -19,13 +19,14 @@ services:
       MLFLOW_TRACKING_URI: postgresql://mlflowuser:mlflowpassword@postgres:5432/mlflowdb
 
   mysql:
-    image: mysql:5.7
+    image: mysql
     restart: always
     environment:
       MYSQL_ROOT_PASSWORD: root-password
       MYSQL_DATABASE: mlflowdb
       MYSQL_USER: mlflowuser
       MYSQL_PASSWORD: mlflowpassword
+    command: mysqld --default-authentication-plugin=mysql_native_password
 
   mlflow-mysql:
     depends_on:

tests/db/run_checks.py

@@ -35,6 +35,12 @@ def run_logging_operations():
         )
         print(mlflow.get_run(run.info.run_id))
 
+    # Ensure the following migration scripts are applied correctly:
+    # - cfd24bdc0731_update_run_status_constraint_with_killed.py
+    # - versions/0a8213491aaa_drop_duplicate_killed_constraint.py
+    client = mlflow.tracking.MlflowClient()
+    client.set_terminated(run_id=run.info.run_id, status="KILLED")
+
 
 def get_db_schema():
     engine = sqlalchemy.create_engine(mlflow.get_tracking_uri())