Skip to content

Commit 9f41b79

Browse files
TimPietruskyTimPietrusky
committed
fix: optimize dockerfile to reduce image size
- Implement multi-stage build to reduce final image size - Selectively copy only required packages from builder stage - Clean up unnecessary CUDA libraries and cache files - Optimize package installation and caching settings
1 parent 3834448 commit 9f41b79

File tree

3 files changed

+128
-37
lines changed

3 files changed

+128
-37
lines changed

‎Dockerfile‎

Lines changed: 65 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -1,47 +1,77 @@
1-
# FROM nvidia/cuda:12.1.0-base-ubuntu22.04
1+
FROM axolotlai/axolotl-cloud:main-latest AS builder
22

3+
# Set environment variables to reduce size
4+
ENV PYTHONDONTWRITEBYTECODE=1 \
5+
PYTHONUNBUFFERED=1 \
6+
PIP_NO_CACHE_DIR=1
37

4-
# RUN ldconfig /usr/local/cuda-12.1/compat/
5-
6-
7-
# # # Set CUDA environment variables
8-
# ENV PATH="/usr/local/cuda/bin:${PATH}"
9-
# ENV LD_LIBRARY_PATH="/usr/local/cuda/lib64:${LD_LIBRARY_PATH}"
10-
11-
# RUN apt-get update -y \
12-
# && apt-get install -y python3-pip git
13-
14-
15-
16-
# # The base image comes with many system dependencies pre-installed to help you get started quickly.
17-
# # Please refer to the base image's Dockerfile for more information before adding additional dependencies.
18-
# # IMPORTANT: The base image overrides the default huggingface cache location.
19-
# # Python dependencies
20-
FROM axolotlai/axolotl-cloud:main-latest
21-
8+
# Install dependencies
229
COPY builder/requirements.txt /requirements.txt
23-
RUN --mount=type=cache,target=/root/.cache/pip \
24-
python3 -m pip install --upgrade pip && \
25-
python3 -m pip install --upgrade -r /requirements.txt
26-
27-
# RUN pip install --upgrade torch --index-url https://download.pytorch.org/whl/cu121
28-
29-
# RUN git clone https://github.com/runpod-workers/axolotl.git && \
30-
# cd axolotl && \
31-
# pip install packaging ninja && \
32-
# pip install --no-build-isolation -e '.[flash-attn,deepspeed]'
10+
RUN pip install --no-cache-dir --upgrade pip && \
11+
pip install --no-cache-dir --upgrade -r /requirements.txt
3312

13+
# Create a new stage with minimal files
14+
FROM axolotlai/axolotl-cloud:main-latest
3415

16+
# Set environment variables
17+
ENV PYTHONDONTWRITEBYTECODE=1 \
18+
PYTHONUNBUFFERED=1 \
19+
PIP_NO_CACHE_DIR=1
20+
21+
# Copy installed packages from builder stage - selectively to save space
22+
COPY --from=builder /usr/local/lib/python3.11/site-packages/runpod /usr/local/lib/python3.11/site-packages/runpod
23+
COPY --from=builder /usr/local/lib/python3.11/site-packages/huggingface_hub /usr/local/lib/python3.11/site-packages/huggingface_hub
24+
COPY --from=builder /usr/local/lib/python3.11/site-packages/typing_extensions* /usr/local/lib/python3.11/site-packages/
25+
COPY --from=builder /usr/local/lib/python3.11/site-packages/pydantic* /usr/local/lib/python3.11/site-packages/
26+
COPY --from=builder /usr/local/lib/python3.11/site-packages/hf_transfer /usr/local/lib/python3.11/site-packages/hf_transfer
27+
COPY --from=builder /usr/local/lib/python3.11/site-packages/numpy /usr/local/lib/python3.11/site-packages/numpy
28+
COPY --from=builder /usr/local/bin/runpod* /usr/local/bin/
29+
30+
# Clean up to save space
31+
RUN apt-get clean && \
32+
rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* && \
33+
# Remove unnecessary CUDA files to save space
34+
rm -rf /usr/local/cuda/samples /usr/local/cuda/doc /usr/local/cuda/extras && \
35+
# Remove unnecessary Python cache files
36+
find /usr/local -type d -name __pycache__ -exec rm -rf {} + 2>/dev/null || true && \
37+
find /root -type d -name __pycache__ -exec rm -rf {} + 2>/dev/null || true
38+
39+
# Clean up NVIDIA CUDA libraries in both potential locations to save space
40+
RUN rm -rf /usr/local/cuda/lib64/libcudnn* \
41+
/usr/local/cuda/lib64/libnvjpeg* \
42+
/usr/local/cuda/lib64/libnvjitlink* \
43+
/usr/local/cuda/lib64/libcufft* \
44+
/usr/local/cuda/lib64/libcurand* \
45+
/usr/local/cuda/lib64/libnvgraph* \
46+
/usr/local/cuda/lib64/libnpp* \
47+
/usr/local/cuda/lib64/libnvrtc* \
48+
/usr/local/cuda/lib64/libnvToolsExt* \
49+
/usr/local/cuda/lib64/libcupti* 2>/dev/null || true && \
50+
# Keep only essential CUDA libraries
51+
mkdir -p /tmp/essential_libs && \
52+
# Save essential libraries before cleanup
53+
cp -a /root/miniconda3/envs/py3.11/lib/python3.11/site-packages/nvidia/cublas/lib/libcublas.so* /tmp/essential_libs/ 2>/dev/null || true && \
54+
cp -a /root/miniconda3/envs/py3.11/lib/python3.11/site-packages/nvidia/cusolver/lib/libcusolver.so* /tmp/essential_libs/ 2>/dev/null || true && \
55+
# Clean up NVIDIA libraries in miniconda environment
56+
rm -rf /root/miniconda3/envs/py3.11/lib/python3.11/site-packages/nvidia/*/lib/* 2>/dev/null || true && \
57+
# Restore essential libraries
58+
mkdir -p /root/miniconda3/envs/py3.11/lib/python3.11/site-packages/nvidia/cublas/lib/ 2>/dev/null || true && \
59+
mkdir -p /root/miniconda3/envs/py3.11/lib/python3.11/site-packages/nvidia/cusolver/lib/ 2>/dev/null || true && \
60+
cp -a /tmp/essential_libs/libcublas.so* /root/miniconda3/envs/py3.11/lib/python3.11/site-packages/nvidia/cublas/lib/ 2>/dev/null || true && \
61+
cp -a /tmp/essential_libs/libcusolver.so* /root/miniconda3/envs/py3.11/lib/python3.11/site-packages/nvidia/cusolver/lib/ 2>/dev/null || true && \
62+
rm -rf /tmp/essential_libs
3563

3664
# Environment settings
3765
ARG BASE_VOLUME="/runpod-volume"
38-
ENV BASE_VOLUME=$BASE_VOLUME
39-
ENV HF_DATASETS_CACHE="${BASE_VOLUME}/huggingface-cache/datasets"
40-
ENV HUGGINGFACE_HUB_CACHE="${BASE_VOLUME}/huggingface-cache/hub"
41-
ENV TRANSFORMERS_CACHE="${BASE_VOLUME}/huggingface-cache/hub"
42-
66+
ENV BASE_VOLUME=$BASE_VOLUME \
67+
HF_DATASETS_CACHE="${BASE_VOLUME}/huggingface-cache/datasets" \
68+
HUGGINGFACE_HUB_CACHE="${BASE_VOLUME}/huggingface-cache/hub" \
69+
TRANSFORMERS_CACHE="${BASE_VOLUME}/huggingface-cache/hub"
4370

4471
# Add src files (Worker Template)
4572
COPY src /src
4673

47-
CMD ["python3", "/src/handler.py"]
74+
# Copy test_input.json if it exists (for RunPod SDK automated testing)
75+
COPY test_input.json /
76+
77+
CMD ["python3", "/src/handler.py"]

‎builder/requirements.txt‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,11 @@
44
# You can also install packages from a git repository, e.g.:
55
# git+https://github.com/runpod/runpod-python.git
66
# To learn more, see https://pip.pypa.io/en/stable/reference/requirements-file-format/
7-
runpod~=1.7.0
7+
runpod~=1.7.13
88
huggingface_hub
99
typing-extensions
1010
pydantic
1111
pydantic-settings
1212
hf-transfer
1313
setuptools
14-
numpy==2.0.0
14+
numpy==2.0.0

‎test_input.json‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
{
2+
"input": {
3+
"user_id": "user",
4+
"model_id": "llama-test",
5+
"run_id": "",
6+
"credentials": {
7+
"wandb_api_key": "",
8+
"hf_token": ""
9+
},
10+
"args": {
11+
"base_model": "NousResearch/Meta-Llama-3-8B",
12+
"model_type": "LlamaForCausalLM",
13+
"tokenizer_type": "AutoTokenizer",
14+
"load_in_8bit": true,
15+
"load_in_4bit": false,
16+
"strict": false,
17+
"datasets": [
18+
{
19+
"path": "mhenrichsen/alpaca_2k_test",
20+
"type": "alpaca"
21+
}
22+
],
23+
"val_set_size": 0.05,
24+
"output_dir": "./outputs/lora-out",
25+
"sequence_len": 4096,
26+
"sample_packing": true,
27+
"eval_sample_packing": false,
28+
"pad_to_sequence_len": true,
29+
"adapter": "lora",
30+
"lora_r": 32,
31+
"lora_alpha": 16,
32+
"lora_dropout": 0.05,
33+
"lora_target_linear": true,
34+
"lora_modules_to_save": [
35+
"embed_tokens",
36+
"lm_head"
37+
],
38+
"gradient_accumulation_steps": 4,
39+
"micro_batch_size": 2,
40+
"num_epochs": 4,
41+
"optimizer": "adamw_bnb_8bit",
42+
"lr_scheduler": "cosine",
43+
"learning_rate": 0.0002,
44+
"train_on_inputs": false,
45+
"group_by_length": false,
46+
"bf16": "auto",
47+
"tf32": false,
48+
"gradient_checkpointing": true,
49+
"logging_steps": 1,
50+
"flash_attention": true,
51+
"warmup_steps": 10,
52+
"evals_per_epoch": 4,
53+
"eval_max_new_tokens": 128,
54+
"saves_per_epoch": 1,
55+
"weight_decay": 0.0,
56+
"special_tokens": {
57+
"pad_token": "<|end_of_text|>"
58+
}
59+
}
60+
}
61+
}

0 commit comments

Comments
 (0)