containers/train/Dockerfile

FROM python:3.9-slim

RUN \
    apt-get update && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*

RUN useradd --create-home app
WORKDIR /home/app

COPY requirements.txt /home/app/
COPY run.py /home/app/

RUN \
    chown app:app /home/app/requirements.txt && \
    chmod 0755 /home/app/requirements.txt && \
    chown app:app /home/app/run.py && \
    chmod 0755 /home/app/run.py

RUN \
    pip install --upgrade pip

# Install PyTorch separately from CPU index
RUN pip install --index-url https://download.pytorch.org/whl/cpu torch>=2.0.0

# Install other requirements from default PyPI
RUN pip install transformers>=4.36.0 pillow>=10.0.0 datasets>=2.14.0

USER app

CMD ["python", "run.py", "worker", "-l", "info"]