How to Efficiently Copy Detectron2 Build Files in a Multi-Stage Docker Build?

Question

I'm working on a application that uses LayoutMV2 model, which uses Facebook AI’s Detectron2 package for its visual backbone. Both of these dependencies, as well as the application itself, require torch.

Both Detectron2 and the LayoutMV2 require torch as well as the application I'm working on itself.

My goal is to have a container image in which the application can run a training using a Nvidia graphics card and CUDA.

Detectron2 does not have pre-built wheels for the newer versions of CUDA and Pytorch, so I need to build Detectron2 from source. See this link for installation instructions.

To allow Detectron2 to make use of CUDA, it has to be build in an environment where CUDA development tools are present. For this the nvidia/cuda:11.7.1-cudnn8-devel-ubuntu20.04 image provided by nvidia can be used. However, my application does not require these development tools, as the necessary CUDA libraries that my application requires are already bundled with PyTorch. This results in a large application image size, as the CUDA development tools are not removed after the Detectron2 build.

The documentation for these images states that the following about the devel tagged images: "These images are particularly useful for multi-stage builds."

Okay, so I can use a multi-stage build to build Detectron2 in the devel image and then copy the necessary files to a smaller image. However, I'm not sure what files I need to copy and where to.

When I build Detectron2 from source using the devel image, the following files are created inside the cloned repository:

====================================================================================================
 189 MB      $ pip install --user -e detectron2_repo # buildkit
====================================================================================================
  48 MB      home/appuser/detectron2_repo/build/temp.linux-x86_64-3.8/home/appuser/detectron2_repo
  25 MB      home/appuser/detectron2_repo/build/lib.linux-x86_64-3.8/detectron2/_C.cpython-38-x86_64-linux-gnu.so
  25 MB      home/appuser/detectron2_repo/detectron2/_C.cpython-38-x86_64-linux-gnu.so
 521 kB      home/appuser/detectron2_repo/build/temp.linux-x86_64-3.8/.ninja_deps

(Created using dlayer)

Would it make sense to copy the detectron2_repo directory to the smaller image? Or should should I build a wheel and copy that to the smaller image? How would I go about doing that?

I would appreciate any guidance on the best approach to take for the multi-stage build and what specific files should be copied to the smaller image.

Answer 1

This is what I have come up with.

The torch.txt and detectron2.txt I generate using pip-compile from pip-tools.

requirements/base.in

# Main project dependencies go in here. 

requirements/torch.in

-c base.txt
torch==1.13.1
torchvision==0.14.1

requirements/detectron2.in

-r torch.txt
-c base.txt
detectron2@git+https://github.com/facebookresearch/[email protected]

The Dockerfile

FROM nvidia/cuda:11.7.1-cudnn8-devel-ubuntu20.04 as detectron2

ENV DEBIAN_FRONTEND=noninteractive

ARG PYTHON_VERSION="3.8"
ARG DETECTRON2_TAG="v0.6"
# must match CUDA version in base image
ENV PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cu117" \
    PIP_NO_CACHE_DIR="1"

# Set FORCE_CUDA because during `docker build` Cuda is not accessible
ARG TORCH_CUDA_ARCH_LIST="Kepler;Kepler+Tesla;Maxwell;Maxwell+Tegra;Pascal;Volta;Turing"

ENV TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST}" \
    FORCE_CUDA="1"

# Install system dependencies
RUN apt-get update -y && \
    apt-get -y install --no-install-recommends ca-certificates \
    software-properties-common \
    git \
    curl \
    wget \
    ninja-build \
    cmake \
    "python${PYTHON_VERSION}" \
    "python${PYTHON_VERSION}-dev" \
    "python${PYTHON_VERSION}-venv" \
    "python3-pip" \
    "python3-opencv" \
    "python-is-python3" \
    && rm -rf /var/lib/apt/lists/*

# Create non-root user
ARG UID=1000
ARG APP_USER=app2
RUN useradd -m --no-log-init --system ${APP_USER}

# Use non-root user
USER $APP_USER
WORKDIR /home/$APP_USER
ENV PATH="/home/$APP_USER/.local/bin:${PATH}"

# Clone detectron2 repository
RUN git clone --depth 1 --branch $DETECTRON2_TAG https://github.com/facebookresearch/detectron2 detectron2_repo

# Create and activate virtual environment
RUN python -m venv detectron2_venv
ENV VIRTUAL_ENV=/home/$APP_USER/detectron2_venv
ENV PATH="$VIRTUAL_ENV/bin:${PATH}"

# Install dependency management tools
RUN pip install pip-tools

# Install dependencies
COPY requirements/torch.txt requirements/
RUN pip-sync requirements/torch.txt

# Brrrr
RUN cd detectron2_repo && python setup.py bdist_wheel

FROM scratch
# TODO: push to registry
COPY --from=detectron2 /home/app2/detectron2_repo/dist/*.whl /