Skip to content

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

  1. Checkout this branch.
  2. Set the following environment variables to .env:
# Instrumentators
AIGW_INSTRUMENTATOR__THREAD_MONITORING_ENABLED=True
AIGW_INSTRUMENTATOR__THREAD_MONITORING_INTERVAL=1
PYTHONASYNCIODEBUG=True
  1. poetry run ai_gateway
  2. 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
}'
  1. While the request is being executed, check http://0.0.0.0:8082/ to see the ai_gateway_threads_count is increased momentarily, which means the CodeParser.from_language_id is executed in a separate thread.
  2. You can also check threads logger in modelgateway_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.
Edited by Shinya Maeda

Merge request reports