Partial tool concurrency (#305)

Author: sidnarayanan

src/aviary/env.py

@@ -31,7 +31,7 @@
     ToolRequestMessage,
     ToolResponseMessage,
 )
-from aviary.utils import format_exc, is_coroutine_callable
+from aviary.utils import ReaderWriterLock, format_exc, is_coroutine_callable
 
 logger = logging.getLogger(__name__)
 
@@ -215,6 +215,7 @@ async def exec_tool_calls(
             Ordered list of ToolResponseMessages, order matches the order of tool calls
                 in the input message.
         """
+        concurrency_lock = ReaderWriterLock()
 
         async def _exec_tool_call(tool_call: ToolCall) -> ToolResponseMessage:
             start = time.monotonic()
@@ -227,6 +228,7 @@ async def _exec_tool_call(tool_call: ToolCall) -> ToolResponseMessage:
                     f"{tool_call.function.name!r} not a valid name in"
                     f" { {t.info.name for t in self.tools} }."
                 ) from exc
+
             # we do a special convenience to make
             # state be optional in the function signature
             need_to_filter = (
@@ -239,25 +241,33 @@ async def _exec_tool_call(tool_call: ToolCall) -> ToolResponseMessage:
                 if need_to_filter
                 else function_kwargs
             )
+
+            concurrency_context = (
+                concurrency_lock.read_lock()
+                if tool.concurrency_safe
+                else concurrency_lock.write_lock()
+            )
+
             tool_exc: Exception | None = None
             try:
-                if is_coroutine_callable(tool._tool_fn):
-                    content = await maybe_wait_for(
-                        tool._tool_fn(
-                            **tool_call.function.arguments, **filtered_kwargs
-                        ),
-                        exec_timeout,
-                    )
-                else:
-                    # If the function is synchronous, run on a thread
-                    content = await maybe_wait_for(
-                        asyncio.to_thread(
-                            tool._tool_fn,
-                            **tool_call.function.arguments,
-                            **filtered_kwargs,
-                        ),
-                        exec_timeout,
-                    )
+                async with concurrency_context:
+                    if is_coroutine_callable(tool._tool_fn):
+                        content = await maybe_wait_for(
+                            tool._tool_fn(
+                                **tool_call.function.arguments, **filtered_kwargs
+                            ),
+                            exec_timeout,
+                        )
+                    else:
+                        # If the function is synchronous, run on a thread
+                        content = await maybe_wait_for(
+                            asyncio.to_thread(
+                                tool._tool_fn,
+                                **tool_call.function.arguments,
+                                **filtered_kwargs,
+                            ),
+                            exec_timeout,
+                        )
             except Exception as exc:
                 if not handle_tool_exc:
                     raise

src/aviary/tools/base.py

@@ -344,6 +344,19 @@ class Tool(BaseModel):
             " serialization, and the validation alias enables deserialization."
         ),
     )
+    concurrency_safe: bool = Field(
+        default=True,
+        # Exclude since we need Tool.model_dump() to conform to OpenAI schema.
+        # Note that this is safe: while we do (de)serialize tools when e.g. passing to
+        # agents, only Environment.exec_tool_calls uses this field. And we never serialize
+        # an env after reset.
+        exclude=True,
+        description=(
+            "Whether the tool is safe to run concurrently with itself and other tools. "
+            "If set to False (not default), then executing this tool will block all "
+            "other tool calls (including concurrency-safe tools)."
+        ),
+    )
 
     def __init__(
         self,
@@ -382,6 +395,7 @@ def from_function(
         docstring_style: DocstringStyle = DocstringStyle.AUTO,
         allow_empty_param_descriptions: bool = False,
         types_in_param_descriptions: bool = False,
+        concurrency_safe: bool = True,
         **formats,
     ) -> "Tool":
         """Hydrate this class via inspection from a free function with a docstring."""
@@ -450,6 +464,7 @@ def from_function(
                 ),
                 parameters=json_schema,
             ),
+            concurrency_safe=concurrency_safe,
         )
 
 

src/aviary/utils.py

@@ -1,11 +1,13 @@
+import asyncio
 import base64
 import inspect
 import io
 import random
 import string
 from ast import literal_eval
 from collections import UserDict
-from collections.abc import Sequence
+from collections.abc import AsyncIterator, Sequence
+from contextlib import asynccontextmanager
 from enum import StrEnum
 from typing import (
     TYPE_CHECKING,
@@ -510,3 +512,50 @@ def format_exc(exc: BaseException) -> str:
             f" {', '.join(repr(e) for e in exc.exceptions)}"
         )
     return repr(exc)
+
+
+class ReaderWriterLock:
+    """An asyncio lock that allows execution of multiple readers or a single writer.
+
+    When a writer is executing, it will block all readers and writers.
+    The main use case here is for concurrency-unsafe tools to block execution
+    of other tool calls, while still allowing concurrency-safe tools to execute
+    in parallel with each other.
+    """
+
+    def __init__(self):
+        self._readers = 0
+        self._writer = False
+        self._lock = asyncio.Lock()
+        self._write_ok = asyncio.Condition(self._lock)
+        self._read_ok = asyncio.Condition(self._lock)
+
+    @asynccontextmanager
+    async def read_lock(self) -> AsyncIterator[None]:
+        """Acquire a read lock. This blocks all writers."""
+        async with self._lock:
+            while self._writer:
+                await self._read_ok.wait()
+            self._readers += 1
+        try:
+            yield
+        finally:
+            async with self._lock:
+                self._readers -= 1
+                if self._readers == 0:
+                    self._write_ok.notify_all()
+
+    @asynccontextmanager
+    async def write_lock(self) -> AsyncIterator[None]:
+        """Acquire a write lock. This blocks all readers and writers."""
+        async with self._lock:
+            while self._writer or self._readers > 0:
+                await self._write_ok.wait()
+            self._writer = True
+        try:
+            yield
+        finally:
+            async with self._lock:
+                self._writer = False
+                self._read_ok.notify_all()
+                self._write_ok.notify_all()

tests/test_tools.py

@@ -4,7 +4,7 @@
 import pickle
 from collections.abc import Callable, Sequence
 from enum import IntEnum, auto
-from typing import Any
+from typing import Any, cast
 from unittest.mock import patch
 
 import pytest
@@ -956,3 +956,54 @@ async def export_frame(self):
         )
         assert response.status_code == 200
         assert response.json()["result"] == "3"
+
+
+@pytest.mark.asyncio
+async def test_mixed_concurrency() -> None:
+    # Counts the number of tools executing concurrently
+    counter = 0
+    counter_lock = asyncio.Lock()
+
+    async def sleep_fn() -> int:
+        """Stub."""
+        nonlocal counter
+        async with counter_lock:
+            counter += 1
+            counter_val = counter
+        await asyncio.sleep(0.5)
+        async with counter_lock:
+            counter -= 1
+        return counter_val
+
+    async def unsafe_sleep_fn() -> int:
+        """Stub."""
+        return await sleep_fn()
+
+    safe_sleep = Tool.from_function(sleep_fn)
+    unsafe_sleep = Tool.from_function(unsafe_sleep_fn, concurrency_safe=False)
+
+    dummy_env = DummyEnv()
+    await dummy_env.reset()
+    dummy_env.tools = [safe_sleep, unsafe_sleep]
+
+    safes = [True, True, True, False, True, False, False, True, True]
+    obs, *_ = await dummy_env.step(
+        ToolRequestMessage(
+            tool_calls=[
+                ToolCall.from_tool(safe_sleep if safe else unsafe_sleep)
+                for safe in safes
+            ]
+        )
+    )
+
+    at_least_one_parallel = False
+    for safe, msg in zip(safes, obs, strict=True):
+        count = int(cast(str, msg.content))
+        if safe:
+            at_least_one_parallel |= count > 1
+        else:
+            assert count == 1, "Expected unsafe tools to block all other tool calls."
+
+    assert at_least_one_parallel, (
+        "Expected at least one safe tool call to run concurrently with another."
+    )