Official MLflow docker image support - slim & full version

.github/workflows/push-images.yml

@@ -44,11 +44,22 @@ jobs:
           tags: |
             type=ref,event=tag
 
-      - name: Build and Push Base Image
+      - name: Build and Push Base Slim Image
         uses: docker/build-push-action@v3
         with:
           context: docker
           push: true
+          tags: ${{ steps.meta.outputs.tags }}-slim
+          platforms: linux/amd64,linux/arm64
+          build-args: |
+            VERSION=${{ steps.meta.outputs.version }}
+
+      - name: Build and Push Base Full Image
+        uses: docker/build-push-action@v3
+        with:
+          context: docker
+          file: docker/Dockerfile.full
+          push: true
           tags: ${{ steps.meta.outputs.tags }}
           platforms: linux/amd64,linux/arm64
           build-args: |

dev/pyproject.py

@@ -89,6 +89,12 @@ def build(skinny: bool) -> None:
                     # Required for exporting metrics from the MLflow server to Prometheus
                     # as part of the MLflow server monitoring add-on
                     "prometheus-flask-exporter",
+                    # Required to use MySQL as the backend store
+                    "PyMySQL",
+                    # Required to use PostgreSQL as the backend store
+                    "psycopg2-binary",
+                    # Required to use SQL Server as the backend store
+                    "pymssql",
                 ],
                 "databricks": [
                     # Required to write model artifacts to unity catalog locations

docker/Dockerfile.full

@@ -0,0 +1,4 @@
+FROM python:3.10-slim-bullseye
+ARG VERSION
+RUN pip install --upgrade pip && \
+    pip install --no-cache mlflow[extras,databricks,gateway,genai,sqlserver,aliyun-oss]==$VERSION

docker/Dockerfile.full.dev

@@ -0,0 +1,6 @@
+FROM python:3.10-slim-bullseye
+
+WORKDIR /app
+ADD . /app
+
+RUN pip install --upgrade pip && pip install --no-cache-dir -e .[extras,databricks,gateway,genai,sqlserver,aliyun-oss]

docker/README.md

@@ -0,0 +1,60 @@
+# mlflow docker image
+
+## Image variants
+
+### mlflow:\<version\>
+
+This image contains mlflow plus all extra dependencies. It is intended to support "out of the box" all mlflow integrations.
+
+Use this image if you just want to quickly have a working instance of mlflow, or if you don't want to bother rolling out your own image.
+
+### mlflow:\<version\>-slim
+
+This image only contains the core mlflow dependencies, so most of the integrations (ie. backend stores databases,
+artifact stores, etc.) will not work.
+
+Use this image if you don't need any integration, or as the base image in case you want to have full control over
+the extra dependencies installed.
+
+# Testing
+
+The [testcontainers](https://testcontainers.com/) package is used to the test the different integrations, for each one there
+is a Docker Compose file. So far the following integrations are tested:
+
+Backend store:
+
+- MySQL
+- PostgreSQL
+- MSSQL
+
+Artifact store:
+
+- Amazon S3 (through the [minio](https://hub.docker.com/r/minio/minio) image)
+- Azure Blob Storage (throug the [Azurite](https://hub.docker.com/_/microsoft-azure-storage-azurite) image)
+- Google Cloud Storage (through the [fake-gcs-server](https://github.com/fsouza/fake-gcs-server) image)
+
+## Executing the tests
+
+Make sure that you install all required dependencies first.
+
+```bash
+# mlflow, from a release
+$ pip install -U mlflow
+# or in editable mode
+$ pip install -e .
+# testcontaienrs & pytest
+$ pip install -U testcontainers<4 pytest
+```
+
+Ensure [Docker](https://www.docker.com/) is installed. You will also need `docker-compose`, make sure to install
+[docker-compose v1](https://docs.docker.com/compose/install/standalone/), since `testcontainers` only supports Compose v2
+from version 4.0 onwards, which is available only for Python 3.9+.
+
+Once you have all the dependencies you can execute the tests (from within `docker` folder):
+
+```bash
+# Execute all tests
+pytest
+# Execute a specific test
+pytest -k 'test_backend_and_artifact_store_integration[docker-compose.aws-test.yaml]'
+```

docker/tests/docker-compose.aws-test.yaml

@@ -0,0 +1,30 @@
+version: "2.4"
+
+services:
+  minio:
+    container_name: mlflow-aws-test-storage
+    image: minio/minio
+    entrypoint: sh
+    command: -c 'mkdir -p /data/mlflow && minio server /data --console-address ":9001"'
+    environment:
+      MINIO_ROOT_USER: minioadmin
+      MINIO_ROOT_PASSWORD: minioadmin
+    expose:
+      - 9000
+      - 9001
+
+  mlflow:
+    container_name: mlflow-aws-test
+    image: mlflow-integration-test
+    build:
+      context: ../..
+      dockerfile: docker/Dockerfile.full.dev
+    command: "mlflow server --artifacts-destination=s3://mlflow/ --host=0.0.0.0 --port=5000"
+    environment:
+      MLFLOW_S3_ENDPOINT_URL: http://minio:9000
+      AWS_ACCESS_KEY_ID: minioadmin
+      AWS_SECRET_ACCESS_KEY: minioadmin
+    ports:
+      - "5000:5000"
+    depends_on:
+      - minio

docker/tests/docker-compose.azure-test.yaml

@@ -0,0 +1,36 @@
+version: "2.4"
+
+services:
+  azurite:
+    image: mcr.microsoft.com/azure-storage/azurite
+    container_name: mlflow-azure-test-storage
+    command: azurite --blobHost 0.0.0.0 --queueHost 0.0.0.0 --loose --skipApiVersionCheck
+    environment:
+      AZURITE_ACCOUNTS: devstoreaccount1:Eby8vdM02xNOcqFlqUwJPLlmEtlCDXJ1OUzFT50uSRZ6IFsuFq2UVErCz4I6tq/K1SZFPTOtr/KBHBeksoGMGw==
+    expose:
+      - 10000
+      - 10001
+
+  azurite-init:
+    image: mcr.microsoft.com/azure-cli
+    container_name: azurite-init
+    command: az storage container create --name mlflow
+    environment:
+      AZURE_STORAGE_CONNECTION_STRING: DefaultEndpointsProtocol=http;AccountName=devstoreaccount1;AccountKey=Eby8vdM02xNOcqFlqUwJPLlmEtlCDXJ1OUzFT50uSRZ6IFsuFq2UVErCz4I6tq/K1SZFPTOtr/KBHBeksoGMGw==;BlobEndpoint=http://azurite:10000/devstoreaccount1;QueueEndpoint=http://azurite:10001/devstoreaccount1
+    depends_on:
+      - azurite
+
+  mlflow:
+    container_name: mlflow-azure-test
+    image: mlflow-integration-test
+    build:
+      context: ../..
+      dockerfile: docker/Dockerfile.full.dev
+    command: "mlflow server --artifacts-destination=wasbs://mlflow@devstoreaccount1.blob.core.windows.net --host=0.0.0.0 --port=5000"
+    environment:
+      AZURE_STORAGE_CONNECTION_STRING: DefaultEndpointsProtocol=http;AccountName=devstoreaccount1;AccountKey=Eby8vdM02xNOcqFlqUwJPLlmEtlCDXJ1OUzFT50uSRZ6IFsuFq2UVErCz4I6tq/K1SZFPTOtr/KBHBeksoGMGw==;BlobEndpoint=http://azurite:10000/devstoreaccount1;QueueEndpoint=http://azurite:10001/devstoreaccount1
+    ports:
+      - "5000:5000"
+    depends_on:
+      - azurite-init
+      - azurite

docker/tests/docker-compose.gcp-test.yaml

@@ -0,0 +1,24 @@
+version: "2.4"
+
+services:
+  gcs:
+    container_name: mlflow-gcp-test-storage
+    image: fsouza/fake-gcs-server
+    entrypoint: sh -c "mkdir -p /data/mlflow && /bin/fake-gcs-server -data /data -scheme http"
+    expose:
+      - 4443
+
+  mlflow:
+    container_name: mlflow-gcp-test
+    image: mlflow-integration-test
+    build:
+      context: ../..
+      dockerfile: docker/Dockerfile.full.dev
+    command: "mlflow server --artifacts-destination=gs://mlflow/ --host=0.0.0.0 --port=5000"
+    environment:
+      GOOGLE_CLOUD_PROJECT: mlflow
+      STORAGE_EMULATOR_HOST: http://gcs:4443
+    ports:
+      - "5000:5000"
+    depends_on:
+      - gcs

docker/tests/docker-compose.mssql-test.yaml

@@ -0,0 +1,23 @@
+version: "2.4"
+
+services:
+  mssql:
+    image: mcr.microsoft.com/mssql/server
+    container_name: mlflow-mssql-test-db
+    environment:
+      MSSQL_SA_PASSWORD: StrongPassword!
+      ACCEPT_EULA: "Y"
+    expose:
+      - 1433
+
+  mlflow:
+    container_name: mlflow-mssql-test
+    image: mlflow-integration-test
+    build:
+      context: ../..
+      dockerfile: docker/Dockerfile.full.dev
+    command: "mlflow server --backend-store-uri=mssql+pymssql://sa:StrongPassword!@mssql:1433 --host=0.0.0.0 --port=5000"
+    ports:
+      - "5000:5000"
+    depends_on:
+      - mssql

docker/tests/docker-compose.mysql-test.yaml

@@ -0,0 +1,26 @@
+version: "2.4"
+
+services:
+  mysql:
+    image: "mysql"
+    container_name: mlflow-mysql-test-db
+    environment:
+      MYSQL_ROOT_PASSWORD: password
+      MYSQL_DATABASE: mlflow
+      MYSQL_USER: mlflow
+      MYSQL_PASSWORD: password
+      MYSQL_TCP_PORT: 3306
+    expose:
+      - 3306
+
+  mlflow:
+    container_name: mlflow-mysql-test
+    image: mlflow-integration-test
+    build:
+      context: ../..
+      dockerfile: docker/Dockerfile.full.dev
+    command: "mlflow server --backend-store-uri=mysql+pymysql://mlflow:password@mysql:3306/mlflow --host=0.0.0.0 --port=5000"
+    ports:
+      - "5000:5000"
+    depends_on:
+      - mysql

docker/tests/docker-compose.postgres-test.yaml

@@ -0,0 +1,29 @@
+version: "2.4"
+
+services:
+  postgres:
+    image: "postgres"
+    container_name: mlflow-pg-test-db
+    environment:
+      POSTGRES_PASSWORD: postgres
+      POSTGRES_USER: postgres
+    expose:
+      - 5432
+
+  mlflow:
+    container_name: mlflow-pg-test
+    image: mlflow-integration-test
+    build:
+      context: ../..
+      dockerfile: docker/Dockerfile.full.dev
+    command: "mlflow server --backend-store-uri=postgresql:// --host=0.0.0.0 --port=5000"
+    environment:
+      PGHOST: postgres
+      PGPORT: 5432
+      PGDATABASE: postgres
+      PGUSER: postgres
+      PGPASSWORD: postgres
+    ports:
+      - "5000:5000"
+    depends_on:
+      - postgres

docker/tests/test_integrations.py

@@ -0,0 +1,54 @@
+import os
+
+import pytest
+from testcontainers.compose import DockerCompose
+
+import mlflow
+from mlflow.tracking.client import MlflowClient
+
+
+class TestModel(mlflow.pyfunc.PythonModel):
+    def load_context(self, context):
+        pass
+
+    def predict(self, context, model_input, params=None):
+        pass
+
+
+_COMPOSE_FILE_LIST = [
+    "docker-compose.aws-test.yaml",
+    "docker-compose.azure-test.yaml",
+    "docker-compose.gcp-test.yaml",
+    "docker-compose.mssql-test.yaml",
+    "docker-compose.mysql-test.yaml",
+    "docker-compose.postgres-test.yaml",
+]
+
+
+def create_experiment_and_register_model(base_url: str):
+    experiment_name = "integration-test-experiment"
+    model_name = "integration-test-model"
+    mlflow.set_tracking_uri(base_url)
+    mlflow.set_experiment(experiment_name)
+    experiment = mlflow.get_experiment_by_name(experiment_name)
+    test_model = TestModel()
+    with mlflow.start_run(experiment_id=experiment.experiment_id) as run:
+        mlflow.log_params({"param": 1})
+        mlflow.log_metric("metric", 1.0)
+        mlflow.pyfunc.log_model("model", python_model=test_model)
+        model_uri = f"runs:/{run.info.run_id}/model"
+        mlflow.register_model(model_uri, model_name, tags={"status": "ready"})
+        client = MlflowClient()
+        client.set_registered_model_alias(model_name, "champion", "1")
+
+
+@pytest.mark.parametrize("compose_file_name", _COMPOSE_FILE_LIST)
+def test_backend_and_artifact_store_integration(compose_file_name):
+    with DockerCompose(
+        filepath=os.path.dirname(os.path.abspath(__file__)), compose_file_name=[compose_file_name]
+    ) as compose:
+        mlflow_host = compose.get_service_host("mlflow", 5000)
+        mlflow_port = compose.get_service_port("mlflow", 5000)
+        base_url = f"http://{mlflow_host}:{mlflow_port}"
+        compose.wait_for(base_url)
+        create_experiment_and_register_model(base_url)

docs/source/docker.rst

@@ -3,6 +3,22 @@ Official MLflow Docker Image
 
 The official MLflow Docker image is available on GitHub Container Registry at https://ghcr.io/mlflow/mlflow.
 
+There are two image variants.
+
+mlflow:\<version\>
+^^^^^^^^^^^^^^^^^^
+This image contains mlflow plus all extra dependencies. It is intended to support "out of the box" all mlflow integrations.
+
+Use this image if you just want to quickly have a working instance of mlflow, or if you don't want to bother rolling out your own image.
+
+mlflow:\<version\>-slim
+^^^^^^^^^^^^^^^^^^^^^^^
+This image only contains the core mlflow dependencies, so most of the integrations (ie. backend stores databases,
+artifact stores, etc.) will not work.
+
+Use this image if you don't need any integration, or as the base image in case you want to have full control over
+the extra dependencies installed.
+
 .. code-block:: shell
 
     export CR_PAT=YOUR_TOKEN
@@ -11,3 +27,5 @@ The official MLflow Docker image is available on GitHub Container Registry at ht
     docker pull ghcr.io/mlflow/mlflow
     # Pull 2.0.1
     docker pull ghcr.io/mlflow/mlflow:v2.0.1
+    # Pull 2.0.1 slim version
+    docker pull ghcr.io/mlflow/mlflow:v2.0.1-slim

pyproject.skinny.toml

@@ -55,6 +55,9 @@ extras = [
   "mlserver-mlflow>=1.2.0,!=1.3.1,<1.4.0",
   "virtualenv",
   "prometheus-flask-exporter",
+  "PyMySQL",
+  "psycopg2-binary",
+  "pymssql",
 ]
 databricks = [
   "azure-storage-file-datalake>12",

pyproject.toml

@@ -72,6 +72,9 @@ extras = [
   "mlserver-mlflow>=1.2.0,!=1.3.1,<1.4.0",
   "virtualenv",
   "prometheus-flask-exporter",
+  "PyMySQL",
+  "psycopg2-binary",
+  "pymssql",
 ]
 databricks = [
   "azure-storage-file-datalake>12",