orign 0.2.3

FROM pytorch/pytorch:2.5.1-cuda12.4-cudnn9-devel

RUN apt-get update && apt-get install -y \
    wget \
    curl \
    build-essential \
    libopenmpi-dev \
    openmpi-bin \
    git

RUN pip install --no-cache-dir --upgrade pip

ARG CACHE_BUST=unknown
RUN pip install ms-swift wandb deepspeed qwen-vl-utils torchvision pyav mpi4py
RUN pip install transformers==4.49.0

RUN apt-get install -y curl unzip

# Install AWS CLI
RUN curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "awscliv2.zip" && \
    unzip awscliv2.zip && \
    ./aws/install && \
    rm -rf awscliv2.zip aws

RUN mkdir -p /workspace/output

RUN wget https://storage.googleapis.com/orign/releases/latest/orign-latest-linux-amd64.tar.gz && \
    tar -xvzf orign-latest-linux-amd64.tar.gz && \
    chmod +x orign && \
    ./orign --version || exit 1

# Install Mountpoint for Amazon S3 dependencies
RUN apt-get update && apt-get install -y \
    fuse3 \
    libfuse3-dev \
    golang \
    ca-certificates

# Install Mountpoint for Amazon S3
ARG S3_VERSION=latest
ARG ARCH=x86_64
RUN wget https://s3.amazonaws.com/mountpoint-s3-release/${S3_VERSION}/${ARCH}/mount-s3.deb && \
    apt-get update && \
    apt-get install -y ./mount-s3.deb && \
    rm mount-s3.deb

# Create S3 mount point and cache directory
RUN mkdir -p /s3 && chmod 755 /s3 && \
    mkdir -p /cache && chmod 777 /cache

# Copy mounthelper.go and build it
WORKDIR /tmp
COPY mounthelper.go .
RUN go build -o /usr/local/bin/mounthelper mounthelper.go && \
    chmod +x /usr/local/bin/mounthelper

# Copy the entrypoint script
COPY entrypoint.sh /usr/local/bin/
RUN chmod +x /usr/local/bin/entrypoint.sh

# Switch back to workspace directory
WORKDIR /workspace

# Set entrypoint to our startup script
ENTRYPOINT ["/usr/local/bin/entrypoint.sh"]
CMD ["bash"]