{% if tracking_mlflow -%}
version: "3.9"
services:
  mlflow:
    image: ghcr.io/mlflow/mlflow:latest
    command: >
      mlflow server --host 0.0.0.0 --port 5000
      --backend-store-uri /mlruns
      {% if "minio" in extras -%} --default-artifact-root s3://artifacts {%- endif %}
    ports: ["5000:5000"]
    environment:
      - MLFLOW_S3_ENDPOINT_URL=http://minio:9000
      - AWS_ACCESS_KEY_ID={{ '${MINIO_ROOT_USER}' }}
      - AWS_SECRET_ACCESS_KEY={{ '${MINIO_ROOT_PASSWORD}' }}
    volumes:
      - ./mlruns:/mlruns
    {% if "minio" in extras -%}
    depends_on: [minio, minio-setup]
    {%- endif %}

  {% if "minio" in extras -%}
  minio:
    image: quay.io/minio/minio:RELEASE.2024-10-02T17-50-41Z
    command: server /data --console-address ":9001"
    environment:
      - MINIO_ROOT_USER={{ '${MINIO_ROOT_USER}' }}
      - MINIO_ROOT_PASSWORD={{ '${MINIO_ROOT_PASSWORD}' }}
    ports: ["9000:9000","9001:9001"]
    volumes:
      - ./minio-data:/data

  minio-setup:
    image: quay.io/minio/mc:RELEASE.2024-10-02T15-30-21Z
    depends_on: [minio]
    entrypoint: ["/bin/sh","-c"]
    command: |
      mc alias set local http://minio:9000 $$MINIO_ROOT_USER $$MINIO_ROOT_PASSWORD && \
      mc mb -p local/artifacts || true
    environment:
      - MINIO_ROOT_USER={{ '${MINIO_ROOT_USER}' }}
      - MINIO_ROOT_PASSWORD={{ '${MINIO_ROOT_PASSWORD}' }}
  {%- endif %}

  app:
    build:
      context: .
      dockerfile: docker/{{ 'cuda' if docker=='cuda' else 'slim' }}.Dockerfile
    environment:
      - MLFLOW_TRACKING_URI=http://mlflow:5000
    {% if docker == 'cuda' -%}
    deploy:
      resources:
        reservations:
          devices:
            - capabilities: [gpu]
    {%- endif %}
    depends_on: [mlflow]
{%- else -%}
{{ SKIP_FILE }}
{%- endif %}
