From 39256c19bbf4e757a528fb7fc647738e8db7ee92 Mon Sep 17 00:00:00 2001
From: Jake Poznanski <jakep@allenai.org>
Date: Wed, 13 Nov 2024 10:25:35 -0800
Subject: [PATCH] Beaker running

---
 pdelfin/beakerpipeline.py           | 6 ++++--
 pyproject.toml                      | 2 +-
 scripts/beaker/Dockerfile-inference | 2 +-
 scripts/build-docker.sh             | 8 +++++++-
 4 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/pdelfin/beakerpipeline.py b/pdelfin/beakerpipeline.py
index 6961c05..bb1c7b7 100644
--- a/pdelfin/beakerpipeline.py
+++ b/pdelfin/beakerpipeline.py
@@ -515,6 +515,8 @@ async def metrics_reporter():
         logger.info("\n" + str(await tracker.get_status_table()))
         await asyncio.sleep(10)
 
+
+
 def submit_beaker_job(args):
     from beaker import (
         Beaker,
@@ -534,12 +536,12 @@ def submit_beaker_job(args):
     
     b = Beaker.from_env(default_workspace=args.beaker_workspace)
     account = b.account.whoami()
-    beaker_image = "ai2/cuda11.8-ubuntu20.04"
+    beaker_image = "jakep/pdelfin-inference"
 
     task_name = f"pdelfin-{os.path.basename(args.workspace.rstrip('/'))}"
     priority = "normal"
 
-    args_list = sum(([f"--{k}", str(v)] if not isinstance(v, bool) else [f"--{k}"] for k, v in vars(args).items() if v is not None), [])
+    args_list = [arg for arg in sys.argv[1:] if arg != "--beaker"]
 
     # Create the experiment spec
     experiment_spec = ExperimentSpec(
diff --git a/pyproject.toml b/pyproject.toml
index df8792d..4fa82ea 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -36,7 +36,7 @@ dependencies = [
   "zstandard",
   "aiohttp",
   "boto3",
-  "torch>=2.5.0",
+  "torch>=2.4.0",
   "transformers>=4.46.2",
 ]
 license = {file = "LICENSE"}
diff --git a/scripts/beaker/Dockerfile-inference b/scripts/beaker/Dockerfile-inference
index a809551..5a8ecd0 100644
--- a/scripts/beaker/Dockerfile-inference
+++ b/scripts/beaker/Dockerfile-inference
@@ -34,7 +34,7 @@ RUN /root/.local/bin/uv pip install --system --no-cache -e .[inference]
 
 # TODO You can remove this and move it into the pyproject.toml once sglang makes a release > 0.35.0
 RUN /root/.local/bin/uv pip install --system git+https://github.com/sgl-project/sglang.git@eff468dd5a3d24646560eb044276585f7a11ac3c#subdirectory=python&egg=sglang[all]
-
+RUN /root/.local/bin/uv pip install --system flashinfer -i https://flashinfer.ai/whl/cu121/torch2.4/
 
 RUN python3 -m pdelfin.beakerpipeline --help
 
diff --git a/scripts/build-docker.sh b/scripts/build-docker.sh
index ca1dbd7..1539bef 100644
--- a/scripts/build-docker.sh
+++ b/scripts/build-docker.sh
@@ -1 +1,7 @@
-docker build -f ./scripts/beaker/Dockerfile-inference  -t pdelfin-inference .
\ No newline at end of file
+#!/bin/bash
+
+set -e
+
+docker build --platform linux/amd64 -f ./scripts/beaker/Dockerfile-inference  -t pdelfin-inference .
+
+beaker image create --workspace ai2/oe-data-pdf --name pdelfin-inference pdelfin-inference
\ No newline at end of file