Skip to content

[Benchmark] Share one ClientSession across benchmark requests #8114

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
kevincheng2 wants to merge 1 commit into from
Closed

[Benchmark] Share one ClientSession across benchmark requests #8114

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 22 additions & 2 deletions benchmarks/benchmark_serving.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
import argparse
import asyncio
import gc
import inspect
import json
import os
import random
Expand All @@ -32,6 +33,7 @@
from datetime import datetime
from typing import Any, Optional

import aiohttp
import numpy as np
import yaml
from backend_request_func import (
Expand Down Expand Up @@ -553,11 +555,27 @@ async def benchmark(
api_url = f"http://{ip_list[0]}{args.endpoint}"
semaphore = asyncio.Semaphore(max_concurrency) if max_concurrency else None

# 复用一个 ClientSession + 长连接池：默认每个请求新建 session 会给每次压测请求
# 建一条新 TCP 连接，长时间大压测（num-prompts 上万）会堆积大量 TIME_WAIT，
# 甚至耗尽本地端口，干扰时延测量。
shared_session = None
if "session" in inspect.signature(request_func).parameters:
conn_limit = max_concurrency * 2 if max_concurrency else 0
shared_session = aiohttp.ClientSession(
trust_env=True,
read_bufsize=10 * 1024 * 1024,
timeout=aiohttp.ClientTimeout(total=6 * 60 * 60, sock_connect=60),
connector=aiohttp.TCPConnector(limit=conn_limit, ttl_dns_cache=600),
)

async def limited_request_func(request_func_input, pbar):
kwargs = {"request_func_input": request_func_input, "pbar": pbar}
if shared_session is not None:
kwargs["session"] = shared_session
if semaphore is None:
return await request_func(request_func_input=request_func_input, pbar=pbar)
return await request_func(**kwargs)
async with semaphore:
return await request_func(request_func_input=request_func_input, pbar=pbar)
return await request_func(**kwargs)

tasks: list[asyncio.Task] = []
benchmark_start_time = time.perf_counter()
Expand Down Expand Up @@ -603,6 +621,8 @@ async def limited_request_func(request_func_input, pbar):
tasks.append(asyncio.create_task(limited_request_func(request_func_input=request_func_input, pbar=pbar)))

outputs: list[RequestFuncOutput] = await asyncio.gather(*tasks)
if shared_session is not None:
await shared_session.close()
else:
# 多ip按DP均分并发
assert max_concurrency, "multi-IP 模式必须指定 max_concurrency"
Expand Down
Loading