Execute CodeParser.from_language_id in a separate thread
This MR is on top of Monitor thread activities (!596 - merged)
What does this merge request do and why?
This MR addresses one of the problems in AI Gateway becomes unresponsive by CPU-bound task (#405 - closed) that CodeParser.from_language_id
(CPU-bound) cloud block the main server event loop. This MR changes to run the treesitter parsing task in a separate thread, so that the server can accept incoming requests while it's working on the CPU-bound task. This archives a similar behavior with coroutine (async
) support of Anthropic/VertexAI clients, which is non-blocking network-IO.
The thread activities can be monitored by Monitor thread activities (!596 - merged), which is an upstream MR.
To be clear, this MR does NOT change the business logic.
Related AI Gateway becomes unresponsive by CPU-bound task (#405 - closed)
How to set up and validate locally
- Checkout this branch.
- Set the following environment variables to
.env
:
# Instrumentators
AIGW_INSTRUMENTATOR__THREAD_MONITORING_ENABLED=True
AIGW_INSTRUMENTATOR__THREAD_MONITORING_INTERVAL=1
PYTHONASYNCIODEBUG=True
poetry run ai_gateway
- Run the following request:
curl -X 'POST' 'http://0.0.0.0:5052/v2/completions' -H 'accept: application/json' -H 'Content-Type: application/json' -d '{
"project_path": "string",
"project_id": 0,
"current_file": {
"file_name": "main.py",
"language_identifier": "python",
"content_above_cursor": "print",
"content_below_cursor": ""
},
"model_provider": "vertex-ai",
"model_name": "code-gecko@002",
"telemetry": [],
"stream": false,
"prompt_version": 1
}'
- While the request is being executed, check
http://0.0.0.0:8082/
to see theai_gateway_threads_count
is increased momentarily, which means theCodeParser.from_language_id
is executed in a separate thread. - You can also check
threads
logger inmodelgateway_debug.log
to see the details of the thread activities.
Further reading
Currently, we constantly run 5 threads for snowplow tracker (daemon mode). This is too much thus we should adjust it. See Adjust snowplow emitter thread count (!594 - merged).
Merge request checklist
-
Tests added for new functionality. If not, please raise an issue to follow up. -
Documentation added/updated, if needed.