Dockerfile.bench

# ── Stage 1: C++ build (re-uses the main Dockerfile builder stage) ────────────
FROM servoflow:latest AS sf_builder

# ── Stage 2: benchmark image ──────────────────────────────────────────────────
# Start from the ServoFlow builder so we have the compiled binaries,
# then add Python + PyTorch + HuggingFace Transformers on top.
FROM nvidia/cuda:12.4.1-devel-ubuntu22.04

ENV DEBIAN_FRONTEND=noninteractive \
    PYTHONUNBUFFERED=1

# System packages
RUN apt-get update && apt-get install -y --no-install-recommends \
        python3 python3-pip python3-dev \
        git wget ca-certificates \
        libgomp1 \
        cmake build-essential \
    && apt-get clean && rm -rf /var/lib/apt/lists/*

RUN ln -sf /usr/bin/python3 /usr/bin/python

# PyTorch + HuggingFace stack (match CUDA 12.4)
RUN pip install --no-cache-dir --upgrade pip && \
    pip install --no-cache-dir \
        torch torchvision --index-url https://download.pytorch.org/whl/cu124 && \
    pip install --no-cache-dir \
        transformers \
        accelerate \
        safetensors \
        sentencepiece \
        huggingface_hub

WORKDIR /workspace/servoflow

# Copy compiled ServoFlow binaries and benchmark scripts from the builder
COPY --from=sf_builder /workspace/servoflow/build/benchmarks ./build/benchmarks
COPY --from=sf_builder /workspace/servoflow/build/tests      ./build/tests
COPY --from=sf_builder /workspace/servoflow/build/libservoflow.so /usr/local/lib/
COPY benchmarks/ ./benchmarks/

RUN ldconfig

ENV LD_LIBRARY_PATH=/usr/local/lib:${LD_LIBRARY_PATH}

CMD ["/bin/bash"]