Repository: mlflow/mlflow

Issue:
Allow alembic>=1.7.5
## What changes are proposed in this pull request?

The change is simply to increase the allowed upper version bound of `alembic` to `1.7.5`. 

This is important because in our organization we install `airflow` and `mlflow` in the same environment. Current version constraint `"alembic<=1.4.1"` is blocking upgrade to latest `airflow==2.2.3`.

Relevant issue: https://github.com/mlflow/mlflow/issues/4215
Similar PR: https://github.com/mlflow/mlflow/pull/4810

## How is this patch tested?
I am aware I could only test this in a limited way, and only for one type of store backend (postgresql), but I did the following:

- Installed `mlflow==0.13.1` locally with a postgresql backend
- Ran some operations against MLflow: create an experiment with a few runs, tags, registered a model
- Installed updated `mlflow` from source, upgraded to `alembic==1.7.5`
- Ran `mlflow db upgrade ${psql_path}`
- Verified data after upgrade

## Does this PR change the documentation?

- [x] No. You can skip the rest of this section.
- [ ] Yes. Make sure the changed pages / sections render correctly by following the steps below.

1. Check the status of the `ci/circleci: build_doc` check. If it's successful, proceed to the
   next step, otherwise fix it.
2. Click `Details` on the right to open the job page of CircleCI.
3. Click the `Artifacts` tab.
4. Click `docs/build/html/index.html`.
5. Find the changed pages / sections and make sure they render correctly.

## Release Notes

### Is this a user-facing change?

- [x] No. You can skip the rest of this section.
- [ ] Yes. Give a description of this change to be included in the release notes for MLflow users.

(Details in 1-2 sentences. You can just refer to another PR with a description if this PR is part of a larger change.)

### What component(s), interfaces, languages, and integrations does this PR affect?
Components 
- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [x] `area/server-infra`: MLflow Tracking server backend
- [ ] `area/tracking`: Tracking Service, tracking client APIs, autologging

Interface 
- [ ] `area/uiux`: Front-end, user experience, plotting, JavaScript, JavaScript dev server
- [ ] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

Language 
- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

Integrations
- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations

<!--
Insert an empty named anchor here to allow jumping to this section with a fragment URL
(e.g. https://github.com/mlflow/mlflow/pull/123#user-content-release-note-category).
Note that GitHub prefixes anchor names in markdown with "user-content-".
-->
<a name="release-note-category"></a>
### How should the PR be classified in the release notes? Choose one:

- [ ] `rn/breaking-change` - The PR will be mentioned in the "Breaking Changes" section
- [ ] `rn/none` - No description will be included. The PR will be mentioned only by the PR number in the "Small Bugfixes and Documentation Updates" section
- [ ] `rn/feature` - A new user-facing feature worth mentioning in the release notes
- [ ] `rn/bug-fix` - A user-facing bug fix worth mentioning in the release notes
- [x] `rn/documentation` - A user-facing documentation change worth mentioning in the release notes

I would figure that updating the db migration tool is quite relevant to mention to users, who might want to take extra precautions during upgrades.


Discussion:

Resolving patch:
diff --git a/mlflow/store/db_migrations/versions/0a8213491aaa_drop_duplicate_killed_constraint.py b/mlflow/store/db_migrations/versions/0a8213491aaa_drop_duplicate_killed_constraint.py
--- a/mlflow/store/db_migrations/versions/0a8213491aaa_drop_duplicate_killed_constraint.py
+++ b/mlflow/store/db_migrations/versions/0a8213491aaa_drop_duplicate_killed_constraint.py
@@ -31,7 +31,10 @@ 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")
+        # For other database backends, the status check constraint is dropped by
+        # cfd24bdc0731_update_run_status_constraint_with_killed.py
+        if op.get_bind().engine.name == "mysql":
+            op.drop_constraint(constraint_name="status", table_name="runs", type_="check")
     except Exception as e:
         _logger.warning(
             "Failed to drop check constraint. Dropping check constraints may not be supported"
diff --git a/mlflow/store/db_migrations/versions/cfd24bdc0731_update_run_status_constraint_with_killed.py b/mlflow/store/db_migrations/versions/cfd24bdc0731_update_run_status_constraint_with_killed.py
--- a/mlflow/store/db_migrations/versions/cfd24bdc0731_update_run_status_constraint_with_killed.py
+++ b/mlflow/store/db_migrations/versions/cfd24bdc0731_update_run_status_constraint_with_killed.py
@@ -5,11 +5,14 @@
 Create Date: 2019-10-11 15:55:10.853449
 
 """
+import alembic
 from alembic import op
+from packaging.version import Version
+from sqlalchemy import CheckConstraint, Enum
+
 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
 
 # revision identifiers, used by Alembic.
 revision = "cfd24bdc0731"
@@ -17,14 +20,15 @@
 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,13 +44,29 @@
 
 
 def upgrade():
-    with op.batch_alter_table("runs", table_args=check_constraint_table_args) as batch_op:
+    # In alembic >= 1.7.0, `table_args` is unnecessary since CHECK constraints are preserved
+    # during migrations.
+    table_args = (
+        [] if Version(alembic.__version__) >= Version("1.7.0") else check_constraint_table_args
+    )
+    with op.batch_alter_table("runs", table_args=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)
+            "status",
+            type_=Enum(
+                *new_run_statuses,
+                create_constraint=True,
+                native_enum=False,
+            ),
+            existing_type=Enum(
+                *old_run_statuses,
+                create_constraint=True,
+                native_enum=False,
+                name="status",
+            ),
         )
 
 
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/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",


Test patch:
diff --git a/tests/db/Dockerfile b/tests/db/Dockerfile
--- a/tests/db/Dockerfile
+++ b/tests/db/Dockerfile
@@ -2,8 +2,7 @@ FROM python:3.6
 
 WORKDIR /tmp/mlflow
 
-COPY dist ./dist
-
-RUN pip install dist/*.whl
 RUN pip install psycopg2 pymysql mysqlclient
+COPY dist ./dist
+RUN pip install dist/mlflow-*.whl
 RUN pip list
diff --git a/tests/db/Dockerfile.mssql b/tests/db/Dockerfile.mssql
--- a/tests/db/Dockerfile.mssql
+++ b/tests/db/Dockerfile.mssql
@@ -2,11 +2,9 @@ FROM python:3.6
 
 WORKDIR /tmp/mlflow
 
-COPY dist ./dist
-
 # apt-get and system utilities
 RUN apt-get update && apt-get install -y \
-	curl apt-transport-https debconf-utils \
+    curl apt-transport-https debconf-utils \
     && rm -rf /var/lib/apt/lists/*
 
 # adding custom MS repository
@@ -16,6 +14,7 @@ RUN curl https://packages.microsoft.com/config/ubuntu/20.04/prod.list > /etc/apt
 # install SQL Server drivers and tools
 RUN apt-get update && ACCEPT_EULA=Y apt-get install -y mssql-tools unixodbc-dev
 
-RUN pip install dist/*.whl
 RUN pip install pyodbc
+COPY dist ./dist
+RUN pip install dist/mlflow-*.whl
 RUN pip list
diff --git a/tests/db/build_wheel.sh b/tests/db/build_wheel.sh
new file mode 100755
--- /dev/null
+++ b/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
diff --git a/tests/db/docker-compose.yml b/tests/db/docker-compose.yml
--- a/tests/db/docker-compose.yml
+++ b/tests/db/docker-compose.yml
@@ -17,15 +17,17 @@ services:
       - .:/tmp/mlflow
     environment:
       MLFLOW_TRACKING_URI: postgresql://mlflowuser:mlflowpassword@postgres:5432/mlflowdb
+    command: python run_checks.py --schema-output schemas/postgres.sql
 
   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:
@@ -36,6 +38,7 @@ services:
       - .:/tmp/mlflow
     environment:
       MLFLOW_TRACKING_URI: mysql://mlflowuser:mlflowpassword@mysql:3306/mlflowdb
+    command: python run_checks.py --schema-output schemas/mysql.sql
 
   mssql:
     image: mcr.microsoft.com/mssql/server
@@ -54,6 +57,7 @@ services:
       - .:/tmp/mlflow
     environment:
       MLFLOW_TRACKING_URI: mssql+pyodbc://mlflowuser:Mlfl*wpassword1@mssql/mlflowdb?driver=ODBC+Driver+17+for+SQL+Server
+    command: bash -ex -c "./init-mssql-db.sh && python run_checks.py --schema-output schemas/mssql.sql"
 
   mlflow-sqlite:
     depends_on:
@@ -64,3 +68,4 @@ services:
       - .:/tmp/mlflow
     environment:
       MLFLOW_TRACKING_URI: sqlite:////tmp/mlflow/mlflowdb
+    command: python run_checks.py --schema-output schemas/sqlite.sql
diff --git a/tests/db/run_checks.py b/tests/db/run_checks.py
--- a/tests/db/run_checks.py
+++ b/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 work correctly:
+    # - cfd24bdc0731_update_run_status_constraint_with_killed.py
+    # - 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())