Add tests to confirm mlflow versions match for all dependencies.

* Update where needed.

Author: superdosh

.github/workflows/tests.yml

@@ -8,7 +8,7 @@ on:
   workflow_dispatch:
     inputs:
       branch:
-        description: 'Branch'
+        description: "Branch"
         required: true
         default: main
 
@@ -27,13 +27,19 @@ jobs:
         with:
           python-version: "3.12"
 
+      - name: Verify MLflow versions match
+        run: ./scripts/check_mlflow_versions.sh
+
       - name: Start MLflow server (no jupyter)
         run: |
           ./start_services.sh --no-jupyter -d
 
       - name: Install poetry
         run: pipx install "poetry == 1.8.5"
 
+      - name: Check poetry lock file
+        run: poetry check --lock
+
       - name: Remove existing virtual environment
         run: |
           poetry env remove python || true
@@ -75,6 +81,5 @@ jobs:
         run: |
           docker exec modelplane-jupyter-1 poetry run python /app/test_notebooks.py
 
-          
       - name: Stop MLflow server
         run: docker compose down

Dockerfile.mlflow

@@ -3,6 +3,5 @@ FROM ghcr.io/mlflow/mlflow:v3.1.1
 # The base image does not include various dependencies that are needed for
 # the MLflow server. We assume a postgres backend, so we need psycopg2.
 # We also need boto3 for S3 support, and google-cloud-storage for GCS support.
-# TODO: better way to install these (maybe using poetry.lock to grab consistent versions?)
-RUN pip install mlflow[auth]==3.1.1 psycopg2-binary==2.9.10 boto3==1.38.31 \
-    google-cloud-storage==3.1.0
+RUN pip install mlflow[auth]==3.1.1 psycopg2-binary==2.9.11 boto3==1.40.42 \
+    google-cloud-storage==3.4.0

poetry.lock

@@ -14,7 +14,7 @@ python = ">=3.10,!=3.12.5,<3.13"
 click = "^8"
 dvc = {extras = ["gs"], version = "^3.60"} 
 modelbench = {git = "https://github.com/mlcommons/modelbench.git", rev = "8bf3a7bedcf2ae4049373a46e73ca9be6961c8c6" }
-mlflow = "3.1.1"
+mlflow = { extras = ["auth"], version = "3.1.1" }
 python-dotenv = "^1"
 requests = "^2"
 prometheus-client = "*" # version controlled by dependencies
@@ -37,6 +37,11 @@ ipython = "<9"
 pytest = "^8"
 pytest-cov = "^6"
 
+[tool.poetry.group.mlflow.dependencies]
+psycopg2-binary = "^2.9.10"
+boto3 = "^1.38.31"
+google-cloud-storage = "^3.1.0"
+
 [tool.poetry.scripts]
 modelplane = "modelplane.cli:cli"
 

pyproject.toml

@@ -0,0 +1,78 @@
+#!/usr/bin/env bash
+set -uo pipefail
+
+pyproject="pyproject.toml"
+dockerfile="Dockerfile.mlflow"
+poetry_lock="poetry.lock"
+
+for file in "${pyproject}" "${dockerfile}" "${poetry_lock}"; do
+  [[ -f "${file}" ]] || { echo "Missing ${file}" >&2; exit 1; }
+done
+
+pyproject_mlflow_version="$(grep '^mlflow[[:space:]]*=' "${pyproject}" | head -n1 | sed -E 's/.*version[[:space:]]*=[[:space:]]*"([^"]+)".*/\1/')"
+if [[ -z "${pyproject_mlflow_version}" ]]; then
+  echo "Unable to determine mlflow version from ${pyproject}" >&2
+  exit 1
+fi
+
+docker_image_version="$(head -n1 "${dockerfile}" | cut -d ':' -f2 | tr -d 'v')"
+if [[ -z "${docker_image_version}" ]]; then
+  echo "Unable to determine base image version from ${dockerfile}" >&2
+  exit 1
+fi
+docker_mlflow_version="$(grep -o 'mlflow\[auth\]==[^[:space:]]*' "${dockerfile}" | head -n1 | cut -d '=' -f3)"
+if [[ -z "${docker_mlflow_version}" ]]; then
+  echo "Unable to determine mlflow version from ${dockerfile}" >&2
+  exit 1
+fi
+
+mismatches=0
+
+if [[ "${pyproject_mlflow_version}" != "${docker_image_version}" ]]; then
+  echo "Mismatch: pyproject mlflow (${pyproject_mlflow_version}) vs Docker base (${docker_image_version})" >&2
+  mismatches=$((mismatches + 1))
+fi
+
+if [[ "${pyproject_mlflow_version}" != "${docker_mlflow_version}" ]]; then
+  echo "Mismatch: pyproject mlflow (${pyproject_mlflow_version}) vs Docker pip install (${docker_mlflow_version})" >&2
+  mismatches=$((mismatches + 1))
+fi
+
+lock_version() {
+  local pkg="$1"
+  local name_line
+  name_line="$(grep -n "^name = \"${pkg}\"$" "${poetry_lock}" | head -n1)"
+  [[ -z "${name_line}" ]] && return 1
+  local line_number="${name_line%%:*}"
+  sed -n "$((line_number + 1))p" "${poetry_lock}" | cut -d '"' -f2
+}
+
+check_dep() {
+  local pkg="$1"
+  local token
+  token="$(grep -o "${pkg}==[^[:space:]]*" "${dockerfile}" | head -n1)"
+  local docker_version="${token##*==}"
+  local lock_version_value="$(lock_version "${pkg}")"
+
+  if [[ -z "${docker_version}" || -z "${lock_version_value}" ]]; then
+    echo "Unable to resolve ${pkg} in Dockerfile or poetry.lock" >&2
+    mismatches=$((mismatches + 1))
+    return
+  fi
+
+  if [[ "${docker_version}" != "${lock_version_value}" ]]; then
+    echo "Mismatch: ${pkg} Docker (${docker_version}) vs poetry.lock (${lock_version_value})" >&2
+    mismatches=$((mismatches + 1))
+  fi
+}
+
+# Other dependencies from the dockerfile
+check_dep "psycopg2-binary"
+check_dep "boto3"
+check_dep "google-cloud-storage"
+
+if [[ ${mismatches} -ne 0 ]]; then
+  exit 1
+fi
+
+echo "Dependency versions are consistent (${pyproject_mlflow_version})."