Skip to content

Commit

Permalink
add Dockerfile, an onstart.sh and some post build steps
Browse files Browse the repository at this point in the history
Add a post-build script that runs the first time you execute it on
a machine with a GPU, since some stuff won't build in the builder
  • Loading branch information
bitplane committed Dec 11, 2024
1 parent ab1b84a commit caae392
Show file tree
Hide file tree
Showing 4 changed files with 129 additions and 0 deletions.
4 changes: 4 additions & 0 deletions .dockerignore
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
.git
.env
*.pyc
Dockerfile
88 changes: 88 additions & 0 deletions Dockerfile
Original file line number Diff line number Diff line change
@@ -0,0 +1,88 @@
FROM pytorch/pytorch:2.4.0-cuda12.4-cudnn9-devel AS builder

# Install build dependencies
RUN apt-get update && \
apt-get install -y ffmpeg build-essential git rdfind

WORKDIR /app
ADD trellis /app/trellis
ADD setup.sh /app/setup.sh
ADD app.py /app/app.py
ADD example.py /app/example.py
ADD extensions /app/extensions
ADD assets /app/assets

# Setup conda
RUN conda config --set always_yes true && \
conda init && \
conda config --add channels defaults

# Use bash shell so we can source activate
SHELL ["/bin/bash", "-c"]


RUN conda install torchvision=0.19.0 \
onnx==1.17.0 \
pytorch=2.4.0 \
pytorch-cuda=12.4 \
--force-reinstall \
-c pytorch \
-c nvidia

# Create a g++ wrapper for JIT, since the include dirs are passed with -i rather than -I for some reason
RUN printf '#!/usr/bin/env bash\nexec /usr/bin/g++ -I/usr/local/cuda/include -I/usr/local/cuda/include/crt "$@"\n' > /usr/local/bin/gxx-wrapper && \
chmod +x /usr/local/bin/gxx-wrapper
ENV CXX=/usr/local/bin/gxx-wrapper


# Run setup.sh - this won't install all the things due to missing GPU in the builder
RUN conda run -n base ./setup.sh --basic --xformers --flash-attn --diffoctreerast --vox2seq --spconv --mipgaussian --kaolin --nvdiffrast --demo

# Now install additional Python packages
# These ones work inside the builder
RUN conda run -n base pip install diso
RUN conda run -n base pip install plyfile utils3d flash_attn spconv-cu120 xformers
RUN conda run -n base pip install kaolin -f https://nvidia-kaolin.s3.us-east-2.amazonaws.com/torch-2.4.0_cu121.html
RUN conda run -n base pip install git+https://github.com/NVlabs/nvdiffrast.git

# Remove downloaded packages from conda and pip
RUN conda clean --all -f -y
RUN pip cache purge

# Deduplicate with rdfind
# This reduces the size of the image by a few hundred megs.
RUN rdfind -makesymlinks true /opt/conda

# Final stage
FROM pytorch/pytorch:2.4.0-cuda12.4-cudnn9-devel AS final

WORKDIR /app
COPY --from=builder /usr/local/bin/gxx-wrapper /usr/local/bin/gxx-wrapper
COPY --from=builder /opt/conda /opt/conda
COPY --from=builder /root /root
COPY --from=builder /app /app

# Reinstall any runtime tools needed
# git and build-essential are needed for post_install.sh script.
# vim and strace are useful for debugging, remove those if you want to.
RUN apt-get update && \
apt-get install -y build-essential \
git \
strace \
vim && \
rm -rf /var/lib/apt/lists/*

# install these last, so we can experiment without excessive build times.
COPY onstart.sh /app/onstart.sh
COPY post_install.sh /app/post_install.sh

ENV PATH=/opt/conda/bin:$PATH

# This script runs the post_install steps

# If you're pushing to a container registry, let this run once, run some
# tests, then do `docker commit` to save the models along with the image.
# This will ensure that it won't fail at runtime due to models being
# unavailable, or network restrictions.
CMD ["/app/onstart.sh"]

12 changes: 12 additions & 0 deletions onstart.sh
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
#!/bin/bash

cd /app

echo "Doing post install steps"
./post_install.sh

export CXX=/usr/local/bin/gxx-wrapper

echo "Launching app"
python3 app.py

25 changes: 25 additions & 0 deletions post_install.sh
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
#!/bin/bash
set -e

# Check if post-install steps have already been run
if [ -f /app/.post_install_done ]; then
echo "Post-install steps already completed."
exit 0
fi

cd /app

echo "Install stuff that couldn't be installed without GPU"
# Run the demo setup
conda run -n base ./setup.sh --mipgaussian --diffoctreerast

echo "Proving it actually works..."

export CXX=/usr/local/bin/gxx-wrapper
python example.py

# Mark completion
touch /app/.post_install_done

echo "Post-install steps completed successfully."

0 comments on commit caae392

Please sign in to comment.