volumes:
  pgdata:
  storage-data:

services:
  postgres:
    image: postgres:15
    container_name: mlflow-postgres
    environment:
      POSTGRES_USER: ${POSTGRES_USER}
      POSTGRES_PASSWORD: ${POSTGRES_PASSWORD}
      POSTGRES_DB: ${POSTGRES_DB}
      PGPORT: ${PGPORT}
    volumes:
      - pgdata:/var/lib/postgresql/data
    ports:
      - ${PGPORT}:${PGPORT}
    healthcheck:
      test: ["CMD-SHELL", "pg_isready -U ${POSTGRES_USER} -d ${POSTGRES_DB} -p ${PGPORT}"]
      interval: 5s
      timeout: 3s
      retries: 10

  storage:
    image: rustfs/rustfs:1.0.0-alpha.83
    container_name: storage
    environment:
      RUSTFS_ADDRESS: :9000
      RUSTFS_SERVER_DOMAINS: storage:9000
      RUSTFS_REGION: ${AWS_DEFAULT_REGION:-us-east-1}
      RUSTFS_ACCESS_KEY: ${AWS_ACCESS_KEY_ID:-s3admin}
      RUSTFS_SECRET_KEY: ${AWS_SECRET_ACCESS_KEY:-s3admin}
      RUSTFS_CONSOLE_ENABLE: ${RUSTFS_CONSOLE_ENABLE:-true}
    ports:
      - "9000:9000"
      - "9001:9001"
    volumes:
      - storage-data:/data
    restart: unless-stopped
    healthcheck:
      test: ["CMD-SHELL", 'curl -s http://127.0.0.1:9000/health | grep -q ''"status":"ok"''']
      interval: 10s
      timeout: 5s
      retries: 5
      start_period: 10s

  create-bucket:
    image: amazon/aws-cli:2.33.25
    container_name: mlflow-create-bucket
    depends_on:
      storage:
        condition: service_healthy
    entrypoint: >
      /bin/sh -c "
        set -e;
        echo 'Waiting for S3 gateway getting ready...';
        if aws --endpoint-url=${MLFLOW_S3_ENDPOINT_URL} s3api head-bucket --bucket ${S3_BUCKET} 2>/dev/null; then
          echo 'Bucket ${S3_BUCKET} already exists. Skipping creation.';
        else
          echo 'Creating bucket ${S3_BUCKET}...';
          aws --endpoint-url=${MLFLOW_S3_ENDPOINT_URL} s3api create-bucket --bucket ${S3_BUCKET} --region ${AWS_DEFAULT_REGION};
        fi
      "
    environment:
      AWS_ACCESS_KEY_ID: ${AWS_ACCESS_KEY_ID}
      AWS_SECRET_ACCESS_KEY: ${AWS_SECRET_ACCESS_KEY}
      AWS_DEFAULT_REGION: ${AWS_DEFAULT_REGION}
      AWS_S3_ADDRESSING_STYLE: path
      MLFLOW_S3_ENDPOINT_URL: ${MLFLOW_S3_ENDPOINT_URL}
      S3_BUCKET: ${S3_BUCKET}
    restart: "no"

  mlflow:
    image: ghcr.io/mlflow/mlflow:${MLFLOW_VERSION}
    container_name: mlflow-server
    depends_on:
      postgres:
        condition: service_healthy
      storage:
        condition: service_healthy
      create-bucket:
        condition: service_completed_successfully
    environment:
      # Backend store URI built from vars
      MLFLOW_BACKEND_STORE_URI: ${MLFLOW_BACKEND_STORE_URI}

      # S3/RustFS settings
      MLFLOW_S3_ENDPOINT_URL: ${MLFLOW_S3_ENDPOINT_URL}
      MLFLOW_ARTIFACTS_DESTINATION: ${MLFLOW_ARTIFACTS_DESTINATION}
      AWS_ACCESS_KEY_ID: ${AWS_ACCESS_KEY_ID}
      AWS_SECRET_ACCESS_KEY: ${AWS_SECRET_ACCESS_KEY}
      AWS_DEFAULT_REGION: ${AWS_DEFAULT_REGION}
      MLFLOW_S3_IGNORE_TLS: "true"

      # Server host/port
      MLFLOW_HOST: ${MLFLOW_HOST}
      MLFLOW_PORT: ${MLFLOW_PORT}
    command:
      - /bin/bash
      - -c
      - |
        pip install --no-cache-dir psycopg2-binary boto3
        mlflow server \
          --backend-store-uri "${MLFLOW_BACKEND_STORE_URI}" \
          --artifacts-destination "${MLFLOW_ARTIFACTS_DESTINATION}" \
          --serve-artifacts \
          --host "${MLFLOW_HOST}" \
          --port "${MLFLOW_PORT}"
    ports:
      - "${MLFLOW_PORT}:${MLFLOW_PORT}"
    healthcheck:
      test:
        [
          "CMD",
          "python",
          "-c",
          "import urllib.request; urllib.request.urlopen('http://localhost:${MLFLOW_PORT}/health')",
        ]
      interval: 10s
      timeout: 5s
      retries: 30

networks:
  default:
    name: mlflow-network