# Build once on CPU; no credentials or model weights are embedded in this image.
# Verified upstream PyTorch image digest, CUDA 12.8 / Torch 2.11.0.
FROM pytorch/pytorch:2.11.0-cuda12.8-cudnn9-devel@sha256:53ab3de62f6101d1e42f9be28623ab7a468a24c070d632f211ed576e30b6abd3 AS build
ENV DEBIAN_FRONTEND=noninteractive TORCH_CUDA_ARCH_LIST=9.0 MAX_JOBS=4
RUN apt-get update && apt-get install -y --no-install-recommends python3-venv \
    && rm -rf /var/lib/apt/lists/* \
    && python -m venv --system-site-packages /opt/flora-venv
ENV PATH=/opt/flora-venv/bin:$PATH
WORKDIR /workspace/app
COPY worker-requirements.txt flora-requirements.txt flora-acceleration-requirements.txt ./
# Keep the reference worker/native dependency overlay order.
RUN python -m pip install --no-cache-dir -r worker-requirements.txt
RUN python -m pip install --no-cache-dir -r flora-requirements.txt
COPY build_causal_conv1d.py ./
ENV MAX_JOBS=1
RUN python build_causal_conv1d.py
RUN python -m pip install --no-cache-dir --no-build-isolation -r flora-acceleration-requirements.txt
RUN python -c "import causal_conv1d; from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule"
COPY *.py *.txt prompt-templates.json pipeline-source.tar.gz pipeline-version.json ./
RUN python bootstrap.py

# Keep the 10-GB compiler toolchain out of each paid Job's image pull.
FROM pytorch/pytorch:2.11.0-cuda12.8-cudnn9-runtime@sha256:eee11b3b3872a8c838e35ef48f08b2d5def2080902c7f666831310ca1a0ef2be
COPY --from=build /opt/flora-venv /opt/flora-venv
COPY --from=build /workspace /workspace
WORKDIR /workspace/app
ENV PATH=/opt/flora-venv/bin:$PATH
ENV PYTHONPATH=/workspace/source/src:/workspace/app PYTHONUNBUFFERED=1 HF_HUB_DISABLE_TELEMETRY=1
ENV FLORA_EXECUTION=accelerated FLORA_BATCH_SIZE=8 FLORA_ENABLE_INFERENCE=1
RUN python -c "import worker, torch, causal_conv1d; from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule; assert torch.__version__.split('+')[0] == '2.11.0'"
# CPU Space builds expose only health information. A GPU Job overrides this CMD
# with python /workspace/app/worker.py and supplies its own scoped HF secret.
EXPOSE 7860
CMD ["python", "flora_image_health.py"]