refactor(ai): 优化华为网关连接池配置与 Token 缓存查找逻辑
This commit is contained in:
1 parent
af92152524
commit
bdec3a6cb0
1 file changed
+26
-18
+26
-18
@@ -43,6 +43,12 @@ UPSTREAM_TIMEOUT_MAX = 300 # 大请求超时 300s
|
||||
RETRY_ON_429 = 2 # 429限流重试次数
|
||||
RETRY_ON_504 = 1 # 504超时重试次数
|
||||
|
||||
# ================= 并发配置 =================
|
||||
POOL_CONNECTIONS = 64 # 连接池大小(须 ≥ Waitress 线程数)
|
||||
POOL_MAXSIZE = 64 # 单主机最大连接数
|
||||
WAITRESS_THREADS = 64 # Waitress 工作线程数
|
||||
SSE_CHUNK_SIZE = 4096 # SSE 流式转发块大小(越小首 token 延迟越低)
|
||||
|
||||
# ================= 日志 =================
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
@@ -184,11 +190,14 @@ def scan_pid_mem(pid):
|
||||
|
||||
def find_token_in_memory():
|
||||
"""在所有进程中扫描 Token"""
|
||||
# 快速路径:缓存有效直接返回(避免每次请求都查环境变量/文件)
|
||||
cached = cache.get()
|
||||
if cached:
|
||||
return cached
|
||||
|
||||
# HUAWEI_TOKEN 环境变量优先级最高
|
||||
env_token = os.environ.get('HUAWEI_TOKEN', '').strip()
|
||||
if env_token and len(env_token) > 200 and not cache.is_blacklisted(env_token):
|
||||
cached = cache.get()
|
||||
if cached != env_token:
|
||||
cache.set(env_token)
|
||||
logger.info("Token 从 HUAWEI_TOKEN 环境变量加载")
|
||||
return env_token
|
||||
@@ -203,8 +212,6 @@ def find_token_in_memory():
|
||||
if line.startswith('HUAWEI_TOKEN='):
|
||||
file_token = line.split('=', 1)[1].strip().strip('"').strip("'")
|
||||
if file_token and len(file_token) > 200 and not cache.is_blacklisted(file_token):
|
||||
cached = cache.get()
|
||||
if cached != file_token:
|
||||
cache.set(file_token)
|
||||
logger.info("Token 从持久化文件加载")
|
||||
return file_token
|
||||
@@ -212,7 +219,7 @@ def find_token_in_memory():
|
||||
except (OSError, IOError):
|
||||
pass
|
||||
|
||||
if cache.is_scan_cooldown() and cache.get():
|
||||
if cache.is_scan_cooldown():
|
||||
return cache.get()
|
||||
|
||||
try:
|
||||
@@ -257,11 +264,12 @@ def find_token_in_memory():
|
||||
|
||||
|
||||
# ================= HTTP 会话池 =================
|
||||
# max_retries=0:禁用 urllib3 自动重试,由 proxy 手动控制重试逻辑(避免双重重试)
|
||||
http_session = requests.Session()
|
||||
adapter = requests.adapters.HTTPAdapter(
|
||||
pool_connections=20,
|
||||
pool_maxsize=20,
|
||||
max_retries=3
|
||||
pool_connections=POOL_CONNECTIONS,
|
||||
pool_maxsize=POOL_MAXSIZE,
|
||||
max_retries=0
|
||||
)
|
||||
http_session.mount('https://', adapter)
|
||||
http_session.mount('http://', adapter)
|
||||
@@ -338,7 +346,7 @@ def _build_response(resp):
|
||||
if 'text/event-stream' in content_type or resp.headers.get('Transfer-Encoding', '') == 'chunked':
|
||||
def sse_stream():
|
||||
try:
|
||||
for chunk in resp.iter_content(chunk_size=16384):
|
||||
for chunk in resp.iter_content(chunk_size=SSE_CHUNK_SIZE):
|
||||
if chunk:
|
||||
yield chunk
|
||||
finally:
|
||||
@@ -355,6 +363,13 @@ def _build_response(resp):
|
||||
content = resp.content
|
||||
resp.close()
|
||||
|
||||
# 记录上游非 200 响应(在读取 content 之后,避免提前消费流)
|
||||
if resp.status_code != 200:
|
||||
try:
|
||||
logger.warning(f"上游返回 {resp.status_code}: {content[:500].decode('utf-8', errors='replace')}")
|
||||
except Exception:
|
||||
logger.warning(f"上游返回 {resp.status_code}")
|
||||
|
||||
# 华为云 ModelArts 错误 → OpenAI 标准格式
|
||||
if resp.status_code >= 400:
|
||||
try:
|
||||
@@ -404,6 +419,7 @@ def proxy(subpath):
|
||||
|
||||
headers['Authorization'] = f'Bearer {real_token}'
|
||||
headers['Host'] = TARGET_HOST
|
||||
headers['Connection'] = 'keep-alive'
|
||||
|
||||
target_url = f'https://{TARGET_HOST}/v2/{subpath}'
|
||||
|
||||
@@ -485,14 +501,6 @@ def proxy(subpath):
|
||||
if resp.status_code != 504:
|
||||
break
|
||||
|
||||
# 记录上游非200响应
|
||||
if resp.status_code != 200:
|
||||
try:
|
||||
_err_peek = resp.content[:500]
|
||||
logger.warning(f"上游返回 {resp.status_code}: {_err_peek.decode('utf-8', errors='replace')}")
|
||||
except Exception:
|
||||
logger.warning(f"上游返回 {resp.status_code}")
|
||||
|
||||
return _build_response(resp)
|
||||
|
||||
except requests.exceptions.Timeout:
|
||||
@@ -530,7 +538,7 @@ def main():
|
||||
try:
|
||||
import waitress
|
||||
logger.info(f"使用 Waitress 启动网关 ({host}:{port})")
|
||||
waitress.serve(app, host=host, port=port, threads=32)
|
||||
waitress.serve(app, host=host, port=port, threads=WAITRESS_THREADS)
|
||||
except ImportError:
|
||||
try:
|
||||
import gunicorn.app.wsgiapp
|
||||
|
||||
Reference in new issue
Block a user