fix: image build failure due to flash attn (#629)

* Add free up disk space to gh runners

Signed-off-by: Dushyant Behl <[email protected]>

* Build flash attn without cache to avoid problems

Signed-off-by: Dushyant Behl <[email protected]>

---------

Signed-off-by: Dushyant Behl <[email protected]>

Author: dushyantbehl

.github/workflows/coverage.yaml

@@ -10,6 +10,8 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
+      - name: "Free up disk space"
+        uses: ./.github/actions/free-up-disk-space
       - name: Set up Python 3.12
         uses: actions/setup-python@v4
         with:

build/Dockerfile

@@ -156,9 +156,10 @@ ENV PIP_NO_BINARY=mamba-ssm,mamba_ssm
 RUN --mount=type=cache,target=/home/${USER}/.cache/pip,uid=${USER_UID} \
     python -m pip install --user wheel && \
     python -m pip install --user "$(head bdist_name)" && \
-    python -m pip install --user "$(head bdist_name)[flash-attn]" && \
     python -m pip install --user --no-build-isolation "$(head bdist_name)[mamba]"
 
+RUN python -m pip install --user --no-build-isolation "$(head bdist_name)[flash-attn]"
+
 # fms_acceleration_peft = PEFT-training, e.g., 4bit QLoRA
 # fms_acceleration_foak = Fused LoRA and triton kernels
 # fms_acceleration_aadp = Padding-Free Flash Attention Computation