49 lines
2.1 KiB
Python
49 lines
2.1 KiB
Python
"""Opt-in configuration of the Gateway event loop's shared blocking-I/O pool."""
|
|
|
|
import asyncio
|
|
import logging
|
|
import os
|
|
import re
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
|
|
from deerflow.config.runtime_performance_config import RuntimePerformanceConfig
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_THREAD_POOL_ENV = "DEER_FLOW_ASYNCIO_THREAD_POOL_SIZE"
|
|
_THREAD_POOL_ALIAS = "DEERFLOW_IO_THREADS"
|
|
|
|
|
|
async def configure_runtime_executor(performance: RuntimePerformanceConfig) -> None:
|
|
"""Install the selected executor once at startup, before background work.
|
|
|
|
The hosting event loop owns shutdown (e.g. ``asyncio.run``/uvicorn). Do not
|
|
shut the pool down from a lifespan hook: other teardown hooks may still
|
|
need it, and scheduling its shutdown inside itself can deadlock.
|
|
"""
|
|
if not performance.asyncio_thread_pool_enabled:
|
|
logger.info("GATEWAY_ASYNCIO_EXECUTOR enabled=false max_workers=python_default source=python_default pid=%d", os.getpid())
|
|
return
|
|
|
|
canonical = os.getenv(_THREAD_POOL_ENV, "").strip()
|
|
alias = os.getenv(_THREAD_POOL_ALIAS, "").strip()
|
|
if canonical:
|
|
raw_size, source = canonical, _THREAD_POOL_ENV
|
|
elif alias:
|
|
raw_size, source = alias, _THREAD_POOL_ALIAS
|
|
else:
|
|
raw_size = str(performance.asyncio_thread_pool_size)
|
|
source = "config.runtime_performance.asyncio_thread_pool_size"
|
|
|
|
# Accept only bounded decimal integers; do not echo untrusted raw env
|
|
# contents in startup exceptions or logs.
|
|
if not re.fullmatch(r"[0-9]{1,4}", raw_size) or not 1 <= int(raw_size) <= 1024:
|
|
raise ValueError(f"{source} must be an integer from 1 to 1024")
|
|
size = int(raw_size)
|
|
if canonical and alias:
|
|
logger.info("GATEWAY_ASYNCIO_EXECUTOR %s ignored because %s takes precedence pid=%d", _THREAD_POOL_ALIAS, _THREAD_POOL_ENV, os.getpid())
|
|
|
|
loop = asyncio.get_running_loop()
|
|
loop.set_default_executor(ThreadPoolExecutor(max_workers=size, thread_name_prefix="gateway-sync"))
|
|
logger.info("GATEWAY_ASYNCIO_EXECUTOR enabled=true max_workers=%d source=%s pid=%d", size, source, os.getpid())
|