Skip to content

Commit 8e65d27

Browse files
junwhajunwha
committed
Implement uv-based base image infra
1 parent 961583f commit 8e65d27

3 files changed

Lines changed: 147 additions & 0 deletions

File tree

‎uv/Dockerfile‎

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
ARG PY_VERSION
2+
ARG DOCKER_CUDA_VERSION=12.4.1
3+
FROM nvidia/cuda:${DOCKER_CUDA_VERSION}-cudnn-devel-ubuntu22.04
4+
5+
ARG PY_VERSION
6+
7+
# For A100, RTX4090, H100, and H200
8+
ENV TORCH_CUDA_ARCH_LIST='7.0 7.5 8.0 8.6 8.9 9.0+PTX'
9+
ENV MAX_JOBS=2
10+
ENV NVCC_THREADS=8
11+
12+
ENV DEBIAN_FRONTEND=noninteractive
13+
ENV CUDA_HOME=/usr/local/cuda
14+
ENV SETUPTOOLS_USE_DISTUTILS=stdlib
15+
16+
RUN apt-get update -y
17+
18+
# Utils
19+
RUN apt-get install -y vim git bzip2 tmux wget tar htop unzip
20+
21+
# SSH
22+
RUN apt-get install -y ssh openssh-server
23+
RUN mkdir -p /var/run/sshd
24+
RUN echo "root:root" | chpasswd
25+
RUN echo "Port 22" >> /etc/ssh/sshd_config
26+
RUN echo "PermitRootLogin yes" >> /etc/ssh/sshd_config
27+
EXPOSE 22
28+
29+
RUN apt-get update -y && apt-get install -y software-properties-common
30+
RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/7fa2af80.pub
31+
RUN add-apt-repository -y "deb https://developer.download.nvidia.com/devtools/repos/ubuntu22.04/$(dpkg --print-architecture) /"
32+
RUN apt-get install -y nsight-systems-2024.6.2
33+
RUN apt-get install -y mesa-utils x11-apps
34+
RUN apt-get install -y freeglut3-dev libglu1-mesa-dev mesa-common-dev libxkbfile-dev libgl1-mesa-glx
35+
36+
ENV LIBGL_DRIVERS_PATH=/usr/lib/x86_64-linux-gnu/dri
37+
ENV LIBGL_ALWAYS_INDIRECT=1
38+
39+
# Build tools
40+
RUN rm /etc/apt/sources.list.d/archive_uri-https_developer_download_nvidia_com_devtools_repos_ubuntu22_04_amd64-jammy.list
41+
RUN apt-get update -y && apt-get install -y gcc-12 g++-12 ninja-build cmake build-essential autoconf libtool automake git
42+
RUN apt-get install -y ffmpeg libsm6 libxext6 libgl1
43+
44+
# Utils
45+
RUN apt-get install -y vim git bzip2 tmux wget tar htop curl
46+
RUN add-apt-repository ppa:deadsnakes/ppa && apt-get install -y python3.12 python3.12-dev python3.12-venv
47+
48+
RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.12 1
49+
RUN update-alternatives --set python3 /usr/bin/python3.12
50+
RUN ln -sf /usr/bin/python3.12-config /usr/bin/python3-config
51+
RUN curl -sS https://bootstrap.pypa.io/get-pip.py | python3.12
52+
53+
RUN curl -LsSf https://astral.sh/uv/0.8.14/install.sh | sh
54+
55+
SHELL ["/bin/bash", "-c"]
56+
ENV PATH=/root/.local/bin:$PATH
57+
58+
RUN uv python pin ${PY_VERSION}
59+
RUN uv pip install --no-cache --system ipython tqdm rich jupyter jupyterlab ipykernel pandas einops safetensors pyyaml requests psutil opencv-python-headless matplotlib seaborn scikit-learn scipy pillow tensorboard h5py
60+
61+
# torch installer
62+
COPY <<EOF /usr/local/bin/install_torch_env
63+
#!/bin/bash
64+
set -e
65+
TORCH_VER=\$1
66+
VISION_VER=\$2
67+
CUDA_VER=\$3
68+
BASE_DIR=/ddiff-base/py\${PY_VERSION}-torch\${TORCH_VER}
69+
mkdir -p "\$BASE_DIR" && cd "\$BASE_DIR"
70+
uv venv --python "\${PY_VERSION}" --system-site-packages
71+
uv pip install --no-cache torch=="\${TORCH_VER}" torchvision=="\${VISION_VER}" torchaudio=="\${TORCH_VER}" --index-url "https://download.pytorch.org/whl/\${CUDA_VER}"
72+
echo "alias uv_init_torch\${TORCH_VER}='ln -sfn \${BASE_DIR}/.venv ./.venv && echo Created uv virtual environment with torch==\${TORCH_VER}'" >> ~/.bashrc
73+
EOF
74+
RUN chmod +x /usr/local/bin/install_torch_env
75+
76+
# pytorch 2.4.1 (251209 Update)
77+
RUN install_torch_env 2.4.1 0.19.1 cu124
78+
79+
# pytorch 2.5.1 (251209 Update)
80+
RUN install_torch_env 2.5.1 0.20.1 cu124
81+
82+
# pytorch 2.6.0 (251209 Update)
83+
RUN install_torch_env 2.6.0 0.21.0 cu124
84+
85+
# pytorch 2.7.1 (251209 Update)
86+
RUN install_torch_env 2.7.1 0.22.1 cu126
87+
88+
# pytorch 2.9.0 (251209 Update)
89+
RUN install_torch_env 2.9.0 0.24.0 cu126

‎uv/README.md‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
# Efficient Docker Storage Management with Pre-built, Torch-Shared uv Images
2+
3+
### Problem Definition
4+
5+
When multiple users on a server repeatedly build their own Docker images, it results in massive duplicated storage usage. Specifically, repeated installations of PyTorch (which exceeds 5GB) can rapidly consume server disk space.
6+
7+
### Solution
8+
9+
We provide a monolithic base image with various PyTorch versions and essential libraries pre-installed. Users can simply build their images on top of this base.
10+
11+
* **Zero-Copy Torch:** PyTorch is already installed in the base layer.
12+
* **Symlink & Inheritance:** Uses `uv` to link existing installations, so your project only takes up space for *additional* packages.
13+
14+
### Usage
15+
16+
### 1. Docker file example
17+
18+
```dockerfile
19+
FROM junwha/ddiff-base:cu12.4.1-py3.10-torch-251205
20+
21+
WORKDIR /<your workspace>
22+
23+
# 1. Initialize venv linked to a specific Torch version
24+
# (Note: This creates a symlink to the base image's venv, thus only one project per torch version is recommended)
25+
RUN uv_init_torch2.5.1
26+
27+
# 2. Install additional packages (e.g., diffusers)
28+
RUN uv pip install -r requirements.txt
29+
```
30+
31+
### 2. Running Commands Inside the container
32+
you can use standard uv commands. The environment is automatically detected via the .venv link.
33+
34+
```uv run python example.py```
35+
36+
# Base Images
37+
### Common Packages
38+
* **Python Libraries**
39+
* **Data & Science:** numpy pandas scipy scikit-learn matplotlib seaborn
40+
* **Tools:** jupyterlab ipykernel tqdm rich
41+
* **CV & Utils:** opencv-python-headless pillow einops safetensors
42+
43+
* **System Libraries**
44+
* **Basic Utils:** git ffmpeg vim git bzip2 tmux wget tar htop curl
45+
* **X11 Utils:** mesa-utils x11-apps freeglut3-dev libglu1-mesa-dev mesa-common-dev libxkbfile-dev libgl1-mesa-glx libgl1
46+
* **Build Tools:** gcc-12 g++-12 ninja-build cmake build-essential
47+
* **Profiler:** Nsight System 2024
48+
## Releases
49+
**Image Tag:** `junwha/ddiff-base:cu12.4.1-py3.10-torch-251205`
50+
* **Python:** 3.10
51+
* **Pre-installed PyTorch Versions:** 2.4.1, 2.5.1, 2.6.0, 2.7.1, 2.9.0
52+
* **Compute Capabilities**: 7.0 7.5 8.0 8.6 8.9 9.0+PTX
53+

‎uv/build_uv.sh‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
PY_VERSION=$1
2+
TAG=junwha/ddiff-base:cu12.4.1-py${PY_VERSION}-torch-$(date +"%y%m%d")
3+
4+
docker build -t $TAG --build-arg PY_VERSION=${PY_VERSION} ./
5+
docker push $TAG

0 commit comments

Comments
 (0)