name: docker

on:
  workflow_dispatch:
  push:
    branches:
      - "main"
    paths:
      - "**/*.py"
      - "pyproject.toml"
      - "docker/**"
      - ".github/workflows/*.yml"
  pull_request:
    branches:
      - "main"
    paths:
      - "**/*.py"
      - "pyproject.toml"
      - "docker/**"
      - ".github/workflows/*.yml"
  release:
    types:
      - published

jobs:
  build:
    strategy:
      fail-fast: false
      matrix:
        include:
          - device: "cuda"
          - device: "npu-a2"
          - device: "npu-a3"

    runs-on: ubuntu-latest

    concurrency:
      group: ${{ github.workflow }}-${{ github.ref }}-${{ matrix.device }}
      cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}

    environment:
      name: docker
      url: https://hub.docker.com/r/hiyouga/llamafactory

    steps:
      - name: Free up disk space
        uses: jlumbroso/free-disk-space@v1.3.1
        with:
          tool-cache: true
          docker-images: false

      - name: Checkout
        uses: actions/checkout@v4

      - name: Get llamafactory version
        id: version
        run: |
          if [ "${{ github.event_name }}" = "release" ]; then
            echo "tag=$(grep -oP 'VERSION = "\K[^"]+' src/llamafactory/extras/env.py)" >> "$GITHUB_OUTPUT"
          else
            echo "tag=latest" >> "$GITHUB_OUTPUT"
          fi

      - name: Set up Docker Buildx
        uses: docker/setup-buildx-action@v3

      - name: Login to Docker Hub
        if: ${{ github.event_name != 'pull_request' }}
        uses: docker/login-action@v3
        with:
          username: ${{ vars.DOCKERHUB_USERNAME }}
          password: ${{ secrets.DOCKERHUB_TOKEN }}

      - name: Login to Quay
        if: ${{ github.event_name != 'pull_request' && startsWith(matrix.device, 'npu') }}
        uses: docker/login-action@v3
        with:
          registry: quay.io
          username: ${{ vars.QUAY_ASCEND_USERNAME }}
          password: ${{ secrets.QUAY_ASCEND_TOKEN }}

      - name: Build and push Docker image (CUDA)
        if: ${{ matrix.device == 'cuda' }}
        uses: docker/build-push-action@v6
        with:
          context: .
          file: ./docker/docker-cuda/Dockerfile
          push: ${{ github.event_name != 'pull_request' }}
          tags: |
            docker.io/hiyouga/llamafactory:${{ steps.version.outputs.tag }}

      - name: Build and push Docker image (NPU-A2)
        if: ${{ matrix.device == 'npu-a2' }}
        uses: docker/build-push-action@v6
        with:
          context: .
          platforms: linux/amd64,linux/arm64
          file: ./docker/docker-npu/Dockerfile
          push: ${{ github.event_name != 'pull_request' }}
          tags: |
            docker.io/hiyouga/llamafactory:${{ steps.version.outputs.tag }}-npu-a2
            quay.io/ascend/llamafactory:${{ steps.version.outputs.tag }}-npu-a2

      - name: Build and push Docker image (NPU-A3)
        if: ${{ matrix.device == 'npu-a3' }}
        uses: docker/build-push-action@v6
        with:
          context: .
          platforms: linux/amd64,linux/arm64
          file: ./docker/docker-npu/Dockerfile
          build-args: |
            BASE_IMAGE=quay.io/ascend/cann:8.3.rc2-a3-ubuntu22.04-py3.11
          push: ${{ github.event_name != 'pull_request' }}
          tags: |
            docker.io/hiyouga/llamafactory:${{ steps.version.outputs.tag }}-npu-a3
            quay.io/ascend/llamafactory:${{ steps.version.outputs.tag }}-npu-a3
