test_e2e.yaml

name: "E2E: kernel-builder init + build + upload + download"

on:
  push:
    branches: [main]
  pull_request:
    branches: [main]
    paths:
      - "kernel-builder/**"
      - "kernels/src/**"
      - "nix-builder/**"
      - "kernels-data/**"
      - ".github/workflows/test_e2e.yaml"
  workflow_dispatch:

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

env:
  HF_ENDPOINT: https://hub-ci.huggingface.co
  HF_TOKEN: ${{ secrets.HF_HUB_CI_TOKEN }}
  E2E_REPO_ID: __DUMMY_KERNELS_USER__/kernels-e2e-${{ github.run_id }}-${{ github.run_attempt }}
  E2E_REPO_NAME: kernels-e2e-${{ github.run_id }}-${{ github.run_attempt }}
  E2E_PKG_NAME: kernels_e2e_${{ github.run_id }}_${{ github.run_attempt }}

jobs:
  init-build-upload:
    name: Init, build, and upload kernel
    runs-on:
      group: aws-highmemory-32-plus-nix
    outputs:
      variant: ${{ steps.variant.outputs.name }}
    steps:
      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
      - uses: DeterminateSystems/nix-installer-action@ef8a148080ab6020fd15196c2084a2eea5ff2d25
        with:
          extra-conf: |
            max-jobs = 8
            cores = 12
            sandbox-fallback = false
      - uses: cachix/cachix-action@3ba601ff5bbb07c7220846facfa2cd81eeee15a1 # v16
        with:
          name: huggingface
          authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"
        env:
          USER: runner

      - name: Init kernel project
        run: |
          cd /tmp
          nix run $GITHUB_WORKSPACE#kernel-builder -- init \
            --name ${{ env.E2E_REPO_ID }} \
            --backends cuda

      - name: Validate scaffold
        run: |
          cd /tmp/${{ env.E2E_REPO_NAME }}
          test -f build.toml
          test -f flake.nix
          test -f torch-ext/${{ env.E2E_PKG_NAME }}/__init__.py
          test -f torch-ext/torch_binding.cpp
          test -f torch-ext/torch_binding.h
          test -f ${{ env.E2E_PKG_NAME }}_cuda/${{ env.E2E_PKG_NAME }}.cu
          test -f tests/test_${{ env.E2E_PKG_NAME }}.py
          test -f example.py
          grep -q 'name = "${{ env.E2E_REPO_NAME }}"' build.toml
          grep -q 'repo-id = "${{ env.E2E_REPO_ID }}"' build.toml
          grep -q 'backend = "cuda"' build.toml

      - name: Patch flake.nix to use local nix-builder
        run: |
          cd /tmp/${{ env.E2E_REPO_NAME }}
          sed -i 's|github:huggingface/kernels|path:'"$GITHUB_WORKSPACE"'|' flake.nix

      - name: Make flake a Git repo
        run: |
          cd /tmp/${{ env.E2E_REPO_NAME }}
          git config --global user.email "bottie@mcbotface.hf.co"
          git config --global user.name "Botty McBotface"
          git init && git add . && git commit -m "e2e test"

      - name: Determine latest variant
        id: variant
        run: |
          cd /tmp/${{ env.E2E_REPO_NAME }}
          # NOTE: Remove Torch 2.14 grep once it is released.
          VARIANT=$(nix run $GITHUB_WORKSPACE#kernel-builder -- list-variants . | grep -v torch214 | tail -1)
          echo "name=$VARIANT" >> $GITHUB_OUTPUT
          echo "Building variant: $VARIANT"

      - name: Build kernel
        run: |
          cd /tmp/${{ env.E2E_REPO_NAME }}
          nix run $GITHUB_WORKSPACE#kernel-builder -- build --variant ${{ steps.variant.outputs.name }} . -L

      - name: Verify build artifacts
        run: |
          cd /tmp/${{ env.E2E_REPO_NAME }}
          VARIANT_DIR=$(ls -d result/torch* | head -1)
          echo "Built variant: $VARIANT_DIR"
          test -f "$VARIANT_DIR/__init__.py"
          test -f "$VARIANT_DIR/metadata.json"
          ls "$VARIANT_DIR"/*.so

      - name: Upload kernel to Hub
        run: |
          nix run $GITHUB_WORKSPACE#kernel-builder -- upload /tmp/${{ env.E2E_REPO_NAME }}
          nix run $GITHUB_WORKSPACE#kernel-builder -- upload /tmp/${{ env.E2E_REPO_NAME }} --repo-type model

  download-and-test:
    name: Download and test kernel via get_kernel
    needs: init-build-upload
    runs-on:
      group: aws-g6-24xlarge
    env:
      UV_PYTHON_PREFERENCE: only-managed
    steps:
      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1

      - name: Install uv and set Python version
        uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0
        with:
          python-version: "3.12"

      - name: Install Python deps
        working-directory: ./kernels
        run: |
          VARIANT="${{ needs.init-build-upload.outputs.variant }}"
          CUDA_TAG=$(echo "$VARIANT" | grep -oP 'cu\d+')
          echo "Installing torch matching variant $VARIANT (CUDA tag: $CUDA_TAG)"
          uv sync --all-extras --dev
          # E2E validates this checkout, so use the local kernels-data binding instead of the released wheel from uv.lock.
          uv pip install --reinstall ../kernels-data/bindings/python
          uv pip install --upgrade torch --index-url https://download.pytorch.org/whl/$CUDA_TAG
          uv run --no-sync python -c "import torch; print(f'torch={torch.__version__}, cuda={torch.version.cuda}, cxx11_abi={torch.compiled_with_cxx11_abi()}')"

      - name: Test get_kernel download and usage
        working-directory: ./kernels
        run: |
          uv run --no-sync python -c "
          import torch
          from kernels import get_kernel

          kernel = get_kernel('${{ env.E2E_REPO_ID }}', trust_remote_code=True, version=1)

          x = torch.randn(1024, 1024, dtype=torch.float32, device='cuda')
          result = kernel.${{ env.E2E_PKG_NAME }}(x)
          expected = x + 1.0
          torch.testing.assert_close(result, expected)
          print('E2E test passed: get_kernel + correctness check')
          "

  cleanup:
    name: Clean up test repos
    needs: [init-build-upload, download-and-test]
    if: always()
    runs-on: ubuntu-latest
    steps:
      - uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0
        with:
          python-version: "3.12"
      - name: Delete test repos
        run: |
          uv run --with huggingface-hub python -c "
          from huggingface_hub import HfApi
          import os

          api = HfApi(endpoint=os.environ['HF_ENDPOINT'])
          repo_id = os.environ['E2E_REPO_ID']
          for repo_type in ['kernel', 'model']:
              api.delete_repo(repo_id, repo_type=repo_type, missing_ok=True)
              print(f'Deleted {repo_type} repo (or it did not exist)')
          "