Merge pull request #4 from mjun0812/feature/add-self-hosted-runner

[WIP] Add self hosted runner

.github/workflows/test-self-hosted.yml

@@ -0,0 +1,100 @@
+name: "Test self-hosted runner"
+on:
+  workflow_dispatch:
+
+jobs:
+  docker:
+    runs-on:
+      - self-hosted
+    timeout-minutes: 10
+    container:
+      image: docker:dind
+    steps:
+      - run: ls
+      - run: pwd
+
+  test-build:
+    runs-on:
+      - self-hosted
+    timeout-minutes: 1000
+    strategy:
+      fail-fast: false
+      matrix:
+        flash-attn-version: ["2.4.3"]
+        python-version: ["3.12"]
+        torch-version: ["2.7.0"]
+        # https://developer.nvidia.com/cuda-toolkit-archive
+        cuda-version: ["12.8.1"]
+    container:
+      image: ubuntu:22.04
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+      - uses: Jimver/cuda-toolkit@master
+        with:
+          cuda: ${{ matrix.cuda-version }}
+          linux-local-args: '["--toolkit"]'
+          method: "network"
+
+      - name: Set CUDA and PyTorch versions
+        run: |
+          echo "MATRIX_CUDA_VERSION=$(echo ${{ matrix.cuda-version }} | awk -F \. {'print $1 $2'})" >> $GITHUB_ENV
+          echo "MATRIX_TORCH_VERSION=$(echo ${{ matrix.torch-version }} | awk -F \. {'print $1 "." $2'})" >> $GITHUB_ENV
+          echo "CACHE_KEY=cuda-ext-${{ matrix.flash-attn-version }}-py${{ matrix.python-version }}-torch${{ matrix.torch-version }}-cuda${{ matrix.cuda-version }}" >> $GITHUB_ENV
+
+      - name: Install build dependencies
+        run: |
+          sudo apt install -y ninja-build clang
+          pip install -U pip setuptools==75.8.0 wheel setuptools packaging psutil
+
+      - name: Install PyTorch ${{ matrix.torch-version }}+cu${{ matrix.cuda-version }}
+        run: |
+          export TORCH_CUDA_VERSION=$(python -c "from os import environ as env; \
+              support_cuda_versions = { \
+                  '2.0': [117, 118], \
+                  '2.1': [118, 121], \
+                  '2.2': [118, 121], \
+                  '2.3': [118, 121], \
+                  '2.4': [118, 121, 124], \
+                  '2.5': [118, 121, 124], \
+                  '2.6': [118, 124, 126], \
+                  '2.7': [118, 126, 128], \
+              }; \
+              target_cuda_versions = support_cuda_versions[env['MATRIX_TORCH_VERSION']]; \
+              cuda_version = int(env['MATRIX_CUDA_VERSION']); \
+              closest_version = min(target_cuda_versions, key=lambda x: abs(x - cuda_version)); \
+              print(closest_version) \
+          ")
+          pip install --no-cache-dir torch==${{ matrix.torch-version }} --index-url https://download.pytorch.org/whl/cu${TORCH_CUDA_VERSION}
+          nvcc --version
+          python -V
+          python -c "import torch; print('PyTorch:', torch.__version__)"
+          python -c "import torch; print('CUDA:', torch.version.cuda)"
+          python -c "from torch.utils import cpp_extension; print(cpp_extension.CUDA_HOME)"
+
+      - name: Checkout flash-attn
+        run: |
+          git clone https://github.com/Dao-AILab/flash-attention.git -b "v${{ matrix.flash-attn-version }}"
+          cd flash-attention
+
+      - name: Build wheels
+        timeout-minutes: 800
+        run: |
+          export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH
+          export LD_LIBRARY_PATH=/usr/local/nvidia/lib64:/usr/local/cuda/lib64:$LD_LIBRARY_PATH
+          export MAX_JOBS=2
+          export NVCC_THREADS=2
+          export FLASH_ATTENTION_FORCE_BUILD="TRUE"
+          cd flash-attention
+          python setup.py bdist_wheel --dist-dir=dist
+          base_wheel_name=$(basename $(ls dist/*.whl | head -n 1))
+          wheel_name=$(echo $base_wheel_name | sed "s/${{ matrix.flash-attn-version }}/${{ matrix.flash-attn-version }}+cu${{ env.MATRIX_CUDA_VERSION }}torch${{ env.MATRIX_TORCH_VERSION }}/")
+          mv dist/$base_wheel_name dist/$wheel_name
+          echo "wheel_name=$wheel_name" >> $GITHUB_ENV
+
+      - name: Install Test
+        run: |
+          pip install flash-attention/dist/${{ env.wheel_name }}
+          python -c "import flash_attn; print(flash_attn.__version__)"

.github/workflows/test.yml

@@ -1,19 +0,0 @@
-name: "Docker and container hooks test"
-on:
-  workflow_dispatch:
-
-jobs:
-  docker:
-    runs-on:
-      - self-hosted
-    timeout-minutes: 10
-    container:
-      image: docker:dind
-    steps:
-      - name: Show docker info
-        run: |
-          docker -v
-          docker info
-          docker ps
-      - run: ls
-      - run: pwd

self-hosted-runner/Dockerfile

@@ -17,7 +17,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends\
     software-properties-common \
     && rm -rf /var/lib/apt/lists/*
 
-# Install Docker cli
+# Install Docker
 RUN install -m 0755 -d /etc/apt/keyrings \
     && curl -fsSL https://download.docker.com/linux/ubuntu/gpg -o /etc/apt/keyrings/docker.asc \
     && chmod a+r /etc/apt/keyrings/docker.asc \

self-hosted-runner/compose.yml

@@ -12,3 +12,5 @@ services:
         TARGET_ARCH: x64
     volumes:
       - /var/run/docker.sock:/var/run/docker.sock
+    group_add:
+      - "${DOCKER_GID}"

self-hosted-runner/env.template

@@ -1 +1,2 @@
-PERSONAL_ACCESS_TOKEN=
+PERSONAL_ACCESS_TOKEN=
+DOCKER_GID=