Improve devbox polling: eliminate client-side sleep, enable HTTP/2

Replace the client-side busy-poll loop (async_poll_until) for
await_running with a new retry_server_poll_until that passes the
remaining time budget to the server's wait_for_status endpoint as
timeout_seconds. This eliminates a 30-second asyncio.sleep() between
each long-poll attempt that was causing devbox startup times to cluster
at 40s and 80s instead of the actual 3-5s provisioning time.

Enable HTTP/2 by default on AsyncHttpxClient so that concurrent
long-poll requests multiplex over fewer TCP connections instead of
requiring one connection per in-flight request.

For the Trajectory test, this dramatically improved the results with 494 concurrent devboxes:
  Before: p50=40.5s, p90=48.1s, max=80.9s
  After:  p50=4.4s,  p90=9.7s,  max=20.6s

Author: jrvb-rl

pyproject.toml

@@ -9,7 +9,7 @@ authors = [
 ]
 
 dependencies = [
-  "httpx>=0.23.0, <1",
+  "httpx[http2]>=0.23.0, <1",
   "pydantic>=2.0, <3",
   "typing-extensions>=4.14, <5",
   "anyio>=3.5.0, <5",

src/runloop_api_client/_base_client.py

@@ -1375,6 +1375,7 @@ def __init__(self, **kwargs: Any) -> None:
         kwargs.setdefault("timeout", DEFAULT_TIMEOUT)
         kwargs.setdefault("limits", DEFAULT_CONNECTION_LIMITS)
         kwargs.setdefault("follow_redirects", True)
+        kwargs.setdefault("http2", True)
         super().__init__(**kwargs)
 
 

src/runloop_api_client/lib/polling_async.py

@@ -58,3 +58,48 @@ async def async_poll_until(
                 raise PollingTimeout(f"Exceeded timeout of {config.timeout_seconds} seconds", last_result)
 
         await asyncio.sleep(config.interval_seconds)
+
+
+async def retry_server_poll_until(
+    retriever: Callable[[], Awaitable[T]],
+    is_terminal: Callable[[T], bool],
+    timeout_seconds: float = 30.0,
+    on_error: Optional[Callable[[Exception], T]] = None,
+) -> T:
+    """
+    Retry a server-side long-poll until a condition is met or max timeout is reached.
+
+    Args:
+        retriever: Async or sync callable that returns the object to check.  This takes should
+          take one argument, which is the remaing time to poll.q
+        is_terminal: Callable that returns True when polling should stop
+        timeout_seconds: Total time to wait.  Must be > 0
+        on_error: Optional error handler that can return a value to continue polling
+                 or re-raise the exception to stop polling
+
+    Returns:
+        The final state of the polled object
+
+    Raises:
+        PollingTimeout: When max attempts or timeout is reached
+    """
+    start_time = time.time()
+    last_result: Union[T, None] = None
+
+    while True:
+        elapsed = time.time() - start_time
+        remaining_time = timeout_seconds - elapsed
+        if remaining_time <= 0:
+            raise PollingTimeout(f"Exceeded timeout of {timeout_seconds} seconds", last_result)
+
+        try:
+            last_result = await retriever(remaining_time)
+        except Exception as e:
+            if on_error is not None:
+                last_result = on_error(e)
+            else:
+                raise
+
+        if is_terminal(last_result):
+            return last_result
+

src/runloop_api_client/resources/devboxes/devboxes.py

@@ -82,7 +82,7 @@
     DiskSnapshotsResourceWithStreamingResponse,
     AsyncDiskSnapshotsResourceWithStreamingResponse,
 )
-from ...lib.polling_async import async_poll_until
+from ...lib.polling_async import async_poll_until, retry_server_poll_until
 from ...types.devbox_view import DevboxView
 from ...types.tunnel_view import TunnelView
 from ...types.shared_params.mount import Mount
@@ -2042,11 +2042,10 @@ async def await_running(
 
         Args:
             id: The ID of the devbox to wait for
-            config: Optional polling configuration
+            polling_config: Optional polling configuration
             extra_headers: Send extra headers
             extra_query: Add additional query parameters to the request
             extra_body: Add additional JSON properties to the request
-            timeout: Override the client-level default timeout for this request, in seconds
 
         Returns:
             The devbox in running state
@@ -2056,13 +2055,14 @@ async def await_running(
             RunloopError: If devbox enters a non-running terminal state
         """
 
-        async def wait_for_devbox_status() -> DevboxView:
+        async def wait_for_devbox_status(remaining_timeout_seconds) -> DevboxView:
             # This wait_for_status endpoint polls the devbox status for 10 seconds until it reaches either running or failure.
             # If it's neither, it will throw an error.
             try:
                 return await self._post(
                     f"/v1/devboxes/{id}/wait_for_status",
-                    body={"statuses": ["running", "failure", "shutdown"]},
+                    body={"statuses": ["running", "failure", "shutdown"],
+                          "timeout_seconds": remaining_timeout_seconds },
                     cast_to=DevboxView,
                 )
             except (APITimeoutError, APIStatusError) as error:
@@ -2077,7 +2077,19 @@ async def wait_for_devbox_status() -> DevboxView:
         def is_done_booting(devbox: DevboxView) -> bool:
             return devbox.status not in DEVBOX_BOOTING_STATES
 
-        devbox = await async_poll_until(wait_for_devbox_status, is_done_booting, polling_config)
+        # calculate the timeout to use.  The PollingConfig doesn't
+        # match the semantics for server-side polling well, so we
+        # instead convert interval*attempts to a total time, and take
+        # the minimum total.
+        config = polling_config
+        if not config:
+            config = PollingConfig()  # use defaults
+
+        timeout = config.interval_seconds * config.max_attempts
+        if config.timeout_seconds is not None and config.timeout_seconds > 0:
+            timeout = min(config.timeout_seconds, timeout)
+
+        devbox = await retry_server_poll_until(wait_for_devbox_status, is_done_booting, timeout)
 
         if devbox.status != "running":
             raise RunloopError(f"Devbox entered non-running terminal state: {devbox.status}")