ARG BASE_IMAGE=nvcr.io/nvidia/pytorch:26.05-py3
FROM ${BASE_IMAGE}

ENV PYTHONDONTWRITEBYTECODE=1 \
    PYTHONUNBUFFERED=1 \
    TORCH_ALLOW_TF32_CUBLAS_OVERRIDE=0 \
    NVIDIA_TF32_OVERRIDE=0 \
    TRITON_CACHE_DIR=/work/cache/triton \
    TORCH_EXTENSIONS_DIR=/work/cache/torch-extensions

WORKDIR /opt/packed-adamw-reference
COPY . .
RUN python3 -m pip install --no-cache-dir --no-deps . \
    && python3 -m pip install --no-cache-dir "pytest>=8,<10" \
    && rm -rf build src/*.egg-info

ENTRYPOINT ["packed-adamw"]
CMD ["doctor"]
