# Pennyroyal (jpezzulli/sglang-rtxpro6000): patched SGLang for Qwen3.8
# Flash-Next + 27B/DFlash2 on a single RTX PRO 6000 Blackwell (SM120, TP=1).
#
# There is no prebuilt Pennyroyal image -- the repo's docker/ dir is
# upstream SGLang's generic material, which the README explicitly says is
# NOT a qualified Pennyroyal recipe.  So this containerizes the native
# BUILD.md source build: SGLang (editable, -e python) + NIXL POSIX plugin
# (for HiCache/NIXL KV persistence), both compiled from source against a
# consistent CUDA 13.3 toolchain.
#
# Heavy build: compiles sglang-kernel, FlashInfer and NIXL from source
# (tens of minutes).  Base is 13.3 to match BUILD.md's qualified CUDA.
FROM nvidia/cuda:13.3.1-devel-ubuntu24.04

ENV DEBIAN_FRONTEND=noninteractive

# gcc-15 / g++-15 are the qualified host compilers (BUILD.md).  Ubuntu 24.04
# ships gcc-13, so pull 15 from the ubuntu-toolchain-r test PPA.  Also the
# NIXL build deps (cmake/meson/ninja/pkg-config/libaio-dev) and Rust.
RUN apt-get update && apt-get install -y --no-install-recommends \
        software-properties-common ca-certificates curl git \
    && add-apt-repository -y ppa:ubuntu-toolchain-r/test \
    && apt-get update && apt-get install -y --no-install-recommends \
        gcc-15 g++-15 \
        python3.12 python3.12-venv python3.12-dev \
        cmake meson ninja-build pkg-config libaio-dev \
    && rm -rf /var/lib/apt/lists/*

# Rust via rustup, NOT apt.  Pennyroyal's rust/Cargo.toml is `edition =
# "2024"` (needs Rust >= 1.85); Ubuntu 24.04's apt rustc is 1.75 and fails
# with "failed to parse manifest".  Qualified table lists Rust 1.97.1.
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs \
        | sh -s -- -y --default-toolchain 1.97.1 --profile minimal
ENV PATH=/root/.cargo/bin:$PATH

# uv (BUILD.md uses it for the pinned python + editable install).
RUN curl -LsSf https://astral.sh/uv/install.sh | env UV_INSTALL_DIR=/usr/local/bin sh

ENV CUDA_HOME=/usr/local/cuda \
    CUDACXX=/usr/local/cuda/bin/nvcc \
    CC=/usr/bin/gcc-15 CXX=/usr/bin/g++-15 CUDAHOSTCXX=/usr/bin/g++-15 \
    TORCH_CUDA_ARCH_LIST=12.0 \
    PENNY_BUILD_JOBS=4

WORKDIR /opt

# --- SGLang (Pennyroyal branch) --------------------------------------------
ARG PENNY_TAG=pennyroyal-v2.5.0
RUN git clone --branch ${PENNY_TAG} --single-branch \
        https://github.com/jpezzulli/sglang-rtxpro6000.git pennyroyal
WORKDIR /opt/pennyroyal
RUN uv python install 3.12.13 \
    && uv venv --python 3.12.13 .venv
ENV VIRTUAL_ENV=/opt/pennyroyal/.venv PATH=/opt/pennyroyal/.venv/bin:$PATH
# --no-build-isolation means every build-time dep must already be present.
# wheel_stub: NVIDIA cuda-tile/cuda-* wheels' build backend.
# pybind11 / patchelf / meson-python / pyyaml: NIXL's meson build resolves
# `dependency('pybind11')` and the meta-wheel packaging against the venv.
RUN uv pip install pip "setuptools>=80.9.0" "setuptools-rust>=1.10" \
        "setuptools-scm>=8.0" wheel build tomlkit wheel_stub \
        pybind11 patchelf meson-python pyyaml pytest
# --- NIXL POSIX plugin (for HiCache/NIXL) ----------------------------------
# Built BEFORE the SGLang editable install: NIXL's meson build spawns an
# isolated `uv build --wheel` for its python bindings, and if SGLang's
# editable `__editable__.sglang*.pth` is already on the venv, that
# sitecustomize hook runs in the sub-build with a broken sys.path and dies
# with a bogus `ModuleNotFoundError: No module named 'traceback'`.  Building
# NIXL first (clean venv) sidesteps it.  asio's meson wrap has a GitHub
# fallback URL, so the retry loop rides out the sourceforge flake.
ARG NIXL_COMMIT=aecbc3846d92c34c7507a58d776e1fda50ff4fba
ENV NIXL_PREFIX=/opt/nvidia/nvda_nixl
RUN git clone https://github.com/ai-dynamo/nixl.git /opt/nixl \
    && cd /opt/nixl && git checkout ${NIXL_COMMIT} \
    && python -m pip install . \
    && ./contrib/tomlutil.py --wheel-name nixl-cu13 pyproject.toml \
    && for i in 1 2 3 4 5; do \
        meson setup build-posix --buildtype=release \
            --prefix="$NIXL_PREFIX" --libdir=lib64 \
            -Denable_plugins=POSIX -Dnixl_cuda_arch_list=120 \
        && ninja -C build-posix -j ${PENNY_BUILD_JOBS} install && break; \
        echo "nixl attempt $i failed, retrying in 15s"; rm -rf build-posix; sleep 15; \
    done \
    && python -m pip install build-posix/src/bindings/python/nixl-meta/nixl-*-py3-none-any.whl
ENV LD_LIBRARY_PATH=$NIXL_PREFIX/lib64

# --- SGLang editable install (LAST, so its .pth can't break NIXL's build) --
# UV_HTTP_TIMEOUT: the default 30 s times out on the big wheels (triton,
# nccl ~200 MB) on a slow link.  Retry loop rides out transient network
# drops -- uv's cache makes each retry resume, not restart.
RUN export MAX_JOBS=$PENNY_BUILD_JOBS \
        CMAKE_BUILD_PARALLEL_LEVEL=$PENNY_BUILD_JOBS \
        CARGO_BUILD_JOBS=$PENNY_BUILD_JOBS \
        FLASHINFER_NINJA_JOBS=$PENNY_BUILD_JOBS \
        FLASHINFER_NVCC_THREADS=1 \
        TORCHINDUCTOR_COMPILE_THREADS=$PENNY_BUILD_JOBS \
        UV_HTTP_TIMEOUT=600 \
    && for i in 1 2 3 4 5; do \
        uv pip install --prerelease=allow --index-strategy unsafe-best-match \
            --extra-index-url https://docs.sglang.ai/whl/cu130/ \
            --no-build-isolation -e python && break; \
        echo "attempt $i failed, retrying in 15s"; sleep 15; \
    done

# libssl-dev: HiCache JIT-compiles a C++ hash extension (hicache_hash_cpp)
# at first cached-request time, which #include's <openssl/sha.h>.  Runtime
# openssl is present but not its headers.  Added as a LATE layer so the
# expensive NIXL + SGLang compile layers above stay cached.
RUN apt-get update && apt-get install -y --no-install-recommends libssl-dev \
    && rm -rf /var/lib/apt/lists/*

# Optional NVMe-PLE reader (ssd_stream, a Rust package).  Streams the 47.7
# GiB PLE table from a prepared on-disk overlay instead of holding it in
# host shmem, which frees ~50 GiB of host RAM.  Built into .ple-nvme;
# activated at launch with PENNY_PLE_BACKEND=nvme + PENNY_PLE_NVME_MODEL.
# Late layer so it does not invalidate the big compile layers.
RUN PYTHON=/opt/pennyroyal/.venv/bin/python \
    bash /opt/pennyroyal/tools/ple_nvme/install.sh

WORKDIR /opt/pennyroyal
# Launchers read TARGET_MODEL / CACHE_BASE / NIXL_STORAGE_BASE from env and
# run sglang; invoke the FR-Spec Flash-Next launcher as the entrypoint.
ENTRYPOINT ["bash", "configs/pennyroyal/serve-flash-next-frspec.sh"]