version: "3.8"

services:
  mlflow:
    image: ghcr.io/mlflow/mlflow:${MLFLOW_VERSION:-latest}
    container_name: mlflow-server
    restart: unless-stopped
    ports:
      - "${MLFLOW_PORT:-5000}:5000"
    environment:
      - MLFLOW_TRACKING_URI=postgresql://${PG_USER:-mlflow}:${PG_PASSWORD:-mlflow}@mlflow-db:5432/${PG_DB:-mlflow}
      - MLFLOW_S3_ENDPOINT_URL=http://mlflow-minio:9000
      - AWS_ACCESS_KEY_ID=${MINIO_ROOT_USER:-mlflow}
      - AWS_SECRET_ACCESS_KEY=${MINIO_ROOT_PASSWORD:-mlflow123}
      - MLFLOW_DEFAULT_ARTIFACT_ROOT=s3://${ARTIFACT_BUCKET:-mlflow-artifacts}/
    command: >
      mlflow server
      --host 0.0.0.0
      --port 5000
      --backend-store-uri postgresql://${PG_USER:-mlflow}:${PG_PASSWORD:-mlflow}@mlflow-db:5432/${PG_DB:-mlflow}
      --default-artifact-root s3://${ARTIFACT_BUCKET:-mlflow-artifacts}/
      --serve-artifacts
    depends_on:
      mlflow-db:
        condition: service_healthy
      mlflow-minio:
        condition: service_started

  mlflow-db:
    image: postgres:16-alpine
    container_name: mlflow-db
    restart: unless-stopped
    environment:
      - POSTGRES_USER=${PG_USER:-mlflow}
      - POSTGRES_PASSWORD=${PG_PASSWORD:-mlflow}
      - POSTGRES_DB=${PG_DB:-mlflow}
    volumes:
      - mlflow_pg_data:/var/lib/postgresql/data
    healthcheck:
      test: ["CMD-SHELL", "pg_isready -U ${PG_USER:-mlflow}"]
      interval: 10s
      timeout: 5s
      retries: 5

  mlflow-minio:
    image: quay.io/minio/minio:latest
    container_name: mlflow-minio
    restart: unless-stopped
    ports:
      - "${MINIO_API_PORT:-9000}:9000"
      - "${MINIO_CONSOLE_PORT:-9001}:9001"
    volumes:
      - mlflow_minio_data:/data
    environment:
      - MINIO_ROOT_USER=${MINIO_ROOT_USER:-mlflow}
      - MINIO_ROOT_PASSWORD=${MINIO_ROOT_PASSWORD:-mlflow123}
    command: server /data --console-address ':9001'

  # Init container to create the default bucket
  mlflow-minio-init:
    image: quay.io/minio/mc:latest
    container_name: mlflow-minio-init
    depends_on:
      - mlflow-minio
    entrypoint: >
      /bin/sh -c "
      sleep 5;
      mc alias set myminio http://mlflow-minio:9000 ${MINIO_ROOT_USER:-mlflow} ${MINIO_ROOT_PASSWORD:-mlflow123};
      mc mb --ignore-existing myminio/${ARTIFACT_BUCKET:-mlflow-artifacts};
      mc anonymous set download myminio/${ARTIFACT_BUCKET:-mlflow-artifacts};
      exit 0;
      "

volumes:
  mlflow_pg_data:
  mlflow_minio_data: