From 7e8492059c5d84a97acabd916b00b44658f18ea2 Mon Sep 17 00:00:00 2001
From: Chris Wilhelm <chrisw@allenai.org>
Date: Thu, 13 Mar 2025 14:10:55 -0700
Subject: [PATCH] wip

---
 pyproject.toml                   |  1 -
 scripts/beaker/Dockerfile-gpu-ci |  1 +
 scripts/beaker/gpu-ci-script.sh  | 15 ++++++++-------
 setup.py                         | 12 ------------
 4 files changed, 9 insertions(+), 20 deletions(-)
 mode change 100644 => 100755 scripts/beaker/gpu-ci-script.sh
 delete mode 100644 setup.py

diff --git a/pyproject.toml b/pyproject.toml
index 2161e08..ba8fbec 100644
--- a/pyproject.toml
+++ b/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 = [
diff --git a/scripts/beaker/Dockerfile-gpu-ci b/scripts/beaker/Dockerfile-gpu-ci
index 3dfb164..d51b592 100644
--- a/scripts/beaker/Dockerfile-gpu-ci
+++ b/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
diff --git a/scripts/beaker/gpu-ci-script.sh b/scripts/beaker/gpu-ci-script.sh
old mode 100644
new mode 100755
index 0e863a4..e68d337
--- a/scripts/beaker/gpu-ci-script.sh
+++ b/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
 
 
 
diff --git a/setup.py b/setup.py
deleted file mode 100644
index abd2f8e..0000000
--- a/setup.py
+++ /dev/null
@@ -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/"],
-)