Docker安装JAX镜像体积过大的原因及优化方法咨询

尝试通过Docker使用pip install安装JAX后，镜像体积膨胀至4.82GB。已使用--no-cache-dir参数绕过缓存安装，但体积仍不合理。

现有Dockerfile #

FROM ubuntu:22.04

WORKDIR /app

RUN apt-get update && apt-get install -y \
    libosmesa6-dev \
    sudo \
    wget \
    curl \
    unzip \
    gcc \
    g++

ENV PATH="/root/miniconda3/bin:${PATH}"
ARG PATH="/root/miniconda3/bin:${PATH}"

RUN mkdir -p ~/miniconda3
RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda3/miniconda.sh
RUN bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3
RUN rm -rf ~/miniconda3/miniconda.sh
RUN ~/miniconda3/bin/conda init bash
RUN conda init

RUN pip install --no-cache-dir --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

构建及镜像信息 #

构建命令：

docker build -t tbd_jax .

镜像大小：

REPOSITORY   TAG       IMAGE ID       CREATED          SIZE
tbd_jax      latest    812292e2264e   7 minutes ago    4.82GB

关键层体积：

SIZE      COMMENT
sha256:812292e2264e4340b7715956824055d7409f9546f8dfa54ccad1da056febf300   8 minutes ago    RUN |1 PATH=/root/miniconda3/bin:/root/miniconda3/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin /bin/sh -c pip install --no-cache-dir --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html # buildkit   3.54GB    buildkit.dockerfile.v0

1. 使用多阶段构建分离依赖与运行环境 #

多阶段构建可以把编译、安装依赖的过程放在临时镜像里，只把最终需要的运行文件复制到目标镜像，大幅减少体积。示例：

# 第一阶段：构建依赖
FROM ubuntu:22.04 as builder

WORKDIR /app

RUN apt-get update && apt-get install -y --no-install-recommends \
    libosmesa6-dev \
    wget \
    curl \
    gcc \
    g++ \
    && apt-get clean \
    && rm -rf /var/lib/apt/lists/*

RUN mkdir -p ~/miniconda3
RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda3/miniconda.sh
RUN bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3
RUN rm -rf ~/miniconda3/miniconda.sh

ENV PATH="/root/miniconda3/bin:${PATH}"
RUN pip install --no-cache-dir --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
RUN conda clean -afy && rm -rf /root/miniconda3/pkgs/*

# 第二阶段：生成轻量运行镜像
FROM ubuntu:22.04

WORKDIR /app

# 从构建阶段复制必要文件
COPY --from=builder /root/miniconda3 /root/miniconda3
COPY --from=builder /usr/lib/x86_64-linux-gnu/libosmesa.so.6 /usr/lib/x86_64-linux-gnu/

ENV PATH="/root/miniconda3/bin:${PATH}"

RUN apt-get clean && rm -rf /var/lib/apt/lists/*

2. 精简系统依赖安装 #

• 移除不需要的包：比如sudo、unzip如果后续运行用不到，直接删除；
• 添加--no-install-recommends参数，避免安装额外推荐依赖；
• 安装后立即清理apt缓存，且和安装命令放在同一个RUN层，避免缓存残留。
  修改后的系统安装命令：

RUN apt-get update && apt-get install -y --no-install-recommends \
    libosmesa6-dev \
    wget \
    curl \
    gcc \
    g++ \
    && apt-get clean \
    && rm -rf /var/lib/apt/lists/*

3. 替换基础镜像为轻量版本 #

• 直接使用conda/miniconda3:latest作为基础镜像，省去手动安装Miniconda的步骤，官方镜像已做体积优化；
• 或者用Python官方slim镜像（如python:3.10-slim）替代Ubuntu+Miniconda，避免Miniconda自带的冗余包：

FROM python:3.10-slim

WORKDIR /app

RUN apt-get update && apt-get install -y --no-install-recommends \
    libosmesa6-dev \
    gcc \
    g++ \
    && apt-get clean \
    && rm -rf /var/lib/apt/lists/*

RUN pip install --no-cache-dir --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

4. 清理conda冗余文件 #

在pip安装完成后，添加conda缓存清理命令，删除下载的包缓存：

RUN conda clean -afy && rm -rf /root/miniconda3/pkgs/*

内容的提问来源于stack exchange，提问作者desert_ranger