5602578b2f
This PR does three things: 1. Executes synchronous tool calls in thread pool allowing for up to 4 + # of CPUs executions in parallel. 2. Makes force quitting via double SIGINT/SIGTERM possible and via single SIGINT/SIGTERM + graceful shutdown timeout expiry possible, even if there are active connections. 3. Sets `timeout_graceful_shutdown` to `ARCADE_UVICORN_TIMEOUT_GRACEFUL_SHUTDOWN` env var if set, else defaults to 15. 4. Disable the worker health check span to reduce noise Tradeoffs: Since this PR introduces executing synchronous tools via `await asyncio.to_thread(func, **func_args)`, this means that there is no way for the thread to be killed until it finishes. The ramifications of this is that the force quitting logic that is also implemented in this PR has to be very harsh `os._exit(1)` just in case there is a sync tool actively executing. This means that `MCPApp` teardown logic will not execute when force quitting is required. Although this was already the case because we weren't previously able to force quit! This tradeoff is justified for now since "parallel" tool executions will relieve us of many worker timeouts that we are seeing in prod. Future work: Minimize/eliminate the need for `os._exit(1)` such that `MCPApp` teardown logic will always execute, even when force quitting. The solution will likely be moving away from `await asyncio.to_thread(func, **func_args)` (while maintaining "parallelism" and then utilize the `TaskTrackerMiddleware` introduced in this PR to cancel all of the active HTTP requests. Resolves PLT-713
99 lines
2.8 KiB
Python
99 lines
2.8 KiB
Python
from arcade_core.schema import (
|
|
ToolCallRequest,
|
|
ToolCallResponse,
|
|
)
|
|
from opentelemetry import trace
|
|
|
|
from arcade_serve.core.common import (
|
|
CatalogResponse,
|
|
HealthCheckResponse,
|
|
RequestData,
|
|
Router,
|
|
Worker,
|
|
WorkerComponent,
|
|
)
|
|
|
|
|
|
class CatalogComponent(WorkerComponent):
|
|
def __init__(self, worker: Worker) -> None:
|
|
self.worker = worker
|
|
|
|
def register(self, router: Router) -> None:
|
|
"""
|
|
Register the catalog route with the router.
|
|
"""
|
|
router.add_route(
|
|
"tools",
|
|
self,
|
|
method="GET",
|
|
response_type=CatalogResponse,
|
|
operation_id="get_catalog",
|
|
description="Get the catalog of tools",
|
|
summary="Get the catalog of tools",
|
|
tags=["Arcade"],
|
|
)
|
|
|
|
async def __call__(self, request: RequestData) -> CatalogResponse:
|
|
"""
|
|
Handle the request to get the catalog.
|
|
"""
|
|
tracer = trace.get_tracer(__name__)
|
|
with tracer.start_as_current_span("Catalog"):
|
|
return self.worker.get_catalog()
|
|
|
|
|
|
class CallToolComponent(WorkerComponent):
|
|
def __init__(self, worker: Worker) -> None:
|
|
self.worker = worker
|
|
|
|
def register(self, router: Router) -> None:
|
|
"""
|
|
Register the call tool route with the router.
|
|
"""
|
|
router.add_route(
|
|
"tools/invoke",
|
|
self,
|
|
method="POST",
|
|
response_type=ToolCallResponse,
|
|
operation_id="call_tool",
|
|
description="Call a tool",
|
|
summary="Call a tool",
|
|
tags=["Arcade"],
|
|
)
|
|
|
|
async def __call__(self, request: RequestData) -> ToolCallResponse:
|
|
"""
|
|
Handle the request to call (invoke) a tool.
|
|
"""
|
|
tracer = trace.get_tracer(__name__)
|
|
with tracer.start_as_current_span("CallTool"):
|
|
call_tool_request_data = request.body_json
|
|
call_tool_request = ToolCallRequest.model_validate(call_tool_request_data)
|
|
return await self.worker.call_tool(call_tool_request)
|
|
|
|
|
|
class HealthCheckComponent(WorkerComponent):
|
|
def __init__(self, worker: Worker) -> None:
|
|
self.worker = worker
|
|
|
|
def register(self, router: Router) -> None:
|
|
"""
|
|
Register the health check route with the router.
|
|
"""
|
|
router.add_route(
|
|
"health",
|
|
self,
|
|
method="GET",
|
|
response_type=HealthCheckResponse,
|
|
operation_id="health_check",
|
|
description="Health check",
|
|
summary="Health check",
|
|
tags=["Arcade"],
|
|
require_auth=False,
|
|
)
|
|
|
|
async def __call__(self, request: RequestData) -> HealthCheckResponse:
|
|
"""
|
|
Handle the request to check the health of the worker.
|
|
"""
|
|
return self.worker.health_check()
|