Beaker running

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(

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"}

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
 

scripts/build-docker.sh

@@ -1 +1,7 @@
-docker build -f ./scripts/beaker/Dockerfile-inference  -t pdelfin-inference .
+#!/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