wip

pyproject.toml

@@ -51,7 +51,6 @@ Changelog = "https://github.com/allenai/olmocr/blob/main/CHANGELOG.md"
 gpu = [
     "sgl-kernel==0.0.3.post1",
     "sglang[all]==0.4.2",
-    "flashinfer"
 ]
 
 dev = [

scripts/beaker/Dockerfile-gpu-ci

@@ -32,5 +32,6 @@ RUN /install.sh && rm /install.sh
 
 
 WORKDIR /root
+COPY gpu-ci-script.sh .
 
 ENV PYTHONUNBUFFERED=1

scripts/beaker/gpu-ci-script.sh

@@ -1,14 +1,15 @@
+#/usr/bin/bash
+
 set -ex
 
-git clone git@github.com:allenai/olmocr.git olmocr \
+git clone https://github.com/allenai/olmocr.git olmocr \
   && cd olmocr \
   && git checkout $GIT_REVISION \
-  && /root/.local/bin/uv pip install --system --no-cache . \
-  && /root/.local/bin/uv pip install --system --no-cache sgl-kernel==0.0.3.post1 --force-reinstall --no-deps \
-  && /root/.local/bin/uv pip install --system --no-cache "sglang[all]==0.4.2" --find-links https://flashinfer.ai/whl/cu124/torch2.4/flashinfer/
-  && /root/.local/bin/uv pip install --system --no-cache pytest \
-  && python -m olmocr.pipeline ./localworkspace --pdfs tests/gnarly_pdfs/*.pdf \
-  && python tests/gnarly_pdfs/evaluate_gnarly_pdfs.py
+  && /root/.local/bin/uv pip install --system --no-cache \
+    .[gpu] \
+    pytest \
+    --find-links https://flashinfer.ai/whl/cu124/torch2.4/flashinfer/ \
+  && pytest tests/gnarly_pdfs/test_evaluate_gnarly_pdfs.py
 
 
 

setup.py

@@ -1,12 +0,0 @@
-from setuptools import setup
-
-setup(
-    extras_require={
-        "gpu": [
-            "sgl-kernel==0.0.3.post1",
-            "sglang[all]==0.4.2",
-            "flashinfer"
-        ],
-    },
-    dependency_links=["https://flashinfer.ai/whl/cu124/torch2.4/flashinfer/"],
-)