test-docker-gpu.yaml

name: docker-test-gpu

on:
  workflow_dispatch:

permissions: read-all

jobs:
  test:
    runs-on: self-hosted
    timeout-minutes: 240
    strategy:
      fail-fast: true
      max-parallel: 1
      matrix:
        backend:
          - llama-cuda
          # - exllama # https://github.com/sozercan/aikit/issues/94
          - exllama2-gptq
          - exllama2-exl2
          - mamba
    steps:
      - name: cleanup workspace
        run: |
          rm -rf ./* || true
          rm -rf ./.??* || true
      - uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # v4.1.1

        # use default docker driver builder with containerd image store for local aikit image
      - run: docker buildx use default

      - name: build aikit
        run: |
          docker buildx build . -t aikit:test \
            --load --provenance=false --progress plain

      - name: build test model
        run: |
          docker buildx build . -t testmodel:test \
            -f test/aikitfile-${{ matrix.backend }}.yaml \
            --load --provenance=false --progress plain

      - name: list images
        run: docker images

      - name: run test model
        run: docker run --name testmodel -d --rm -p 8080:8080 --gpus all testmodel:test

      - name: run test
        if: matrix.backend != 'mamba'
        run: |
          result=$(curl --fail --retry 10 --retry-all-errors http://127.0.0.1:8080/v1/chat/completions -H "Content-Type: application/json" -d '{
            "model": "llama-2-7b-chat",
            "messages": [{"role": "user", "content": "explain kubernetes in a sentence"}]
          }')
          echo $result

          choices=$(echo "$result" | jq '.choices')
          if [ -z "$choices" ]; then
            exit 1
          fi

      - name: run test
        if: matrix.backend == 'mamba'
        run: |
          result=$(curl --fail --retry 10 --retry-all-errors http://127.0.0.1:8080/v1/chat/completions -H "Content-Type: application/json" -d '{
            "model": "mamba-chat",
            "messages": [{"role": "user", "content": "explain kubernetes in a sentence"}]
          }')
          echo $result

          choices=$(echo "$result" | jq '.choices')
          if [ -z "$choices" ]; then
            exit 1
          fi

      - name: save logs
        if: always()
        run: docker logs testmodel > /tmp/docker-${{ matrix.backend }}.log

      - run: docker stop testmodel
        if: always()

      - run: docker system prune -a -f --volumes
        if: always()

      - name: publish test artifacts
        if: always()
        uses: actions/upload-artifact@5d5d22a31266ced268874388b861e4b58bb5c2f3 # v4.3.1
        with:
          name: test-${{ matrix.backend }}
          path: |
            /tmp/*.log