FROM nvidia/cuda:13.1.1-cudnn-devel-ubuntu24.04

ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHON_VERSION=3.12.7
ENV PATH=/usr/local/bin:$PATH

RUN apt update && apt install -y --no-install-recommends build-essential libssl-dev zlib1g-dev \
    libbz2-dev libreadline-dev libsqlite3-dev curl git libncursesw5-dev xz-utils tk-dev libxml2-dev \
    libxmlsec1-dev libffi-dev liblzma-dev \
    && apt clean && rm -rf /var/lib/apt/lists/*

RUN curl -O https://www.python.org/ftp/python/${PYTHON_VERSION}/Python-${PYTHON_VERSION}.tgz && \
    tar -xzf Python-${PYTHON_VERSION}.tgz && \
    cd Python-${PYTHON_VERSION} && \
    ./configure --enable-optimizations && \
    make -j$(nproc) && \
    make altinstall && \
    cd .. && \
    rm -rf Python-${PYTHON_VERSION} Python-${PYTHON_VERSION}.tgz

RUN ln -s /usr/local/bin/python3.12 /usr/local/bin/python && \
    ln -s /usr/local/bin/pip3.12 /usr/local/bin/pip

COPY requirements.txt /modded-nanogpt/requirements.txt
WORKDIR /modded-nanogpt

RUN python -m pip install --upgrade pip && \
    pip install -r requirements.txt

# [ANVIL2] The certified stack: pinned stable torch==2.10 from the cu128 wheel index, on a
# CUDA-13 base (the attention kernel links libcudart.so.13; a 12.x base cannot run it).
# The stock repo installed an unpinned nightly here; a nightly loads (the kernel is stable-ABI)
# but its numerics are unpinned and have produced NaNs in reproduction attempts. Do not substitute.
RUN pip install torch==2.10 --index-url https://download.pytorch.org/whl/cu128

CMD ["bash"]
ENTRYPOINT []