diff --git a/ais_bench/benchmark/models/api_models/base_api.py b/ais_bench/benchmark/models/api_models/base_api.py index db5eef7f..5b769c8e 100644 --- a/ais_bench/benchmark/models/api_models/base_api.py +++ b/ais_bench/benchmark/models/api_models/base_api.py @@ -1,6 +1,6 @@ import sys +import os import json -import urllib.parse import warnings import asyncio import os.path as osp @@ -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 @@ -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 +_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 + 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 @@ -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. @@ -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}" # For IPv6 literals, wrap in brackets when constructing the URL. host = self.host_ip @@ -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: + 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() @@ -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 @@ -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. @@ -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, @@ -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. @@ -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 @@ -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