Dockerfile.bench

# syntax=docker/dockerfile:1.7
#
# Runtime bench image (code only, ~1-2 GB).
# Repo cache (~20 GB) is populated at provision time by bake_bench_cache.py
# running inside this image and mounted at /cache/contextbench_repos via a
# shared host bind-mount on the CCX63 server (see bench-sweep.yml cloud-init).
#
# Layers:
#   1. python base + system deps             (rare changes)
#   2. python deps (requirements-bench.txt)  (weekly changes)
#   3. diffctx wheel (cross-compiled)        (per-commit changes)
#   4. benchmarks code                       (changes most often)

ARG PYTHON_VERSION=3.12
ARG RUST_VERSION=1.92.0

# Accepted for build_bench_image.sh compatibility; the repo cache is baked at
# provision time, not into this image.
ARG BAKE_DATASET=full
ARG BAKE_LIMIT=0
ARG BAKE_PARALLELISM=4

# ============================================================================
# Stage A — Rust cross-compiler (runs on BUILDPLATFORM, targets TARGETPLATFORM)
# ============================================================================
FROM --platform=$BUILDPLATFORM python:${PYTHON_VERSION}-slim-bookworm AS rust-builder
ARG TARGETPLATFORM
ARG RUST_VERSION

RUN apt-get update && apt-get install -y --no-install-recommends \
    git build-essential pkg-config libssl-dev curl ca-certificates \
    gcc-aarch64-linux-gnu libc6-dev-arm64-cross \
    gcc-x86-64-linux-gnu libc6-dev-amd64-cross \
    && rm -rf /var/lib/apt/lists/*

ENV RUSTUP_HOME=/usr/local/rustup CARGO_HOME=/usr/local/cargo PATH=/usr/local/cargo/bin:$PATH
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
    sh -s -- -y --default-toolchain ${RUST_VERSION} --profile minimal

# Pick Rust target triple + cross-linker from TARGETPLATFORM
RUN set -eux; \
    case "${TARGETPLATFORM}" in \
      linux/amd64) TGT=x86_64-unknown-linux-gnu;  LINKER=x86_64-linux-gnu-gcc ;; \
      linux/arm64) TGT=aarch64-unknown-linux-gnu; LINKER=aarch64-linux-gnu-gcc ;; \
      *) echo "Unsupported TARGETPLATFORM=${TARGETPLATFORM}" >&2; exit 1 ;; \
    esac; \
    echo "$TGT" > /target_triple; \
    echo "$LINKER" > /linker; \
    rustup target add "$TGT"

ENV VIRTUAL_ENV=/opt/venv
RUN python3 -m venv "$VIRTUAL_ENV"
ENV PATH="$VIRTUAL_ENV/bin:$PATH"
RUN pip install --no-cache-dir "maturin>=1.10,<1.11"

# Maturin builds the top-level package from the repo root: python source in
# src/, Rust crate in diffctx/ (see [tool.maturin] in pyproject.toml).
WORKDIR /build
COPY pyproject.toml README.md LICENSE ./
COPY src ./src
COPY diffctx/Cargo.toml diffctx/Cargo.lock ./diffctx/
COPY diffctx/src ./diffctx/src
# Cargo.toml declares [[test]] targets; cargo refuses to parse the manifest
# if their files are absent.
COPY diffctx/tests ./diffctx/tests

RUN set -eux; \
    TGT=$(cat /target_triple); \
    LINKER=$(cat /linker); \
    LINKER_VAR="CARGO_TARGET_$(echo "$TGT" | tr '[:lower:]-' '[:upper:]_')_LINKER"; \
    export "${LINKER_VAR}=${LINKER}"; \
    export PYO3_USE_ABI3_FORWARD_COMPATIBILITY=1; \
    maturin build --release --target "$TGT" --out /wheels

# ============================================================================
# Stage B — Runtime
# ============================================================================
FROM python:${PYTHON_VERSION}-slim-bookworm AS runtime

RUN apt-get update && apt-get install -y --no-install-recommends \
    git ca-certificates \
    && rm -rf /var/lib/apt/lists/*

ENV VIRTUAL_ENV=/opt/venv
RUN python3 -m venv "$VIRTUAL_ENV"
ENV PATH="$VIRTUAL_ENV/bin:$PATH"

WORKDIR /app

# Layer 2: Python deps (heavy, slow — install BEFORE source).
# uv is required by the Aider baseline (`uv tool run --from aider-chat ...`);
# without it, every aider cell crashes on first request with
# "RuntimeError: `uv` not found on PATH".
COPY requirements-bench.txt ./
RUN pip install --no-cache-dir -r requirements-bench.txt uv

# Layer 3: diffctx wheel with the compiled _diffctx extension
COPY --from=rust-builder /wheels/*.whl /tmp/wheels/
RUN pip install --no-cache-dir /tmp/wheels/*.whl && rm -rf /tmp/wheels

# Layer 4: Benchmarks code (changes most often)
COPY benchmarks ./benchmarks
COPY scripts/bake_bench_cache.py ./scripts/bake_bench_cache.py

ENV CB_REPOS_DIR=/cache/contextbench_repos
ENV PYTHONUNBUFFERED=1

LABEL org.opencontainers.image.source="https://github.com/nikolay-e/diffctx"
LABEL org.opencontainers.image.description="diffctx benchmark runtime image (repo cache mounted separately)"
LABEL org.opencontainers.image.licenses="Apache-2.0"

ENTRYPOINT ["python", "-m", "benchmarks"]
CMD ["cb", "--help"]