# docker-compose for biopesticide-AI.
#
# Usage:
#   docker compose up --build         # launches Ollama + the Gradio UI on http://localhost:7860
#   docker compose run --rm bioai python -m bioai train --epochs 10
#   docker compose run --rm bioai python -m bioai design --user-text "Brown planthopper in rice paddy"
#
# The Ollama service is optional -- if you already run Ollama on the host,
# comment out the `ollama` service below and set OLLAMA_HOST=http://host.docker.internal:11434
# in the bioai environment block.
#
# To run on an AMD GPU host, also pass:
#   devices:
#     - /dev/kfd:/dev/kfd
#     - /dev/dri:/dev/dri
# and `group_add: [video]` and `ipc: host` (see Dockerfile comments).

services:
  ollama:
    image: ollama/ollama:latest
    container_name: bioai-ollama
    ports:
      - "11434:11434"
    volumes:
      # Persist the pulled models (Llama 3.2 3B is ~2 GB) across container rebuilds.
      - ./.ollama_models:/root/.ollama
    # Pull the model on first run, then keep the server alive.
    # The entrypoint script waits for the server, pulls the model, then tails logs.
    entrypoint: ["/bin/sh", "-c"]
    command:
      - |
        ollama serve &
        SERVER_PID=$$!
        echo "[ollama] waiting for server to be ready..."
        for i in $$(seq 1 30); do
          if ollama list >/dev/null 2>&1; then break; fi
          sleep 1
        done
        if ! ollama list | grep -q llama3.2:3b; then
          echo "[ollama] pulling llama3.2:3b (first run only, ~2 GB)..."
          ollama pull llama3.2:3b
        fi
        echo "[ollama] ready. models:"
        ollama list
        wait $$SERVER_PID
    restart: unless-stopped
    # For AMD GPU passthrough on the Ollama container (uncomment on a ROCm host):
    # devices:
    #   - /dev/kfd:/dev/kfd
    #   - /dev/dri:/dev/dri
    # group_add:
    #   - video

  bioai:
    build: .
    image: biopesticide-ai:latest
    container_name: bioai
    working_dir: /app
    depends_on:
      - ollama
    environment:
      # Point the bioai container at the Ollama service in the same compose network.
      - OLLAMA_HOST=http://ollama:11434
      - OLLAMA_MODEL=llama3.2:3b
      - PYTHONUNBUFFERED=1
    ports:
      - "7860:7860"
    volumes:
      - ./data:/app/data
      - ./checkpoints:/app/checkpoints
      # Persist the Ollama response cache across runs so demo calls don't
      # re-spend compute on identical prompts.
      - ./.ollama_cache:/app/.ollama_cache
    # Default command matches the Dockerfile CMD (launches the FastAPI web UI).
    # Override via `docker compose run bioai <command>` for training / design.
    command: ["python", "-m", "bioai", "web", "--host", "0.0.0.0", "--port", "7860"]
    # For AMD GPU passthrough (uncomment on a ROCm host):
    # devices:
    #   - /dev/kfd:/dev/kfd
    #   - /dev/dri:/dev/dri
    # group_add:
    #   - video
    # ipc: host
    # mem_limit: 64g