# syntax=docker/dockerfile:1.7
FROM python:3.12-slim

ENV PYTHONUNBUFFERED=1 \
    PIP_NO_CACHE_DIR=1 \
    PIP_DISABLE_PIP_VERSION_CHECK=1 \
    UV_HTTP_TIMEOUT=600 \
    UV_CONCURRENT_DOWNLOADS=1 \
    UV_LINK_MODE=copy \
    PIP_DEFAULT_TIMEOUT=600

WORKDIR /app

RUN apt-get update && \
    apt-get install -y --no-install-recommends libgl1 libglib2.0-0 git && \
    rm -rf /var/lib/apt/lists/*

RUN pip install --upgrade pip setuptools wheel uv && \
    printf '%s\n' \
        'pynput>=0.0.0; sys_platform == "never"' \
        'evdev>=0.0.0; sys_platform == "never"' \
        'transformers==4.53.2' \
        'huggingface-hub>=0.30.0,<1.0' \
        'torch==2.8.0' \
        'torchvision==0.23.0' \
        > /tmp/overrides.txt

# Combined JAX + PyTorch + LeRobot dependency set, pinned to the versions the
# ARX5 openpi runtime was validated against (minus the ARX5-only packages).
# LeRobot pulls feetech motor support for the SO-101. Keep it one layer so
# Docker doesn't upload every intermediate Torch/CUDA replacement layer.
RUN --mount=type=cache,target=/root/.cache/uv \
    uv pip install --system \
        --override /tmp/overrides.txt \
        --extra-index-url https://download.pytorch.org/whl/cu129 \
        --index-strategy unsafe-best-match \
        'transformers[vision]==4.53.2' \
        accelerate diffusers einops 'huggingface_hub>=0.30.0,<1.0' numpy pandas pyarrow scipy \
        opencv-python-headless safetensors \
        augmax dm-tree equinox flatbuffers flax==0.10.2 'fsspec[gcs]>=2024.6.0' \
        imageio 'jaxlib==0.5.3' 'jax[cuda12]==0.5.3' jaxtyping==0.2.36 ml_collections==1.0.0 \
        'numpydantic>=1.6.6' orbax-checkpoint==0.11.13 'sentencepiece>=0.2.0' \
        'tqdm-loggable>=0.2' 'tyro>=0.9.5' 'wandb>=0.19.1' 'filelock>=3.16.1' \
        beartype==0.19.0 'treescope>=0.1.7' 'rich>=14.0.0' 'polars>=1.30.0' chex \
        'lerobot[feetech]==0.5.1' \
        torch==2.8.0 torchvision==0.23.0

COPY core /app/core
COPY runtime /app/runtime

RUN --mount=type=cache,target=/root/.cache/uv \
    uv pip install --system --override /tmp/overrides.txt ./core && \
    python -m pip install --no-deps ./runtime

# openpi from the pravsels fork (main). Submodules are required (the README
# clones with --recurse-submodules). Install both the client and the main
# package with --no-deps; their dependencies are pinned explicitly above.
RUN git clone --recurse-submodules --depth 1 \
        https://github.com/pravsels/openpi /app/openpi

RUN python -m pip install --no-deps --ignore-requires-python \
        /app/openpi/packages/openpi-client \
        /app/openpi && \
    python -c "import torch, torchvision; print(torch.__version__, torchvision.__version__)"

COPY openpi_eval_runtime.py /app/openpi_eval_runtime.py

# JAX/XLA GPU memory: don't preallocate the whole device, and cap JAX at 30% so
# it shares the GPU with other processes (e.g. concurrent evals / the cell).
ENV XLA_PYTHON_CLIENT_MEM_FRACTION=0.30 \
    XLA_PYTHON_CLIENT_PREALLOCATE=false

CMD ["armnet-runtime", "/app/openpi_eval_runtime.py"]
