# SageMaker: stage 1 recipe. Believed to be the source of the prebuilt
# cosmos-predict2-base ECR image that Dockerfile_271-2stage FROMs (the
# ingredients line up), but there is no script in this repo that builds,
# tags, and pushes that image -- it was done manually by the original
# author. There is also no provenance label on the ECR image, so the link
# is convention only. Not invoked per-run.

# Use latest Pytorch 2.7, Python 3.12 image to align with the original docker
ARG AWS_REGION="us-west-2"
FROM 763104351884.dkr.ecr.${AWS_REGION}.amazonaws.com/pytorch-training:2.7.1-gpu-py312-cu128-ubuntu22.04-sagemaker

############# Hot fixes to install things properly
############# To be updated when SageMaker base image becomes more stable
# Set PYTHON and PATH explicitly
ENV PATH="/usr/local/bin:$PATH" \
    PYTHONPATH="/usr/local/lib/python3.12/site-packages"
RUN rm /usr/local/lib/python3.12/site-packages/typing.py
RUN pip install --upgrade setuptools wheel ninja

# Install apex 0.1
RUN git clone https://github.com/NVIDIA/apex /opt/apex && cd /opt/apex/ && \
    NVCC_APPEND_FLAGS="--threads=4" \
    python setup.py install --cpp_ext --cuda_ext --parallel 8



# Install Flash Attention 3
# NOTE: SM docker installs under site-packages, and --no-build-isolation is needed
RUN MAX_JOBS=$(( $(nproc) / 4 )) pip install git+https://github.com/Dao-AILab/flash-attention.git@27f501d#subdirectory=hopper --no-build-isolation
COPY cosmos_predict2/utils/flash_attn_3/flash_attn_interface.py /usr/local/lib/python3.12/site-packages/flash_attn_3/flash_attn_interface.py

# pin the exact version of transformer_engine
# RUN pip install transformer_engine==2.2.0+c55e425
RUN pip install git+https://github.com/NVIDIA/TransformerEngine.git@c55e425 --no-build-isolation
COPY cosmos_predict2/utils/flash_attn_3/te_attn.diff /tmp/te_attn.diff
# RUN patch /usr/local/lib/python3.12/site-packages/transformer_engine/pytorch/attention.py /tmp/te_attn.diff

RUN pip install "git+https://github.com/facebookresearch/pytorch3d.git@stable" --no-build-isolation


################## from install.sh
RUN pip install lovely_tensors \
                wandb==0.16.6 \
                lpips

RUN curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "awscliv2.zip" && \
    unzip awscliv2.zip && ./aws/install

RUN pip install "git+https://github.com/openai/CLIP.git"

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