File size: 2,112 Bytes
0839907 | 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 | FROM nvcr.io/nvidia/pytorch:25.06-py3
# Defines the architecture (amd64 or arm64) automatically during build
ARG TARGETARCH
ENV PYTHONDONTWRITEBYTECODE=1
ENV PYTHONUNBUFFERED=1
ENV DEBIAN_FRONTEND=noninteractive
# Set the timezone
ENV TZ=America/Los_Angeles
RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime && echo $TZ > /etc/timezone
# Install system dependencies
RUN apt-get update --fix-missing && \
apt-get install -y --no-install-recommends \
build-essential \
ninja-build \
cmake \
wget \
ca-certificates \
git \
curl \
unzip \
python3-tk \
&& rm -rf /var/lib/apt/lists/* \
&& cmake --version \
&& echo "✅ All system dependencies installed successfully"
# Install s5cmd with architecture logic
RUN S5CMD_VERSION="2.1.0-beta.1" && \
if [ "$TARGETARCH" = "arm64" ]; then \
S5CMD_ARCH="Linux-arm64"; \
else \
S5CMD_ARCH="Linux-64bit"; \
fi && \
wget "https://github.com/peak/s5cmd/releases/download/v${S5CMD_VERSION}/s5cmd_${S5CMD_VERSION}_${S5CMD_ARCH}.tar.gz" && \
tar -xf "s5cmd_${S5CMD_VERSION}_${S5CMD_ARCH}.tar.gz" && \
install s5cmd /usr/local/bin/ && \
rm s5cmd "s5cmd_${S5CMD_VERSION}_${S5CMD_ARCH}.tar.gz"
# Set build optimization flags
ENV MAX_JOBS=4
ENV TORCH_CUDA_ARCH_LIST="8.0;8.6;9.0;10.0"
# Upgrade pip
RUN pip install --upgrade pip setuptools wheel
# numpy<2.0.0 is required for the preinstalled numba and thinc
RUN pip install wandb[media] \
diffusers==0.35.1 \
transformers==4.49.0 \
accelerate \
safetensors \
scipy \
einops \
jupyter \
hydra-core \
imageio \
tqdm \
webdataset \
av \
loguru \
boto3 \
timm \
ftfy \
opencv-python-headless \
sentencepiece \
"numpy<2.0.0"
RUN pip uninstall -y apex
RUN pip install --index-url=https://sc-hw-artf.nvidia.com/artifactory/api/pypi/hwinf-mlwfo-pypi/simple --upgrade one-logger-utils
WORKDIR /workspace
# Entry point
RUN (printf '#!/bin/bash\nexec \"$@\"\n' >> /entry.sh) && chmod a+x /entry.sh
ENTRYPOINT ["/entry.sh"]
|