FROM nvidia/cuda:12.2.2-cudnn8-runtime-ubuntu22.04
COPY --from=ghcr.io/astral-sh/uv:0.5.1 /uv /uvx /bin/
WORKDIR /app

RUN apt-get update && apt-get install -y --no-install-recommends \
    python3.11 python3.11-venv python3.11-dev \
    git git-lfs linux-headers-generic build-essential clang && \
    git lfs install && rm -rf /var/lib/apt/lists/*

ENV UV_LINK_MODE=copy UV_PROJECT_ENVIRONMENT=/.venv UV_PYTHON_DOWNLOADS=never
RUN uv venv --python python3.11 $UV_PROJECT_ENVIRONMENT

COPY pyproject.toml uv.lock LICENSE README.md ./
COPY packages/ packages/
RUN --mount=type=cache,target=/root/.cache/uv \
    GIT_LFS_SKIP_SMUDGE=1 uv sync --frozen --no-install-project --no-dev

COPY src/ src/
COPY scripts/ scripts/
COPY examples/ examples/
RUN --mount=type=cache,target=/root/.cache/uv \
    GIT_LFS_SKIP_SMUDGE=1 uv sync --frozen --no-dev

# Checkpoint base
ARG C=checkpoints/acot_icra_simulation_challenge_reasoning_to_action/jubail_h200_4gpu/6000
ARG D=${C}/params/ocdbt.process_0/d

# Small checkpoint files (one layer)
COPY ${C}/assets ${C}/assets
COPY ${C}/_CHECKPOINT_METADATA ${C}/_CHECKPOINT_METADATA
COPY ${C}/params/array_metadatas ${C}/params/array_metadatas
COPY ${C}/params/d ${C}/params/d
COPY ${C}/params/manifest.ocdbt ${C}/params/manifest.ocdbt
COPY ${C}/params/_METADATA ${C}/params/_METADATA
COPY ${C}/params/_sharding ${C}/params/_sharding
COPY ${C}/params/ocdbt.process_0/manifest.ocdbt ${C}/params/ocdbt.process_0/manifest.ocdbt

# Medium files (~910M + 60M + tiny files) - one layer
COPY ${D}/621357ee5ce63185e775b10737f613cd ${D}/621357ee5ce63185e775b10737f613cd
COPY ${D}/af8bb8e6f022b1aa60cf898d57ce2786 ${D}/af8bb8e6f022b1aa60cf898d57ce2786
COPY ${D}/194726844871162a68fd3706fe5df322 ${D}/
COPY ${D}/3b8f638975b6cd0bfe076a67be0bda33 ${D}/
COPY ${D}/3f4cfeae67046e5f1bd969fbf71c9eea ${D}/
COPY ${D}/670a778fa8fb810afa51a5bbf7d47236 ${D}/
COPY ${D}/97700388192cfdf53731d9d98084d23b ${D}/
COPY ${D}/9f6f221cd61ff8aaad9f8934ac1be593 ${D}/
COPY ${D}/b0a266dda268c4e59baaad58ee232f34 ${D}/
COPY ${D}/ba3c7716a6f6ca414cf75ff226a91102 ${D}/
COPY ${D}/d79abe13eed95cab9c9fc2734d009a7e ${D}/

# Big file 1 chunks (2.1G -> 5x500M) - each a separate layer
COPY ${D}/split/224973dd319b9f0cd7268554fc312684_00 ${D}/split/224973dd319b9f0cd7268554fc312684_00
COPY ${D}/split/224973dd319b9f0cd7268554fc312684_01 ${D}/split/224973dd319b9f0cd7268554fc312684_01
COPY ${D}/split/224973dd319b9f0cd7268554fc312684_02 ${D}/split/224973dd319b9f0cd7268554fc312684_02
COPY ${D}/split/224973dd319b9f0cd7268554fc312684_03 ${D}/split/224973dd319b9f0cd7268554fc312684_03
COPY ${D}/split/224973dd319b9f0cd7268554fc312684_04 ${D}/split/224973dd319b9f0cd7268554fc312684_04

# Big file 2 chunks (2.1G -> 5x500M)
COPY ${D}/split/896b601176d4e8371cc79fc6271a7a27_00 ${D}/split/896b601176d4e8371cc79fc6271a7a27_00
COPY ${D}/split/896b601176d4e8371cc79fc6271a7a27_01 ${D}/split/896b601176d4e8371cc79fc6271a7a27_01
COPY ${D}/split/896b601176d4e8371cc79fc6271a7a27_02 ${D}/split/896b601176d4e8371cc79fc6271a7a27_02
COPY ${D}/split/896b601176d4e8371cc79fc6271a7a27_03 ${D}/split/896b601176d4e8371cc79fc6271a7a27_03
COPY ${D}/split/896b601176d4e8371cc79fc6271a7a27_04 ${D}/split/896b601176d4e8371cc79fc6271a7a27_04

# Big file 3 chunks (2.5G -> 5x500M)
COPY ${D}/split/deb79b36487c5a15fa2fd64b812d6a00_00 ${D}/split/deb79b36487c5a15fa2fd64b812d6a00_00
COPY ${D}/split/deb79b36487c5a15fa2fd64b812d6a00_01 ${D}/split/deb79b36487c5a15fa2fd64b812d6a00_01
COPY ${D}/split/deb79b36487c5a15fa2fd64b812d6a00_02 ${D}/split/deb79b36487c5a15fa2fd64b812d6a00_02
COPY ${D}/split/deb79b36487c5a15fa2fd64b812d6a00_03 ${D}/split/deb79b36487c5a15fa2fd64b812d6a00_03
COPY ${D}/split/deb79b36487c5a15fa2fd64b812d6a00_04 ${D}/split/deb79b36487c5a15fa2fd64b812d6a00_04

# Reassemble original files from chunks and remove chunks
RUN cd ${D} && \
    for name in 224973dd319b9f0cd7268554fc312684 896b601176d4e8371cc79fc6271a7a27 deb79b36487c5a15fa2fd64b812d6a00; do \
        cat split/${name}_* > ${name} && rm split/${name}_*; \
    done && rmdir split

# Symlinks for default G2SIM path
RUN ln -s jubail_h200_4gpu checkpoints/acot_icra_simulation_challenge_reasoning_to_action/exp_name && \
    ln -s 6000 checkpoints/acot_icra_simulation_challenge_reasoning_to_action/jubail_h200_4gpu/30000

ENV PYTHONPATH="/app/src"
CMD ["bash", "./scripts/server.sh", "0", "8999"]