# Use PyTorch 2.9 + Python 3.12 to avoid py313 runtime linkage issues with torchcodec.
ARG AWS_REGION="us-west-2"
FROM public.ecr.aws/deep-learning-containers/pytorch-training:2.9.0-gpu-py312-cu130-ubuntu22.04-sagemaker-v1.9
############# 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 -f /usr/local/lib/python3.12/site-packages/typing.py
RUN pip install --upgrade setuptools wheel ninja
RUN pip install --no-cache-dir uv


# Install runtime system deps early so torchcodec linkage can be tested first.
RUN apt-get update && apt-get install -y curl tar unzip ffmpeg && \
    rm -rf /var/lib/apt/lists/*

ENV SAGEMAKER_SUBMIT_DIRECTORY /opt/ml/code

WORKDIR ${SAGEMAKER_SUBMIT_DIRECTORY}
COPY pyproject.toml uv.lock ${SAGEMAKER_SUBMIT_DIRECTORY}/
COPY externals/webdataset ${SAGEMAKER_SUBMIT_DIRECTORY}/externals/webdataset
COPY externals/camviz_internal/pyproject.toml ${SAGEMAKER_SUBMIT_DIRECTORY}/externals/camviz_internal/pyproject.toml

# torchcodec is not required for webbing job now due to imcompatibility with torch 2.9.0, but we install it here to avoid linkage issues with torch 2.9.0 in the future when we need to use torchcodec again.
RUN cd ${SAGEMAKER_SUBMIT_DIRECTORY} && \
    uv pip install --system --index-strategy unsafe-best-match --extra-index-url https://download.pytorch.org/whl/cu130 \
        -r pyproject.toml "torch==2.9.0" "torchvision==0.24.0" "torchcodec==0.8.1"
ENV LD_LIBRARY_PATH=/usr/local/lib/python3.12/site-packages/torch/lib:/usr/local/lib:/usr/lib/x86_64-linux-gnu:${LD_LIBRARY_PATH}
RUN ldconfig

# 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

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

# Install s5cmd into /tmp and prepend to PATH (safe for SageMaker runtime)
RUN curl -L https://github.com/peak/s5cmd/releases/download/v2.3.0/s5cmd_2.3.0_Linux-64bit.tar.gz | tar -xz -C /usr/local/bin && \
    chmod +x /usr/local/bin/s5cmd

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

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

################## from original Dockerfile
RUN pip install --no-cache-dir \
    "boto3<1.35.0" \
    "botocore<1.35.0" \
    "s3fs==2024.6.0" \
    "aiobotocore==2.13.0"
RUN pip install --no-cache-dir "ffmpeg-python" "zarr"
RUN python -c "import s3fs, boto3, ffmpeg, zarr; print('✅ s3fs, boto3, ffmpeg, zarr loaded OK')"


# RUN apt-get update && apt-get install -y curl tar && \
#     curl -L https://github.com/peak/s5cmd/releases/download/v2.3.0/s5cmd_2.3.0_Linux-64bit.tar.gz \
#     | tar -xz -C /usr/local/bin && \
#     mv /usr/local/bin/s5cmd_2.3.0_Linux-64bit/s5cmd /usr/local/bin/s5cmd && \
#     rm -rf /usr/local/bin/s5cmd_2.3.0_Linux-64bit && \
#     chmod +x /usr/local/bin/s5cmd


################## other installs, envs
COPY . ${SAGEMAKER_SUBMIT_DIRECTORY}
RUN cd ${SAGEMAKER_SUBMIT_DIRECTORY} && \
    uv pip install --system --no-deps .

ENV PYTHONPATH=/opt/ml/code/externals/cv_datasets/:${PYTHONPATH}
ENV PYTHONPATH=/opt/ml/code/externals/webdataset/:${PYTHONPATH}
ENV PYTHONPATH=/opt/ml/code/externals/camviz/:${PYTHONPATH}

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