# Nanochat with GPU support (CUDA + PyTorch). Use to get a clean NCCL/CUDA environment.
# Build: docker build -t nanochat:gpu .
# Run:   docker run --gpus all -it --rm -v "$(pwd)":/workspace/nanochat -w /workspace/nanochat nanochat:gpu bash
# Test NCCL: docker run --gpus all -it --rm -v "$(pwd)":/workspace/nanochat -w /workspace/nanochat --shm-size=1g nanochat:gpu bash scripts/try_nccl_8gpu.sh

FROM nvidia/cuda:12.4.0-cudnn-runtime-ubuntu22.04

ENV DEBIAN_FRONTEND=noninteractive
RUN apt-get update && apt-get install -y --no-install-recommends \
    python3 \
    python3-pip \
    python3-venv \
    git \
    curl \
    && rm -rf /var/lib/apt/lists/*

# Install uv (install script puts it in /root/.local/bin when run as root)
RUN curl -LsSf https://astral.sh/uv/install.sh | sh
ENV PATH="/root/.local/bin:$PATH"

WORKDIR /workspace/nanochat

# Copy project (context should be repo root)
COPY pyproject.toml uv.lock ./
COPY nanochat ./nanochat
COPY scripts ./scripts
COPY runs ./runs
COPY tasks ./tasks
COPY docs ./docs

# Create venv and install with GPU extra (PyTorch cu128); venv in /opt so mount doesn't override it
RUN uv venv /opt/venv && UV_PROJECT_ENVIRONMENT=/opt/venv uv sync --extra gpu

ENV PATH="/opt/venv/bin:$PATH" VIRTUAL_ENV=/opt/venv

# Default: bash so you can run torchrun / scripts manually
CMD ["bash"]