# syntax=docker/dockerfile:1.7-labs
# SPDX-License-Identifier: Apache-2.0
# arbi-serve — slim image (~4.4 GB; vs ~10-12 GB for full vLLM image).
#
# Design notes (this is the only image we ship; the prior single-stage
# full variant was removed). vs a naive full cudnn-runtime + system-site
# image:
#   * runtime base = `nvidia/cuda:VERSION-runtime-ubuntu24.04` instead
#     of `cudnn-runtime`. Torch wheels bundle libcudnn under
#     `torch/lib/libcudnn_*.so`; the system cudnn from the base image
#     is unused at inference time.
#   * /opt/venv layout instead of system site-packages. Single COPY
#     from builder; no PEP-668 `--break-system-packages` workaround.
#   * Builder pre-bakes xgrammar's torch cpp_extension into
#     /opt/cache-baked/. The arbi_serve._prebake_loader monkey-patch
#     short-circuits torch.utils.cpp_extension.load via
#     torch.ops.load_library so the runtime image needs no
#     ninja/g++/nvcc for xgrammar.
#   * Strip torch's `tests/`, `examples/`, `dist-info/RECORD` from
#     site-packages. ~200-400 MB savings depending on torch release.
#   * Final image carries no nvcc, no headers, no ninja, no git, no
#     gcc.
#
# What works in the slim image:
#   * /v1/models, /v1/chat/completions API surface
#   * Tokenizer + xgrammar structured generation
#   * Multi-group routing (engine_params override per request)
#
# Attention backend (torch-2.12 stack):
#   * DEFAULT is paged_kv:bf16-turbo (FA4 CuTeDSL prefill + turbo split-K
#     decode), which is FA2-free. flash-attn (FlashAttention-2) is DROPPED —
#     no torch-2.12 wheel exists and source-building it is forbidden by policy;
#     a measured A/B showed dropping it costs nothing on sm_89 head_dim<=256
#     bf16 decode. The bf16-fa / bf16-fa-persistent backends remain code-only
#     and lazy-fail with a hint pointing at bf16-turbo.
#
#   * --kv-cache-dtype tkv — runtime stage now ships nvcc + g++ +
#     ninja so torch.cpp_extension.load_inline can JIT per-arch
#     kernels for any (head_dim, num_kv_heads) shape on first request.
#     WARMUP_MODEL still recommended for production (skips the ~30s
#     first-request JIT) but no longer required.
#
# Build (context is just arbi-serve/ — turbo-attn installs from PyPI at
# the uv.lock-pinned version, no sibling checkout needed):
#
#   docker build -f arbi-serve/Dockerfile -t arbi-serve:slim arbi-serve/
#
# (turbo-attn was published to PyPI; the old parent-dir context + local
# turbo-attn COPY were removed in favour of a normal locked PyPI install.)
#
# Optional kernel-warmup model (so kernel cache is pre-baked for
# this exact arch + bit width — saves ~30 s on first request):
#   docker build -f arbi-serve/Dockerfile \
#       --build-arg WARMUP_MODEL=/models/Qwen3.8-27B-exl3-4.0bpw \
#       --build-arg WARMUP_K_BITS=4 --build-arg WARMUP_V_BITS=4 \
#       -t arbi-serve:slim .
#
# A group-2 VQ (vq2) KV-codec deployment must ALSO bake the vq2 kernel
# family — it is a different kernel, not a flag on the scalar one, so a
# scalar-only image cannot serve it under TKV_NO_JIT=1:
#   docker build -f arbi-serve/Dockerfile \
#       --build-arg TKV_BAKE_BITS=2,3,4 --build-arg TKV_BAKE_VQ=both \
#       -t arbi-serve:vq2 .

# CUDA Toolkit 13.1. Drives BOTH the builder (-cudnn-devel) and runtime
# (-base) stages.
#
# NOTE: this pin was originally taken for the sm89 e4m3 warp-level MMA that
# the (now removed) split-D fp8 prefill path needed — 12.9/13.0 refused it.
# That requirement is retired. The pin now rests on the coordinated
# cutlass-dsl 4.6.x / quack / turbo-attn upgrade documented below; nobody has
# re-tested whether a lower toolkit would serve, so it is kept as-is rather
# than lowered on inference.
ARG CUDA_VERSION=13.1.1
ARG UBUNTU_VERSION=ubuntu24.04
ARG PYTHON_VERSION=3.12
# Single GPU arch by default (sm_89 = RTX 4090 / Ada Lovelace). Build a
# matrix of `arbi-serve:slim-sm89`, `arbi-serve:slim-sm90` etc. by
# overriding this. Stripping per-arch SASS from libtorch_cuda.so + tkv
# extensions saves several hundred MB per image.
ARG GPU_ARCH=8.9

# ──────────────────────────────────────────────────────────────────
# Stage 1: builder — full CUDA dev toolkit, compile + warm kernels.
# ──────────────────────────────────────────────────────────────────
FROM nvidia/cuda:${CUDA_VERSION}-cudnn-devel-${UBUNTU_VERSION} AS builder

ARG PYTHON_VERSION
ARG GPU_ARCH

# CuTeDSL ships as a FAMILY: the frontend (nvidia-cutlass-dsl) plus
# libs-base/libs-core/libs-cu12/libs-cu13. They must all be the SAME version —
# a frontend newer than libs-cu13 fails _get_baked_dso_digest() with "Missing or
# stale DSO version module", and forcing past that SIGSEGVs. Both install sites
# below read this ARG so they cannot drift.
ARG CUTLASS_DSL_VERSION=4.6.2
ENV DEBIAN_FRONTEND=noninteractive \
    PIP_NO_CACHE_DIR=1 \
    PYTHONDONTWRITEBYTECODE=1 \
    PYTHONUNBUFFERED=1 \
    # Two-part peak-disk control for the torch/CUDA install (it ENOSPC'd lean
    # build runners — "failed to copy ... No space left on device"):
    #   * UV_NO_CACHE=1 — don't retain a uv download/extract cache. Otherwise uv
    #     holds each wheel TWICE (once in /root/.cache/uv/archive-v0, once in the
    #     venv), so torch peaked ~30GB. No-cache installs straight to the venv.
    #   * UV_LINK_MODE=hardlink — uv installs by HARDLINKING the unpacked wheel
    #     into the venv instead of COPYing it. The unpack dir and /opt/venv are
    #     the same overlayfs upper layer during a RUN, so the link is free and
    #     the venv shares inodes with the unpack (~1x peak) rather than holding a
    #     second physical copy (~2x). This roughly halves the remaining torch +
    #     nvidia-* install peak — the difference between fitting and ENOSPC on a
    #     small build VM. (uv silently falls back to copy if it ever can't link,
    #     e.g. a cross-device TMPDIR, so it's never worse than copy.)
    # docker layer caching still keeps unchanged RUN steps fast; mirrors the
    # PIP_NO_CACHE_DIR=1 above.
    UV_NO_CACHE=1 \
    UV_LINK_MODE=hardlink

RUN apt-get update && apt-get install -y --no-install-recommends \
        python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv \
        ninja-build curl ca-certificates \
        git build-essential cmake \
    && rm -rf /var/lib/apt/lists/*

# Use uv for fast resolves + parallel installs. Falls back to pip if
# uv install fails on this platform.
RUN curl -LsSf https://astral.sh/uv/install.sh | sh && \
    mv /root/.local/bin/uv /usr/local/bin/uv

# Venv at /opt/venv — single dir to copy into runtime.
RUN python${PYTHON_VERSION} -m venv /opt/venv
ENV PATH="/opt/venv/bin:${PATH}" \
    VIRTUAL_ENV="/opt/venv" \
    UV_CACHE_DIR="/root/.cache/uv" \
    UV_LINK_MODE="copy"

# PyTorch 2.12.1+cu130 from the stable channel. cu130 is the NEWEST 13.x torch
# wheel PyTorch ships — there is no cu131 index (download.pytorch.org/whl/cu131
# 403s), so torch stays on the CUDA-13.0 runtime. That is minor-compatible under
# the 13.1.1 toolkit base (13.0 runtime libs load fine on a 13.1 toolkit) and is
# NOT a sub-13.x confound; it also does not gate the fp8 MMA — CuTeDSL's own cu13
# toolkit JIT-compiles that kernel, independent of torch's bundled runtime.
# Ships triton 3.7.1. cu130 dropped the cu128 wheels; the ABI lines up with the
# source-built causal-conv1d / mamba-ssm below.
RUN --mount=type=cache,target=/root/.cache/uv \
    uv pip install torch==2.12.1 --index-url https://download.pytorch.org/whl/cu130

# torchvision — multimodal (vision) input path: the bicubic-with-antialias
# resize that matches Qwen-VL's fast image processor bit-for-bit + the
# on-GPU resize/normalize/patchify pipeline. torchvision tracks the torch
# ABI, so pin to the torch-2.12 pair (0.27.x <-> torch 2.12.x) and pull from
# the SAME cu130 index — it reuses the torch installed above instead of
# re-resolving a CPU build (which would also drag in a mismatched torch).
RUN --mount=type=cache,target=/root/.cache/uv \
    uv pip install torchvision==0.27.1 --index-url https://download.pytorch.org/whl/cu130

# arbi-serve runtime deps. Non-editable so /workspace/arbi-serve is
# also throwaway after this stage.
#
# LAYER ORDER (cache discipline): copy ONLY pyproject.toml + uv.lock here —
# everything from here down to the kernel bakes depends solely on the locked
# dependency set, NOT on arbi_serve source. The source tree (arbi_serve/) is
# copied much later, just before the self-install + exl3 bake that actually
# need it. A code-only merge (the common CI trigger) therefore re-runs only
# that short tail and hits cache on the whole dependency + bake block above —
# instead of busting turbo-attn + the ~20-pkg install + FLA + flash-attn +
# the xgrammar/TKV bakes on every commit, which is what copying the source
# here (above the installs) used to force.
# The lock-vs-installed assertion rides in THIS copy, not a layer of its own,
# so it busts exactly when uv.lock does. The same script runs against a
# PUBLISHED image in the GPU lanes' preflight, so the build gate and the CI
# gate can never diverge in how they compare a version to the lock.
COPY arbi-serve/pyproject.toml arbi-serve/uv.lock arbi-serve/README.md \
     arbi-serve/scripts/ci/assert_image_backend_matches_lock.py \
     arbi-serve/scripts/packaging/patch_xgrammar_optional_transformers.py \
     /workspace/arbi-serve/
# Pin EVERY subsequent PyPI install to the committed uv.lock so the image
# is byte-for-byte reproducible across builds. Floating constraints
# (`xgrammar>=0.2`, its declared but unused Transformers dependency, …) are the
# root cause of "green in one build, red in the next" — the resolved
# version drifts and a cross-checked oracle (transformers GELU) or a
# kernel signature (xgrammar CPU bitmask) moves underneath the tests.
# `uv export` flattens uv.lock to a constraints file; the `grep -v ' @ '`
# drops any non-PyPI (path/git/url) lines so the constraints carry only
# registry pins. turbo-attn is now a normal PyPI pin and stays in.
#
# A `-c` constraints file cannot express a direct reference, and a dropped
# direct reference does not fail the install that needs it -- it resolves the
# same NAME off PyPI at whatever version the index offers, which for a package
# whose fork carries the same `__version__` is indistinguishable downstream.
# The direct references are therefore split into their own `-r` file and the
# `grep` that extracts them is the gate: no match, no image.
RUN cd /workspace/arbi-serve && \
    uv export --frozen --extra dev --no-hashes --no-emit-project 2>/dev/null \
        > /tmp/arbi-lock-export.txt && \
    grep -v ' @ ' /tmp/arbi-lock-export.txt > /etc/arbi-lock-constraints.txt && \
    grep '^exllamav3 @ ' /tmp/arbi-lock-export.txt > /etc/arbi-exllamav3-req.txt && \
    (grep '^turbo-attn @ ' /tmp/arbi-lock-export.txt > /etc/arbi-turbo-attn-req.txt || \
        : > /etc/arbi-turbo-attn-req.txt) && \
    rm -f /tmp/arbi-lock-export.txt && \
    echo "locked constraints: $(grep -c '==' /etc/arbi-lock-constraints.txt) pins" && \
    echo "exllamav3 direct reference: $(cat /etc/arbi-exllamav3-req.txt)" && \
    echo "turbo-attn direct reference: $(cat /etc/arbi-turbo-attn-req.txt || echo '(none — PyPI pin)')"

# turbo-attn (tkv codec + runtime) at the uv.lock-pinned version, from PyPI or
# from the lock's direct git reference — see the WHICHEVER note below.
# Pure-Python wheel (py3-none-any) — no compile; the per-arch CUDA kernels
# JIT at runtime (or get pre-baked below by bake_grid). --no-deps: turbo-attn's
# only base dep is torch (already pinned + installed above). We do NOT install
# turbo-attn's flash-attn extra: FA2 is retired on torch 2.12 (see the flash-attn
# note below) and the default path is FA4/turbo, so --no-deps also guarantees uv
# can't pull flash-attn from PyPI source here. Pinned via the lock constraints.
# NB: the package is "turbo-attn" (import name `tkv`); do NOT confuse with the
# unrelated upstream "turboquant" vLLM project.
# WHICHEVER the lock resolved. `grep -v ' @ '` strips direct references from the
# constraints file, so a git-sourced turbo-attn is absent from it and installing
# by NAME would silently take the newest PyPI release instead — which is exactly
# what happened when the lock moved to a git rev: every image build failed the
# backend-match assert below with "lock pins 0.49.1.dev58+gcc67ca56, image
# carries 0.49.0", and `:test-latest` could not refresh. Mirror the exllamav3
# handling: install the direct reference when the lock has one, the constrained
# PyPI name when it does not.
# The direct reference resolves from a PRIVATE repo while turbo-attn's release
# lane is stalled (turbo-attn#908), so the clone needs a credential. Read from a
# BuildKit secret and applied through GIT_CONFIG_* for the length of THIS command
# only: a secret mount is not committed to a layer and the rewrite never reaches
# the image, so the token cannot be recovered from `docker history` or the
# filesystem. Same GIT_CONFIG_* form the CPU-venv CI step uses for the same repo.
# No secret is not an error -- a public source (exllamav3) and a PyPI-pinned lock
# both clone fine without one, and the install below fails loudly if it needed it.
RUN --mount=type=cache,target=/root/.cache/uv \
    --mount=type=secret,id=gh_token \
    if [ -s /run/secrets/gh_token ]; then \
        export GIT_CONFIG_COUNT=1 && \
        export GIT_CONFIG_KEY_0="url.https://x-access-token:$(cat /run/secrets/gh_token)@github.com/.insteadOf" && \
        export GIT_CONFIG_VALUE_0="https://github.com/"; \
    fi && \
    if [ -s /etc/arbi-turbo-attn-req.txt ]; then \
        uv pip install --no-deps -r /etc/arbi-turbo-attn-req.txt; \
    else \
        uv pip install --no-deps -c /etc/arbi-lock-constraints.txt turbo-attn; \
    fi

# The image's turbo-attn MUST equal the uv.lock pin, asserted where it is
# installed. arbi_serve.runtime.backend_version_guard raises on the same
# mismatch at engine boot, i.e. in every consumer of a published image.
RUN python /workspace/arbi-serve/assert_image_backend_matches_lock.py \
        --lock /workspace/arbi-serve/uv.lock --image "(builder)" --exact \
        --from-env turbo-attn

# arbi-serve's third-party runtime deps (NOT arbi-serve itself — that
# self-install moved below the source COPY, so this whole layer stays cached
# across code-only changes). All pinned by the lock constraints.
#
# nvidia-cutlass-dsl carries the [cu13] EXTRA on its requirement line (not just
# the lock): `-c` constraints pin the version but ignore extras, so a bare
# `nvidia-cutlass-dsl` resolves to libs-base (CUDA 12.9). [cu13] pulls
# libs-cu13 (Toolkit 13.x), which is what the pin above keeps.
#
# cutlass-dsl 4.6.x cu13: the coordinated cutlass + quack + turbo-attn upgrade is
# now unblocked. 4.6.x removed ``cute.core.ThrMma`` (moved to ``cute.ThrMma``);
# turbo-attn >=0.28.2 uses the 4.6 compat in ``tkv/kernels/cute/_fa/utils.py`` and
# quack-kernels >=0.6 dropped the ``cute.core.ThrMma`` annotation, so the FA4 cute
# import surface is green on 4.6. At 4.6.1 the [cu13] wheels read CUDA 13 cleanly.
RUN --mount=type=cache,target=/root/.cache/uv \
    uv pip install -c /etc/arbi-lock-constraints.txt \
        fastapi 'uvicorn[standard]' 'pydantic>=2' \
        'orjson>=3.9' \
        'pyzmq>=25' 'msgspec>=0.19' \
        tokenizers safetensors jinja2 numpy typing-extensions \
        httpx attrs python-dateutil \
        'apache-tvm-ffi>=0.1.12' \
        loguru \
        'cryptography>=41' \
        "nvidia-cutlass-dsl[cu13]==${CUTLASS_DSL_VERSION}" \
        'quack-kernels>=0.6' \
        'opentelemetry-api>=1.36.0' 'opentelemetry-sdk>=1.36.0' \
        'opentelemetry-exporter-otlp-proto-http>=1.36.0' \
        'opentelemetry-exporter-otlp-proto-grpc>=1.36.0' \
        'opentelemetry-exporter-prometheus>=0.58b0' \
        'opentelemetry-instrumentation-fastapi>=0.58b0' \
        'opentelemetry-instrumentation-logging>=0.58b0' \
        'grpcio>=1.78.1' \
        'nvidia-ml-py>=12.0' \
        'psutil>=5.9' \
        'Pillow>=10.0' && \
    uv pip install --no-deps -c /etc/arbi-lock-constraints.txt 'xgrammar>=0.2' && \
    python /workspace/arbi-serve/patch_xgrammar_optional_transformers.py --verify

# Force the cu13 CuTeDSL flavor to WIN on disk. The [cu13] extra ADDS
# nvidia-cutlass-dsl-libs-cu13 but does NOT remove nvidia-cutlass-dsl-libs-base;
# both ship a ``version_info.py`` to the SAME path, so whichever lands LAST wins.
# Reinstalling libs-cu13 LAST pins CUDA_VERSION to the Toolkit-13 build. At 4.6.1
# the cu13 wheels already read CUDA 13; this keeps that guaranteed. --no-deps:
# touch only this wheel, re-resolve nothing. The assert fails the build LOUD if
# the resolved payload is not Toolkit-13 (e.g. a stray libs-base shadow).
RUN --mount=type=cache,target=/root/.cache/uv \
    uv pip install --force-reinstall --no-deps "nvidia-cutlass-dsl-libs-cu13==${CUTLASS_DSL_VERSION}" && \
    python -c "from cutlass.base_dsl import version_info as v; cv=v.CUDA_VERSION; assert cv.major==13, cv; print('cutlass version_info =', cv)" && \
    python -c "\
import importlib.metadata as md; \
fam={ (d.metadata['Name'] or ''): d.version for d in md.distributions() if (d.metadata['Name'] or '').startswith('nvidia-cutlass-dsl')}; \
vs=set(fam.values()); \
print('cutlass family:', fam); \
assert len(vs)==1, 'SPLIT CuTeDSL INSTALL — the frontend and the cu<N> native libs are at different versions. Compiling through that mismatch SIGSEGVs. Family: %r' % fam"

# Gated Delta Net kernels — `fla-core` provides
# the production GDN forward (chunk + fused-recurrent) used by the
# Qwen 3.5 / Qwen 3.6 hybrid models. Pure-Python wheel (no CUDA
# compile here), MIT-licensed. Install the core distribution directly:
# the `flash-linear-attention` metapackage adds Transformers, which the
# serving runtime neither imports nor needs. Add at the end so prior layers
# stay cached when only the FLA pin moves.
RUN --mount=type=cache,target=/root/.cache/uv \
    uv pip install -c /etc/arbi-lock-constraints.txt 'fla-core>=0.5'

# flash-attn (FlashAttention-2) is DROPPED on the torch-2.12 stack. There is no
# prebuilt wheel for torch 2.12 (Dao-AILab tops out at cu13torch2.10) and source-
# building it is the one thing project policy forbids. It is also unnecessary: the
# DEFAULT serve path is paged_kv:bf16-turbo (FA4 CuTeDSL prefill + turbo split-K
# decode), which is FA2-free, and a measured A/B showed dropping FA2 costs nothing
# on sm_89 head_dim<=256 bf16 decode (within run-to-run noise). The exl3,
# vision-tower and embed/rerank paths were all moved off FA2; the bf16-fa /
# bf16-fa-persistent backends remain code-only and lazy-fail with a hint pointing
# at bf16-turbo. (Was: a ~30s prebuilt-wheel download here.)

# onnxscript — required by ``torch.onnx.__init__`` (it imports onnxscript
# eagerly). We never EXPORT onnx, but torch.compile's Dynamo tracer does
# ``import torch.onnx`` while building its tracing-state handler table
# (it reads ``torch.onnx.is_in_onnx_export``). Without onnxscript that
# import raises ``ModuleNotFoundError: No module named 'onnxscript'``,
# surfacing as ``InternalTorchDynamoError`` on the FIRST compiled forward
# — i.e. every served model with compile-on (the default) fails its
# first decode step. onnxscript is a small pure-Python wheel; install it
# rather than forcing every deployment onto --no-compile-on.
RUN --mount=type=cache,target=/root/.cache/uv \
    uv pip install -c /etc/arbi-lock-constraints.txt onnxscript
ENV TORCH_CUDA_ARCH_LIST="${GPU_ARCH}"

# Hybrid-arch kernels (Mamba / Jamba / Qwen3-Next / LFM2 ShortConv) AND
# the EXL3 weight backend — both are FIRST-CLASS supported surfaces, so
# they install UNCONDITIONALLY (no opt-in arg). Serving a hybrid or
# EXL3-quantized model on an image that silently lacks these is a footgun:
# the runtime fails closed on multi-token Mamba prefill, LFM2/Jamba
# short-conv falls back to a slow per-row Python reference path, and EXL3
# checkpoints won't load at all. causal-conv1d + mamba-ssm have NO torch-2.12
# wheel, so they are BUILT FROM SOURCE here (the -devel builder has nvcc + the
# CUDA 13 headers); the compiled site-packages are copied into the runtime image.
#
# mamba_ssm/__init__.py is then trimmed to a version stub so it does NOT
# eagerly import ``MambaLMHeadModel`` (→ ``transformers.generation.*`` +
# ``torch._dynamo``, both stripped from slim). arbi-serve only consumes
# ``mamba_ssm.ops.{selective_scan_interface,triton.selective_state_update}``.
# The trim path is resolved via find_spec (NOT by importing mamba_ssm,
# whose un-trimmed __init__ would import transformers).
#
#   * causal-conv1d (v1.6.2.post1) / mamba-ssm (v2.3.2.post1) — BUILT FROM
#     SOURCE against torch 2.12 (CAUSAL_CONV1D_FORCE_BUILD / MAMBA_FORCE_BUILD,
#     arch pinned to ${GPU_ARCH} via the setup.py patch below, bounded MAX_JOBS
#     so the compile doesn't saturate the build host). --no-deps; we add only
#     the lightweight ops deps (einops + packaging; torch brings triton) and
#     skip the heavy transformers / tilelang / tvm transitives. Source-building
#     these is permitted — only flash-attn was ever off-limits, and flash-attn
#     is now dropped entirely.
#   * exllamav3 — --no-deps weights-only trellis-kernel surface; its grammar
#     and architecture-layer dependencies are unused. Version comes from the
#     lock constraints, NEVER a literal here: a hardcoded `==` cannot be
#     reached by `uv lock`, so a pyproject/uv.lock bump silently left the image
#     behind, below the shim's EXLLAMAV3_MIN_VERSION floor, and NO EXL3
#     checkpoint could load in it. rich + marisa-trie are the two declared deps
#     we DO install and take the constraints for the same reason. The assert
#     below fails the build if any of them drifts from the lock.
# (Qwen3.5 / Qwen3.6 / Qwen3.8 GDN do NOT need causal-conv1d — they ride
# fla-core, installed above.)
# MAX_JOBS / NVCC_THREADS bound the CUDA-extension compile fan-out: causal-conv1d
# and mamba-ssm setup.py default to one ninja job per core, which saturates a
# shared build host. 8 jobs x 2 nvcc threads keeps the build fast without pinning
# every core.
#
# ARCH PIN: both setup.py files IGNORE TORCH_CUDA_ARCH_LIST — they hardcode a
# -gencode sweep over every arch the toolkit supports (9 arches on CUDA 13;
# observed compiling sm_87/100/103/110/121 cubins this image can never run,
# ~6x the nvcc time of the one arch it ships for) and even CLEAR
# TORCH_CUDA_ARCH_LIST on a torch/toolkit major mismatch. Torch's
# cpp_extension defers to explicit -gencode flags, so the env pin above never
# reaches these two builds. Patch the pinned-tag setup.py to emit ONLY the
# ${GPU_ARCH} target(s): each hardcoded pair becomes `pass` (keeps the
# version-guard if blocks syntactically valid) and cc_flag is seeded with the
# target arch LIST. ${GPU_ARCH} is a ';'-separated list ("8.9" single-arch, or
# "8.9;12.0" for a fat Ada+Blackwell image), so the pin emits one SASS cubin
# per arch PLUS a +PTX on the highest for forward-compat — mirroring how
# TORCH_CUDA_ARCH_LIST fans the torch-native exts (xgrammar/tkv/exllamav3)
# across the same list. One unified image then carries sm_89 AND sm_120 cubins
# for causal-conv1d + mamba, so a single fleet tag serves 4090s and 5090s.
# Fail-loud asserts catch an upstream layout change at a future tag bump —
# re-verify the patch then, rather than silently sweeping again. Same
# pinned-tag surgery pattern as the mamba_ssm __init__ trim below.
COPY <<'PIN_ARCH_PY' /tmp/pin-cuda-arch.py
import re
import sys

# arch_spec: ';'-separated SM codes (GPU_ARCH with dots stripped), e.g. "89"
# (single-arch) or "89;120" (fat Ada+Blackwell). Emit one SASS cubin
# (-gencode arch=compute_XX,code=sm_XX) per arch PLUS a single +PTX
# (code=compute_XX) on the HIGHEST arch for forward-compat on future GPUs the
# image was not built for — the same fan-out torch applies for the exts that
# honor TORCH_CUDA_ARCH_LIST. These two setup.py files ignore that env, so we
# splice the equivalent -gencode LIST in by hand.
path, arch_spec = sys.argv[1], sys.argv[2]
archs = [a for a in arch_spec.split(";") if a]
assert archs, f"no arch codes parsed from {arch_spec!r}"
gencode: list[str] = []
for a in archs:
    gencode += ["-gencode", f"arch=compute_{a},code=sm_{a}"]
top = max(archs, key=int)
gencode += ["-gencode", f"arch=compute_{top},code=compute_{top}"]
cc_flag_init = "cc_flag = [" + ", ".join(f'"{f}"' for f in gencode) + "]"

src = open(path).read()
src, n = re.subn(
    r'cc_flag\.append\("-gencode"\)\n\s*cc_flag\.append\("arch=compute_\d+,code=sm_\d+"\)',
    "pass",
    src,
)
assert n >= 6, (
    f"{path}: hardcoded -gencode sweep not found (matched {n}) — "
    "upstream setup.py layout changed; re-verify this patch"
)
# Function replacement (not an f-string) so a literal '[' / '\\' in the
# assembled -gencode list is never read as a regex backreference.
src, n = re.subn(
    r"cc_flag = \[\]",
    lambda _m: cc_flag_init,
    src,
    count=1,
)
assert n == 1, (
    f"{path}: cc_flag initializer not found — "
    "upstream setup.py layout changed; re-verify this patch"
)
open(path, "w").write(src)
print(f"pinned {path} to {cc_flag_init}")
PIN_ARCH_PY
RUN --mount=type=cache,target=/root/.cache/uv \
    MAX_JOBS=8 NVCC_THREADS=2 \
    uv pip install -c /etc/arbi-lock-constraints.txt einops packaging wheel setuptools psutil && \
    ARCH_SM=$(printf '%s' "${GPU_ARCH}" | tr -d .) && \
    git clone --depth 1 -b v1.6.2.post1 https://github.com/Dao-AILab/causal-conv1d.git /tmp/cc1d && \
    python /tmp/pin-cuda-arch.py /tmp/cc1d/setup.py "$ARCH_SM" && \
    MAX_JOBS=8 NVCC_THREADS=2 CAUSAL_CONV1D_FORCE_BUILD=TRUE \
        uv pip install --no-build-isolation --no-deps /tmp/cc1d && \
    rm -rf /tmp/cc1d && \
    git clone --depth 1 -b v2.3.2.post1 https://github.com/state-spaces/mamba.git /tmp/mamba && \
    python /tmp/pin-cuda-arch.py /tmp/mamba/setup.py "$ARCH_SM" && \
    MAX_JOBS=8 NVCC_THREADS=2 MAMBA_FORCE_BUILD=TRUE \
        uv pip install --no-build-isolation --no-deps /tmp/mamba && \
    rm -rf /tmp/mamba && \
    SP=$(python -c "import importlib.util as u; print(u.find_spec('mamba_ssm').submodule_search_locations[0])") && \
    printf '__version__ = "stub-for-arbi-serve"\n' > "$SP/__init__.py" && \
    EXLLAMA_NOCOMPILE=1 uv pip install --no-deps -c /etc/arbi-lock-constraints.txt -r /etc/arbi-exllamav3-req.txt && \
    uv pip install -c /etc/arbi-lock-constraints.txt rich marisa-trie

# The EXL3 family MUST equal the uv.lock pins, for the same reason turbo-attn
# does: below the shim's version floor every EXL3 load is refused outright, and
# a native family split across versions ends in SIGSEGV rather than a clean
# import error. Transformers is deliberately absent from the runtime image;
# reference/parity tests install it in the test stage below.
RUN python /workspace/arbi-serve/assert_image_backend_matches_lock.py \
        --lock /workspace/arbi-serve/uv.lock --image "(builder)" --exact \
        --from-env marisa-trie --from-env rich \
        --behaviour-critical runtime

# EVERY package a version skew can make a GREEN GATE WRONG ON, asserted here
# because this is the first moment the image and the lock exist together. The
# CPU lane resolves its own venv from uv.lock and never sees an image, so
# "test in the image container" — right for GPU behaviour — does not reach it:
# the gate can be green on code the image behaves differently on, and red on
# code the image is fine with. That is not hypothetical for this set; the
# membership rule and what each one costs are in
# scripts/ci/assert_image_backend_matches_lock.py's BEHAVIOUR_CRITICAL.

# Serving has no Hugging Face model-stack dependency. xgrammar declares and
# imports Transformers for its optional from_huggingface adapter, while
# arbi-serve supplies tokenizer vocabulary + metadata directly. The image
# therefore installs xgrammar --no-deps and applies the exact-source patch
# above. Prove both halves: Transformers stayed absent and the patched xgrammar
# package really imports. An upstream source drift fails the patch step first.
RUN python -c "import importlib.util as u; assert u.find_spec('transformers') is None, 'Transformers leaked into the production runtime'; import xgrammar; print('xgrammar imports without Transformers')"

# Bake `uv` into the venv so in-container `uv pip install` works in BOTH
# the shipped slim runtime (operators add optional extras at run time) and
# the `test` stage below. Test/dev tooling itself is deliberately NOT
# installed in the builder: the runtime stage's catch-all site-packages
# COPY would sweep it into the shipped image. It lives in the `test` stage
# (FROM runtime) at the end of this file, so the slim image stays lean and
# the test/dev image is the slim image plus one thin pure-Python layer.
RUN --mount=type=cache,target=/root/.cache/uv \
    uv pip install uv

# Optional kernel cache warmup. Provide --build-arg WARMUP_MODEL=...
# to bake `~/.cache/torch_extensions/` for that model's specific
# (gqa, vec_size, k_bits, v_bits) tuple. Skipped when WARMUP_MODEL is
# empty (default).
ARG WARMUP_MODEL=""
ARG WARMUP_K_BITS=4
ARG WARMUP_V_BITS=4
# Which sides serve the group-2 VQ (vq2) KV codec: auto|none|k|v|both.
# vq2 compiles a DIFFERENT kernel per side than the scalar Lloyd codec (the
# split-K .so name carries a _vqk1/_vqv1 tag), so a scalar-only bake leaves a
# vq2 deployment with no kernel to load — which under the production
# TKV_NO_JIT=1 posture is a hard boot failure, not a slow path.
#
# "auto" (default) reads each side's kind from the calibration bundle at
# TKV_CALIBRATION_FILE. That is the right answer whenever a bundle is
# readable at build time; when one is not, name the kinds explicitly. vq2
# serves widths 2, 3 and 4 only.
ARG WARMUP_VQ=auto
# Pre-baked caches land in /opt/cache-baked/, NOT /cache/, so user
# bind-mounts on /cache don't mask the pre-baked content. Two consumers:
#  * exl3 / xgrammar (torch ``cpp_extension.load``): resolved directly out
#    of /opt/cache-baked by the arbi_serve._prebake_loader short-circuit.
#  * tkv split-K decode/MTP (torch ``cpp_extension.load_inline``): NOT
#    covered by that short-circuit — it resolves under $TORCH_EXTENSIONS_DIR
#    (= /cache/torch_extensions at runtime). bootstrap._seed_prebaked_torch_
#    extensions() seeds the missing kernel dirs from here into /cache on
#    boot, so the runtime warm is a 0-JIT cache HIT.
ENV CACHE_HOME=/opt/cache-baked
RUN mkdir -p ${CACHE_HOME}/torch_extensions ${CACHE_HOME}/cute-dsl ${CACHE_HOME}/fa4-cute
# CUTE_DSL_CACHE_DIR is the production cache (tkv arbi-prefill cute_dsl_cache);
# the FLASH_ATTENTION_* pair only feeds the retired FA4 kernel, which loads in
# parity-test environments — kept so those runs stay cached too.
ENV TORCH_EXTENSIONS_DIR=${CACHE_HOME}/torch_extensions \
    CUTE_DSL_CACHE_DIR=${CACHE_HOME}/cute-dsl \
    FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1 \
    FLASH_ATTENTION_CUTE_DSL_CACHE_DIR=${CACHE_HOME}/fa4-cute

# xgrammar JITs apply_token_bitmask_inplace_cuda on first import via
# torch.utils.cpp_extension.load — needs ninja + g++ which the runtime
# stage drops. Bake the .so here so the runtime image hits the torch
# extension cache and skips the build path. Unconditional — without
# this, the runtime image fails on first request with
# "Ninja is required to load C++ extensions".
#
# TORCH_CUDA_ARCH_LIST is required because the build container has no
# GPU, so torch can't probe arch at runtime. Covers Turing through
# Hopper; adjust if you need older or Blackwell.
ENV TORCH_CUDA_ARCH_LIST="${GPU_ARCH}"
RUN python -c "import xgrammar"

# NB: the exllamav3_ext (EXL3 trellis kernel) bake needs arbi_serve installed
# (it triggers via the exl3 shim), so it lives below the source COPY at the
# end of this stage — see "EXL3 trellis kernel pre-bake" there. The xgrammar
# bake above and the tkv bakes below need NO arbi_serve source, so they stay
# here in the cached dependency block.

# --vq is passed ONLY when it departs from the default. turbo-attn installs
# from the uv.lock PyPI pin, and the flag postdates some pinned versions — an
# unconditional flag would break every default build against an older wheel
# for a knob that build is not using. Omitting it also keeps the default
# build's command line byte-identical to before the flag existed.
RUN if [ -n "${WARMUP_MODEL}" ]; then \
        vq_arg=""; \
        [ "${WARMUP_VQ}" = "auto" ] || vq_arg="--vq ${WARMUP_VQ}"; \
        python -m tkv.kernels.precompile \
            --model "${WARMUP_MODEL}" \
            --k-bits ${WARMUP_K_BITS} \
            --v-bits ${WARMUP_V_BITS} \
            ${vq_arg}; \
    else \
        echo "WARMUP_MODEL not set; skipping tkv kernel pre-bake. First request will JIT tkv kernels (~30s)."; \
    fi

# ── Model-agnostic tkv decode kernel pre-bake (GPU-free) ──────────────
# Bake the EXACT split-K decode + compress-store kernels the served config
# requests at warm_tkv_decode_autotune, so the runtime warm is a 0-JIT
# cache HIT instead of minutes of boot-time nvcc (py-spy caught a 27B TP2
# boot stalled in warm_autotune_layers -> compile_decode_and_cs ->
# _compile_unified_splitk). Unlike the WARMUP_MODEL path this needs NO
# model and NO GPU: TORCH_CUDA_ARCH_LIST drives -gencode so nvcc emits
# sm_<arch> cubins device-less. The arch is known (${GPU_ARCH}) and the
# serving shapes are known, so a precise (not a giant grid) bake is right.
#
# TKV_BAKE_SHAPES is head_dim:num_kv_heads:num_q_heads in PER-RANK terms
# (bake_grid derives gqa = num_q_heads // num_kv_heads directly — no TP
# awareness), so for a TP=N deployment divide the model's total head counts
# by N. Default 256:2:12 = Qwen3.x-27B-AWQ (head_dim 256, total n_kv 4 /
# n_q 24) at TP2 → per-rank n_kv 2, n_q 12 → gqa 6 (NOT the 256:2:8 gqa-4
# single-GPU lead shape baked in turbo-attn's own image).
#
# MTP-verify / LSE-emitting decode kernels are NOT baked here: the PyPI
# bake_grid CLI exposes only --shapes / --bits (plus --vq on newer
# turbo-attn) and deliberately leaves MTP + smart-mix to lazy JIT into the persistent
# /cache/tkv volume (its docstring: "MTP falls back to lazy JIT … to keep
# the build short"). So an MTP deployment JITs the verify kernel on its
# first spec step (~30 s), then it's a warm cache hit on the mounted volume
# for every later boot. WARMUP_MODEL still pre-bakes a specific model's full
# (incl. MTP) kernel set when given. Set TKV_BAKE=0 to skip (pure lazy-JIT
# image).
ARG TKV_BAKE=1
ARG TKV_BAKE_SHAPES=256:2:12
ARG TKV_BAKE_BITS=4
# Codec family for the model-agnostic bake: none|k|v|both. Unlike the
# WARMUP_* bake above there is no bundle to read here (this stage runs with
# no model and no calibration mounted), so the family is stated, not derived
# — hence "none" rather than "auto". vq2 serves widths 2, 3 and 4 only, and
# the bake refuses a width outside that set rather than emitting a kernel
# whose pair index cannot address its codebook.
#
# A vq2 image wants BOTH families: this grid's scalar kernels still serve the
# bypass arm and any scalar-calibrated layer, so set TKV_BAKE_VQ=both and
# keep a second build (or a second bits list) for the scalar rung rather than
# trading one for the other.
ARG TKV_BAKE_VQ=none
# --vq only when non-default, for the same PyPI-pin reason as the WARMUP bake.
RUN if [ "${TKV_BAKE}" = "1" ]; then \
        vq_arg=""; \
        [ "${TKV_BAKE_VQ}" = "none" ] || vq_arg="--vq ${TKV_BAKE_VQ}"; \
        TORCH_CUDA_ARCH_LIST="${GPU_ARCH}" \
        python -m tkv.kernels.bake_grid \
            --shapes "${TKV_BAKE_SHAPES}" \
            --bits "${TKV_BAKE_BITS}" \
            ${vq_arg}; \
    else \
        echo "TKV_BAKE=0 — skipping model-agnostic tkv kernel bake."; \
    fi

# ── Heavy, arbi_serve-INDEPENDENT venv post-processing — MOVED ABOVE the
#    arbi_serve source COPY so a code-only merge (the common CI trigger)
#    cache-HITS all of it. Every step below only touches torch / nvidia /
#    triton / cuda / xgrammar / tkv in the venv — never arbi_serve — so
#    ordering them here is pure cache discipline, not a behaviour change.
#    Previously they sat BELOW the source install, so a one-line code commit
#    re-ran the whole-venv strip + the ~6 GB dir `mv` AND re-EXPORTED those
#    multi-GB layers to the registry cache under zstd:19 (~20 min for a
#    code-only change — torch was cached, its DERIVED layers were not).
#    The GPU-kernel bakes genuinely need arbi_serve, so they stay in the
#    tail — but only the exl3 trellis one is keyed on the WHOLE package; see
#    the bake-inputs COPY block there. ──────────────────────────────────────

# Strip torch test trees + huge dist-info to shrink the venv copy.
RUN find /opt/venv/lib -type d \( -name 'tests' -o -name 'test' -o -name 'examples' \) \
        -prune -exec rm -rf {} + 2>/dev/null || true && \
    find /opt/venv/lib -type d -name '__pycache__' -prune -exec rm -rf {} + && \
    find /opt/venv/lib -name 'RECORD' -path '*/torch-*.dist-info/*' -delete 2>/dev/null || true

# Drop the AMD Triton backend; we run on NVIDIA only.
# Also strip the `amd` line from triton's entry_points.txt — anything
# that walks `importlib.metadata.entry_points(group='triton.backends')`
# (e.g. ``mamba_ssm.ops.triton.selective_state_update``) would otherwise
# try to import the deleted ``triton.backends.amd`` and crash.
RUN set -e; \
    SP=$(python -c 'import triton, os; print(os.path.dirname(triton.__file__))'); \
    test -n "$SP"; \
    rm -rf "$SP/backends/amd"; \
    DI=$(python -c 'import importlib.metadata as m; print(m.distribution("triton")._path)'); \
    sed -i '/^amd = triton.backends.amd$/d' "$DI/entry_points.txt"

# CUDA library prune. The cu13x torch wheel pulls ~2.6 GB of nvidia-*
# pip packages under site-packages/nvidia/ — most are eagerly opened
# at `import torch` and can't be dropped without breaking the image.
#
# DO NOT TOUCH — eagerly opened by torch._C at `import torch` (per
# `LD_DEBUG=files python -c 'import torch'` trace), or lazily dlopen'd
# in the inference hot path:
#   libcufft (278 MB)         eager NEEDED
#   libcusparse (160 MB)      eager NEEDED
#   libcusolver (150 MB)      eager NEEDED
#   libcusolverMg (115 MB)    lazy dlopen via torch.linalg.eigh path
#   libcurand, libcublas{Lt}, libnvrtc, libnvJitLink, libcudart,
#   libcudnn, libcupti, libcufile, libnvtx — all eager NEEDED
#   nvidia/cusparselt (224 MB), nvidia/nccl (187 MB),
#   nvidia/nvshmem (75 MB)   — all eager NEEDED
#
# SAFE-TO-DROP (verified via the same trace — never opened, no production
# code path lazy-dlopens them):
#   libnvperf_host   33 MB  — Nsight PerfWorks host runtime; ncu profiles
#                              from the HOST against a containerized
#                              process inject their own copy. In-container
#                              profiling would need this lib re-added.
#   libnvperf_target  6 MB  — same family, target-side
#   libnvblas        1 MB   — Fortran BLAS API; LLM stack uses cublas
#   libcufftw        1 MB   — FFTW3 bridge (libcufft is NEEDED, but
#                              libcufftw is only the FFTW3 wrapper)
#   *.a static libs  3 MB   — never dlopen'd at runtime
#
# Net saving: ~45 MB. The big nvidia-* wins are blocked by torch's
# NEEDED list — no realistic path to drop them short of source-building
# torch with USE_CUSPARSELT=0 / USE_NCCL=0 / USE_NVSHMEM=0 etc.
RUN set -e; \
    SP=$(python -c 'import torch, os; print(os.path.dirname(os.path.dirname(torch.__file__)))'); \
    test -n "$SP"; \
    if [ -d "$SP/nvidia/cu12/lib" ]; then \
        cd "$SP/nvidia/cu12/lib"; \
    elif [ -d "$SP/nvidia/cu13/lib" ]; then \
        cd "$SP/nvidia/cu13/lib"; \
    else \
        cd "$SP/nvidia"; \
    fi && \
    find . -name 'libnvperf_host.so*' -delete && \
    find . -name 'libnvperf_target.so*' -delete && \
    find . -name 'libnvblas.so*' -delete && \
    find . -name 'libcufftw.so*' -delete && \
    find . -name 'libcudadevrt.a' -delete && \
    find . -name 'libcudart_static.a' -delete

# strip --strip-unneeded every shared lib in the venv. Removes debug
# symbols + unused dynamic-export entries while keeping every external
# dlsym the runtime needs. Skip files that fail (some .so are linker
# scripts or compressed). Net: ~150-400 MB on the giant torch / triton
# blobs; verified safe across torch 2.x release wheels.
RUN find /opt/venv -name '*.so' -o -name '*.so.*' 2>/dev/null \
    | xargs -r -n 1 -P "$(nproc)" strip --strip-unneeded 2>/dev/null || true

# NB on libtorch_cuda.so: PyTorch wheels embed kernels for every (sm_70,
# sm_75, sm_80, sm_86, sm_89, sm_90, sm_100, sm_120) — about 8 GPU
# generations, 7 of 8 dead weight on a single-arch deployment. cuobjdump
# can't surgically replace fatbin sections in a .so reliably, so the
# source-build path (TORCH_CUDA_ARCH_LIST= build env) is the canonical fix;
# for prebuilt wheels we settle for the strip pass above.
#
# The runtime arch record used to be written HERE, from ${GPU_ARCH}. It is
# not any more: a record that restates the build's REQUEST is a positive
# claim nothing verified, and the step that verified it (the coverage
# assertion below) is separate, so any build that skipped it produced a
# record that could be FALSE — which is the one thing
# arbi_serve.runtime.built_arch's design cannot survive. The record is now
# DERIVED from the baked cubins, in the tail after every bake. See #2196.

# Split the big, STABLE site-packages dirs OUT of the venv so the runtime
# stage can COPY them as their own cache layers while a plain catch-all
# COPY of the remaining /opt/venv/lib does NOT re-copy them. This replaces
# a `COPY --exclude` (a dockerfile:1.7-labs feature) that silently
# regressed in newer BuildKit and duplicated ~6 GB of nvidia/torch/triton
# into the shipped image — a `mv` in the builder is robust and needs no
# labs feature. arbi_serve is deliberately NOT in this list: it is
# installed + split out in its own small mv below the source COPY, so a
# code-only change moves only its ~11 MB layer, never these multi-GB dirs.
# Runs after every builder step that imports torch/triton/xgrammar/tkv; the
# kernel bakes below read them via PYTHONPATH=/opt/venv-big, not the venv.
RUN cd /opt/venv/lib/python${PYTHON_VERSION}/site-packages && \
    mkdir -p /opt/venv-big && \
    for d in nvidia cuda torch triton xgrammar tkv; do \
        [ -e "$d" ] && mv "$d" /opt/venv-big/ || true; \
    done

# ── bake inputs: the kernel sources, and the package surface the bakes
#    IMPORT to reach them, put in place BEFORE the whole-package COPY below.
#
# WHY. Every kernel pre-bake below is an nvcc run whose output is a function of
# its own sources, ${GPU_ARCH} and the pinned torch — and of nothing else in
# this package. Sitting after a COPY of the whole tree they were keyed on the
# whole tree instead, so a docstring edit anywhere in arbi_serve re-ran all of
# them. They are therefore given exactly what they import, and nothing else.
#
# INTO SITE-PACKAGES, at the path the installed package will occupy, NOT a tree
# of their own. ``Path(__file__).resolve().parent`` is how the vendored csrc
# dirs are found, so the source path reaches nvcc as ``-I`` and is embedded in
# the debug line tables the kernels compile with: baking from any other
# directory yields a different .so for the same sources. The install below
# overwrites every one of these files with byte-identical content (this set is
# a subset of that one), so nothing here survives it but the layer boundary.
#
# The set is a CLOSURE, not a guess, and tests/test_dockerfile_bake_inputs.py
# is what keeps it one: it materialises this very COPY list, runs each bake's
# entry point against it with ``cpp_extension.load`` stubbed, and fails when a
# module resolves from outside the list or when the recorded compile — name,
# source digests, every flag — differs from the whole-tree one. An import added
# to a bake-reachable module therefore fails on the PR gate, not thirty minutes
# into a merge-to-main image build. Widen the list and it goes green.
#
# The exl3 TRELLIS bake is deliberately NOT fed from here: reaching it runs
# ``runtime.built_arch``'s device refusal, which imports ``arbi_serve.runtime``
# — the engine core. It stays below the full COPY, where its inputs honestly
# are the whole package.
ENV SITE_PACKAGES=/opt/venv/lib/python${PYTHON_VERSION}/site-packages
# The tkv bake FIRST, and keyed on ONE file, because it needs exactly one:
# tkv_prebake.py imports nothing from ``arbi_serve`` (its own docstring says
# so, and the contract test holds it to that) and is run as a script, so the
# package surface the three bakes below need must not sit above it — a layer
# is invalidated by every COPY before it, not only by the ones it reads.
COPY arbi-serve/arbi_serve/tkv_prebake.py ${SITE_PACKAGES}/arbi_serve/
# tkv turbo_attn_int8 prefill kernel pre-bake + tkv source fingerprint.
# TKV_PREFILL_INT8 routes the codec prefill through a hand-written
# CUDA kernel that tkv JITs with ``cpp_extension.load_inline``. A cold build
# is nvcc, minutes long, and it holds the extension baton while it runs —
# and unlike the split-K decode grid above, no bake_grid entry reaches it.
# So it is asked for exactly the way a serving dispatch asks: through the
# launcher's own ``_module``, which owns the cache name, the tile shape and
# the -D define list. GPU-free like the other bakes (TORCH_CUDA_ARCH_LIST
# drives -gencode). Run as the SCRIPT its own ``__main__`` guard already
# makes it: the module's docstring records that it imports nothing from
# ``arbi_serve`` so the image can bake before the install, and ``python -m``
# would nonetheless execute the package ``__init__`` and key this bake on
# the whole package.
#
# The SAME step then stamps EVERY tkv extension dir under ${CACHE_HOME}
# with a digest of the tkv sources it was built from — the decode /
# MTP-verify / compress-store dirs baked earlier included. That record is
# what lets the boot seed (`cli.bootstrap._seed_prebaked_torch_extensions`)
# refuse to carry a bake across to a differently-sourced tkv: on that path
# the extension NAME is the entire cache key, and a TKV_NO_JIT=1 boot will
# not rebuild over a wrong hit. So the stamp runs even when the bake does
# not, which is the case for every turbo-attn pin that has no such kernel
# (it prints why and leaves the exit code 0).
#
# NEVER FATAL, for the same reason the exl3 int8 GEMM bake below is not: a
# failed bake costs back the JIT this removes and nothing else, and a boot
# that pays that nvcc says so on its own timeline. Failing every image build
# over it would be the larger outage. A
# failure prints differently from an absence, so the two are never confused,
# and the stamp runs either way.
#
# Widths track the model-agnostic bake above (TKV_BAKE_BITS / TKV_BAKE_VQ)
# and are resolved through tkv's own bw-combo resolver, so a deployment that
# widens one widens both. TKV_TURBO_ATTN_INT8_BAKE=0 stamps without baking.
ARG TKV_TURBO_ATTN_INT8_BAKE=1
RUN TORCH_CUDA_ARCH_LIST="${GPU_ARCH}" PYTHONDONTWRITEBYTECODE=1 \
    PYTHONPATH=/opt/venv-big \
    python -c "import runpy, sys; sys.argv = ['tkv_prebake'] + sys.argv[1:]; \
runpy.run_path('${SITE_PACKAGES}/arbi_serve/tkv_prebake.py', run_name='__main__')" \
        --bits "${TKV_BAKE_BITS}" \
        --vq "${TKV_BAKE_VQ}" \
        $([ "${TKV_TURBO_ATTN_INT8_BAKE}" = "1" ] || echo --stamp-only)

COPY arbi-serve/arbi_serve/weight_quant   ${SITE_PACKAGES}/arbi_serve/weight_quant
COPY arbi-serve/arbi_serve/_custom_ops    ${SITE_PACKAGES}/arbi_serve/_custom_ops
COPY arbi-serve/arbi_serve/server/metrics ${SITE_PACKAGES}/arbi_serve/server/metrics
COPY arbi-serve/arbi_serve/server/__init__.py ${SITE_PACKAGES}/arbi_serve/server/
COPY arbi-serve/arbi_serve/adapters/__init__.py \
     arbi-serve/arbi_serve/adapters/adapter.py \
     ${SITE_PACKAGES}/arbi_serve/adapters/
COPY arbi-serve/arbi_serve/distributed/__init__.py \
     arbi-serve/arbi_serve/distributed/gather_dest.py \
     arbi-serve/arbi_serve/distributed/parallel_state.py \
     ${SITE_PACKAGES}/arbi_serve/distributed/
COPY arbi-serve/arbi_serve/engine/__init__.py \
     arbi-serve/arbi_serve/engine/boot_clock.py \
     ${SITE_PACKAGES}/arbi_serve/engine/
COPY arbi-serve/arbi_serve/kernels/__init__.py \
     arbi-serve/arbi_serve/kernels/fp4_format.py \
     ${SITE_PACKAGES}/arbi_serve/kernels/
COPY arbi-serve/arbi_serve/models/__init__.py \
     arbi-serve/arbi_serve/models/_batch_invariant.py \
     arbi-serve/arbi_serve/models/linear.py \
     ${SITE_PACKAGES}/arbi_serve/models/
COPY arbi-serve/arbi_serve/__init__.py \
     arbi-serve/arbi_serve/_compile_cache_env.py \
     arbi-serve/arbi_serve/_cpp_ext_locks.py \
     arbi-serve/arbi_serve/_ext_cache_key.py \
     arbi-serve/arbi_serve/_import_hooks.py \
     arbi-serve/arbi_serve/_jit_baton.py \
     arbi-serve/arbi_serve/_libcudart.py \
     arbi-serve/arbi_serve/_prebake_loader.py \
     arbi-serve/arbi_serve/_runtime_shims.py \
     arbi-serve/arbi_serve/_units.py \
     arbi-serve/arbi_serve/cache_paths.py \
     arbi-serve/arbi_serve/config.py \
     arbi-serve/arbi_serve/config_groups.py \
     arbi-serve/arbi_serve/config_server.py \
     arbi-serve/arbi_serve/flag_truth.py \
     arbi-serve/arbi_serve/runtime_flags.py \
     arbi-serve/arbi_serve/runtime_flags_fields_backend.py \
     arbi-serve/arbi_serve/runtime_flags_fields_core.py \
     arbi-serve/arbi_serve/runtime_flags_fields_spec.py \
     arbi-serve/arbi_serve/runtime_flags_parsers.py \
     ${SITE_PACKAGES}/arbi_serve/

# fp4-Marlin (NVFP4 W4A16) kernel pre-bake — the Blackwell fast path for
# weight-only NVFP4 layers (Qwen3.6-27B-NVFP4 MLP + lm_head). The vendored
# stable-ABI extension (arbi_serve/weight_quant/nvfp4/marlin/csrc) JIT-builds
# via torch.utils.cpp_extension.load; bake it HERE (GPU-free — nvcc emits the
# sm_120 cubin from TORCH_CUDA_ARCH_LIST) into the same
# /opt/cache-baked/torch_extensions/<name>/<name>.so layout, so the runtime
# load() short-circuits through arbi_serve._prebake_loader with no compile.
#
# GATED on a Blackwell GPU_ARCH: the fe2m1f kernel is compiled compute_120 /
# sm_120, and the marlin forward is gated cc>=(12,0) — on a pure sm_89/sm_90
# image the extension is dead weight (weight-only NVFP4 falls back to the
# arch-agnostic Triton kernel), so we skip the bake there. Forces
# TORCH_CUDA_ARCH_LIST=12.0 for the bake regardless of the image's GPU_ARCH.
# The match is CONTAINS-12 (``*12.*``), not prefix: a unified GPU_ARCH like
# "8.9;12.0" carries Blackwell in the list, so the bake must fire for it too
# (the old ``12.*`` prefix only matched a Blackwell-FIRST arch and silently
# skipped the fp4 bake for the fat image). Same PYTHONPATH=/opt/venv-big as
# the exl3 bakes, to reach the moved torch/nvidia during compile.
# PYTHONDONTWRITEBYTECODE, on this bake and every other one fed from the
# block above, because they run BEFORE the install that overwrites their
# .py: a .pyc written now is stale the moment it lands, and would ride the
# arbi_serve layer out of this stage.
RUN case "${GPU_ARCH}" in \
      *12.*) TORCH_CUDA_ARCH_LIST=12.0 PYTHONDONTWRITEBYTECODE=1 \
            PYTHONPATH=/opt/venv-big python -c \
            "from arbi_serve.weight_quant.nvfp4.marlin.loader import _ensure_loaded; _ensure_loaded()" ;; \
      *) echo "GPU_ARCH=${GPU_ARCH} has no Blackwell (sm_120) arch; skipping fp4-Marlin NVFP4 bake." ;; \
    esac

# AWQ-Marlin (INT4 W4A16 + W4A8) kernel pre-bake — the repack + GEMM
# extension every AWQ load path JITs on first touch
# (arbi_serve_awq_marlin_v2: awq_marlin_repack / gptq_marlin_repack /
# marlin_gemm / marlin_int4_fp8_preprocess). Baking it makes
# ON-boots of the head-quant modes (and every AWQ boot on a cold /cache)
# 0-JIT deterministic: the runtime load() short-circuits through
# arbi_serve._prebake_loader to this .so and never enters the multi-minute
# nvcc build that the two TP ranks previously raced in shared
# /cache/torch_extensions ("invalid ELF header", 2/6 ON-arm boots on the
# 2026-07-16 lm-head-quant A/B). GPU-free: TORCH_CUDA_ARCH_LIST drives
# -gencode so nvcc emits the cubins device-less (same pattern as the exl3 /
# fp4 bake above). Calls cpp_extension.load directly with the loader's own
# _build_args() — _ensure_loaded() itself refuses on a CUDA-less host, which
# is a RUNTIME guard, not a build-time one. The loader is loaded from its
# FILE rather than imported: ``arbi_serve.weight_quant.awq.marlin`` runs
# three package __init__ chains that reach the engine core, none of which
# _build_args() uses, and every one of which would put the engine's own
# source back into this nvcc run's inputs. Same file, same function, same
# flags — asserted, not assumed, by tests/test_dockerfile_bake_inputs.py.
#
# MAX_JOBS=1: marlin's template TUs are multi-GB in nvcc; serial compile is
# the safe floor on any builder (the same host-OOM guard the runtime loader
# applies — froze two dev boxes on 2026-05-27). Build intermediates are
# pruned so the baked layer carries only the .so torch needs to cache-hit.
RUN MAX_JOBS=1 PYTHONDONTWRITEBYTECODE=1 PYTHONPATH=/opt/venv-big python -c "\
import importlib.util as u, sys, torch.utils.cpp_extension as ce; \
s = u.spec_from_file_location('_awq_marlin_loader', \
    '${SITE_PACKAGES}/arbi_serve/weight_quant/awq/marlin/loader.py'); \
m = u.module_from_spec(s); sys.modules[s.name] = m; s.loader.exec_module(m); \
ce.load(**m._build_args())" \
    && MARLIN_DIR="$(dirname "$(find ${CACHE_HOME}/torch_extensions \
         -name 'arbi_serve_awq_marlin_v*.so' -print -quit)")" \
    && test -n "${MARLIN_DIR}" \
    && find "${MARLIN_DIR}" -name '*.o' -delete \
    && ls -la "${MARLIN_DIR}"

# EXL3 int8 prefill GEMM pre-bake (#1861) — the prefill fast path on a part
# that has to dequantise the EXL3 format in software. Same reasoning as the exl3 trellis bake and the AWQ one just before
# it, with one difference that makes it MORE necessary rather than less:
# this kernel's instantiation matrix is (8 trellis widths) x (2 codebooks) x
# (the shape ladder), so an unbaked first boot pays tens of minutes of nvcc
# INSIDE the boot seam the KV sizer waits behind — and on the slim runtime
# stage, which has no nvcc at all, it cannot build, so the leg would refuse
# every call and the flag would be inert in exactly the image production runs.
#
# GPU-free: TORCH_CUDA_ARCH_LIST drives nvcc's -gencode and the loader emits no
# -arch of its own when the list is set (see int8_kernel._arch_flags), which is
# the whole reason that function reads the env before the device.
#
# WHICH ARCH IS int8_kernel.int8_bake_arch's, not this file's. The gate used to
# be a literal "contains 8.9" and that is narrower than its own reason: the
# kernel exists because a part with no native path for the exl3 format has to
# dequantise it in software, which is every Ampere and Ada part, while
# ``exl3_int8_gemm`` arms on all of them by default — so a GPU_ARCH=8.6 image
# shipped the leg armed with nothing baked, and its first armed TP boot met
# nvcc inside the seam every rank waits at. One arch, not a set, because torch
# keys its JIT cache on the extension NAME: a second single-arch bake under
# that name overwrites the first. The policy is Python so the CPU lane can
# check it (tests/test_exl3_int8_cold_build_under_tp.py).
#
# NEVER FATAL. A failed bake costs the JIT this removes; it must not fail an
# image build, because every other weights backend in the image is unaffected.
RUN I8_ARCH="$(PYTHONDONTWRITEBYTECODE=1 PYTHONPATH=/opt/venv-big python -c \
        "from arbi_serve.weight_quant.exl3.int8_kernel import int8_bake_arch; \
         print(int8_bake_arch('${GPU_ARCH}') or '')")" \
    && if [ -z "${I8_ARCH}" ]; then \
         echo "GPU_ARCH=${GPU_ARCH} names no single arch the exl3 int8 GEMM serves; skipping its bake."; \
       else \
         TORCH_CUDA_ARCH_LIST="${I8_ARCH}" PYTHONDONTWRITEBYTECODE=1 \
         PYTHONPATH=/opt/venv-big python -c \
            "from arbi_serve.weight_quant.exl3.int8_kernel import _load; \
             import sys; sys.exit(0 if _load() is not None else 1)" \
         && I8_DIR="$(dirname "$(find ${CACHE_HOME}/torch_extensions \
              -name 'arbi_serve_exl3_i8_gemm_v*.so' -print -quit)")" \
         && test -n "${I8_DIR}" \
         && find "${I8_DIR}" -name '*.o' -delete \
         && ls -la "${I8_DIR}" \
         || echo "exl3 int8 GEMM bake FAILED; the leg will JIT on first armed boot"; \
       fi

# ── arbi_serve source: copied after the bakes that do NOT depend on it, so a
#    code-only change re-runs only this short tail (self-install + the exl3
#    trellis bake + the arch record + an ~11 MB arbi_serve mv) and cache-HITS
#    the entire dependency block, the xgrammar/tkv bakes, the whole-venv
#    strip/prune/big-dir mv, AND the four nvcc bakes above. Non-editable
#    install so /workspace/arbi-serve stays throwaway after this stage. It
#    overwrites the bake-input subset already in site-packages with the same
#    bytes, and adds the rest of the package. ─────────────────────────────────
COPY arbi-serve/arbi_serve /workspace/arbi-serve/arbi_serve
# The generated admin client is a second top-level package in the arbi-serve
# wheel. Include it in the build context so hatchling vendors it alongside
# arbi_serve; `--no-deps` is intentional because this stage already has the
# locked runtime environment.
COPY arbi-serve/client/arbi_serve_client /workspace/arbi-serve/client/arbi_serve_client
RUN --mount=type=cache,target=/root/.cache/uv \
    uv pip install --no-deps /workspace/arbi-serve

# EXL3 trellis kernel pre-bake. exllamav3_ext JITs on first EXL3 load via
# torch.utils.cpp_extension.load(is_python_module=True). The runtime stage
# (-base, no nvcc / CUDA dev headers) cannot compile it — so a stock slim
# image fails to load ANY EXL3 checkpoint with "cusparse.h: No such file".
# Bake the .so HERE (the -devel builder has the headers) into the same
# /opt/cache-baked/torch_extensions/<name>/<name>.so layout xgrammar uses;
# arbi_serve._prebake_loader short-circuits the runtime load() straight to
# this prebaked module. Trigger via the exl3 shim (which stubs exllamav3's
# grammar-filter transitives before importing the kernel) — so this MUST run
# after the arbi-serve self-install above. Unconditional: EXL3 is a
# first-class weights backend (mirrors the install above).
# PYTHONPATH=/opt/venv-big makes torch/nvidia/triton (moved out of
# site-packages above) importable for the bake; arbi_serve is still in
# site-packages (installed just above), so the ``from arbi_serve…`` trigger
# resolves both halves.
RUN PYTHONPATH=/opt/venv-big python -c "from arbi_serve.weight_quant.exl3.shim import LinearEXL3; _ = LinearEXL3"

# ARCH COVERAGE. ${GPU_ARCH} is a promise about which cards this image serves,
# and two separate mechanisms have to keep it — TORCH_CUDA_ARCH_LIST for the
# torch-native extensions and the pin-cuda-arch.py -gencode splice for
# causal-conv1d/mamba. A build that silently drops one arch produces an image
# whose first kernel launch on that card dies with "no kernel image is
# available for execution on the device", minutes into a boot. So check the
# ARTIFACT, not the intent: cuobjdump the just-baked trellis kernel and fail
# the build when a listed arch has no cubin in it. FATAL, unlike the bakes: a
# missing bake costs back a JIT, a missing arch produces an image that cannot
# run at all on a card it advertises.
#
# This is the check that the REQUEST was honoured, and that is now its whole
# job: the runtime record arbi_serve.runtime.built_arch refuses against is no
# longer written from ${GPU_ARCH} but derived from these same cubins, in the
# tail below, so the runtime no longer depends on this step having run.
COPY arbi-serve/scripts/ci/assert_baked_arch_coverage.py /tmp/assert-arch-coverage.py
RUN python /tmp/assert-arch-coverage.py "${GPU_ARCH}" \
        "${CACHE_HOME}/torch_extensions/exllamav3_ext/exllamav3_ext.so"

# The EXL3 MoE expert-parallel dispatch hands exl3_mgemm this rank's expert
# range; the exllamav3 revision uv.lock pins is the one whose filtering path
# keeps slot positions. The registry wheel reports the SAME __version__, so
# the lock-vs-image version assertion above cannot tell the two apart -- this
# reads the capability out of the header the extension compiles from, and
# fails the BUILD rather than the first EP boot.
RUN PYTHONPATH=/opt/venv-big python -c "from arbi_serve.weight_quant.exl3.shim import assert_mgemm_expert_range_supported as a; print('mgemm expert-range slots:', a(context='image build'))"

# RUNTIME ARCH RECORD, derived from the artifact. arbi_serve.runtime.built_arch
# refuses a card whose arch this record excludes, before a weight is read, so
# the record's whole worth is that a POSITIVE entry is true. Writing it from
# ${GPU_ARCH} made it a restatement of the build's REQUEST, true only when the
# coverage assertion above had also run — a separate step, and nothing at
# runtime knew whether it had (#2196). So the record is read off the binaries
# instead, by the same cuobjdump pass the assertion takes: the archs EVERY
# arch-fanned baked extension carries. No ${GPU_ARCH} on this command line, by
# design — the record cannot restate what it is never told.
#
# LAST, after every bake above, because it describes all of them; the
# deliberately single-arch extensions (nvfp4-marlin, exl3-i8-gemm) are excluded
# inside the script, where the reason for excluding them is written down.
# FATAL: a build that cannot derive a record must not ship a silent one.
RUN python /tmp/assert-arch-coverage.py \
        --record /opt/venv/share/arbi-serve.gpu-arch \
        "${CACHE_HOME}/torch_extensions"

# THE BAKED CACHE IS READ-ONLY CONTENT, so every uid that can run this image
# has to be able to READ it. The bakes above run as root and one of them
# publishes a PRIVATE directory: tkv's content-addressed include snapshot is
# staged with `tempfile.mkdtemp` (0700 by construction, umask-independent) and
# renamed into place, so the leaf lands 0700 root:root beneath 0755 parents —
# a tree whose root passes a permission check while the boot still dies, in
# ninja, with "Permission denied" on a header inside it. Only a deployment
# whose /cache is already warm escapes it, because then no boot reaches for
# the bake at all; every FRESH one does.
#
# MODE, NOT OWNERSHIP. A chown to ${ARBI_UID} would fix the default account
# and break `docker run --user "$(id -u):$(id -g)"`, which
# scripts/docker-entrypoint.sh supports on purpose. `a+rX` is readable by
# whoever runs it, and `X` grants search on directories without making a
# non-executable file executable.
#
# HERE: after the last step that writes under ${CACHE_HOME}, and in the
# BUILDER — `COPY --from` carries modes, so the ship image gains no second
# copy of the ~1 GB of baked kernels.
# tests/test_baked_cache_is_world_readable.py holds both halves of this.
RUN chmod -R a+rX ${CACHE_HOME}

# Split arbi_serve (+ its dist-info) into its OWN small /opt/venv-big layer,
# mirroring the stable-dir mv above. LAST + separate so a code-only change
# re-exports only this ~11 MB layer, never the multi-GB torch/nvidia dirs.
# The dist-info is moved out too — its RECORD carries per-file hashes that
# change on every code edit, so leaving it in site-packages would bust the
# 665 MB catch-all /opt/venv/lib COPY in the runtime stage on every commit.
# The runtime stage COPYs this pkg dir's contents back into site-packages.
RUN cd /opt/venv/lib/python${PYTHON_VERSION}/site-packages && \
    mkdir -p /opt/venv-big/arbi_serve_pkg && \
    mv arbi_serve /opt/venv-big/arbi_serve_pkg/ && \
    mv arbi_serve-*.dist-info /opt/venv-big/arbi_serve_pkg/

# ──────────────────────────────────────────────────────────────────
# Stage 2: runtime — CUDA *base* image (libcudart only). The PyTorch
# wheel bundles libcublas / libcudnn / libcusparse / libcufft inside
# its own torch/lib/ tree; the cuda-runtime image's copies are
# duplicates we never resolve. Net: ~800 MB - 1.5 GB cut vs the
# -runtime base image.
# ──────────────────────────────────────────────────────────────────
FROM nvidia/cuda:${CUDA_VERSION}-base-${UBUNTU_VERSION} AS runtime

ARG PYTHON_VERSION
ENV DEBIAN_FRONTEND=noninteractive \
    PYTHONDONTWRITEBYTECODE=1 \
    PYTHONUNBUFFERED=1 \
    PATH="/opt/venv/bin:${PATH}" \
    VIRTUAL_ENV="/opt/venv"

# Runtime stage carries:
#   * python3.12      — interpreter
#   * cuda-nvcc       — needed by tkv's torch.cpp_extension.load_inline
#                        (per-arch CUDA module compile on first request
#                        for unfamiliar (head_dim, num_kv_heads) shapes
#                        the WARMUP_MODEL builder didn't pre-bake)
#   * cuda-cudart-dev — cudart.h headers nvcc needs at compile
#   * cuda-cccl       — Thrust / CUB headers nvcc needs (~150 MB)
#   * g++             — host compiler nvcc invokes
#   * ninja-build     — build system torch.cpp_extension uses
#
# CUDA math-library dev headers. AWQ Marlin AND exllamav3 (EXL3) both
# JIT-compile their kernels on first cold load, and every torch CUDA
# extension pulls ``ATen/cuda/CUDAContext.h`` → ``cusparse.h`` (+ cublas
# / cusolver). Without these the cold JIT dies with
# ``fatal error: cusparse.h: No such file or directory`` — so a fresh
# container can't serve any AWQ/EXL3 checkpoint. Both are first-class
# surfaces, so the headers stay in (~3.5 GB). A warmed kernel cache
# avoids the cold JIT at runtime, but the image must still be able to
# build it.
RUN apt-get update && apt-get install -y --no-install-recommends \
        python${PYTHON_VERSION} \
        python${PYTHON_VERSION}-dev \
        cuda-nvcc-13-0 \
        cuda-cudart-dev-13-0 \
        cuda-cccl-13-0 \
        libcusparse-dev-13-0 \
        libcublas-dev-13-0 \
        libcusolver-dev-13-0 \
        libcurand-dev-13-0 \
        g++ \
        ninja-build \
        curl ca-certificates \
        iproute2 \
    && rm -rf /var/lib/apt/lists/* \
    && mkdir -p /cache/cute-dsl /cache/fa4-cute /cache/torch_extensions /cache/triton \
                /cache/hf /cache/xdg /cache/tkv \
                /cache/arbi-serve/xgrammar /cache/xgrammar

# ``/cache/arbi-serve/xgrammar`` AND ``/cache/xgrammar``: the grammar cache is
# the one managed cache with no env var of its own. It derives from the BOOT
# CACHE (``boot_cache_subdir("xgrammar")`` via ``_ARBI_CACHE_SUBDIRS``), so it
# moved when nine ``*_dir`` flags collapsed into one root — and the image was
# still creating only the OLD path, which ``ARBI_ASSERT_WARM_CACHE`` then read
# as a warm→cold regression on any host that had served structured output.
# The old path stays beside the new one so a volume written by a pre-collapse
# image still resolves: the entries are content-hash-named grammars carrying no
# build key (unlike their ``inductor-<key>`` siblings), so they remain usable
# across the move.

# Cap the CUDA-extension parallel compile at 2 jobs by default —
# nvcc + g++ from torch.cpp_extension.load fans out unbounded
# (one job per .cu translation unit) which OOMs / overloads
# small dev boxes on first model load. Operators with bigger
# build hosts can override at run time.
# CUDA_HOME points at cuda-13.0, NOT the `/usr/local/cuda` -> cuda-13.1
# symlink. The 13.1.1-base runtime image ships the cuda-13.1 tree with NO
# headers (`cuda.h` absent), while the cu130 torch stack drops a complete
# header tree under /usr/local/cuda-13.0. The boot-time cumem extension
# (arbi_serve_cumem_ext) JIT-compiles against `cuda.h`, so with CUDA_HOME
# on the header-less 13.1 tree it dies `fatal error: cuda.h: No such file`.
# Pointing CUDA_HOME at cuda-13.0 gives the JIT its headers with NO runtime
# override (the 13.0 runtime is minor-compatible under the 13.1 toolkit and
# does not gate any served kernel; nvcc-driven paths are all pre-baked at
# build time in the -devel builder, so runtime CUDA_HOME feeds only the
# header-only g++ JIT).
ENV CUDA_HOME=/usr/local/cuda-13.0 \
    PATH="/usr/local/cuda/bin:${PATH}" \
    MAX_JOBS=2

# Persistent kernel + grammar caches (mount as a docker volume).
# CUTE_DSL_CACHE_DIR = production arbi-prefill CuTeDSL JIT cache; the
# FLASH_ATTENTION_* pair only feeds the retired FA4 kernel (parity tests).
ENV CUTE_DSL_CACHE_DIR=/cache/cute-dsl \
    FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1 \
    FLASH_ATTENTION_CUTE_DSL_CACHE_DIR=/cache/fa4-cute \
    TORCH_EXTENSIONS_DIR=/cache/torch_extensions \
    TRITON_HOME=/cache/triton \
    HF_HOME=/cache/hf \
    XDG_CACHE_HOME=/cache/xdg \
    TKV_CACHE_DIR=/cache/tkv
# ^ The xgrammar cache used to be baked here too. It is not any more: every
# cache a BOOT writes derives from ONE root (arbi_serve/cache_paths.py),
# which is /cache/arbi-serve whenever /cache is writable — exactly what this
# image's volume makes it. A baked per-cache path is a second place for the
# tree to be half-configured.
# ^ TKV_CACHE_DIR is the tkv split-K autotune table dir (the ~2-3 min cold
# sweep). It MUST be baked here, not left to tkv's default. tkv defaults its
# cache to ``$XDG_CACHE_HOME/tkv`` (= /cache/xdg/tkv above) — already on the
# volume — BUT docker/compose.yaml separately sets TKV_CACHE_DIR=/cache/tkv,
# so a compose-launched server warmed /cache/tkv while a bare ``docker run``
# (benchmarks, ad-hoc) fell back to /cache/xdg/tkv. Same volume, DIFFERENT
# dir → the autotune table was written twice and the two copies diverged, so a
# bench run re-ran the autotune the compose boot had already cached. Pinning it
# here makes EVERY launch path share ONE table → autotune is paid once per
# shape, ever, across compose / docker run / serve_bench. (Compose's identical
# override is now a harmless no-op.)

# Universal serving defaults: offline HF (models live on local disk), and
# loopback gloo so torch.distributed doesn't bind a routable interface.
# NOTE: TORCHINDUCTOR_COMPILE_THREADS is intentionally NOT pinned here.
# The old value 1 ("idle-thread hygiene") serialized the cold Inductor
# compile/codegen passes (~2-4x slower one-time cache population). Inductor's
# compile workers only spawn DURING compile and go idle right after — that
# idle cost is acceptable for a much faster cold build. We honor the code
# default (8, set in arbi_serve/compile/runtime.py::configure_inductor_caches),
# which operators can still override with their own env value.
ENV HF_HUB_OFFLINE=1 \
    GLOO_SOCKET_IFNAME=lo

# CUDA caching-allocator: expandable_segments MUST be OFF — it is
# mutually exclusive with torch.cuda.MemPool (pytorch#147851), and the
# engine routes weights/KV/recurrent/arena + the serving hot path
# through ~8 cuMem-backed MemPools. ON throws "MemPool doesn't support
# expandable_segments" on any unwrapped pool use. Set at the image layer
# (before first cuda init) on BOTH the new (torch>=2.9) and legacy var
# names. Mirrors arbi_serve/__init__.py. (Was True for a since-removed
# slack-reclaim scheme; cuMem pools own VA directly so that's moot.)
ENV PYTORCH_ALLOC_CONF=expandable_segments:False \
    PYTORCH_CUDA_ALLOC_CONF=expandable_segments:False

# LAYERED COPY — ordered from least- to most-frequently-changed so
# incremental pulls only re-download the diff. A code-only update only
# refreshes ~50-100 MB (arbi_serve layer); the ~3 GB nvidia-* + torch
# layers stay cached on the puller.
#
# Ordering rationale (estimated change frequency):
#   1. nvidia/*           — torch wheel updates (rare, weeks-months)
#   2. cuda/, cuda_*       — same cadence as nvidia
#   3. torch/             — torch wheel updates (rare)
#   4. triton/            — torch upgrade or independent triton bump
#   5. xgrammar/          — independent SDK version bump (rare)
#   6. tkv/              — turboquant commits (frequent — daily-weekly)
#   7. arbi_serve/        — every commit (most frequent — hourly-daily)
#   8. /opt/venv/{bin,lib/python3.12/site-packages-rest, …, share}
#                         — small Python stdlib + venv plumbing (rare)
#   9. /opt/cache-baked/  — only changes when xgrammar / tkv kernel
#                            sources change OR GPU_ARCH bumps (rare)
#
# An empty /opt/venv mkdir comes first so subsequent paths land where
# torch / Python expects.
RUN mkdir -p /opt/venv/lib/python${PYTHON_VERSION}/site-packages

# Big, stable dirs first (rarely change → cached on pullers). These were
# `mv`d to /opt/venv-big in the builder's last step, so copying them here
# and the catch-all /opt/venv/lib below cannot double-store them.
COPY --from=builder /opt/venv-big/nvidia \
    /opt/venv/lib/python${PYTHON_VERSION}/site-packages/nvidia
COPY --from=builder /opt/venv-big/cuda \
    /opt/venv/lib/python${PYTHON_VERSION}/site-packages/cuda
COPY --from=builder /opt/venv-big/torch \
    /opt/venv/lib/python${PYTHON_VERSION}/site-packages/torch
COPY --from=builder /opt/venv-big/triton \
    /opt/venv/lib/python${PYTHON_VERSION}/site-packages/triton
COPY --from=builder /opt/venv-big/xgrammar \
    /opt/venv/lib/python${PYTHON_VERSION}/site-packages/xgrammar
COPY --from=builder /opt/venv-big/tkv \
    /opt/venv/lib/python${PYTHON_VERSION}/site-packages/tkv
# arbi_serve + its dist-info (moved together into /opt/venv-big/arbi_serve_pkg
# in the builder). Copying the dir's CONTENTS lands both back in site-packages,
# exactly where the catch-all /opt/venv/lib COPY used to carry them — but as a
# small, isolated layer that a code-only change can refresh without touching the
# multi-GB dirs above or the 665 MB /opt/venv/lib layer below.
COPY --from=builder /opt/venv-big/arbi_serve_pkg/ \
    /opt/venv/lib/python${PYTHON_VERSION}/site-packages/

# Everything else in the venv — Python interpreter (bin/, lib/python3.12/
# {config-*, lib-dynload}), small site-packages (fastapi, pydantic,
# uvicorn, tokenizers, safetensors, jinja2, loguru, opentelemetry-*,
# numpy, sympy, networkx, hf_xet, etc), pyvenv.cfg, share/.
# The big dirs above were moved out of /opt/venv/lib in the builder, so
# this catch-all copies the remainder with NO overlap — no `--exclude`
# (labs feature) and no double-storage.
COPY --from=builder /opt/venv/bin /opt/venv/bin
COPY --from=builder /opt/venv/lib /opt/venv/lib
COPY --from=builder /opt/venv/pyvenv.cfg /opt/venv/pyvenv.cfg
COPY --from=builder /opt/venv/share /opt/venv/share

# The unified wheel's client package is part of the production artifact too.
# Exercise it only after the manually assembled dependency closure and all
# split site-packages layers have landed, so no transitive dependency from the
# test stage can mask an omitted base requirement.
RUN python -c "import importlib.util as u; assert u.find_spec('transformers') is None; import arbi_serve_client, xgrammar; print('runtime client + structured generation imports OK')"

# Bytecode for everything the boot imports. PYTHONDONTWRITEBYTECODE=1 above
# stops the RUNTIME from writing .pyc files but not from reading them, and a
# venv installed under it ships none — so every boot compiled every imported
# source (torch alone is over two thousand modules) before it could start.
# One layer, after the last venv copy, so it is rebuilt with the code it
# covers; a stale .pyc is ignored by source mtime, never served.
RUN python -m compileall -q -j 0 /opt/venv/lib/python${PYTHON_VERSION}/site-packages

# Pre-baked extension .so files load via torch.ops.load_library through
# the arbi_serve._prebake_loader monkey-patch. Last because it's the
# most-likely-to-change-on-arch-bump and small.
COPY --from=builder /opt/cache-baked /opt/cache-baked

# Put the baked exllamav3 kernel on the IMPORT PATH, not only behind the
# prebake short-circuit.
#
# exllamav3's own `ext.py` decides whether to JIT with
# `try: import exllamav3_ext / except ModuleNotFoundError: build_jit = True`,
# and it imports torch first, so a module it can import is one it will not
# compile. The `_prebake_loader` monkey-patch only redirects a caller that
# imported arbi_serve FIRST; anything else — a test probe, a fresh
# subprocess, any tool that imports exllamav3 on its own — misses it and
# runs ninja + nvcc over a dozen compilation units. Cold that is ~20
# minutes, spent silently with the calling python near-idle, which reads as
# a hang rather than a build. The CPU lane, which has no GPU and wants no
# nvcc, paid it on every run.
#
# One symlink makes the guard find the binary the image already built:
# `tests/test_exl3_fork_surface.py` goes 362s (2 of 4 probes skipped on the
# timeout) -> 23s, all 4 passing. The prebake short-circuit stays as it is;
# this only removes the reason anything else has to compile.
RUN set -eu;     so=/opt/cache-baked/torch_extensions/exllamav3_ext/exllamav3_ext.so;     sp="$(/opt/venv/bin/python -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')";     if [ -f "$so" ]; then         ln -sfn "$so" "$sp/exllamav3_ext.so";         echo "exllamav3_ext linked onto the import path: $sp/exllamav3_ext.so";     else         echo "no baked exllamav3_ext at $so - importers will JIT-compile it" >&2;         exit 1;     fi

# Container launcher shim. Single-GPU: exec's `python -m arbi_serve`
# (unchanged). ARBI_TP_SIZE>1: exec's `torchrun --nproc-per-node=N -m
# arbi_serve --tp-size N` so the container serves tensor-parallel through
# `docker compose up` (the server reads WORLD_SIZE/RANK from torchrun —
# a bare python launch can't do TP>1). See scripts/docker-entrypoint.sh.
COPY arbi-serve/scripts/docker-entrypoint.sh /usr/local/bin/arbi-entrypoint.sh
RUN chmod +x /usr/local/bin/arbi-entrypoint.sh

# Stamp the git SHA this image was built from. A stale image (baked
# arbi_serve older than the repo tip) is the single most expensive bench
# footgun — see scripts/build-slim.sh (passes --build-arg ARBI_GIT_SHA)
# and the bench harness preflight, which reads this and FAILS LOUD if it
# != the repo HEAD. The engine also logs it at boot.
ARG ARBI_GIT_SHA=unknown
RUN echo "${ARBI_GIT_SHA}" > /etc/arbi-build-sha

# The exllamav3 fork revision this image resolved. It is provenance the
# arbi-serve SHA cannot stand in for: the fork carries the EXL3 GEMM kernel
# map, so the pin decides which shape table (and which reconstruct surface)
# the image actually holds, and two different pins both report version
# `1.4.2` — a version-keyed staleness check is blind by construction.
#
# ASSERTED, not merely recorded: a label that can disagree with the venv is
# worse than none, because it is trusted. uv/pip write the resolved commit
# into the dist-info `direct_url.json`, which is the installed truth; the
# build fails if the declared ARG does not match it.
ARG ARBI_EXL3_REV=unknown
RUN set -eu; \
    installed="$(/opt/venv/bin/python -c 'import glob, json; h = sorted(glob.glob("/opt/venv/lib/python*/site-packages/exllamav3-*.dist-info/direct_url.json")); print(json.load(open(h[0]))["vcs_info"]["commit_id"] if h else "")')"; \
    if [ -z "$installed" ]; then \
        echo "provenance: exllamav3 has no dist-info direct_url.json - cannot record a rev" >&2; \
        exit 1; \
    fi; \
    if [ "$ARBI_EXL3_REV" != "unknown" ] && [ "$ARBI_EXL3_REV" != "$installed" ]; then \
        echo "provenance: declared ARBI_EXL3_REV=$ARBI_EXL3_REV != installed $installed" >&2; \
        exit 1; \
    fi; \
    echo "$installed" > /etc/arbi-exllamav3-rev

# Provenance as LABELs so `docker inspect` answers "what is in this image"
# without pulling and running it — the question a stale-image investigation
# has to answer first, and the one a registry tag alone cannot.
LABEL org.opencontainers.image.source="https://github.com/arbicity/arbi-serve" \
      org.opencontainers.image.revision="${ARBI_GIT_SHA}" \
      city.arbi.exllamav3.revision="${ARBI_EXL3_REV}"

# ── The serving uid ───────────────────────────────────────────────────
# The engine runs as `arbi` (uid/gid 10001), NOT root.
#
# WHAT ROOT WAS ACTUALLY BUYING, MEASURED, NOT ASSUMED. Docker's default
# capability set does not include CAP_IPC_LOCK, CAP_SYS_ADMIN,
# CAP_SYS_RESOURCE or CAP_SYS_PTRACE — a uid-0 container here holds
# CapEff=0x00000000a80425fb, and inside it `mlock(64 MiB)` fails ENOMEM
# against an 8 MiB RLIMIT_MEMLOCK exactly as it does for any other uid.
# So the gigabytes this engine page-locks were never reaching the host
# through `mlock`: `cudaHostAlloc` pins through the NVIDIA driver's own
# get_user_pages path, which RLIMIT_MEMLOCK does not govern. Root bought
# no pinned bytes. It bought CAP_DAC_OVERRIDE — the ability to write a
# directory it does not own — and that is a permissions bug wearing a
# privilege costume. Correct ownership replaces it below.
#
# The device nodes need nothing either: the nvidia runtime exposes
# /dev/nvidia* as crw-rw-rw-, so any uid opens them (nvidia-caps, the
# 0600 MIG surface, is unused on these cards).
#
# A REAL user, not a bare numeric uid. `getpass.getuser()` and
# `os.path.expanduser("~")` both consult the passwd database, and torch
# names its Inductor scratch `/tmp/torchinductor_<username>` — a uid with
# no passwd entry turns those into a KeyError at import time, which reads
# as a torch bug rather than a missing account.
#
# 10001 is above the distro's dynamic range and clear of the base image's
# own `ubuntu` (1000), so a rebuild on a newer base cannot silently alias
# the serving account onto a pre-existing one.
ARG ARBI_UID=10001
ARG ARBI_GID=10001
RUN set -eu; \
    groupadd -g "${ARBI_GID}" arbi; \
    useradd -u "${ARBI_UID}" -g "${ARBI_GID}" -m -d /home/arbi -s /bin/bash arbi

# HOME is the engine's own, not root's. Every `~`-rooted cache in the
# tree (`_cpp_ext_locks`, `_compile_cache_env`, `cli/kv_resolve`,
# `runtime/_cumem_shim`, the flat-weight and shape-pin caches) resolves
# through this, so leaving it at /root would point a non-root process at
# a directory it cannot write and send each of those to its fallback —
# silently, since they all degrade rather than refuse.
ENV HOME=/home/arbi

# THE MOUNT POINTS ARE CREATED AND OWNED HERE, and that is what makes the
# non-root recipe work with no privileged step on a fresh deployment:
# Docker seeds an EMPTY named volume from the image path it is mounted
# over, and it copies the ownership and mode along with the content. So a
# first `docker compose up` on a clean host gets /cache and the two home
# caches already owned by uid 10001. (Volumes that a PRIOR root container
# already populated keep their root ownership — Docker only seeds an
# empty one — which is the one-time migration
# `scripts/migrate-cache-volumes-to-nonroot.sh` performs, and which the
# entrypoint's write probe refuses to boot without.)
#
# The request-path tiers are included by name. They are derived at runtime
# under the RUNTIME CACHE (`_blob_vault.runtime_cache_tier`, called from
# engine/active.py for `savepoints` and `prefix_tier`), which
# resolves to the boot cache's local `runtime-cache` leaf here; created 0700 with
# 0600 blobs, and a tier that cannot create its directory answers None
# and switches ITSELF OFF with a log line — a silent capability loss, not
# a crash. Pre-creating them owned means that path is never taken.
RUN set -eu; \
    mkdir -p \
        /cache/cute-dsl /cache/fa4-cute /cache/torch_extensions /cache/triton \
        /cache/hf /cache/xdg /cache/xgrammar /cache/tkv /cache/pycache \
        /cache/arbi-serve/budget-cache \
        /cache/arbi-serve/runtime-cache/savepoints \
        /cache/arbi-serve/runtime-cache/prefix_tier \
        /home/arbi/.cache/arbi-serve /home/arbi/.cache/torch_extensions; \
    chown -R "${ARBI_UID}:${ARBI_GID}" /cache /home/arbi; \
    chmod 0700 /cache/arbi-serve/runtime-cache/savepoints \
        /cache/arbi-serve/runtime-cache/prefix_tier

USER arbi

EXPOSE 8000
ENTRYPOINT ["/usr/local/bin/arbi-entrypoint.sh"]
CMD ["--help"]

# ── Test / dev image ──────────────────────────────────────────────────
# The slim `runtime` stage above is what we SHIP (`docker build --target
# runtime`). This `test` stage is the slim image plus the run-the-suite +
# dev tooling in one thin pure-Python layer. CI tests THIS image and ships
# `runtime`; because `test` is `FROM runtime`, every byte of the shipped
# image is exercised under test — the only difference is additive,
# server-inert test packages. Use this image for local development and to
# run the suite / benches inside the container.
#
# It is the LAST stage, so a bare `docker build` (no --target) yields this
# fuller dev/test image; pass `--target runtime` to build the slim ship
# artifact.
#
#   * pytest + pytest-asyncio (asyncio_mode=auto) + pytest-xdist (CPU lane
#     runs `-n auto`) + pytest-timeout — the test runner.
#   * httpx / pytest-httpx — FastAPI TestClient + mocked-HTTP endpoint tests
#     (httpx is also used by the serving stack).
#   * docker — tests/integration/test_observability_stack.py drives the
#     otel/grafana stack via the docker SDK.
#   * peft — test-only LoRA *reference oracle* (cross-checks our BGMV in
#     test_lora_adapter.py). The runtime never imports it; baked in here so
#     the test image is self-sufficient (no gate-time install). Pinned via
#     the lock constraints so it cannot drag in a different transformers.
#   * node (apt) — the admin console's JS tests EXECUTE the shipped module
#     rather than asserting over its text; without the binary they skip and a
#     green run covers none of the console. See the RUN line below.
#   * openai — tests/test_admin_console.py et al. exercise the standalone
#     admin console (arbi_serve/admin_console/), whose OpenAI-compat proxy
#     path imports the official `openai` SDK directly (the `ui` extra in
#     pyproject.toml). Baked in here for the same self-sufficiency reason
#     as peft.
# Transformers is a reference oracle for tests only. It is intentionally
# installed here and absent from the runtime/slim stage.
FROM runtime AS test
# `runtime` ends as uid 10001; the layers below install packages into
# /opt/venv, which that uid does not own. Root for the BUILD only — the
# stage hands the image back to `arbi` at its end, so `dev` (an alias of
# this stage) serves under the same uid as `slim`. A test image that ran
# as root would test a permission surface the ship image does not have.
USER root
# `node` — the admin console's JavaScript is TESTED by executing it.
#
# The console ships as an ES module the server serves verbatim, so the only
# way to assert what `renderConfig` (and its hundred siblings) DO, rather than
# what their source text says, is to run them. ~101 tests shell out to `node`
# for exactly that. Without the binary every one of them SKIPS, the suite
# exits 0, and a green run says nothing whatsoever about the console — a
# runtime fault there ships under a passing check. tests/_node_gate.py turns
# that silence into a failure; this line is what keeps the failure from being
# the normal state of the image.
#
# In the `test` stage ONLY: the shipped `runtime`/`slim` image never executes
# the console (a browser does), so it pays nothing. Cost here, measured on
# :test-latest: +120 MiB uncompressed on a 15.0 GiB image (+0.8 %), +30 MB
# compressed to pull. Ubuntu's `nodejs` (v18) is enough — the harness runs
# plain scripts and `--input-type=module`, no npm, no package.json, no
# registry fetch at test time.
RUN apt-get update && apt-get install -y --no-install-recommends \
        nodejs \
    && rm -rf /var/lib/apt/lists/*
COPY --from=builder /etc/arbi-lock-constraints.txt /etc/arbi-lock-constraints.txt
RUN --mount=type=cache,target=/root/.cache/uv \
    UV_LINK_MODE=copy uv pip install -c /etc/arbi-lock-constraints.txt \
        pytest pytest-asyncio pytest-xdist pytest-timeout httpx pytest-httpx docker peft openai attrs transformers
# Transformers is the reference oracle a parity test compares against, so a
# version the CPU gate did not resolve is a gate comparing against a different
# oracle. It cannot be asserted in `runtime` — it is deliberately absent there
# — so it is asserted in the stage that has it.
COPY --from=builder /workspace/arbi-serve/assert_image_backend_matches_lock.py \
     /workspace/arbi-serve/uv.lock /tmp/lockcheck/
RUN python /tmp/lockcheck/assert_image_backend_matches_lock.py \
        --lock /tmp/lockcheck/uv.lock --image "(test)" --exact \
        --behaviour-critical test
# attrs is the runtime dependency of the vendored arbi_serve_client package.
# The `audio` extra (silero-vad, s3tokenizer, onnxruntime, torchaudio) is
# deliberately left OUT of `runtime`/slim to keep the ship image minimal —
# but this IS the comprehensive test/dev image, so it belongs here: without
# it, every audio-extra-gated test (VAD backends beyond energy, Smart Turn
# v3 EOU detection, the diarizer, s3tokenizer voice cloning) only ever
# verifies its own clean-skip path in CI, never real dependency-backed
# behavior. Same `-c` lock-pinned pattern as the line above.
RUN --mount=type=cache,target=/root/.cache/uv \
    UV_LINK_MODE=copy uv pip install -c /etc/arbi-lock-constraints.txt \
        silero-vad s3tokenizer onnxruntime torchaudio
# The test/dev image adds packages the runtime precompile never saw.
RUN python -m compileall -q -j 0 /opt/venv/lib/python${PYTHON_VERSION}/site-packages

# Back to the serving uid — see the `USER root` note at the top of this stage.
USER arbi

# ─────────────────────────────────────────────────────────────────────
# User-facing build targets — call from `docker build --target …` or compose:
#   --target slim  → the minimal production SHIP image (alias of `runtime`).
#   --target dev   → slim + test/dev tooling for local iteration (alias of
#                    `test`); compose's `dev` profile bind-mounts live source
#                    over the baked package (PYTHONPATH=/work) so edits apply
#                    without a rebuild.
# `runtime` / `test` are retained as the internal/CI stage names
# (scripts/build-slim.sh + .github/workflows/ci.yml reference them — keep in
# sync). `slim`/`dev` are pure aliases (no extra layers, no size cost).
# ─────────────────────────────────────────────────────────────────────
FROM runtime AS slim
FROM test AS dev

# ──────────────────────────────────────────────────────────────────
# console — the admin UI as a thin HTTP client
# ──────────────────────────────────────────────────────────────────
# The console talks to arbi-serve over the NETWORK and imports only the
# pure-Python parts of this package (runtime_flags, playground_workloads).
# No CUDA, no engine, no model ever loaded here: it does not belong in the
# serving image, which carries the whole GPU stack to serve a web page.
#
# arbi_serve_client is NOT optional. Every /proxy/* admin route imports it, so
# the console copies the same generated source vendored into the unified wheel.
# That keeps the thin console image aligned with the admin routes in this exact
# checkout without installing the local GPU server dependency set.
FROM python:${PYTHON_VERSION}-slim AS console

WORKDIR /app

COPY arbi-serve/client/arbi_serve_client /app/arbi_serve_client
RUN pip install --no-cache-dir \
        fastapi "uvicorn[standard]" httpx websockets openai attrs \
        opentelemetry-api opentelemetry-sdk prometheus-client pydantic pyyaml \
 && rm -rf /root/.cache/pip

# The Evals tab's harness (the `evals` extra: lm-eval). Installed here rather
# than left to the operator so the committed recipe produces a console whose
# every tab works; without it the tab renders the catalogue and reports the
# install line, which is correct but is not a product.
#
# THE TRADEOFF, plainly: lm-eval depends on torch, and torch plus the dataset
# stack (datasets / pyarrow / pandas / scipy / scikit-learn) is by far the
# largest thing in this image — it dominates the ~200 MB the console would
# otherwise be. `scripts/build-slim.sh`-style size checks should expect that.
# torch comes from PyTorch's CPU wheel index because the default PyPI wheel
# bundles the CUDA runtime (several more GB of nvidia-* wheels) for a process
# that never touches a GPU: the console is an HTTP client, and the harness it
# drives scores a REMOTE endpoint, so nothing here runs a model locally.
# torch is installed FIRST so lm-eval's own resolution finds it already
# satisfied and never reaches for the CUDA build.
RUN pip install --no-cache-dir --index-url https://download.pytorch.org/whl/cpu torch \
 && pip install --no-cache-dir lm-eval \
 && rm -rf /root/.cache/pip

COPY arbi-serve/arbi_serve /app/arbi_serve

# Same non-root account, SAME uid as the serving image — not merely
# "also not root". The console shares the `arbi-serve-compile-cache`
# volume with the engine (it is where its server registry and its
# self-signed key land, the key written 0600), so a console on a
# different uid could not read back what the engine's uid wrote, and
# vice versa. Port 8899 is unprivileged, so nothing here wants uid 0.
ARG ARBI_UID=10001
ARG ARBI_GID=10001
RUN set -eu; \
    groupadd -g "${ARBI_GID}" arbi; \
    useradd -u "${ARBI_UID}" -g "${ARBI_GID}" -m -d /home/arbi -s /bin/bash arbi; \
    mkdir -p /cache; \
    chown -R "${ARBI_UID}:${ARBI_GID}" /cache /home/arbi /app
ENV HOME=/home/arbi
USER arbi

EXPOSE 8899

# --host 0.0.0.0: the console's own default is 127.0.0.1, unreachable from
# outside the container. Add --self-signed-tls to serve HTTPS (the Audio
# tab's mic capture needs a secure context).
#
# NO --engine and NO --admin-token-file here: a console that runs beside an
# engine is told BOTH by its deployment (the committed recipes pass them), and
# a default for either would be a default that is right on one deployment and
# silently wrong on the next -- an address that connects to nothing, or a
# credential file that is not the engine's. Unset, the console dials loopback
# with no credential, which is correct for a console on the engine's own host
# and is the state a `pip install arbi-serve[ui]` console starts in.
ENTRYPOINT ["python", "-m", "arbi_serve.admin_console", "--host", "0.0.0.0", "--port", "8899"]
