FROM nvidia/cuda:12.9.1-devel-ubuntu24.04
ENV DEBIAN_FRONTEND=noninteractive
RUN apt-get update && apt-get install -y --no-install-recommends python3 python3-dev python3-venv git ca-certificates build-essential
RUN mkdir -p /home/dev/pretrain_suffix_min/vendor && python3 -m venv /home/dev/venvs/main
ENV PATH=/home/dev/venvs/main/bin:$PATH
ENV HF_HOME=/home/dev/.cache/huggingface
ENV TRITON_CACHE_DIR=/home/dev/.cache/triton
ENV TORCHINDUCTOR_CACHE_DIR=/home/dev/.cache/torchinductor
RUN pip install --no-cache-dir setuptools wheel packaging
RUN pip install --no-cache-dir torch==2.13.0+cu129 --index-url https://download.pytorch.org/whl/cu129
RUN git clone https://github.com/fla-org/flash-linear-attention /home/dev/pretrain_suffix_min/vendor/flash-linear-attention && git -C /home/dev/pretrain_suffix_min/vendor/flash-linear-attention checkout 7b5138a92133132cfc5f80455c9c0920d4fe0278
RUN git clone https://github.com/EleutherAI/lm-evaluation-harness /home/dev/pretrain_suffix_min/vendor/lm-evaluation-harness && git -C /home/dev/pretrain_suffix_min/vendor/lm-evaluation-harness checkout d6de81643928d653435c431bae19945d41d32520
RUN pip install --no-cache-dir --no-build-isolation -e /home/dev/pretrain_suffix_min/vendor/flash-linear-attention -e /home/dev/pretrain_suffix_min/vendor/lm-evaluation-harness transformers==4.57.6 datasets==5.0.1 numpy==2.5.3 pytest==9.1.1 matplotlib==3.11.2
# Keep upstream source and its license; build against this image's exact torch/CUDA.
ENV MAX_JOBS=8
ENV MAMBA_FORCE_BUILD=TRUE
ENV CAUSAL_CONV1D_FORCE_BUILD=TRUE
ENV TORCH_CUDA_ARCH_LIST=9.0
RUN pip install --no-cache-dir ninja einops && pip install --no-cache-dir --no-build-isolation causal-conv1d==1.5.2
RUN git clone https://github.com/state-spaces/mamba /home/dev/pretrain_mamba2_outer_vendor && git -C /home/dev/pretrain_mamba2_outer_vendor checkout d7b1ceb3c367ec9022925e812f507bcf706937c6
RUN pip install --no-cache-dir --no-build-isolation -e /home/dev/pretrain_mamba2_outer_vendor
RUN mkdir -p /home/dev/pretrain_mamba2_outer
COPY source.tar.gz /home/dev/source.tar.gz
RUN tar -xzf /home/dev/source.tar.gz -C /home/dev/pretrain_mamba2_outer
WORKDIR /home/dev/pretrain_mamba2_outer
ENV PYTHONPATH=/home/dev/pretrain_mamba2_outer
ENV FLA_CI_ENV=0
ENV CUBLAS_WORKSPACE_CONFIG=:4096:8
CMD ["python", "-u", "-m", "pretraining_mamba2_outer.debug_gate", "--out", "/home/dev/mamba2-debug"]