Dockerfile.cuda

# NVIDIA CUDA base image with Python 3.11
FROM nvidia/cuda:12.2.0-devel-ubuntu22.04

# Set environment variables
ENV DEBIAN_FRONTEND=noninteractive
ENV CUDA_HOME=/usr/local/cuda
ENV PATH=${CUDA_HOME}/bin:${PATH}
ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}

# Install system dependencies
RUN apt-get update && apt-get install -y \
    python3.11 \
    python3.11-dev \
    python3.11-venv \
    python3-pip \
    build-essential \
    cmake \
    git \
    wget \
    curl \
    && rm -rf /var/lib/apt/lists/*

# Set Python 3.11 as default
RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.11 1 && \
    update-alternatives --install /usr/bin/python python /usr/bin/python3.11 1

# Upgrade pip
RUN python3 -m pip install --upgrade pip setuptools wheel

# Create working directory
WORKDIR /app

# Copy project files
COPY pyproject.toml ./
COPY pocket_agent_cli/ ./pocket_agent_cli/
COPY README.md ./

# Install llama-cpp-python with CUDA support
# CMAKE_ARGS ensures CUDA compilation
ENV CMAKE_ARGS="-DLLAMA_CUDA=on"
ENV FORCE_CMAKE=1

# Install the package with CUDA-enabled llama-cpp-python
RUN pip install --no-cache-dir llama-cpp-python --force-reinstall

# Install the rest of the dependencies
RUN pip install --no-cache-dir -e .

# Verify CUDA installation
RUN python3 -c "import llama_cpp; print('llama-cpp-python installed successfully')"

# Set up entrypoint
ENTRYPOINT ["pocket-agent"]
CMD ["--help"]