# SageMaker: stage 2 (active path, used by Basile via run_sm.sh "271-2stage").
# FROMs cosmos-predict2-base:latest, which is built by Dockerfile_base alongside.

# Use latest Pytorch 2.7, Python 3.12 image to align with the original docker
ARG AWS_REGION="us-west-2"
FROM 124224456861.dkr.ecr.us-west-2.amazonaws.com/cosmos-predict2-base:latest

# Install the dependencies from requirements-docker.txt
COPY ./requirements-docker.txt /requirements.txt
RUN pip install --no-cache-dir -r /requirements.txt

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

################## other installs, envs
ENV PYTHONPATH=/opt/ml/code/externals/AnyData/:${PYTHONPATH}
ENV PYTHONPATH=/opt/ml/code/externals/AnyData/externals/webdataset/:${PYTHONPATH}
# NOTE(bvh): vidar etc removed, only inactive code paths import from these libraries now
# 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 SAGEMAKER_SUBMIT_DIRECTORY /opt/ml/code

COPY . ${SAGEMAKER_SUBMIT_DIRECTORY}

################## AnyData (replaces vidar/cv_webdatasets)
RUN pip install --no-cache-dir -e ${SAGEMAKER_SUBMIT_DIRECTORY}/externals/AnyData/ \
 && pip install --no-cache-dir -e ${SAGEMAKER_SUBMIT_DIRECTORY}/externals/AnyData/externals/webdataset/ \
 && pip install --no-cache-dir numpy==1.26.4
# ^ for compatibility

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