fix flash-attn version

Dockerfile

@@ -34,8 +34,9 @@ ENV HOME=/home/user \
 
 RUN pip3 install torch==2.3.1 torchvision==0.18.1 torchaudio==2.3.1 --index-url https://download.pytorch.org/whl/cu121
 RUN pip3 install --no-cache-dir --upgrade -r /code/requirements.txt
-ARG DYNAMIC_PARAMS=$(python3 /code/get_flash_attn.py)
-RUN pip3 install https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.2.post1/${DYNAMIC_PARAMS}
+RUN python /code/get_flash_attn.py > url.txt
+RUN source ./url.txt
+RUN pip3 install https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.2.post1/${FLASH_ATTN_VER} --no-cache-dir
 
 # Set the working directory to the user's home directory
 WORKDIR $HOME/app

get_flash_attn.py

@@ -54,6 +54,6 @@ def generate_flash_attn_filename(flash_attn_version="2.7.2.post1"):
 if __name__ == "__main__":
     try:
         filename = generate_flash_attn_filename()
-        print(filename)
+        print(f"FLASH_ATTN_VER={filename}")
     except Exception as e:
         print("Error generating filename:", e)