docker-publish.yml

name: Docker Publish

# This workflow uses actions that are not certified by GitHub.
# They are provided by a third-party and are governed by
# separate terms of service, privacy policy, and support
# documentation.

on:
  release:
    types: [ "published" ]
  pull_request:
    branches: [ "main" ]

env:
  # Use docker.io for Docker Hub if empty
  REGISTRY: ghcr.io
  JAX_VERSION : 0.4.23
  JAX_CONDA_IMAGE_NAME : "yfukai/conda-jax"
  # github.repository as <account>/<repo>
  IMAGE_NAME: ${{ github.repository }}

jobs:
  build:

    runs-on: ubuntu-latest
    permissions:
      contents: read
      packages: write
      # This is used to complete the identity challenge
      # with sigstore/fulcio when running outside of PRs.
      id-token: write

    steps:
      - name: Checkout repository
        uses: actions/checkout@v3

      # Install the cosign tool except on PR
      # https://github.com/sigstore/cosign-installer
      - name: Install Cosign
        uses: sigstore/cosign-installer@v3.1.1
        with:
          cosign-release: 'v2.2.1'

      # Workaround: https://github.com/docker/build-push-action/issues/461
      - name: Setup Docker buildx
        uses: docker/setup-buildx-action@79abd3f86f79a9d68a23c75a09a9a85889262adf

      # Login against a Docker registry except on PR
      # https://github.com/docker/login-action
      - name: Log into registry ${{ env.REGISTRY }}
        uses: docker/login-action@28218f9b04b4f3f62068d7b6ce6ca5b26e35336c
        with:
          registry: ${{ env.REGISTRY }}
          username: ${{ github.actor }}
          password: ${{ secrets.GITHUB_TOKEN }}

      # Extract metadata (tags, labels) for Docker
      # https://github.com/docker/metadata-action
      - name: Extract Docker metadata (JAX conda CPU)
        id: meta_jax_conda
        uses: docker/metadata-action@v5
        with:
          images: ${{ env.REGISTRY }}/${{ env.JAX_CONDA_IMAGE_NAME }}

      - name: Extract first tag (JAX conda CPU)
        id: extract_first_tag_cpu
        run: |
          # Extracting the first tag
          tags="${{ steps.meta_jax_conda.outputs.tags }}"
          first_line=$(echo "$tags" | head -n 1)
          # Setting the output
          echo "meta_jax_conda_first_tag=$first_line" >> $GITHUB_OUTPUT

      - name: Extract Docker metadata (JAX conda GPU)
        id: meta_jax_conda_cuda
        uses: docker/metadata-action@v5
        with:
          images: ${{ env.REGISTRY }}/${{ env.JAX_CONDA_IMAGE_NAME }}
          flavor: |
            latest=auto
            suffix=-cuda,onlatest=true

      - name: Extract first tag (JAX conda GPU)
        id: extract_first_tag_gpu
        run: |
          # Extracting the first tag
          tags="${{ steps.meta_jax_conda_cuda.outputs.tags }}"
          first_line=$(echo "$tags" | head -n 1)
          # Setting the output
          echo "meta_jax_conda_cuda_first_tag=$first_line" >> $GITHUB_OUTPUT

      - name: Extract Docker metadata (CPU)
        id: meta
        uses: docker/metadata-action@v5
        with:
          images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}

      - name: Extract Docker metadata (GPU)
        id: meta_cuda
        uses: docker/metadata-action@v5
        with:
          images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
          flavor: |
            latest=auto
            suffix=-cuda,onlatest=true

      - name: Set up QEMU
        uses: docker/setup-qemu-action@v2

      # Build and push Docker image with Buildx
      # https://github.com/docker/build-push-action
      - name: Build and push Docker image (JAX conda CPU)
        id: build-and-push-jax-conda-cpu
        uses: docker/build-push-action@v5
        with:
          context: jax_image/
          platforms: linux/amd64,linux/arm64/v8
          push: true
          tags: ${{ steps.meta_jax_conda.outputs.tags }}
          labels: ${{ steps.meta_jax_conda.outputs.labels }}
          build-args: |
            JAX_VERSION=${{ env.JAX_VERSION }}

      - name: Build and push Docker image (JAX conda GPU)
        id: build-and-push-jax-conda-gpu
        uses: docker/build-push-action@v5
        with:
          context: jax_image/
          platforms: linux/amd64
          push: true
          tags: ${{ steps.meta_jax_conda_cuda.outputs.tags }}
          labels: ${{ steps.meta_jax_conda_cuda.outputs.labels }}
          build-args: |
            JAX_VERSION=${{ env.JAX_VERSION }}
            JAX_VERSION_EXTRA=cuda11_local
            BASE_IMAGE=nvidia/cuda:11.8.0-cudnn8-devel-ubuntu20.04

      - name: Build and push Docker image (CPU)
        id: build-and-push-cpu
        uses: docker/build-push-action@v5
        with:
          context: .
          platforms: linux/amd64,linux/arm64/v8
          push: true
          tags: ${{ steps.meta.outputs.tags }}
          labels: ${{ steps.meta.outputs.labels }}
          build-args: |
            BASE_IMAGE=${{ steps.extract_first_tag_cpu.outputs.meta_jax_conda_first_tag }}

      - name: Build and push Docker image (GPU)
        id: build-and-push-gpu
        uses: docker/build-push-action@v5
        with:
          context: .
          platforms: linux/amd64
          push: true
          tags: ${{ steps.meta_cuda.outputs.tags }}
          labels: ${{ steps.meta_cuda.outputs.labels }}
          build-args: |
            BASE_IMAGE=${{ steps.extract_first_tag_gpu.outputs.meta_jax_conda_cuda_first_tag }}