From a83a0da65f2edae97ae8ea37fe6803979f6c180f Mon Sep 17 00:00:00 2001 From: Jake Poznanski Date: Tue, 10 Jun 2025 21:56:05 +0000 Subject: [PATCH] Cleanup of vllm perf branch with @amanr --- Dockerfile | 57 +++++++++++++++++++++++++--------------------- olmocr/pipeline.py | 2 -- 2 files changed, 31 insertions(+), 28 deletions(-) diff --git a/Dockerfile b/Dockerfile index 2ce2b18..d07fca9 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,60 +1,65 @@ ARG CUDA_VERSION=12.8.1 FROM --platform=linux/amd64 nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.04 + # Needs to be repeated below the FROM, or else it's not picked up ARG PYTHON_VERSION=3.12 ARG CUDA_VERSION=12.8.1 + # Set environment variable to prevent interactive prompts ENV DEBIAN_FRONTEND=noninteractive + # From original VLLM dockerfile https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile # Install Python and other dependencies RUN echo 'tzdata tzdata/Areas select America' | debconf-set-selections \ -&& echo 'tzdata tzdata/Zones/America select Los_Angeles' | debconf-set-selections \ -&& apt-get update -y \ -&& apt-get install -y ccache software-properties-common git curl sudo python3-apt \ -&& for i in 1 2 3; do \ -add-apt-repository -y ppa:deadsnakes/ppa && break || \ -{ echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \ -done \ -&& apt-get update -y \ -&& apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv + && echo 'tzdata tzdata/Zones/America select Los_Angeles' | debconf-set-selections \ + && apt-get update -y \ + && apt-get install -y ccache software-properties-common git curl sudo python3-apt \ + && for i in 1 2 3; do \ + add-apt-repository -y ppa:deadsnakes/ppa && break || \ + { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \ + done \ + && apt-get update -y \ + && apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv # olmOCR Specific Installs - Install fonts BEFORE changing Python version RUN echo "ttf-mscorefonts-installer msttcorefonts/accepted-mscorefonts-eula select true" | debconf-set-selections && \ -apt-get update -y && \ -apt-get install -y --no-install-recommends poppler-utils fonts-crosextra-caladea fonts-crosextra-carlito gsfonts lcdf-typetools ttf-mscorefonts-installer + apt-get update -y && \ + apt-get install -y --no-install-recommends poppler-utils fonts-crosextra-caladea fonts-crosextra-carlito gsfonts lcdf-typetools ttf-mscorefonts-installer # Now update Python alternatives RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \ -&& update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \ -&& update-alternatives --install /usr/bin/python python /usr/bin/python${PYTHON_VERSION} 1 \ -&& update-alternatives --set python /usr/bin/python${PYTHON_VERSION} \ -&& ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \ -&& curl -sS https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION} \ -&& python3 --version && python3 -m pip --version + && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \ + && update-alternatives --install /usr/bin/python python /usr/bin/python${PYTHON_VERSION} 1 \ + && update-alternatives --set python /usr/bin/python${PYTHON_VERSION} \ + && ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \ + && curl -sS https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION} \ + && python3 --version && python3 -m pip --version # Install uv for faster pip installs -RUN --mount=type=cache,target=/root/.cache/uv \ -python3 -m pip install uv +RUN --mount=type=cache,target=/root/.cache/uv python3 -m pip install uv # Install some helper utilities for things like the benchmark RUN apt-get update -y && apt-get install -y --no-install-recommends \ -git \ -git-lfs \ -curl \ -wget \ -unzip + git \ + git-lfs \ + curl \ + wget \ + unzip + ENV PYTHONUNBUFFERED=1 WORKDIR /root COPY pyproject.toml pyproject.toml COPY olmocr/version.py olmocr/version.py + # Needed to resolve setuptools dependencies ENV UV_INDEX_STRATEGY="unsafe-best-match" -RUN uv pip install --system --no-cache -e . --extra-index-url https://download.pytorch.org/whl/cu128 -RUN uv pip install --system --no-cache ".[gpu]" --extra-index-url https://download.pytorch.org/whl/cu128 +RUN uv pip install --system --no-cache -e ".[gpu]" --extra-index-url https://download.pytorch.org/whl/cu128 RUN uv pip install --system https://download.pytorch.org/whl/cu128/flashinfer/flashinfer_python-0.2.5%2Bcu128torch2.7-cp38-abi3-linux_x86_64.whl RUN uv pip install --system --no-cache ".[bench]" RUN playwright install-deps RUN playwright install chromium + COPY olmocr olmocr COPY scripts scripts + RUN python3 -m olmocr.pipeline --help \ No newline at end of file diff --git a/olmocr/pipeline.py b/olmocr/pipeline.py index def4c4f..0899d77 100644 --- a/olmocr/pipeline.py +++ b/olmocr/pipeline.py @@ -574,7 +574,6 @@ async def vllm_server_task(model_name_or_path, args, semaphore): "--disable-log-requests", "--uvicorn-log-level", "warning", "--served-model-name", "Qwen/Qwen2-VL-7B-Instruct", - "--dtype", args.dtype, ] cmd.extend(mem_fraction_arg) @@ -987,7 +986,6 @@ async def main(): help="List of paths where you can find the model to convert this pdf. You can specify several different paths here, and the script will try to use the one which is fastest to access", default="allenai/olmOCR-7B-0225-preview", ) - parser.add_argument("--dtype", type=str, default="bfloat16", help="Data type for model weights") parser.add_argument("--model_max_context", type=int, default="8192", help="Maximum context length that the model was fine tuned under") parser.add_argument("--model_chat_template", type=str, default="qwen2-vl", help="Chat template to pass to vllm server") parser.add_argument("--target_longest_image_dim", type=int, help="Dimension on longest side to use for rendering the pdf pages", default=1024)