Dockerfile

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

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

WORKDIR /app

# Install basic Python packages first
RUN pip3 install --no-cache-dir \
    packaging \
    setuptools \
    wheel \
    numpy \
    torch==2.4.0

# Install CUDA toolkit
ENV CUDA_HOME=/usr/local/cuda
ENV PATH=${CUDA_HOME}/bin:${PATH}
ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}

# Install dependencies in order
COPY requirements.txt .
RUN pip3 install --no-cache-dir \
    transformers==4.43.2 \
    accelerate \
    peft \
    datasets \
    sentencepiece \
    protobuf \
    tiktoken \
    scipy \
    gradio \
    cn2an>=0.5.22 \
    langdetect>=1.0.9 \
    openai \
    tqdm \
    && pip3 install --no-cache-dir flash-attn --no-build-isolation \
    && pip3 install --no-cache-dir vllm==0.5.5 vllm-flash-attn

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

# Install LLaMA Factory
RUN pip3 install --no-cache-dir llamafactory

# 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"]