# SageMaker: 1-stage standalone alternative to (Dockerfile_base + Dockerfile_271-2stage).
# Self-contained build from the AWS DLC; no dependency on cosmos-predict2-base.
# Not used by Basile's current flow (which uses 271-2stage). Adds FA4 (--pre).
# Kept for reference; moved to misc/ to avoid confusion with the active path.

# 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

################## from original Dockerfile

# Install the dependencies from requirements-docker.txt
COPY ./requirements-docker.txt /requirements.txt
RUN pip install --no-cache-dir -r /requirements.txt
# 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

# Install Flash Attention 4
# NOTE(dc): remove "--pre" when the pip release becomes stable
RUN pip install --pre flash-attn-4

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

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


################## other installs, envs
ENV PYTHONPATH=/opt/ml/code/externals/vidar/externals/cv_datasets/:${PYTHONPATH}
ENV PYTHONPATH=/opt/ml/code/externals/vidar/externals/webdataset/:${PYTHONPATH}
ENV PYTHONPATH=/opt/ml/code/externals/vidar/:${PYTHONPATH}
ENV PYTHONPATH=/opt/ml/code/externals/cv_datasets/:${PYTHONPATH}
ENV PYTHONPATH=/opt/ml/code/externals/vggt/:${PYTHONPATH}

ENV SAGEMAKER_SUBMIT_DIRECTORY /opt/ml/code

COPY . ${SAGEMAKER_SUBMIT_DIRECTORY}

## Set other environment variables
ENV TORCH_HOME="/opt/ml/code/torchhub"
RUN git config --global --add safe.directory ${SAGEMAKER_SUBMIT_DIRECTORY}