# Use the official NVIDIA CUDA image with cuDNN support
FROM nvidia/cuda:11.3.1-cudnn8-runtime-ubuntu20.04

# Set the working directory in the container
WORKDIR /app

# Install Python, Rust, Cargo, and other dependencies
RUN apt-get update && apt-get install -y \
    python3 \
    python3-pip \
    curl \
    && curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y \
    && . "$HOME/.cargo/env" \
    && rm -rf /var/lib/apt/lists/*

# Copy the requirements.txt file into the container at /app
COPY requirements.txt .

# Install the dependencies
RUN . "$HOME/.cargo/env" && pip3 install --no-cache-dir -r requirements.txt

# Copy the rest of the application code into the container at /app
COPY . .

# Run the training script
CMD ["python3", "train.py"]