【WIP】 control Prefill con in aisbench - #396
Conversation
Added prefill pipeline functionality with environment variable control and logging. Updated various functions to support new prefill behavior and improved error handling.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a prefill pipeline semaphore mechanism in base_api.py to allow overlapping the decode phase of one request with the prefill phase of the next, controlled via environment variables. It also simplifies custom URL parsing and cleans up formatting. The review feedback highlights critical issues: using a single global asyncio.Semaphore can cause RuntimeError across different event loops, removing the trailing slash enforcement in _get_base_url introduces a regression for custom URLs with path prefixes, and a potential semaphore leak exists in stream_infer if an exception occurs before entering the try block.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| _prefill_semaphore: Optional[asyncio.Semaphore] = None | ||
| _prefill_sem_lock: Optional[asyncio.Lock] = None |
There was a problem hiding this comment.
Using a single global asyncio.Semaphore and asyncio.Lock can cause RuntimeError (e.g., attached to a different loop) when running multiple benchmarks or unit tests sequentially, as they get bound to the first active event loop.
Instead, we can use a dictionary mapping the current running event loop to its corresponding semaphore. This also completely eliminates the need for a lock, as the semaphore creation is synchronous and safe from race conditions within the same event loop.
| _prefill_semaphore: Optional[asyncio.Semaphore] = None | |
| _prefill_sem_lock: Optional[asyncio.Lock] = None | |
| _prefill_semaphores: Dict[asyncio.AbstractEventLoop, asyncio.Semaphore] = {} |
| async def _get_prefill_semaphore() -> Optional[asyncio.Semaphore]: | ||
| """惰性创建全局 prefill 信号量（绑定到当前事件循环）。""" | ||
| global _prefill_semaphore, _prefill_sem_lock | ||
| if not _prefill_pipeline_enabled(): | ||
| return None | ||
| if _prefill_sem_lock is None: | ||
| _prefill_sem_lock = asyncio.Lock() | ||
| async with _prefill_sem_lock: | ||
| if _prefill_semaphore is None: | ||
| try: | ||
| conc = int(os.environ.get("AISBENCH_PREFILL_CONCURRENCY", "1")) | ||
| except ValueError: | ||
| conc = 1 | ||
| if conc < 1: | ||
| conc = 1 | ||
| _prefill_semaphore = asyncio.Semaphore(conc) | ||
| return _prefill_semaphore |
There was a problem hiding this comment.
Update _get_prefill_semaphore to use the loop-safe dictionary mapping. This avoids binding a single global semaphore to a closed or different event loop, and removes the redundant lock since the initialization is synchronous.
async def _get_prefill_semaphore() -> Optional[asyncio.Semaphore]:
"""惰性创建全局 prefill 信号量（绑定到当前事件循环）。"""
global _prefill_semaphores
if not _prefill_pipeline_enabled():
return None
loop = asyncio.get_running_loop()
if loop not in _prefill_semaphores:
try:
conc = int(os.environ.get("AISBENCH_PREFILL_CONCURRENCY", "1"))
except ValueError:
conc = 1
if conc < 1:
conc = 1
_prefill_semaphores[loop] = asyncio.Semaphore(conc)
return _prefill_semaphores[loop]| def _get_base_url(self) -> str: | ||
| protocol = "https" if self.enable_ssl else "http" | ||
| clean_url = self.url.strip() if isinstance(self.url, str) else "" | ||
| if clean_url: | ||
| self.logger.info( | ||
| f"Using custom URL: [{clean_url}], [host_ip: {self.host_ip}] and [host_port: {self.host_port}] will be ignored" | ||
| ) | ||
| if self.url: | ||
| self.logger.info(f"Using custom URL: [{self.url}], [host_ip: {self.host_ip}] and [host_port: {self.host_port}] will be ignored") | ||
| # Check if URL already contains protocol | ||
| if clean_url.startswith("http://") or clean_url.startswith("https://"): | ||
| url = clean_url | ||
| else: | ||
| url = f"{protocol}://{clean_url}" | ||
| # Ensure trailing slash on path to avoid urljoin dropping the last path segment. | ||
| # Use urlparse/urlunparse to safely handle URLs with query strings or fragments. | ||
| parsed = urllib.parse.urlparse(url) | ||
| if not parsed.path or parsed.path.endswith("/"): | ||
| return url | ||
| return urllib.parse.urlunparse(parsed._replace(path=parsed.path + "/")) | ||
| if self.url.startswith("http://") or self.url.startswith("https://"): | ||
| return self.url | ||
| return f"{protocol}://{self.url}" |
There was a problem hiding this comment.
Removing the trailing slash enforcement logic introduces a regression. If a user provides a custom URL with a path prefix (e.g., http://localhost:8080/v1), urllib.parse.urljoin in subclasses (like VLLMCustomAPI or MindieStreamApi) will drop the last path segment (e.g., /v1), resulting in incorrect endpoint URLs.
We should restore the trailing slash enforcement to ensure path prefixes are preserved.
def _get_base_url(self) -> str:
protocol = "https" if self.enable_ssl else "http"
if self.url:
self.logger.info(f"Using custom URL: [{self.url}], [host_ip: {self.host_ip}] and [host_port: {self.host_port}] will be ignored")
# Check if URL already contains protocol
if self.url.startswith("http://") or self.url.startswith("https://"):
url = self.url
else:
url = f"{protocol}://{self.url}"
# Ensure trailing slash on path to avoid urljoin dropping the last path segment.
import urllib.parse
parsed = urllib.parse.urlparse(url)
if not parsed.path or parsed.path.endswith("/"):
return url
return urllib.parse.urlunparse(parsed._replace(path=parsed.path + "/"))| if prefill_sem is not None: | ||
| await prefill_sem.acquire() | ||
| _prefill_inflight += 1 | ||
| _prefill_log("acquire") | ||
| # 计时起点放在 acquire() 之后：TTFT 只反映“真正发出请求 -> 收到首 token” | ||
| # 的服务端耗时，不包含客户端等待 prefill 槽位的排队时间（与 benchmark_req.py 口径一致）。 | ||
| await output.record_time_point() | ||
| async with self.session.post( | ||
| url=self.url, json=request_body, headers=self.headers | ||
| ) as response: | ||
| if response.status == 200: | ||
| async for raw_chunk in self.iter_lines(response.content): | ||
| chunk = raw_chunk.strip() | ||
| if not chunk: | ||
| continue | ||
| chunk = chunk.decode("utf-8") | ||
| if chunk.startswith(":"): | ||
| continue | ||
| chunk = chunk.removeprefix("data:").strip() | ||
| if chunk == "[DONE]": | ||
| break | ||
| await output.record_time_point() | ||
| try: | ||
| data = json.loads(chunk) | ||
| except json.JSONDecodeError as e: | ||
| output.success = False | ||
| output.error_info = f"Unexpected response format: {raw_chunk}. Please check if server is working correctly." | ||
| raise AISBenchValueError( | ||
| MODEL_CODES.PARSE_TEXT_RSP_INVALID_FORMAT, | ||
| f"Unexpected response format. Please check 'error_info' in ***_failed.jsonl for more information.", | ||
| ) | ||
| await self.parse_stream_response(data, output) | ||
| output.success = True | ||
| else: | ||
| output.error_info = response.reason | ||
| output.success = False | ||
| try: |
There was a problem hiding this comment.
To prevent potential semaphore leaks and deadlocks, the try...finally block should be entered immediately after acquiring the semaphore. If an exception or cancellation occurs during await output.record_time_point(), the acquired semaphore would otherwise be leaked because the try block has not yet been entered.
Additionally, wrapping the state updates (_prefill_inflight += 1 and logging) inside the try block ensures that any unexpected failure there is also safely recovered by the finally block.
if prefill_sem is not None:
await prefill_sem.acquire()
try:
if prefill_sem is not None:
_prefill_inflight += 1
_prefill_log("acquire")
# 计时起点放在 acquire() 之后：TTFT 只反映“真正发出请求 -> 收到首 token”
# 的服务端耗时，不包含客户端等待 prefill 槽位的排队时间（与 benchmark_req.py 口径一致）。
await output.record_time_point()1 participant
改了什么
对齐
benchmark_req.py的测试方法，新增 “prefill 完成即放行下一个请求” 的流水线机制：asyncio.Semaphore(n)。stream_infer中：发请求前acquire()，流式收到第一个数据块（首 token，prefill 完成）时立即release()，让下一个请求的 prefill 立刻开始；
finally兜底释放防死锁。record_time_point()放在acquire()之后：TTFT 不包含客户端排队等待时间（与 benchmark_req.py 口径一致）。
环境变量开关（默认关闭，不影响原版行为）
Thanks for your contribution; we appreciate it a lot. The following instructions will make your pull request healthier and help you get feedback more easily. If you do not understand some items, don't worry, just make the pull request and seek help from maintainers.
感谢您的贡献，我们非常重视。以下说明将使您的拉取请求更健康，更易于获得反馈。如果您不理解某些项目，请不要担心，只需提交拉取请求并从维护人员那里寻求帮助即可。
PR Type / PR类型
Related Issue | 关联 Issue
Fixes #(issue ID / issue 编号) / Relates to #(issue ID / issue 编号)
🔍 Motivation / 变更动机
Please describe the motivation of this PR and the goal you want to achieve through this PR.
请描述您的拉取请求的动机和您希望通过此拉取请求实现的目标。
📝 Modification / 修改内容
Please briefly describe what modification is made in this PR.
请简要描述此拉取请求中进行的修改。
📐 Associated Test Results / 关联测试结果
Please provide links to the related test results, such as CI pipelines, test reports, etc.
请提供相关测试结果的链接，例如 CI 管道、测试报告等。
Does the modification introduce changes that break the backward compatibility of the downstream repositories? If so, please describe how it breaks the compatibility and how the downstream projects should modify their code to keep compatibility with this PR.
是否引入了会破坏下游存储库向后兼容性的更改？如果是，请描述它如何破坏兼容性，以及下游项目应该如何修改其代码以保持与此 PR 的兼容性。
If the modification introduces performance degradation, please describe the impact of the performance degradation and the expected performance improvement.
如果引入了性能下降，请描述性能下降的影响和预期的性能改进。
🌟 Use cases (Optional) / 使用案例（可选）
If this PR introduces a new feature, it is better to list some use cases here and update the documentation.
如果此拉取请求引入了新功能，最好在此处列出一些用例并更新文档。
✅ Checklist / 检查列表
Before PR:
After PR:
👥 Collaboration Info / 协作信息
🌟 Useful CI Command / 实用的CI命令
/gemini review/gemini summary/gemini help/readthedocs build