From 44845fe80b98adaec8b69107a7b6ac23ef6a6321 Mon Sep 17 00:00:00 2001 From: Keith Zeto Date: Fri, 11 Sep 2026 13:00:32 -0700 Subject: [PATCH] Run Python SDK handlers concurrently, not inline on the dispatch stream The dispatch loop called each handler inline: for response in stub.Dispatch(...): result = handler_to_invoke(ctx) self.agent_stub.ReportInvocation(invoke_info) so the next message was not read until the current handler returned. One slow handler stalled every invocation behind it. There is no threading anywhere else in the SDK, so the limit was exactly one at a time. Webhooks make that visible. The agent accepts a webhook, returns 2xx, queues it, sends it, and starts a timeout timer. While the SDK is busy on the previous invocation the queued ones time out, and the agent logs "invocation error:" or "context deadline exceeded". The sender was told 2xx, so the webhook is simply lost. Paychex hit this on xld_deploy_sync, which syncs deployments from XL Deploy and is slow enough to fall behind. Run handlers on a bounded thread pool instead, which is what the Go SDK already does with a goroutine per invocation. A semaphore stops the loop reading once every worker is busy, so the agent sees its own queue fill rather than this process growing without bound. Reporting now happens on the worker, so a failure there can no longer propagate into the dispatch loop and force a reconnect. It is caught and logged, otherwise the future would swallow it and the invocation would look like it vanished. Two test fixes come along with it. The mock server was not held anywhere, so it could be garbage collected mid-test and stop listening, and a MagicMock cannot serve a streaming RPC at all, which is why no test had ever exercised dispatch. The new test uses a real servicer and fails against the serial version. Co-Authored-By: Claude Opus 5 (1M context) --- sdks/python/README.md | 15 +++ sdks/python/cortex_axon/axon_client.py | 131 +++++++++++++++++-------- sdks/python/tests/test_axon_agent.py | 89 ++++++++++++++++- 3 files changed, 191 insertions(+), 44 deletions(-) diff --git a/sdks/python/README.md b/sdks/python/README.md index b82b532..ba821ae 100644 --- a/sdks/python/README.md +++ b/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. diff --git a/sdks/python/cortex_axon/axon_client.py b/sdks/python/cortex_axon/axon_client.py index b7361d8..68de034 100644 --- a/sdks/python/cortex_axon/axon_client.py +++ b/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 ): diff --git a/sdks/python/tests/test_axon_agent.py b/sdks/python/tests/test_axon_agent.py index 3723f5a..7ae7b2b 100644 --- a/sdks/python/tests/test_axon_agent.py +++ b/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)