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)