build_image_cuda.yml

name: Build Docker Images - CUDA

on:
  # Manual trigger with configurable options
  workflow_dispatch:
    inputs:
      task:
        description: 'Task to build'
        required: true
        type: choice
        options:
          - train
          - inference
          - all
        default: 'train'
      target:
        description: 'Build target stage (dev includes dev tools, release is production)'
        required: true
        type: choice
        options:
          - dev
          - release
        default: 'dev'
      push:
        description: 'Push image to registry'
        type: boolean
        default: true
      no_cache:
        description: 'Build without Docker cache'
        type: boolean
        default: false
      pkg_mgr:
        description: 'Package manager to use'
        required: true
        type: choice
        options:
          - conda
          - uv
        default: 'conda'
      build_all_tasks:
        description: 'Build all tasks (train, inference, all) - overrides task selection'
        type: boolean
        default: false
      tar_dir:
        description: 'Directory to store image tar files (default: from .github/configs/cuda.yml)'
        type: string
        default: ''
      runs_on:
        required: false
        type: string
        description: Override runs_on. Falls back to platform config if not set.
        default: '["self-hosted", "Linux", "X64", "nvidia-0", "gpus-8"]'
      container_volumes:
        required: false
        type: string
        description: Override container_volumes. Falls back to platform config if not set.
        default: >-
          ["/home/flagscale_cicd/docker/docker_build/docker_data:/home/gitlab-runner/data",
          "/home/flagscale_cicd/docker/docker_build/docker_tokenizers:/home/gitlab-runner/tokenizers"]

  # Trigger on PRs that modify docker-related files (build + test only, no push)
  pull_request:
    branches: [main]
    paths:
      - 'docker/cuda/**'
      - 'docker/build.sh'
      - 'tools/install/**'
      - 'requirements/**'
      - '.github/workflows/build_image_cuda.yml'

permissions:
  contents: read

concurrency:
  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
  cancel-in-progress: true

env:
  # Local registry used by CI runners (see .github/configs/cuda.yml)
  REGISTRY: localhost:5000
  # Default build versions (keep in sync with docker/build.sh)
  CUDA_VERSION: '12.8.1'
  UBUNTU_VERSION: '22.04'
  PYTHON_VERSION: '3.12'
  UV_VERSION: '0.7.2'
  PKG_MGR: ${{ inputs.pkg_mgr || 'conda' }}

jobs:
  # ---------------------------------------------------------------------------
  # Prepare: compute build matrix and parameters based on trigger type
  # ---------------------------------------------------------------------------
  prepare:
    name: Prepare build matrix
    runs-on: ubuntu-latest
    outputs:
      matrix: ${{ steps.set-matrix.outputs.matrix }}
      target: ${{ steps.params.outputs.target }}
      push: ${{ steps.params.outputs.push }}
      no_cache: ${{ steps.params.outputs.no_cache }}
      tar_dir: ${{ steps.config.outputs.tar_dir }}
    steps:
      - name: Checkout code
        uses: actions/checkout@v4
        with:
          sparse-checkout: .github/configs

      - name: Load tar_dir from platform config
        id: config
        run: |
          TAR_DIR=$(grep '^tar_dir:' .github/configs/cuda.yml | awk '{print $2}')
          # Allow workflow_dispatch input to override config
          OVERRIDE="${{ inputs.tar_dir }}"
          if [ -n "$OVERRIDE" ]; then
            TAR_DIR="$OVERRIDE"
          fi
          echo "tar_dir=${TAR_DIR}" >> $GITHUB_OUTPUT

      - name: Determine build matrix
        id: set-matrix
        run: |
          EVENT="${{ github.event_name }}"

          if [ "$EVENT" = "workflow_dispatch" ] && [ "${{ inputs.build_all_tasks }}" != "true" ]; then
            # Manual trigger: build selected task only
            echo 'matrix={"task":["${{ inputs.task }}"]}' >> $GITHUB_OUTPUT
          else
            # PR or build_all_tasks=true: build all tasks
            echo 'matrix={"task":["train","inference","all"]}' >> $GITHUB_OUTPUT
          fi

      - name: Set build parameters
        id: params
        run: |
          EVENT="${{ github.event_name }}"

          if [ "$EVENT" = "pull_request" ]; then
            # PR: always build dev images, push to local registry, use cache
            echo "target=dev" >> $GITHUB_OUTPUT
            echo "push=true" >> $GITHUB_OUTPUT
            echo "no_cache=false" >> $GITHUB_OUTPUT
          else
            # workflow_dispatch: use user-provided inputs
            echo "target=${{ inputs.target || 'dev' }}" >> $GITHUB_OUTPUT
            echo "push=${{ inputs.push }}" >> $GITHUB_OUTPUT
            echo "no_cache=${{ inputs.no_cache }}" >> $GITHUB_OUTPUT
          fi

  # ---------------------------------------------------------------------------
  # Build: build and push Docker images (matrix across tasks)
  # ---------------------------------------------------------------------------
  build:
    name: Build ${{ matrix.task }}
    needs: prepare
    runs-on: [self-hosted, Linux, X64, nvidia-0, gpus-8]
    env:
      HTTP_PROXY: ${{ vars.HTTP_PROXY }}
      HTTPS_PROXY: ${{ vars.HTTPS_PROXY }}
      NO_PROXY: ${{ vars.NO_PROXY }}
      # Note: vars.* may be empty for fork PRs; proxy is detected from runner env in the 'proxy' step
    strategy:
      fail-fast: false
      matrix: ${{ fromJson(needs.prepare.outputs.matrix) }}
    outputs:
      train_tag: ${{ steps.export.outputs.train_tag }}
      inference_tag: ${{ steps.export.outputs.inference_tag }}
      all_tag: ${{ steps.export.outputs.all_tag }}
      train_tar: ${{ steps.export.outputs.train_tar }}
      inference_tar: ${{ steps.export.outputs.inference_tar }}
      all_tar: ${{ steps.export.outputs.all_tar }}
    steps:
      - name: Clean workspace
        run: sudo rm -rf "$GITHUB_WORKSPACE"/* "$GITHUB_WORKSPACE"/.[!.]* 2>/dev/null || true

      - name: Checkout code
        uses: actions/checkout@v4

      - name: Set up Docker Buildx
        uses: docker/setup-buildx-action@v3
        with:
          # Use docker driver to avoid pulling moby/buildkit from Docker Hub
          driver: docker

      - name: Compute build metadata
        id: meta
        run: |
          set -euo pipefail

          TASK="${{ matrix.task }}"
          TARGET="${{ needs.prepare.outputs.target }}"
          CUDA_VERSION="${{ env.CUDA_VERSION }}"
          PYTHON_VERSION="${{ env.PYTHON_VERSION }}"
          UBUNTU_VERSION="${{ env.UBUNTU_VERSION }}"

          CUDA_MAJOR=$(echo "$CUDA_VERSION" | cut -d. -f1)
          CUDA_MINOR=$(echo "$CUDA_VERSION" | cut -d. -f2)
          TIMESTAMP=$(date +%Y%m%d%H%M%S)

          # Image naming follows docker/build.sh convention:
          #   flagscale-<task>:<target>-cu<major><minor>-py<version>[-<timestamp>]
          IMAGE_NAME="flagscale-${TASK}"
          TAG="${TARGET}-cu${CUDA_MAJOR}${CUDA_MINOR}-py${PYTHON_VERSION}-${TIMESTAMP}"

          BUILD_TAG="${{ env.REGISTRY }}/${IMAGE_NAME}:${TAG}"

          # Derived build arguments
          BASE_IMAGE="nvidia/cuda:${CUDA_VERSION}-devel-ubuntu${UBUNTU_VERSION}"
          PYTORCH_INDEX="https://download.pytorch.org/whl/cu${CUDA_MAJOR}${CUDA_MINOR}"

          echo "image_name=${IMAGE_NAME}" >> $GITHUB_OUTPUT
          echo "build_tag=${BUILD_TAG}" >> $GITHUB_OUTPUT
          echo "base_image=${BASE_IMAGE}" >> $GITHUB_OUTPUT
          echo "pytorch_index=${PYTORCH_INDEX}" >> $GITHUB_OUTPUT

          # Job summary
          {
            echo "### Build: ${IMAGE_NAME}"
            echo ""
            echo "| Parameter | Value |"
            echo "|---|---|"
            echo "| Task | \`${TASK}\` |"
            echo "| Target | \`${TARGET}\` |"
            echo "| CUDA | \`${CUDA_VERSION}\` |"
            echo "| Python | \`${PYTHON_VERSION}\` |"
            echo "| Dockerfile | \`docker/cuda/Dockerfile.${TASK}\` |"
            echo "| Build Tag | \`${BUILD_TAG}\` |"
          } >> $GITHUB_STEP_SUMMARY

      - name: Detect proxy from runner environment
        id: proxy
        run: |
          # Read proxy from runner environment (works for both fork PRs and direct pushes)
          HTTP_PROXY_VAL="${http_proxy:-${HTTP_PROXY:-}}"
          HTTPS_PROXY_VAL="${https_proxy:-${HTTPS_PROXY:-}}"
          NO_PROXY_VAL="${no_proxy:-${NO_PROXY:-}}"
          echo "http_proxy=${HTTP_PROXY_VAL}" >> $GITHUB_OUTPUT
          echo "https_proxy=${HTTPS_PROXY_VAL}" >> $GITHUB_OUTPUT
          echo "no_proxy=${NO_PROXY_VAL}" >> $GITHUB_OUTPUT
          echo "Detected proxies: HTTP=${HTTP_PROXY_VAL} HTTPS=${HTTPS_PROXY_VAL} NO_PROXY=${NO_PROXY_VAL}"

      - name: Build Docker image
        uses: docker/build-push-action@v6
        with:
          context: .
          file: docker/cuda/Dockerfile.${{ matrix.task }}
          target: ${{ needs.prepare.outputs.target }}
          load: true
          tags: ${{ steps.meta.outputs.build_tag }}
          build-args: |
            BASE_IMAGE=${{ steps.meta.outputs.base_image }}
            CUDA_VERSION=${{ env.CUDA_VERSION }}
            UBUNTU_VERSION=${{ env.UBUNTU_VERSION }}
            PYTHON_VERSION=${{ env.PYTHON_VERSION }}
            UV_VERSION=${{ env.UV_VERSION }}
            PYTORCH_INDEX=${{ steps.meta.outputs.pytorch_index }}
            PKG_MGR=${{ env.PKG_MGR }}
            HTTP_PROXY=${{ steps.proxy.outputs.http_proxy }}
            HTTPS_PROXY=${{ steps.proxy.outputs.https_proxy }}
            NO_PROXY=${{ steps.proxy.outputs.no_proxy }}
            http_proxy=${{ steps.proxy.outputs.http_proxy }}
            https_proxy=${{ steps.proxy.outputs.https_proxy }}
            no_proxy=${{ steps.proxy.outputs.no_proxy }}
          no-cache: ${{ needs.prepare.outputs.no_cache == 'true' }}

      - name: Save image as tar
        run: |
          TAR_DIR="${{ needs.prepare.outputs.tar_dir }}"
          sudo mkdir -p "$TAR_DIR"
          IMAGE_TAG="${{ steps.meta.outputs.build_tag }}"
          PR_NUMBER="${{ github.event.pull_request.number }}"
          # Use PR number for deterministic lookup after merge; fall back to short SHA for manual dispatch
          if [ -n "$PR_NUMBER" ]; then
            SUFFIX="pr${PR_NUMBER}"
          else
            SHA="${{ github.sha }}"
            SUFFIX="${SHA:0:7}"
          fi
          TAR_NAME=$(echo "$IMAGE_TAG" | tr '/: ' '---')-${SUFFIX}.tar
          sudo docker save "$IMAGE_TAG" -o "${TAR_DIR}/${TAR_NAME}"
          echo "tar_path=${TAR_DIR}/${TAR_NAME}" >> $GITHUB_OUTPUT
        id: save_tar

      - name: Export image tag for config update
        id: export
        run: |
          TASK="${{ matrix.task }}"
          echo "${TASK}_tag=${{ steps.meta.outputs.build_tag }}" >> $GITHUB_OUTPUT
          echo "${TASK}_tar=${{ steps.save_tar.outputs.tar_path }}" >> $GITHUB_OUTPUT

      - name: Print build result
        if: success()
        run: |
          echo "" >> $GITHUB_STEP_SUMMARY
          echo "**Result:** Built successfully" >> $GITHUB_STEP_SUMMARY
          echo "**Pushed:** ${{ needs.prepare.outputs.push }}" >> $GITHUB_STEP_SUMMARY

  # ---------------------------------------------------------------------------
  # Summary: verify all builds completed
  # ---------------------------------------------------------------------------
  summary:
    name: Build summary
    needs: build
    runs-on: ubuntu-latest
    if: always()
    steps:
      - name: Verify build results
        run: |
          if [ "${{ needs.build.result }}" != "success" ]; then
            echo "::error::One or more image builds failed"
            exit 1
          fi
          echo "All Docker images built successfully!"

  # ---------------------------------------------------------------------------
  # Load images from tar and push to local registry before running tests
  # ---------------------------------------------------------------------------
  load_images:
    name: Load and push images
    needs: ['build', 'summary']
    runs-on: [self-hosted, Linux, X64, nvidia-0, gpus-8]
    steps:
      - name: Load train image from tar and push
        run: |
          TAR="${{ needs.build.outputs.train_tar }}"
          TAG="${{ needs.build.outputs.train_tag }}"
          if [ -f "$TAR" ]; then
            echo "Loading $TAR"
            sudo docker load -i "$TAR"
            sudo docker push "$TAG"
          else
            echo "::warning::Train image tar not found: $TAR, skipping load"
          fi

      - name: Load inference image from tar and push
        run: |
          TAR="${{ needs.build.outputs.inference_tar }}"
          TAG="${{ needs.build.outputs.inference_tag }}"
          if [ -f "$TAR" ]; then
            echo "Loading $TAR"
            sudo docker load -i "$TAR"
            sudo docker push "$TAG"
          else
            echo "::warning::Inference image tar not found: $TAR, skipping load"
          fi

  # ---------------------------------------------------------------------------
  # Run CUDA tests after build succeeds
  # ---------------------------------------------------------------------------
  run_cuda_tests:
    name: Run CUDA tests
    needs: ['prepare', 'build', 'load_images']
    uses: ./.github/workflows/all_tests_common.yml
    with:
      platform: cuda
      ci_train_image: ${{ needs.build.outputs.train_tag }}
      ci_inference_image: ${{ needs.build.outputs.inference_tag }}
      runs_on: >-
        ${{ inputs.runs_on ||
        '["self-hosted", "Linux", "X64", "nvidia-0", "gpus-8"]' }}
      container_volumes: >-
        ${{ inputs.container_volumes ||
        '["/home/flagscale_cicd/docker/docker_build/docker_data:/home/gitlab-runner/data",
        "/home/flagscale_cicd/docker/docker_build/docker_tokenizers:/home/gitlab-runner/tokenizers"]' }}