FROM nvidia/cuda:12.1.0-runtime-ubuntu22.04

# Install system dependencies
RUN apt-get update && apt-get install -y \
    git \
    python3.10 \
    python3-pip \
    && rm -rf /var/lib/apt/lists/*

# Set working directory
WORKDIR /app

# Install Python packages
COPY requirements.txt .
RUN pip3 install --no-cache-dir -r requirements.txt

# Install custom FastChat version
RUN git clone -b self-lengthen https://github.com/quanshr/FastChat.git && \
    cd FastChat && \
    pip3 install -e ".[model_worker,webui]"

# Copy project files
COPY . .

# Set environment variables
ENV CUDA_VISIBLE_DEVICES=0
ENV WORLD_SIZE=1
ENV RANK=0
ENV MASTER_ADDR=localhost
ENV MASTER_PORT=29500

# Create startup script
RUN echo '#!/bin/bash\n\
cd /app/qwen\n\
bash run.sh --base_model=$MODEL_PATH --instruct_count=$INSTRUCT_COUNT --max_iter=$MAX_ITER\n\
python collect_data.py' > /app/start.sh && \
chmod +x /app/start.sh

# Command to run
ENTRYPOINT ["/app/start.sh"]