[Bug]: Default Docker Image Does not Support Latest Flash Attention (2.8.0)

[Bihan]

Steps to reproduce

1. Create Task as below:

type: task
name: test-task
nvcc: true
python: 3.12

commands:
  - uv pip install torch numpy
  - uv pip install flash-attn --no-build-isolation
  - python -c "import flash_attn.flash_attn_interface"

resources:
  gpu: 1

Actual behaviour

Traceback (most recent call last):
File "", line 1, in
File "/workflow/.venv/lib/python3.12/site-packages/flash_attn/init.py", line 3, in
from flash_attn.flash_attn_interface import (
File "/workflow/.venv/lib/python3.12/site-packages/flash_attn/flash_attn_interface.py", line 15, in
import flash_attn_2_cuda as flash_attn_gpu
ImportError: /usr/lib/x86_64-linux-gnu/libc.so.6: version `GLIBC_2.32' not found (required by /workflow/.venv/lib/python3.12/site-packages/flash_attn_2_cuda.cpython-312-x86_64-linux-gnu.so)

Expected behaviour

No Error

dstack version

Commit 5a5b134

Server logs

Additional information

The issue appears to stem from the fact that our current base image is based on Ubuntu 20.04, which ships with GLIBC 2.31. However, FlashAttention v2.8.0 is compiled against GLIBC 2.32, leading to a runtime ImportError.

Verification with Ubuntu 22
To validate whether Ubuntu 22 solves the issue, I ran the same steps using the nvidia/cuda:12.9.0-devel-ubuntu22.04 image (which includes GLIBC ≥ 2.35), and FlashAttention v2.8.0 imported successfully without errors. I had used GCP backend.

Implications of the issue

1. Update training examples by installing previous version of Flash Attention (2.7.0). Currently it automatically installs latest Flash Attention (2.8.0) and task fails.
2. Identify backends with Ubuntu 20.04 as VM Image. Note: GCP is using Ubuntu 22.04.5 LTS as VM Image.
3. If some VM backends are using Ubuntu 20.04 as VM Image; then we need to find out whether using Ubuntu 20.04 as host VM Image and using Ubuntu 22.04.5 LTS as container image throws error.
4. Find out whether upgrading default image to Ubuntu 22 supports infiniband without additional changes.

[r4victor]

@Bihan #2799 upgraded default docker images to ubuntu 22. Please check if it's fixed.