ci.yml

name: CI

on:
  push:
    branches:
      - main
      - "pull-request/[0-9]+"
    tags:
      - '**'
  # `merge_group` fires when a PR enters GitHub's merge queue. GitHub creates
  # a temporary `gh-readonly-queue/main/pr-<n>-<sha>` ref containing the PR
  # rebased on top of `main` + any earlier queued PRs and runs the required
  # checks against it. The job names below must match the required-check
  # names configured in branch protection.
  merge_group:
    branches: [main]

jobs:
  cpu:
    runs-on: linux-amd64-cpu8
    container:
      image: nvidia/cuda:13.2.1-cudnn-devel-ubuntu24.04
    env:
      UV_PROJECT_ENVIRONMENT: /tmp/flashdreams-venv
      UV_LINK_MODE: copy
      # Exercise the dependency branch used by current Python releases.
      UV_PYTHON: "3.12"
    steps:
      - name: Get PR info
        if: startsWith(github.ref_name, 'pull-request/')
        id: get-pr-info
        uses: nv-gha-runners/get-pr-info@main

      - name: Install lint system dependencies
        run: |
          apt-get update -qq
          DEBIAN_FRONTEND=noninteractive apt-get install -y -qq --no-install-recommends \
            python3 python3-dev python3-venv \
            git curl ca-certificates

      - name: Checkout
        uses: actions/checkout@v4

      - name: Setup proxy cache
        uses: nv-gha-runners/setup-proxy-cache@main

      - name: Setup uv
        uses: astral-sh/setup-uv@v6
        with:
          enable-cache: true
          cache-suffix: "cpu"

      - name: Install lint dependencies
        run: |
          # ty needs third-party packages to resolve imports. Skip packages
          # that cannot be installed on CPU-only runners:
          # - transformer-engine-torch: source-only, requires GPU arch to compile
          uv venv --clear
          uv sync --locked --extra dev --group lint \
            --no-install-package transformer-engine-torch

      - name: Run linter checks
        run: |
          # Container runs as root but the workspace is owned by the runner
          # uid -- git refuses to operate without safe.directory.
          git config --global --add safe.directory "$(pwd)"
          uv run --group lint pre-commit run -a

      - name: Install remaining system dependencies
        run: |
          DEBIAN_FRONTEND=noninteractive apt-get install -y -qq --no-install-recommends \
            ffmpeg \
            gcc g++ ninja-build \
            libnccl-dev \
            unzip
          rm -rf /var/lib/apt/lists/*

      - name: Build flashdreams wheel
        run: uv build --wheel --package flashdreams

      - name: Run CPU unit tests
        run: uv run --no-sync pytest -m ci_cpu -v

      # Documentation is built and deployed by `.github/workflows/doc.yml`.

  gpu:
    runs-on: linux-amd64-gpu-rtxpro6000-latest-2
    timeout-minutes: 120
    container:
      image: nvidia/cuda:13.2.1-cudnn-devel-ubuntu24.04
      options: --gpus all
    env:
      # Place the venv outside the workspace so uv does not interfere with
      # the checkout directory owned by the runner.
      UV_PROJECT_ENVIRONMENT: /tmp/flashdreams-venv
      UV_LINK_MODE: copy
      # imgui-bundle publishes Linux ARM64 wheels for Python 3.11+, but not
      # 3.10. Keep GB300 CI on a wheel-supported interpreter and fail instead
      # of attempting its large, platform-sensitive native source build.
      UV_PYTHON: "3.11"
      UV_NO_BUILD_PACKAGE: imgui-bundle
      # Limit parallel CUDA compilation jobs to avoid OOM (each nvcc job
      # uses ~9GB peak memory).
      MAX_JOBS: 8
      OMNIDREAMS_QUALITY_REFERENCE_REPO: ${{ vars.OMNIDREAMS_QUALITY_REFERENCE_REPO || 'nvidia/omni-dreams-samples' }}
      OMNIDREAMS_QUALITY_REFERENCE_PATH: ${{ vars.OMNIDREAMS_QUALITY_REFERENCE_PATH || 'data/quality_references/single_view/omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae/239560dc-33d1-11ef-9720-00044bcbccac/reference_compare_region.mp4' }}
      OMNIDREAMS_QUALITY_REFERENCE_CLIP: /tmp/omnidreams_quality_reference/reference_compare_region.mp4
      OMNIDREAMS_QUALITY_ARTIFACT_DIR: /tmp/omnidreams_quality_artifacts
    steps:
      - name: Get PR info
        if: startsWith(github.ref_name, 'pull-request/')
        id: get-pr-info
        uses: nv-gha-runners/get-pr-info@main

      - name: Detect GPU architecture
        id: gpu-arch
        run: |
          nvidia-smi
          # Get compute capability (e.g. "12.0") and strip the dot for nvcc
          # arch format (e.g. "120").
          compute_cap=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null | head -1 | tr -d '[:space:]')
          arch=$(echo "${compute_cap}" | tr -d '.')
          echo "arch=${arch}" >> "$GITHUB_OUTPUT"
          echo "Detected GPU compute capability: ${compute_cap} -> sm_${arch}"

      - name: Checkout
        uses: actions/checkout@v4

      - name: Install system dependencies
        run: |
          apt-get update -qq
          DEBIAN_FRONTEND=noninteractive apt-get install -y -qq --no-install-recommends \
            python3 python3-dev python3-venv \
            ffmpeg \
            gcc g++ ninja-build \
            libgl1-mesa-dev \
            libx11-dev libxcursor-dev libxext-dev libxi-dev libxinerama-dev libxrandr-dev \
            libnccl-dev \
            curl git ca-certificates unzip
          rm -rf /var/lib/apt/lists/*

      - name: Setup proxy cache
        uses: nv-gha-runners/setup-proxy-cache@main

      - name: Setup uv
        uses: astral-sh/setup-uv@v6
        with:
          enable-cache: true
          cache-suffix: "gpu-sm${{ steps.gpu-arch.outputs.arch }}"
          # Disable default prune --ci which removes pre-built wheels but also
          # invalidates git-sourced builds (timestamps change on cache restore).
          # Instead we selectively trim the cache in a post step below.
          prune-cache: false

      - name: Install dependencies
        env:
          NVTE_CUDA_ARCHS: ${{ steps.gpu-arch.outputs.arch }}
        run: |
          uv venv --clear
          uv sync --locked --extra dev

      - name: Verify GPU availability
        run: |
          echo "Verifying GPU access..."
          nvidia-smi

      - name: Download Omnidreams quality reference
        env:
          HF_TOKEN: ${{ secrets.HF_TOKEN }}
        run: |
          mkdir -p "$(dirname "${OMNIDREAMS_QUALITY_REFERENCE_CLIP}")"
          if [ -n "${HF_TOKEN:-}" ]; then
            uv run --no-sync hf download \
              "${OMNIDREAMS_QUALITY_REFERENCE_REPO}" \
              "${OMNIDREAMS_QUALITY_REFERENCE_PATH}" \
              --repo-type dataset \
              --local-dir "$(dirname "${OMNIDREAMS_QUALITY_REFERENCE_CLIP}")" \
              --token "${HF_TOKEN}"
          else
            uv run --no-sync hf download \
              "${OMNIDREAMS_QUALITY_REFERENCE_REPO}" \
              "${OMNIDREAMS_QUALITY_REFERENCE_PATH}" \
              --repo-type dataset \
              --local-dir "$(dirname "${OMNIDREAMS_QUALITY_REFERENCE_CLIP}")"
          fi
          cp \
            "$(dirname "${OMNIDREAMS_QUALITY_REFERENCE_CLIP}")/${OMNIDREAMS_QUALITY_REFERENCE_PATH}" \
            "${OMNIDREAMS_QUALITY_REFERENCE_CLIP}"
          test -s "${OMNIDREAMS_QUALITY_REFERENCE_CLIP}"

      - name: Run GPU unit tests
        env:
          HF_TOKEN: ${{ secrets.HF_TOKEN }}
          FLASHDREAMS_OMNIDREAMS_QUALITY_REFERENCE_CLIP: ${{ env.OMNIDREAMS_QUALITY_REFERENCE_CLIP }}
          FLASHDREAMS_OMNIDREAMS_QUALITY_EXAMPLE_DATA: "1"
          FLASHDREAMS_OMNIDREAMS_QUALITY_TOTAL_BLOCKS: "4"
          FLASHDREAMS_OMNIDREAMS_QUALITY_ARTIFACT_DIR: ${{ env.OMNIDREAMS_QUALITY_ARTIFACT_DIR }}
        run: uv run --no-sync pytest -m ci_gpu -v

      - name: Upload Omnidreams quality artifacts
        if: always()
        uses: actions/upload-artifact@v4
        with:
          name: omnidreams-quality-artifacts
          path: ${{ env.OMNIDREAMS_QUALITY_ARTIFACT_DIR }}
          if-no-files-found: ignore

      - name: Trim uv cache for upload
        if: always()
        run: |
          cache_dir="${UV_CACHE_DIR:-/github/home/.cache/uv}"
          echo "=== Cache size before trim ==="
          du -sh "${cache_dir}" || true
          du -sh "${cache_dir}"/*/ 2>/dev/null || true

          # Remove pre-built (downloaded) wheels -- these are re-downloaded
          # quickly from the proxy cache on each run.
          rm -rf "${cache_dir}/wheels-v6"

          # Remove unzipped wheel archives -- uv can re-extract from .whl
          # files in sdists-v9/ on next run.
          rm -rf "${cache_dir}/archive-v0"

          # Remove build artifacts from git checkouts (build/, *.egg-info,
          # __pycache__) that bloat the cache but aren't needed for reuse.
          find "${cache_dir}/git-v0/checkouts" \
            \( -name "build" -o -name "*.egg-info" -o -name "__pycache__" \) \
            -type d -exec rm -rf {} + 2>/dev/null || true

          # Remove editable installs -- these are in-repo workspace members
          # that get rebuilt instantly from source on every sync.
          rm -rf "${cache_dir}/sdists-v9/editable"

          echo ""
          echo "=== Cache size after trim ==="
          du -sh "${cache_dir}" || true
          du -sh "${cache_dir}"/*/ 2>/dev/null || true
          echo ""
          echo "=== Cached built wheels (sdists-v9) ==="
          find "${cache_dir}/sdists-v9" -name "*.whl" -exec ls -lh {} \; 2>/dev/null || true

  # ===========================================================================
  # Publish to PyPI -- runs only on main after cpu + gpu jobs pass.
  # ===========================================================================
  publish-pypi:
    if: github.ref == 'refs/heads/main' && github.event_name == 'push'
    needs: [cpu, gpu]
    runs-on: linux-amd64-cpu8
    steps:
      - name: Checkout
        uses: actions/checkout@v4

      - name: Setup uv
        uses: astral-sh/setup-uv@v6
        with:
          enable-cache: true

      - name: Build wheel
        run: uv build --wheel --package flashdreams

      - name: Previously published check
        id: previously_published
        run: |
          version=$(sed -n 's/^__version__ = "\(.*\)"/\1/p' flashdreams/flashdreams/_version.py)
          echo "version=$version"
          if curl -sSf "https://pypi.org/simple/flashdreams/" \
             | grep -q "flashdreams-${version}-"; then
            echo "Version $version already exists on PyPI -- skipping publish"
            echo "skip=true" >> "$GITHUB_OUTPUT"
          else
            echo "Version $version not found on PyPI -- will publish"
            echo "skip=false" >> "$GITHUB_OUTPUT"
          fi

      - name: Publish to PyPI
        if: steps.previously_published.outputs.skip != 'true'
        run: >
          uv publish
          --publish-url https://upload.pypi.org/legacy/
          --check-url https://pypi.org/simple/
          --token ${{ secrets.PYPI_API_TOKEN }}
          dist/*