Merge pull request #138 from allenai/favyen/20250527-beaker-upgrad

Use Beaker queues instead of Google Cloud Pub/Sub for the worker system.

Author: favyen2

.github/workflows/build_test.yaml

@@ -130,8 +130,9 @@ jobs:
           echo "Executing integration tests in ${TEST_DIRS}"
           docker compose -f docker-compose.yaml run \
             -e CI="true" \
-            -e AWS_ACCESS_KEY_ID=${{ secrets.AWS_ACCESS_KEY_ID }} \
-            -e AWS_SECRET_ACCESS_KEY=${{ secrets.AWS_SECRET_ACCESS_KEY }} \
+            -e AWS_ACCESS_KEY_ID="${AWS_ACCESS_KEY_ID}" \
+            -e AWS_SECRET_ACCESS_KEY="${AWS_SECRET_ACCESS_KEY}" \
+            -e BEAKER_TOKEN="${BEAKER_TOKEN}" \
             -v ${{env.GOOGLE_GHA_CREDS_PATH}}:/tmp/gcp-credentials.json:ro \
             -e GOOGLE_APPLICATION_CREDENTIALS=/tmp/gcp-credentials.json \
             -e RSLP_BUCKET=rslearn-eai \
@@ -143,8 +144,12 @@ jobs:
             -e TEST_PUBSUB_PROJECT=earthsystem-dev-c3po \
             -e TEST_PUBSUB_TOPIC=rslearn_projects_test_topic \
             -e TEST_PUBSUB_SUBSCRIPTION=rslearn_projects_test_subscription \
+            -e TEST_QUEUE_NAME=favyen/rslearn-projects-test-queue \
             test pytest ${TEST_DIRS} -vv
-
+        env:
+          AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
+          AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
+          BEAKER_TOKEN: ${{ secrets.BEAKER_TOKEN_2 }}
 
       - name: Clean up
         if: always()

helios.Dockerfile

@@ -1,4 +1,4 @@
-FROM pytorch/pytorch:2.5.0-cuda11.8-cudnn9-runtime@sha256:d15e9803095e462e351f097fb1f5e7cdaa4f5e855d7ff6d6f36ec4c2aa2938ea
+FROM pytorch/pytorch:2.7.0-cuda11.8-cudnn9-runtime
 
 RUN apt update
 RUN apt install -y libpq-dev ffmpeg libsm6 libxext6 git wget

one_off_projects/convert_satlas_webmercator_to_rslearn/lib/__init__.py

@@ -109,7 +109,7 @@ def convert_window(
 
     layer_name = "label"
     layer_dir = window.get_layer_dir(layer_name)
-    GeojsonVectorFormat().encode_vector(layer_dir, dst_projection, features)
+    GeojsonVectorFormat().encode_vector(layer_dir, features)
     window.mark_layer_completed(layer_name)
 
     # (3) Write mask corresponding to old window projected onto new window.

requirements.txt

@@ -1,4 +1,4 @@
-beaker-py>=1.32,<2.0
+beaker-py>=2.0
 fastapi>=0.115
 google-cloud-bigtable>=2.18
 google-cloud-pubsub>=2.18

rslp/common/beaker_launcher.py

@@ -3,8 +3,8 @@
 import uuid
 from datetime import datetime
 
-from beaker import Beaker, EnvVar, ExperimentSpec, ImageSource
-from beaker.exceptions import ImageNotFound
+from beaker import Beaker, BeakerEnvVar, BeakerExperimentSpec, BeakerImageSource
+from beaker.exceptions import BeakerImageNotFound
 
 from rslp.log_utils import get_logger
 from rslp.utils.beaker import (
@@ -45,7 +45,7 @@ def launch_job(
     gpu_count: int = 0,
     shared_memory: str | None = None,
     priority: str = DEFAULT_PRIORITY,
-    task_specific_env_vars: list[EnvVar] = [],
+    task_specific_env_vars: list[BeakerEnvVar] = [],
     budget: str = DEFAULT_BUDGET,
     workspace: str = DEFAULT_WORKSPACE,
     preemptible: bool = True,
@@ -80,8 +80,7 @@ def launch_job(
 
     logger.info("Starting Beaker client...")
     logger.info(f"Workspace: {workspace}")
-    beaker = Beaker.from_env(default_workspace=workspace)
-    with beaker.session():
+    with Beaker.from_env(default_workspace=workspace) as beaker:
         logger.info("Getting base env vars...")
         base_env_vars = get_base_env_vars()
         logger.info("Generating task name...")
@@ -92,8 +91,8 @@ def launch_job(
         try:
             beaker.image.get(image)
             logger.info(f"Image already exists: {image}")
-            image_source = ImageSource(beaker=image)
-        except ImageNotFound:
+            image_source = BeakerImageSource(beaker=image)
+        except BeakerImageNotFound:
             logger.info(f"Uploading image: {image}")
             # Handle image upload
             image_source = upload_image(image, workspace, beaker)
@@ -102,7 +101,7 @@ def launch_job(
         logger.info("Creating experiment spec...")
         datasets = [create_gcp_credentials_mount()]
         datasets += [weka_mount.to_data_mount() for weka_mount in weka_mounts]
-        experiment_spec = ExperimentSpec.new(
+        experiment_spec = BeakerExperimentSpec.new(
             budget=budget,
             task_name=unique_task_name,
             beaker_image=image_source.beaker,

rslp/common/beaker_train.py

@@ -4,7 +4,13 @@
 import shutil
 import uuid
 
-from beaker import Beaker, Constraints, EnvVar, ExperimentSpec, Priority, TaskResources
+from beaker import (
+    Beaker,
+    BeakerConstraints,
+    BeakerEnvVar,
+    BeakerExperimentSpec,
+    BeakerTaskResources,
+)
 
 from rslp import launcher_lib
 from rslp.utils.beaker import (
@@ -31,6 +37,7 @@ def beaker_train(
     project_id: str | None = None,
     experiment_id: str | None = None,
     extra_args: list[str] = [],
+    priority: str = "high",
 ) -> None:
     """Launch training for the specified config on Beaker.
 
@@ -51,6 +58,7 @@ def beaker_train(
         project_id: override the project ID.
         experiment_id: override the experiment ID.
         extra_args: extra arguments to pass in the Beaker job.
+        priority: the priority to assign to the Beaker job.
     """
     hparams_configs_dir = None
 
@@ -80,41 +88,39 @@ def beaker_train(
     if hparams_configs_dir is not None:
         shutil.rmtree(hparams_configs_dir)
 
-    beaker = Beaker.from_env(default_workspace=workspace)
-
-    for run_id, config_path in config_paths.items():
-        with beaker.session():
+    with Beaker.from_env(default_workspace=workspace) as beaker:
+        for run_id, config_path in config_paths.items():
             env_vars = get_base_env_vars()
             env_vars.extend(
                 [
-                    EnvVar(
+                    BeakerEnvVar(
                         name="RSLP_PROJECT",  # nosec
                         value=project_id,
                     ),
-                    EnvVar(
+                    BeakerEnvVar(
                         name="RSLP_EXPERIMENT",
                         value=experiment_id,
                     ),
-                    EnvVar(
+                    BeakerEnvVar(
                         name="RSLP_RUN_ID",
                         value=run_id,
                     ),
                 ]
             )
             if username:
                 env_vars.append(
-                    EnvVar(
+                    BeakerEnvVar(
                         name="WANDB_USERNAME",
                         value=username,
                     )
                 )
             datasets = [create_gcp_credentials_mount()]
             datasets += [weka_mount.to_data_mount() for weka_mount in weka_mounts]
-            spec = ExperimentSpec.new(
+            spec = BeakerExperimentSpec.new(
                 budget=DEFAULT_BUDGET,
                 description=f"{project_id}/{experiment_id}/{run_id}",
                 beaker_image=image_name,
-                priority=Priority.high,
+                priority=priority,
                 command=["python", "-m", "rslp.docker_entrypoint"],
                 arguments=[
                     "model",
@@ -130,13 +136,17 @@ def beaker_train(
                     project_id,
                 ]
                 + extra_args,
-                constraints=Constraints(
+                constraints=BeakerConstraints(
                     cluster=cluster,
                 ),
                 preemptible=True,
                 datasets=datasets,
                 env_vars=env_vars,
-                resources=TaskResources(gpu_count=gpus, shared_memory=shared_memory),
+                resources=BeakerTaskResources(
+                    gpu_count=gpus, shared_memory=shared_memory
+                ),
             )
             unique_id = str(uuid.uuid4())[0:8]
-            beaker.experiment.create(f"{project_id}_{experiment_id}_{unique_id}", spec)
+            beaker.experiment.create(
+                name=f"{project_id}_{experiment_id}_{unique_id}", spec=spec
+            )