Skip to content

【WIP】 control Prefill con in aisbench #396

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

Open
vasede wants to merge 2 commits into
base: master
Choose a base branch
from
Open

【WIP】 control Prefill con in aisbench #396

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
44e1da2
Implement prefill pipeline feature in BaseAPIModel
vasede Jul 14, 2026
ad5a208
Fix string quotes and formatting in base_api.py
vasede Jul 14, 2026
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
215 changes: 133 additions & 82 deletions ais_bench/benchmark/models/api_models/base_api.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import sys
import os
import json
import urllib.parse
import warnings
import asyncio
import os.path as osp
Expand All @@ -15,19 +15,12 @@
from ais_bench.benchmark.utils.logging.logger import AISLogger
from ais_bench.benchmark.utils.logging.error_codes import MODEL_CODES
from ais_bench.benchmark.utils.logging.exceptions import (
AISBenchNotImplementedError,
AISBenchValueError,
AISBenchKeyError,
AISBenchTypeError,
AISBenchRuntimeError,
AISBenchImplementationError,
)
AISBenchNotImplementedError, AISBenchValueError, AISBenchKeyError,
AISBenchTypeError, AISBenchRuntimeError, AISBenchImplementationError)
from ais_bench.benchmark.utils.prompt import PromptList
from ais_bench.benchmark.models import BaseModel
from ais_bench.benchmark.models.output import Output
from ais_bench.benchmark.openicl.icl_inferencer.output_handler.ppl_inferencer_output_handler import (
PPLRequestOutput,
)
from ais_bench.benchmark.openicl.icl_inferencer.output_handler.ppl_inferencer_output_handler import PPLRequestOutput
from ais_bench.benchmark.utils.logging.error_codes import ICLI_CODES
from ais_bench.benchmark.global_consts import REQUEST_TIME_OUT

Expand All @@ -36,12 +29,67 @@

PromptType = Union[PromptList, str]

# ---------------------------------------------------------------------------
# Prefill 流水线信号量（自定义功能，对齐 benchmark_req.py 的测试方法）
# 目的：当某个请求完成 prefill（流式收到第一个 token，进入 decode）时，
# 立即释放槽位，让下一个请求的 prefill 马上开始，从而让前一个请求的 decode
# 与后一个请求的 prefill 在服务端重叠。
# 通过环境变量开启，默认关闭，不影响 ais_bench 原有行为：
# AISBENCH_PREFILL_PIPELINE=1 开启该功能
# AISBENCH_PREFILL_CONCURRENCY=1 同时处于 prefill 阶段的最大请求数（默认 1）
# AISBENCH_PREFILL_OBSERVE=1 在 stderr 打印实时 prefill/decode 并发数
# ---------------------------------------------------------------------------
_prefill_semaphore: Optional[asyncio.Semaphore] = None
_prefill_sem_lock: Optional[asyncio.Lock] = None
Comment on lines +42 to +43

@gemini-code-assist gemini-code-assist Bot Jul 14, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

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.

Suggested change
_prefill_semaphore: Optional[asyncio.Semaphore] = None
_prefill_sem_lock: Optional[asyncio.Lock] = None
_prefill_semaphores: Dict[asyncio.AbstractEventLoop, asyncio.Semaphore] = {}

_prefill_inflight: int = 0
_decode_inflight: int = 0


def _prefill_pipeline_enabled() -> bool:
return os.environ.get("AISBENCH_PREFILL_PIPELINE", "0") in ("1", "true", "True")


def _prefill_observe_enabled() -> bool:
return os.environ.get("AISBENCH_PREFILL_OBSERVE", "0") in ("1", "true", "True")


def _prefill_log(tag: str) -> None:
"""打印当前 prefill / decode 并发数到 stderr（带前缀便于 grep）。"""
if _prefill_observe_enabled():
try:
conc = os.environ.get("AISBENCH_PREFILL_CONCURRENCY", "1")
sys.stderr.write(
f"[PREFILL_PIPE] {tag} | prefill_inflight={_prefill_inflight} "
f"decode_inflight={_decode_inflight} (prefill_cap={conc})\n"
)
sys.stderr.flush()
except Exception:
pass


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
Comment on lines +70 to +86

@gemini-code-assist gemini-code-assist Bot Jul 14, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

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]



class MockAISLogger(AISLogger):
"""Mock logger for API model. Because model will init in task for warmup and init in infer process,
so we mock the logger to avoid the print each log in init process twice
"""

def info(self, msg, *args, **kwargs):
pass

Expand All @@ -51,7 +99,6 @@ def debug(self, msg, *args, **kwargs):
def warning(self, msg, *args, **kwargs):
pass


class BaseAPIModel(BaseModel):
"""Base class for API model wrapper.

Expand Down Expand Up @@ -114,27 +161,17 @@ def _get_url(self) -> str:
raise AISBenchNotImplementedError(
MODEL_CODES.UNKNOWN_ERROR,
f"{self.__class__.__name__} does not supported"
" to be called in base classes",
" to be called in base classes"
)

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}"
Comment on lines 167 to +174

@gemini-code-assist gemini-code-assist Bot Jul 14, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

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 + "/"))


# For IPv6 literals, wrap in brackets when constructing the URL.
host = self.host_ip
Expand Down Expand Up @@ -280,36 +317,63 @@ async def generate(
return output

async def stream_infer(self, request_body: dict, output: Output):
global _prefill_inflight, _decode_inflight
prefill_sem = await _get_prefill_semaphore()
prefill_released = False
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:
Comment on lines +323 to +330

@gemini-code-assist gemini-code-assist Bot Jul 14, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

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()

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)
# prefill 完成（收到第一个数据块，进入 decode）：立即释放槽位，
# 让下一个请求的 prefill 马上开始。
if prefill_sem is not None and not prefill_released:
prefill_sem.release()
prefill_released = True
_prefill_inflight -= 1
_decode_inflight += 1
_prefill_log("release(first_token)")
output.success = True
else:
output.error_info = response.reason
output.success = False
finally:
# 兜底：若请求无 token 返回 / 出错 / 被取消，确保槽位被释放，避免死锁。
if prefill_sem is not None and not prefill_released:
prefill_sem.release()
prefill_released = True
_prefill_inflight -= 1
_prefill_log("release(finally)")
elif prefill_sem is not None:
_decode_inflight -= 1

async def text_infer(self, request_body, output: Output):
await output.record_time_point()
Expand All @@ -326,7 +390,7 @@ async def text_infer(self, request_body, output: Output):
output.error_info = f"Unexpected response format: {raw_data}. Please check if server is working correctly."
raise AISBenchValueError(
MODEL_CODES.PARSE_TEXT_RSP_INVALID_FORMAT,
f"Unexpected response format. Please check ***_details.jsonl for more information.",
f"Unexpected response format. Please check ***_details.jsonl for more information."
)
await self.parse_text_response(data, output)
output.success = True
Expand All @@ -340,8 +404,8 @@ async def get_ppl(
max_out_len: int,
output: PPLRequestOutput,
session: aiohttp.ClientSession = None,
**args,
):
**args
):
"""Compute perplexity for a given prompt via the remote API.
Args:
input_data: Prompt text or list structure the backend expects.
Expand Down Expand Up @@ -408,29 +472,15 @@ async def get_ppl(
if close_session:
await self.session.close()

async def get_ppl_request_body(
self, input_data: PromptType, max_out_len: int, output: PPLRequestOutput, **args
):
raise AISBenchNotImplementedError(
ICLI_CODES.IMPLEMENTATION_ERROR_PPL_METHOD_NOT_IMPLEMENTED,
f"PPL is not supported for this model.",
)
async def get_ppl_request_body(self, input_data:PromptType, max_out_len: int, output: PPLRequestOutput, **args):
raise AISBenchNotImplementedError(ICLI_CODES.IMPLEMENTATION_ERROR_PPL_METHOD_NOT_IMPLEMENTED, f"PPL is not supported for this model.")

def get_prompt_logprobs(self, data: dict):
raise AISBenchNotImplementedError(
ICLI_CODES.IMPLEMENTATION_ERROR_PPL_METHOD_NOT_IMPLEMENTED,
f"PPL is not supported for this model.",
)
raise AISBenchNotImplementedError(ICLI_CODES.IMPLEMENTATION_ERROR_PPL_METHOD_NOT_IMPLEMENTED, f"PPL is not supported for this model.")

def _calc_ppl(self, prompt_logprobs: list):
logprobs = [
list(item.values())[0]["logprob"]
for item in prompt_logprobs
if item is not None
]
tokenids = [
list(item.keys())[0] for item in prompt_logprobs if item is not None
]
logprobs = [list(item.values())[0]['logprob'] for item in prompt_logprobs if item is not None]
tokenids = [list(item.keys())[0] for item in prompt_logprobs if item is not None]
if len(tokenids) == 0:
raise AISBenchImplementationError(
ICLI_CODES.PPL_COMPUTE_ERROR_NO_VALID_TOKENS,
Expand All @@ -439,7 +489,6 @@ def _calc_ppl(self, prompt_logprobs: list):
loss = -sum(logprobs) / len(tokenids)
return loss


class APITemplateParser:
"""Intermidate prompt template parser, specifically for API models.

Expand Down Expand Up @@ -521,10 +570,12 @@ def parse_template(self, prompt_template: PromptType, mode: str) -> PromptType:
f"Parsing mode must be 'ppl' or 'gen', but got {mode}",
)


if isinstance(prompt_template, str):
return prompt_template

if self.meta_template:

prompt = PromptList()
# Whether to keep generating the prompt
generate = True
Expand Down Expand Up @@ -759,6 +810,6 @@ def _role2api_role(
else:
raise AISBenchValueError(
MODEL_CODES.INVALID_PROMPT_CONTENT,
"Invalid prompt content: without 'prompt' or 'prompt_mm' param!",
"Invalid prompt content: without 'prompt' or 'prompt_mm' param!"
)
return res, True
Loading