####
# Runtime Image 
####
# ARG AWS_REGION=us-east-1
# FROM 763104351884.dkr.ecr.${AWS_REGION}.amazonaws.com/base:12.9.1-gpu-py312-cu129-ubuntu22.04-ec2
FROM ubuntu:22.04
# COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/
COPY --from=ghcr.io/astral-sh/uv:0.9.25 /uv /uvx /bin/



ENV UV_LINK_MODE=copy
ENV UV_PROJECT_ENVIRONMENT=/.venv
ENV PATH="/.venv/bin:$PATH"

# Before running pip install
RUN apt-get update && apt-get install -y \
    libgl1 \
    libglib2.0-0 \
    build-essential git \
 && rm -rf /var/lib/apt/lists/*


COPY sagemaker/requirements.txt /opt/ml/code/requirements.txt
RUN uv venv --python 3.11 $UV_PROJECT_ENVIRONMENT
RUN uv pip install sagemaker-training
RUN uv pip install -r /opt/ml/code/requirements.txt
RUN uv pip install pytest
RUN uv cache clean

COPY src/openpi /opt/ml/code/openpi
WORKDIR /opt/ml/code/openpi
ENV PYTHONPATH=/opt/ml/code/openpi/src:/opt/ml/code/openpi/packages/openpi-client/src
COPY sagemaker/train.sh /opt/ml/code/train.sh

RUN apt-get update && apt-get install -y curl unzip && rm -rf /var/lib/apt/lists/*
RUN curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "awscliv2.zip"
RUN unzip awscliv2.zip
RUN ./aws/install
RUN rm -rf ./aws awscliv2.zip

CMD ["/bin/bash"]
# ENTRYPOINT ["/bin/bash", "/opt/ml/code/sagemaker/launch.sh"]