|
| 1 | +FROM nvidia/cuda:12.9.1-cudnn-devel-ubuntu22.04 |
| 2 | + |
| 3 | +ENV DEBIAN_FRONTEND=noninteractive |
| 4 | + |
| 5 | +SHELL ["/bin/bash", "-c"] |
| 6 | + |
| 7 | +WORKDIR /FastVideo |
| 8 | + |
| 9 | +RUN apt-get update && apt-get install -y --no-install-recommends \ |
| 10 | + wget \ |
| 11 | + git \ |
| 12 | + ca-certificates \ |
| 13 | + openssh-server \ |
| 14 | + zsh \ |
| 15 | + vim \ |
| 16 | + curl \ |
| 17 | + gcc-11 \ |
| 18 | + g++-11 \ |
| 19 | + clang-11 \ |
| 20 | + && rm -rf /var/lib/apt/lists/* |
| 21 | + |
| 22 | +# Set up C++20 compilers for ThunderKittens |
| 23 | +RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11 |
| 24 | + |
| 25 | +# Set CUDA environment variables |
| 26 | +ENV CUDA_HOME=/usr/local/cuda-12.9 |
| 27 | +ENV PATH=${CUDA_HOME}/bin:${PATH} |
| 28 | +ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH |
| 29 | + |
| 30 | +# Install uv and source its environment |
| 31 | +RUN curl -LsSf https://astral.sh/uv/install.sh | sh && \ |
| 32 | + echo 'source $HOME/.local/bin/env' >> /root/.bashrc |
| 33 | + |
| 34 | +# Copy just the pyproject.toml first to leverage Docker cache |
| 35 | +COPY pyproject.toml ./ |
| 36 | + |
| 37 | +# Create a dummy README to satisfy the installation |
| 38 | +RUN echo "# Placeholder" > README.md |
| 39 | + |
| 40 | +# Create and activate virtual environment with specific Python version and seed |
| 41 | +RUN source $HOME/.local/bin/env && \ |
| 42 | + uv venv --python 3.12 --seed /opt/venv && \ |
| 43 | + source /opt/venv/bin/activate && \ |
| 44 | + uv pip install --no-cache-dir --upgrade pip && \ |
| 45 | + uv pip install --no-cache-dir .[dev] && \ |
| 46 | + uv pip install --no-cache-dir flash-attn==2.8.3 --no-build-isolation |
| 47 | + |
| 48 | +COPY . . |
| 49 | + |
| 50 | +# Install dependencies using uv and set up shell configuration |
| 51 | +RUN source $HOME/.local/bin/env && \ |
| 52 | + source /opt/venv/bin/activate && \ |
| 53 | + uv pip install --no-cache-dir -e .[dev] && \ |
| 54 | + git config --unset-all http.https://github.com/.extraheader || true && \ |
| 55 | + echo 'source /opt/venv/bin/activate' >> /root/.bashrc && \ |
| 56 | + echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile |
| 57 | + |
| 58 | +# Install STA (Sliding Tile Attention) |
| 59 | +RUN source $HOME/.local/bin/env && \ |
| 60 | + source /opt/venv/bin/activate && \ |
| 61 | + cd csrc/attn && \ |
| 62 | + git submodule update --init --recursive && \ |
| 63 | + python setup_sta.py install |
| 64 | + |
| 65 | +# Install VSA |
| 66 | +RUN source $HOME/.local/bin/env && \ |
| 67 | + source /opt/venv/bin/activate && \ |
| 68 | + cd csrc/attn && \ |
| 69 | + git submodule update --init --recursive && \ |
| 70 | + python setup_vsa.py install |
| 71 | + |
| 72 | +EXPOSE 22 |
0 commit comments