FSDP-Multi-GPU-Training/Dockerfile at main · codeamt/FSDP-Multi-GPU-Training · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
# Training image with CUDA, awscli v2, s5cmd, uv-managed env, and PyTorch CUDA wheels (no conda)
FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04

# System deps and tools (install Python 3.10 explicitly)
RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
    python3.10 python3.10-venv python3-pip \
    git curl ca-certificates unzip jq \
    && ln -sf /usr/bin/python3.10 /usr/bin/python \
    && ln -sf /usr/bin/pip3 /usr/bin/pip \
    && rm -rf /var/lib/apt/lists/*


## Install uv
RUN curl -LsSf https://astral.sh/uv/install.sh | sh -s --
ENV PATH="/root/.local/bin:${PATH}"


# Create and register dedicated venv for the project
RUN uv venv /opt/venv
ENV VIRTUAL_ENV=/opt/venv
ENV PATH="${VIRTUAL_ENV}/bin:${PATH}"
ENV PYTHONNOUSERSITE=1

# Install base Python deps into /opt/venv using uv pip
RUN uv pip install --upgrade pip
RUN uv pip install \
    omegaconf hydra-core boto3 s3fs datasets transformers accelerate wandb

## Install awscli v2 (official bundle)
RUN curl -Ls https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip -o /tmp/awscliv2.zip && \
    unzip -q /tmp/awscliv2.zip -d /tmp && \
    /tmp/aws/install && rm -rf /tmp/aws /tmp/awscliv2.zip

# Install s5cmd (fast S3 sync) with arch detection
ARG S5CMD_VERSION=2.3.0
RUN set -eux; \
    arch=$(uname -m); \
    if [ "$arch" = "x86_64" ] || [ "$arch" = "amd64" ]; then \
      FILE="s5cmd_${S5CMD_VERSION}_Linux-64bit.tar.gz"; \
    elif [ "$arch" = "aarch64" ] || [ "$arch" = "arm64" ]; then \
      FILE="s5cmd_${S5CMD_VERSION}_Linux-arm64.tar.gz"; \
    else \
      echo "Unsupported architecture: $arch"; exit 1; \
    fi; \
    curl -fL -o /tmp/s5cmd.tgz "https://github.com/peak/s5cmd/releases/download/v${S5CMD_VERSION}/$FILE"; \
    tar -xzf /tmp/s5cmd.tgz -C /usr/local/bin s5cmd; \
    chmod +x /usr/local/bin/s5cmd; \
    rm -f /tmp/s5cmd.tgz

# Workdir
WORKDIR /workspace

# Copy dependency manifests for better caching
COPY pyproject.toml /workspace/pyproject.toml
# If present, include lockfile for reproducibility
COPY uv.lock /workspace/uv.lock

# Set workdir and prefer the already-created /opt/venv on PATH for all subsequent steps
WORKDIR /workspace
RUN which python && python -c 'import sys; print(sys.executable)'

# Install NumPy compatible with torch 2.2.x, then PyTorch CUDA 12.1 wheels
RUN uv pip install --no-cache-dir 'numpy<2'
RUN uv pip install --no-cache-dir --extra-index-url https://download.pytorch.org/whl/cu121 \
    torch==2.3.1 torchvision==0.18.1 torchaudio==2.3.1 && \
    python -c 'import sys, torch; print(sys.executable); print(torch.__version__, "cuda=", torch.cuda.is_available())'

# Copy the rest of the repo
COPY . /workspace

# Default environment variables (can be overridden by SageMaker)
ENV PYTHONUNBUFFERED=1 \
    TRANSFORMERS_NO_ADVISORY_WARNINGS=1 \
    HF_HUB_DISABLE_TELEMETRY=1 \
    PYTHONPATH=/workspace/src

# Provide an entrypoint that can fetch secrets; orchestrators may override
COPY docker/entrypoint.sh /usr/local/bin/entrypoint.sh
RUN chmod +x /usr/local/bin/entrypoint.sh
ENTRYPOINT ["/usr/local/bin/entrypoint.sh"]
CMD ["bash"]