name: Build and push docker image to internal registry

on:
  workflow_dispatch:
  push:
    branches:
      - 'main'
    tags:
      - 'v*'
  pull_request:
    paths:
      - ".github/workflows/build.yaml"
      - "server/**"
      - "proto/**"
      - "router/**"
      - "launcher/**"
      - "Cargo.lock"
      - "rust-toolchain.toml"
      - "Dockerfile"
    branches:
      - 'main'

concurrency:
  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
  cancel-in-progress: true

jobs:
  build-and-push-image:
    runs-on: large-runner
    permissions:
      contents: write
      packages: write
      # This is used to complete the identity challenge
      # with sigstore/fulcio when running outside of PRs.
      id-token: write
      security-events: write
    steps:
      - name: Checkout repository
        uses: actions/checkout@v3
      - name: Initialize Docker Buildx
        uses: docker/setup-buildx-action@v2.0.0
        with:
          install: true
      - name: Inject slug/short variables
        uses: rlespinasse/github-slug-action@v4.4.1
      - name: Install cosign
        if: github.event_name != 'pull_request'
        uses: sigstore/cosign-installer@f3c664df7af409cb4873aa5068053ba9d61a57b6 #v2.6.0
        with:
          cosign-release: 'v1.13.1'
      - name: Tailscale
        uses: tailscale/github-action@v1
        with:
          authkey: ${{ secrets.TAILSCALE_AUTHKEY }}
      - name: Login to GitHub Container Registry
        if: github.event_name != 'pull_request'
        uses: docker/login-action@v2
        with:
          registry: ghcr.io
          username: ${{ github.actor }}
          password: ${{ secrets.GITHUB_TOKEN }}
      - name: Login to internal Container Registry
        uses: docker/login-action@v2.1.0
        with:
          username: ${{ secrets.TAILSCALE_DOCKER_USERNAME }}
          password: ${{ secrets.TAILSCALE_DOCKER_PASSWORD }}
          registry: registry.internal.huggingface.tech
      - name: Login to Azure Container Registry
        if: github.event_name != 'pull_request'
        uses: docker/login-action@v2.1.0
        with:
          username: ${{ secrets.AZURE_DOCKER_USERNAME }}
          password: ${{ secrets.AZURE_DOCKER_PASSWORD }}
          registry: db4c2190dd824d1f950f5d1555fbadf0.azurecr.io
      - name: Extract metadata (tags, labels) for Docker
        id: meta
        uses: docker/metadata-action@v4.3.0
        with:
          flavor: |
            latest=auto
          images: |
            registry.internal.huggingface.tech/api-inference/community/text-generation-inference
            ghcr.io/huggingface/text-generation-inference
            db4c2190dd824d1f950f5d1555fbadf0.azurecr.io/text-generation-inference
          tags: |
            type=semver,pattern={{version}}
            type=semver,pattern={{major}}.{{minor}}
            type=raw,value=latest,enable=${{ github.ref == format('refs/heads/{0}', github.event.repository.default_branch) }}
            type=raw,value=sha-${{ env.GITHUB_SHA_SHORT }}
      - name: Build and push Docker image
        id: build-and-push
        uses: docker/build-push-action@v4
        with:
          context: .
          file: Dockerfile
          push: ${{ github.event_name != 'pull_request' }}
          platforms: 'linux/amd64'
          build-args: |
            GIT_SHA=${{ env.GITHUB_SHA }}
            DOCKER_LABEL=sha-${{ env.GITHUB_SHA_SHORT }}
          tags: ${{ steps.meta.outputs.tags }}
          labels: ${{ steps.meta.outputs.labels }}
          cache-from: type=registry,ref=registry.internal.huggingface.tech/api-inference/community/text-generation-inference:cache,mode=max
          cache-to: type=registry,ref=registry.internal.huggingface.tech/api-inference/community/text-generation-inference:cache,mode=max
      # Sign the resulting Docker image digest except on PRs.
      # This will only write to the public Rekor transparency log when the Docker
      # repository is public to avoid leaking data.
      - name: Sign the published Docker image
        if: ${{ github.event_name != 'pull_request' }}
        env:
          COSIGN_EXPERIMENTAL: "true"
        # This step uses the identity token to provision an ephemeral certificate
        # against the sigstore community Fulcio instance.
        run: echo "${{ steps.meta.outputs.tags }}" | xargs -I {} cosign sign {}@${{ steps.build-and-push.outputs.digest }}
      - name: Run Trivy in GitHub SBOM mode and submit results to Dependency Graph
        uses: aquasecurity/trivy-action@master
        if: ${{ github.event_name != 'pull_request' }}
        with:
          image-ref: 'ghcr.io/huggingface/text-generation-inference:sha-${{ env.GITHUB_SHA_SHORT }}'
          format: 'github'
          output: 'dependency-results.sbom.json'
          github-pat: ${{ secrets.GITHUB_TOKEN }}
          scanners: 'vuln'
      - name: Run Trivy vulnerability scanner
        uses: aquasecurity/trivy-action@master
        if: ${{ github.event_name != 'pull_request' }}
        with:
          image-ref: 'ghcr.io/huggingface/text-generation-inference:sha-${{ env.GITHUB_SHA_SHORT }}'
          format: 'sarif'
          output: 'trivy-results.sarif'
          severity: 'CRITICAL'
          scanners: 'vuln'
      - name: Upload Trivy scan results to GitHub Security tab
        uses: github/codeql-action/upload-sarif@v2
        if: ${{ github.event_name != 'pull_request' }}
        with:
          sarif_file: 'trivy-results.sarif'

  build-and-push-sagemaker-image:
    needs:
      - build-and-push-image
    runs-on: ubuntu-latest
    steps:
      - name: Checkout repository
        uses: actions/checkout@v3
      - name: Initialize Docker Buildx
        uses: docker/setup-buildx-action@v2.0.0
        with:
          install: true
      - name: Inject slug/short variables
        uses: rlespinasse/github-slug-action@v4.4.1
      - name: Tailscale
        uses: tailscale/github-action@v1
        with:
          authkey: ${{ secrets.TAILSCALE_AUTHKEY }}
      - name: Login to internal Container Registry
        uses: docker/login-action@v2.1.0
        with:
          username: ${{ secrets.TAILSCALE_DOCKER_USERNAME }}
          password: ${{ secrets.TAILSCALE_DOCKER_PASSWORD }}
          registry: registry.internal.huggingface.tech
      - name: Extract metadata (tags, labels) for Docker
        id: meta
        uses: docker/metadata-action@v4.3.0
        with:
          flavor: |
            latest=auto
          images: |
            registry.internal.huggingface.tech/api-inference/community/text-generation-inference/sagemaker
          tags: |
            type=semver,pattern={{version}}
            type=semver,pattern={{major}}.{{minor}}
            type=raw,value=latest,enable=${{ github.ref == format('refs/heads/{0}', github.event.repository.default_branch) }}
            type=raw,value=sha-${{ env.GITHUB_SHA_SHORT }}
      - name: Build and push Docker image
        uses: docker/build-push-action@v2
        with:
          context: .
          file: Dockerfile
          push: ${{ github.event_name != 'pull_request' }}
          platforms: 'linux/amd64'
          build-args: |
            GIT_SHA={{ env.GITHUB_SHA }}
          target: sagemaker
          tags: ${{ steps.meta.outputs.tags }}
          labels: ${{ steps.meta.outputs.labels }}
          cache-from: type=registry,ref=registry.internal.huggingface.tech/api-inference/community/text-generation-inference:cache,mode=max