deepscholar: patched the Gemma startup hang

Author: donovanr

src/olmo_eval/common/beaker_status.py

@@ -14,8 +14,9 @@
 import os
 import time
 from collections.abc import Callable
+from threading import Lock, Thread
 
-from beaker import Beaker, BeakerWorkload
+from beaker import Beaker, BeakerExperiment, BeakerWorkload
 from beaker.exceptions import BeakerConfigurationError
 
 DEFAULT_MIN_INTERVAL = 10.0
@@ -35,14 +36,23 @@ class BeakerStatusReporter:
     def __init__(self, min_interval: float = DEFAULT_MIN_INTERVAL) -> None:
         self.min_interval = min_interval
         self._git_suffix = _git_suffix()
-        self._workload: BeakerWorkload | None = None
+        workload_id = os.environ.get("BEAKER_WORKLOAD_ID")
+        self._workload = (
+            BeakerWorkload(experiment=BeakerExperiment(id=workload_id)) if workload_id else None
+        )
+        self._lock = Lock()
+        self._update_in_flight = False
         self._last_update: float = float("-inf")
+        self._client: Beaker | None = None
+        if self._workload is None:
+            return
+
         try:
-            self._client: Beaker | None = Beaker.from_env()
+            self._client = Beaker.from_env()
         except BeakerConfigurationError:
-            self._client = None
             return
-        self._workload = self._client.workload.get(os.environ["BEAKER_WORKLOAD_ID"])
+        except Exception as error:
+            logger.warning("Beaker status reporting disabled during setup: %s", error)
 
     def update(self, message: str, force: bool = False) -> None:
         """Push a status message to the Beaker workload description.
@@ -54,12 +64,41 @@ def update(self, message: str, force: bool = False) -> None:
             return
 
         now = time.monotonic()
-        if not force and now - self._last_update < self.min_interval:
-            return
+        with self._lock:
+            if not force and now - self._last_update < self.min_interval:
+                return
+            # Status is cosmetic. Never queue more work behind a slow Beaker API
+            # request, and never let that request block model startup or evaluation.
+            if self._update_in_flight:
+                return
+            self._update_in_flight = True
+            self._last_update = now
 
+        client = self._client
+        workload = self._workload
         full_message = f"{message} {self._git_suffix}"
-        self._client.workload.update(self._workload, description=full_message)
-        self._last_update = now
+
+        def send_update() -> None:
+            try:
+                client.workload.update(workload, description=full_message)
+            except Exception as error:
+                logger.warning("Beaker status reporting disabled after update failure: %s", error)
+                self._client = None
+            finally:
+                with self._lock:
+                    self._update_in_flight = False
+
+        try:
+            Thread(
+                target=send_update,
+                name="beaker-status-update",
+                daemon=True,
+            ).start()
+        except Exception as error:
+            with self._lock:
+                self._update_in_flight = False
+            self._client = None
+            logger.warning("Beaker status reporting disabled after thread failure: %s", error)
 
     def progress_callback(self, label: str, units: str = "items/sec") -> Callable[..., None]:
         """Return a ``(count, total, *, force=False)`` callback bound to a fresh start time.

tests/core/test_beaker_status.py

@@ -6,6 +6,18 @@
 from olmo_eval.common import beaker_status
 
 
+class InlineThread:
+    """Run a thread target synchronously to keep status reporter tests deterministic."""
+
+    def __init__(self, *, target, name: str, daemon: bool) -> None:
+        self.target = target
+        self.name = name
+        self.daemon = daemon
+
+    def start(self) -> None:
+        self.target()
+
+
 class BeakerStatusReporterTest(unittest.TestCase):
     def test_disabled_when_beaker_config_missing(self) -> None:
         with (
@@ -27,15 +39,15 @@ def test_throttles_updates_within_interval(self) -> None:
             "GIT_BRANCH": "main",
         }
         fake_client = mock.MagicMock()
-        fake_workload = mock.MagicMock()
-        fake_client.workload.get.return_value = fake_workload
         with (
             mock.patch.dict("os.environ", env, clear=True),
             mock.patch.object(beaker_status.Beaker, "from_env", return_value=fake_client),
+            mock.patch.object(beaker_status, "Thread", InlineThread),
         ):
             reporter = beaker_status.BeakerStatusReporter(min_interval=60.0)
 
             self.assertIsNotNone(reporter._client)
+            fake_client.workload.get.assert_not_called()
 
             with mock.patch("time.monotonic", side_effect=[0.0, 1.0, 61.0]):
                 reporter.update("first")
@@ -44,30 +56,36 @@ def test_throttles_updates_within_interval(self) -> None:
 
         self.assertEqual(fake_client.workload.update.call_count, 2)
         suffix = "git_commit: abc123 git_branch: main"
-        fake_client.workload.update.assert_any_call(fake_workload, description=f"first {suffix}")
-        fake_client.workload.update.assert_any_call(fake_workload, description=f"third {suffix}")
+        workload = beaker_status.BeakerWorkload(
+            experiment=beaker_status.BeakerExperiment(id="wl_123")
+        )
+        fake_client.workload.update.assert_any_call(workload, description=f"first {suffix}")
+        fake_client.workload.update.assert_any_call(workload, description=f"third {suffix}")
 
     def test_git_suffix_uses_unknown_when_env_missing(self) -> None:
         env = {"BEAKER_WORKLOAD_ID": "wl_123"}
         fake_client = mock.MagicMock()
-        fake_workload = mock.MagicMock()
-        fake_client.workload.get.return_value = fake_workload
         with (
             mock.patch.dict("os.environ", env, clear=True),
             mock.patch.object(beaker_status.Beaker, "from_env", return_value=fake_client),
+            mock.patch.object(beaker_status, "Thread", InlineThread),
         ):
             reporter = beaker_status.BeakerStatusReporter(min_interval=0.0)
             reporter.update("hello")
 
+        workload = beaker_status.BeakerWorkload(
+            experiment=beaker_status.BeakerExperiment(id="wl_123")
+        )
         fake_client.workload.update.assert_called_once_with(
-            fake_workload, description="hello git_commit: unknown git_branch: unknown"
+            workload, description="hello git_commit: unknown git_branch: unknown"
         )
 
     def test_force_bypasses_throttle(self) -> None:
         fake_client = mock.MagicMock()
         with (
             mock.patch.dict("os.environ", {"BEAKER_WORKLOAD_ID": "wl_xyz"}, clear=True),
             mock.patch.object(beaker_status.Beaker, "from_env", return_value=fake_client),
+            mock.patch.object(beaker_status, "Thread", InlineThread),
         ):
             reporter = beaker_status.BeakerStatusReporter(min_interval=60.0)
 
@@ -77,6 +95,35 @@ def test_force_bypasses_throttle(self) -> None:
 
         self.assertEqual(fake_client.workload.update.call_count, 2)
 
+    def test_update_failure_is_nonfatal_and_disables_reporting(self) -> None:
+        fake_client = mock.MagicMock()
+        fake_client.workload.update.side_effect = RuntimeError("API unavailable")
+        with (
+            mock.patch.dict("os.environ", {"BEAKER_WORKLOAD_ID": "wl_xyz"}, clear=True),
+            mock.patch.object(beaker_status.Beaker, "from_env", return_value=fake_client),
+            mock.patch.object(beaker_status, "Thread", InlineThread),
+        ):
+            reporter = beaker_status.BeakerStatusReporter()
+            reporter.update("starting")
+
+        self.assertIsNone(reporter._client)
+
+    def test_update_starts_a_daemon_thread(self) -> None:
+        fake_client = mock.MagicMock()
+        fake_thread = mock.MagicMock()
+        with (
+            mock.patch.dict("os.environ", {"BEAKER_WORKLOAD_ID": "wl_xyz"}, clear=True),
+            mock.patch.object(beaker_status.Beaker, "from_env", return_value=fake_client),
+            mock.patch.object(beaker_status, "Thread", return_value=fake_thread) as thread_class,
+        ):
+            reporter = beaker_status.BeakerStatusReporter()
+            reporter.update("starting")
+
+        thread_class.assert_called_once()
+        self.assertEqual(thread_class.call_args.kwargs["name"], "beaker-status-update")
+        self.assertTrue(thread_class.call_args.kwargs["daemon"])
+        fake_thread.start.assert_called_once_with()
+
 
 if __name__ == "__main__":
     unittest.main()