Run Python SDK handlers concurrently, not inline on the dispatch stream

sdks/python/README.md

@@ -91,3 +91,18 @@ To invoke the webhook handler, you can just send a POST request to `http://local
 
 
 
+
+## Handler concurrency
+
+Handlers run on a thread pool, so a slow handler does not delay the ones
+behind it. Eight run at once by default. Once every worker is busy, the SDK
+stops reading new invocations until one finishes.
+
+Write handlers to be thread safe, because two invocations of the same handler
+can run at the same time. To change the limit:
+
+```python
+client = AxonClient(scope=globals(), max_concurrent_handlers=16)
+```
+
+Set it to `1` for the old behaviour, where handlers run one at a time.

sdks/python/cortex_axon/axon_client.py

@@ -1,7 +1,9 @@
 import json
 import os
 import queue
+import threading
 import traceback
+from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime
 from time import sleep
 from typing import Any, Optional
@@ -23,6 +25,8 @@
     find_annotated_methods,
 )
 
+DEFAULT_MAX_CONCURRENT_HANDLERS = 8
+
 
 # TODO Add docstrings
 class CortexRequestException(IOError):
@@ -147,7 +151,11 @@ def __init__(
         handlers: Optional[list[CortexAnnotation]] = None,
         scope=None,
         max_receive_message_size: Optional[int] = None,
+        max_concurrent_handlers: int = DEFAULT_MAX_CONCURRENT_HANDLERS,
     ):
+        if max_concurrent_handlers < 1:
+            raise ValueError("max_concurrent_handlers must be at least 1")
+
         self.id = str(datetime.now().timestamp())
         self.agent_hostport = f"{agent_host}:{agent_port}"
 
@@ -178,6 +186,16 @@ def __init__(
 
         self.handlers = handlers or find_annotated_methods(scope)
 
+        # Handlers run on this pool rather than inline on the dispatch stream,
+        # so a slow one does not hold back every invocation behind it. The
+        # semaphore stops reading once every worker is busy, which pushes
+        # backpressure onto the agent instead of queueing without limit here.
+        self._executor = ThreadPoolExecutor(
+            max_workers=max_concurrent_handlers,
+            thread_name_prefix="axon-handler",
+        )
+        self._handler_slots = threading.BoundedSemaphore(max_concurrent_handlers)
+
     def _get_stub(self):
         if not self.agent_stub:
             self.agent_channel = grpc.insecure_channel(
@@ -216,6 +234,75 @@ def _find_handler(self, name):
     def _handler_context(self, args):
         return HandlerContext(self.cortex_stub, args)
 
+    def _submit_handler(self, handler_to_invoke, invoke):
+        # Blocks once every worker is busy. That is deliberate: the agent then
+        # sees its own queue fill rather than this process growing without
+        # bound, and the dispatch stream stays readable up to that point.
+        self._handler_slots.acquire()
+        stub = self.agent_stub
+        try:
+            future = self._executor.submit(
+                self._invoke_handler, handler_to_invoke, invoke, stub
+            )
+        except BaseException:
+            self._handler_slots.release()
+            raise
+        future.add_done_callback(lambda _: self._handler_slots.release())
+
+    def _invoke_handler(self, handler_to_invoke, invoke, stub):
+        print(
+            f"Dispatch handler START {invoke.handler_name} (id={invoke.invocation_id}) reason={invoke.reason}"
+        )
+        handler_err = None
+        handler_result = None
+        result = None
+        now = datetime.now()
+        ctx = self._handler_context(invoke.args)
+        duration_ms = 0
+
+        try:
+            result = handler_to_invoke(ctx)
+            duration_ms = int(
+                (datetime.now() - now).total_seconds() * 1000
+            )
+            if result:
+                handler_result = cortex_axon_agent_pb2.InvokeResult(
+                    value=str(result)
+                )
+            print(f"Dispatch handler SUCCESS {invoke.handler_name} (id={invoke.invocation_id}) in {round((datetime.now() - now).total_seconds() * 1000, 1)}ms. Returned result={result is not None}")
+        except Exception as e:
+            print(f"Dispatch handler ERROR {invoke.handler_name} (id={invoke.invocation_id}) in {round((datetime.now() - now).total_seconds() * 1000, 1)}ms. Error={e}")
+            traceback.print_exc()
+            handler_err = common_pb2.Error(
+                code="unexpected", message="Error calling handler: " + str(e)
+            )
+
+        invoke_info = cortex_axon_agent_pb2.ReportInvocationRequest(
+            handler_invoke=invoke,
+            start_client_timestamp=now,
+            duration_ms=duration_ms,
+            result=handler_result,
+            error=handler_err,
+            logs=ctx.logs,
+        )
+
+        # A failure here must not escape the worker: the dispatch loop is no
+        # longer in the call path, so an exception would be swallowed by the
+        # future and the invocation would look like it vanished.
+        try:
+            report_response = stub.ReportInvocation(invoke_info)
+        except Exception as e:
+            print(
+                f"Error reporting invocation {invoke.invocation_id}: {e}"
+            )
+            traceback.print_exc()
+            return
+
+        if report_response.error and report_response.error.code:
+            print(
+                f"Error reporting invocation: {report_response.error}"
+            )
+
     def run(self):
         registered = False
         while True:
@@ -235,54 +322,12 @@ def run(self):
 
                     if response.type == cortex_axon_agent_pb2.DISPATCH_MESSAGE_INVOKE:
                         invoke = response.invoke
-                        args = invoke.args
                         handler_to_invoke = self._find_handler(
                             invoke.handler_name)
                         if not handler_to_invoke:
                             print(f"Unknown handler: {response.handler_name}")
                             continue
-                        print(
-                            f"Dispatch handler START {invoke.handler_name} (id={invoke.invocation_id}) reason={invoke.reason}"
-                        )
-                        handler_err = None
-                        handler_result = None
-                        result = None
-                        now = datetime.now()
-                        ctx = self._handler_context(args)
-                        duration_ms = 0
-
-                        try:
-                            result = handler_to_invoke(ctx)
-                            duration_ms = int(
-                                (datetime.now() - now).total_seconds() * 1000
-                            )
-                            if result:
-                                handler_result = cortex_axon_agent_pb2.InvokeResult(
-                                    value=str(result)
-                                )
-                            print(f"Dispatch handler SUCCESS {invoke.handler_name} (id={invoke.invocation_id}) in {round((datetime.now() - now).total_seconds() * 1000, 1)}ms. Returned result={result is not None}")
-                        except Exception as e:
-                            print(f"Dispatch handler ERROR {invoke.handler_name} (id={invoke.invocation_id}) in {round((datetime.now() - now).total_seconds() * 1000, 1)}ms. Error={e}")
-                            traceback.print_exc()
-                            handler_err = common_pb2.Error(
-                                code="unexpected", message="Error calling handler: " + str(e)
-                            )
-
-                        invoke_info = cortex_axon_agent_pb2.ReportInvocationRequest(
-                            handler_invoke=invoke,
-                            start_client_timestamp=now,
-                            duration_ms=duration_ms,
-                            result=handler_result,
-                            error=handler_err,
-                            logs=ctx.logs,
-                        )
-                        report_response = self.agent_stub.ReportInvocation(
-                            invoke_info)
-
-                        if report_response.error and report_response.error.code:
-                            print(
-                                f"Error reporting invocation: {report_response.error}"
-                            )
+                        self._submit_handler(handler_to_invoke, invoke)
                     elif (
                         response.type == cortex_axon_agent_pb2.DISPATCH_MESSAGE_WORK_COMPLETED
                     ):

sdks/python/tests/test_axon_agent.py

@@ -1,6 +1,8 @@
 import queue
 import threading
+from concurrent.futures import ThreadPoolExecutor
 
+import grpc
 from utils import (
     mock_agent,  # noqa: F401
     mock_cortex_api,  # noqa: F401
@@ -9,7 +11,7 @@
 
 from cortex_axon import AxonClient
 from cortex_axon.handler import CortexHandler, CortexScheduled, CortexWebhook
-from generated import common_pb2, cortex_axon_agent_pb2
+from generated import common_pb2, cortex_axon_agent_pb2, cortex_axon_agent_pb2_grpc
 
 
 def register_handler(request, assert_queue):
@@ -164,3 +166,88 @@ def my_invoke_handler():
         "handler_name": "my_invoke_handler",
         "options":  [{'invoke': {}}] # this is empty because INVOKE is default
     }
+
+
+class ConcurrencyAgent(cortex_axon_agent_pb2_grpc.AxonAgentServicer):
+    """A real servicer: MagicMock cannot serve a streaming RPC."""
+
+    def __init__(self, handler_name, invoke_count, keep_stream_open):
+        self.handler_name = handler_name
+        self.invoke_count = invoke_count
+        self.keep_stream_open = keep_stream_open
+
+    def RegisterHandler(self, request, context):
+        return cortex_axon_agent_pb2.RegisterHandlerResponse(
+            error=common_pb2.Error(code="0", message="OK"),
+            id=request.dispatch_id,
+        )
+
+    def Dispatch(self, request_iterator, context):
+        request = next(request_iterator)
+        for i in range(self.invoke_count):
+            yield cortex_axon_agent_pb2.DispatchMessage(
+                type=cortex_axon_agent_pb2.DISPATCH_MESSAGE_INVOKE,
+                invoke=cortex_axon_agent_pb2.DispatchHandlerInvoke(
+                    invocation_id=str(i),
+                    dispatch_id=request.dispatch_id,
+                    handler_id=self.handler_name,
+                    handler_name=self.handler_name,
+                ),
+            )
+        self.keep_stream_open.wait(timeout=10)
+
+    def ReportInvocation(self, request, context):
+        return cortex_axon_agent_pb2.ReportInvocationResponse()
+
+
+def test_handlers_run_concurrently():
+    """A slow handler must not hold back the invocations behind it.
+
+    Handlers used to run inline on the dispatch stream loop, so invocation
+    two was not read until invocation one returned. Under webhook load that
+    let the agent time every queued invocation out while the client worked
+    through them one at a time.
+    """
+    handler_count = 4
+    # handler_count workers plus this thread: the barrier only trips if every
+    # handler is inside the call at the same moment.
+    all_started = threading.Barrier(handler_count + 1)
+    release = threading.Event()
+
+    def my_concurrent_webhook(ctx):
+        all_started.wait(timeout=10)
+        release.wait(timeout=10)
+
+    webhook = CortexWebhook(id="my-webhook")
+    webhook.set_func(my_concurrent_webhook)
+
+    server = grpc.server(ThreadPoolExecutor(max_workers=handler_count + 2))
+    port = server.add_insecure_port("127.0.0.1:0")
+    cortex_axon_agent_pb2_grpc.add_AxonAgentServicer_to_server(
+        ConcurrencyAgent("my_concurrent_webhook", handler_count, release),
+        server,
+    )
+    server.start()
+
+    try:
+        client = AxonClient(
+            agent_host="127.0.0.1",
+            agent_port=port,
+            handlers=[webhook],
+            scope=globals(),
+            max_concurrent_handlers=handler_count,
+        )
+
+        client_thread = threading.Thread(target=client.run, daemon=True)
+        client_thread.start()
+
+        try:
+            all_started.wait(timeout=10)
+        except threading.BrokenBarrierError:
+            raise AssertionError(
+                "handlers did not run concurrently: fewer than "
+                f"{handler_count} were in flight at once"
+            ) from None
+    finally:
+        release.set()
+        server.stop(grace=None)