feat(test): add AISBench prefix test tool #1030
Open
Potterluo wants to merge 2 commits into
Open
Conversation
- Add aisbench_utils module with data selection, dataset generation, and API configuration capabilities - Implement API config generation supporting streaming and text test types - Add data selector supporting GSM8K dataset and random token mode - Add multi-prefix dataset generation with variable-length pattern support - Introduce configuration dataclasses to unify AISBench test parameters - Add result parsing and prefix hit rate calculation functions - Add AISBench prefix cache test configuration options in the configuration file
- Add aisbench_utils module with data selection, dataset generation, and API configuration capabilities - Implement API config generation supporting streaming and text test types - Add data selector supporting GSM8K dataset and random token mode - Add multi-prefix dataset generation with variable-length pattern support - Introduce configuration dataclasses to unify AISBench test parameters - Add result parsing and prefix hit rate calculation functions - Add AISBench prefix cache test configuration options in the configuration file
ygwpz
reviewed
Jun 17, 2026
| if result.returncode != 0 or not result.stdout.strip(): | ||
| return {}, {} | ||
|
|
||
| lines = result.stdout.strip().split("\n") |
Contributor
There was a problem hiding this comment.
shell=True with interpolated URLs in subprocess is a potential security risk. While ip_address and port should be validated, consider using subprocess.run(["curl", "-s", url], ...) without shell=True for safer execution.
ygwpz
reviewed
Jun 17, 2026
| except OSError as e: | ||
| if e.errno == errno.EEXIST: | ||
| os.remove(link_name) | ||
| os.symlink(target, link_name) |
Contributor
There was a problem hiding this comment.
💡 Suggestion: The sys.platform == "win32" check handles Windows, but other Unix-like systems (macOS) might behave differently. Consider using os.name == "nt" or a more robust platform detection.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Pull Request Description
via: https://github.com/rayn-zzz/aisbench_auto_tools_prefix
Summary
This PR introduces a new test tool module
aisbench_utilsto enable prefix caching performance testing for our inference service. It provides a flexible framework for generating test datasets, configuring API requests, and parsing results with hit rate metrics.Key Changes
Module structure: Created
aisbench_utilscontaining:data_selector.py– supports GSM8K (grade-school math) dataset and random token generation modes.dataset_generator.py– generates multi-prefix datasets with configurable variable-length patterns.api_config.py– builds API request configurations for both streaming and non‑streaming text tests.config_dataclasses.py– dataclasses for unified parameter management (e.g., batch size, prefix lengths, concurrency).result_parser.py– parses test outputs and computes prefix hit rates.Configuration: Added new YAML/JSON config options under
aisbenchsection, covering:Testing: Included unit tests for data selection, dataset generation, and hit rate calculation (see
/tests/test_aisbench_utils.py).Motivation
Prefix caching is critical for reducing latency in LLM serving. This tool allows us to systematically benchmark cache hit rates under varying prefix distributions, helping optimize our caching policies and deployment configurations.
How to Test
python examples/run_aisbench.py --config configs/aisbench_example.yamlAdditional Notes