feat(benchmark script): add similar to vllm --ready-check-timeout-sec parameter (#15466)

This commit is contained in:
almaslof
2026-03-03 23:44:38 +02:00
committed by GitHub
parent ac2819c81f
commit b0f26698f5
2 changed files with 35 additions and 0 deletions

View File

@@ -125,6 +125,27 @@ def get_request_headers() -> Dict[str, str]:
return headers
def wait_for_endpoint(url: str, timeout_sec: int = 60) -> bool:
"""Wait for the server to become ready by polling the given URL."""
print(f"Waiting up to {timeout_sec}s for {url} to become ready...")
start_time = time.perf_counter()
headers = get_auth_headers()
while True:
try:
response = requests.get(url, headers=headers, timeout=5)
if response.status_code == 200:
elapsed = time.perf_counter() - start_time
print(f"Server ready in {elapsed:.1f}s.")
return True
except requests.exceptions.RequestException:
pass
elapsed = time.perf_counter() - start_time
if elapsed >= timeout_sec:
print(f"Server did not become ready within {timeout_sec}s timeout.")
return False
time.sleep(1)
# trt llm does not support ignore_eos
# https://github.com/triton-inference-server/tensorrtllm_backend/issues/505
async def async_request_trt_llm(
@@ -1663,6 +1684,13 @@ def run_benchmark(args_: argparse.Namespace):
f"http://{args.host}:{args.port}" if args.base_url is None else args.base_url
)
# Wait for server to be ready
if args.ready_check_timeout_sec > 0:
health_url = model_url if args.backend not in ("trt", "gserver") else base_url
if not wait_for_endpoint(health_url, args.ready_check_timeout_sec):
print(f"Server at {health_url} is not ready. Exiting.")
sys.exit(1)
# Get model name
if args.model is None:
if args.backend == "truss":
@@ -1788,6 +1816,12 @@ if __name__ == "__main__":
type=int,
help="If not set, the default port is configured according to its default value for different LLM Inference Engines.",
)
parser.add_argument(
"--ready-check-timeout-sec",
type=int,
default=60,
help="Maximum time in seconds to wait for the server to be ready before benchmarking. Set to 0 to skip. Default: 60.",
)
parser.add_argument(
"--dataset-name",
type=str,

View File

@@ -1090,6 +1090,7 @@ def get_benchmark_args(
gsp_num_turns=gsp_num_turns,
header=header,
max_concurrency=max_concurrency,
ready_check_timeout_sec=0,
)