FROM nvidia/cuda:13.0.2-devel-ubuntu24.04

ENV DEBIAN_FRONTEND=noninteractive

RUN apt-get update \
    && apt-get install -y --no-install-recommends \
      build-essential \
      git \
      ninja-build \
      python3 \
      python3-dev \
      python3-pip \
      python3-venv \
    && rm -rf /var/lib/apt/lists/*

RUN python3 -m venv /opt/flashkda \
    && /opt/flashkda/bin/python -m pip install --upgrade \
      pip \
      setuptools \
      wheel \
      ninja \
      numpy \
      pytest \
      --index-url https://pypi.tuna.tsinghua.edu.cn/simple \
    && /opt/flashkda/bin/python -m pip install \
      torch==2.11.0 \
      --index-url https://download.pytorch.org/whl/cu130

ENV PATH=/opt/flashkda/bin:${PATH}
ENV CUDA_HOME=/usr/local/cuda