FROM nvidia/cuda:12.2.0-base-ubuntu22.04

# Work in a non-interactive environment
ENV DEBIAN_FRONTEND=noninteractive

# Basic system dependencies
RUN apt-get update && \
    apt-get install -y --no-install-recommends \
        python3 \
        python3-pip \
        python3-dev \
        ca-certificates \
        curl \
        git \
    && rm -rf /var/lib/apt/lists/*

# Set locale and working directory
ENV LANG C.UTF-8
WORKDIR /workspace

# Install Python dependencies first (cache-friendly)
COPY requirements.txt .
RUN python3 -m pip install --upgrade pip
RUN python3 -m pip install -r requirements.txt

# Copy the repository contents
COPY . .

# Optional: reduce memory usage of JAX if needed (commented by default)
# ENV XLA_PYTHON_CLIENT_PREALLOCATE=false

# Default command to run the project
CMD ["python3", "run.py"]

# See https://docs.docker.com/engine/reference/builder/#dockerignore
.git
.gitignore
__pycache__/
.pyc
*.pyo
*.pyd
venv/
env/
build/
dist/
*.egg-info/
CHECKPOINTS/
checkpoints/

Summary (under 1500 chars):

- Failing step: [6/7] RUN python3 -m pip install -r requirements.txt
  - Exit code: 1
  - Error (exact messages):
    - ERROR: Could not find a version that satisfies the requirement jaxlib==0.4.25+cuda12.cudnn89; extra == "cuda12-pip" (from jax[cuda12-pip]) (from versions: 0.4.17, 0.4.18, 0.4.19, 0.4.20, 0.4.21, 0.4.22, 0.4.23, 0.4.24, 0.4.25, 0.4.26, 0.4.27, 0.4.28, 0.4.29, 0.4.30, 0.4.31, 0.4.33, 0.4.34, 0.4.35, 0.4.36, 0.4.38, 0.5.0, 0.5.1, 0.5.3, 0.6.0, 0.6.1, 0.6.2)
    - ERROR: No matching distribution found for jaxlib==0.4.25+cuda12.cudnn89; extra == "cuda12-pip"
- Missing packages/files mentioned:
  - The CUDA-specific wheel jaxlib==0.4.25+cuda12.cudnn89 required by jax[cuda12-pip] could not be found for the current environment.
- Version mismatch info:
  - The log shows many jaxlib entries listed as requiring Python >=3.11, while the base image is using Python 3.10 (seen in the pip install context and removed/ignored versions due to Python version constraints). This Python version mismatch likely prevents the requested cuda12-pip variant from resolving.
  - Context: requirements.txt pulls jax[cuda12-pip], which demands a jaxlib variant incompatible with Python 3.10 in this container.