[feat][combination]AISBench支持多组合条件下批量用例的性能测试

ais_bench/benchmark/cli/argument_parser.py

@@ -2,11 +2,15 @@
 import json
 from ais_bench.benchmark.cli.utils import (
     get_current_time_str,
+    validate_batch_size,
     validate_max_workers,
     validate_max_workers_per_gpu,
     validate_num_prompts,
     validate_num_warmups,
-    validate_pressure_time
+    validate_prefix_len,
+    validate_pressure_time,
+    validate_request_count,
+    validate_request_size,
 )
 from ais_bench.benchmark.cli.utils import DEFAULT_PRESSURE_TIME
 
@@ -21,6 +25,7 @@ def __init__(self):
         self._custom_dataset_parser()
         self._agent_parser()
         self._api_model_parser()
+        self._synthetic_dataset_parser()
 
     def parse_args(self):
         args = self.parser.parse_args()
@@ -115,8 +120,11 @@ def _base_parser(self):
         parser.add_argument(
             '--num-prompts',
             help='Num Prompts, Specify the number of prompts to infer and evaluate. '
-            'Must be integer >= 1, If not provided, all prompts will be inferred and evaluated. ',
+            'Must be integer >= 1, If not provided, all prompts will be inferred and evaluated. '
+            'Accepts 1..N integers; with multiple values each one pairs with the '
+            'corresponding --batch-size value (or a single --batch-size broadcasts).',
             type=validate_num_prompts,
+            nargs='+',
             default=None
         )
         parser.add_argument(
@@ -424,8 +432,11 @@ def _api_model_parser(self):
                             help='Endpoint url override')
         parser.add_argument('--max-out-len', type=int, default=None,
                             help='Max output length override')
-        parser.add_argument('--batch-size', type=int, default=None,
-                            help='Batch size override in the range (0, 100000]')
+        parser.add_argument('--batch-size', type=validate_batch_size, nargs='+',
+                            default=None,
+                            help='Batch size override in the range (0, 100000]. '
+                            'Accepts 1..N integers; multiple values expand into one '
+                            'serial sweep combination per concurrency level.')
         parser.add_argument(
             '--trust-remote-code',
             action=argparse.BooleanOptionalAction,
@@ -434,4 +445,42 @@ def _api_model_parser(self):
         )
         parser.add_argument('--generation-kwargs', type=json.loads, default=None,
                             help='Generation kwargs override as a JSON object, '
-                            'e.g. \'{"temperature": 0.01, "ignore_eos": false}\'')
+                            'e.g. \'{"temperature": 0.01, "ignore_eos": false}\'. '
+                            'The `stream` key is not allowed here, use --stream instead.')
+        parser.add_argument('--model-abbr', type=str, default=None,
+                            help='Model abbr override; each sweep combination appends -bs{N}')
+        parser.add_argument(
+            '--stream',
+            action=argparse.BooleanOptionalAction,
+            default=None,
+            help='Stream override (--stream / --no-stream); only streaming can measure TTFT/TPOT'
+        )
+        parser.add_argument(
+            '--use-timestamp',
+            action=argparse.BooleanOptionalAction,
+            default=None,
+            help='use_timestamp override (--use-timestamp / --no-use-timestamp); '
+            'request_rate is ignored when enabled'
+        )
+
+    def _synthetic_dataset_parser(self):
+        """这些参数用于覆盖随机合成数据集配置，均为单值，1:1 映射配置字段。"""
+        parser = self.parser.add_argument_group('synthetic_dataset_args')
+        parser.add_argument('--request-count', type=validate_request_count, default=None,
+                            help='Override RequestCount of the single random synthetic dataset '
+                            '(sample pool size, 1 <= v <= 2^20)')
+        parser.add_argument('--tokenid-config-request-size', type=validate_request_size,
+                            default=None,
+                            help='Override TokenIdConfig.RequestSize of the synthetic dataset '
+                            '(total input length in tokens, 1 <= v <= 2^20)')
+        parser.add_argument('--tokenid-config-prefix-len', type=validate_prefix_len,
+                            default=None,
+                            help='Override TokenIdConfig.PrefixLen of the synthetic dataset '
+                            '(absolute token count, 0 <= v <= RequestSize)')
+        parser.add_argument('--string-config-input', type=json.loads, default=None,
+                            help='Override StringConfig.Input as a JSON object, '
+                            'e.g. \'{"Method": "uniform", "Params": {"MinValue": 1, "MaxValue": 1024}}\'')
+        parser.add_argument('--string-config-output', type=json.loads, default=None,
+                            help='Override StringConfig.Output as a JSON object')
+        parser.add_argument('--dataset-abbr', type=str, default=None,
+                            help='Dataset abbr override (no suffix is appended)')

ais_bench/benchmark/cli/config_manager.py

@@ -1,6 +1,7 @@
 import os
 import os.path as osp
 import inspect
+import json
 import tabulate
 from mmengine.config import Config
 
@@ -11,6 +12,11 @@
 from ais_bench.benchmark.utils.logging.exceptions import CommandError, AISBenchConfigError
 from ais_bench.benchmark.cli.utils import fill_model_path_if_datasets_need, fill_test_range_use_num_prompts, recur_convert_config_type
 from ais_bench.benchmark.utils.response_anomaly import ResponseAnomalyCoordinator
+from ais_bench.benchmark.cli.sweep import (
+    MANIFEST_FILENAME,
+    _normalize_sweep_arg,
+    expand_sweep_combinations,
+)
 
 RESPONSE_ANOMALY_TOP_LOGPROBS = 20
 
@@ -113,12 +119,54 @@ def load_config(self, workflow):
         self.cfg = self._get_config_from_arg()
         self._init_response_anomaly_config()
         self._update_and_init_work_dir()
+        self._expand_sweep_configs()
         self._inject_response_anomaly_payload_storage()
         self._fill_dataset_configs()
         self._update_cfg_of_workflow(workflow)
         self._dump_and_reload_config()
         return self.cfg
 
+    def _expand_sweep_configs(self):
+        """按 `--batch-size`/`--num-prompts` 展开多组合配置并落盘 manifest。"""
+        cli_args = self.cfg.get('cli_args') or {}
+        result = expand_sweep_combinations(
+            self.cfg.get('models') or [],
+            self.cfg.get('datasets') or [],
+            cli_args,
+        )
+        self.cfg['models'] = result.models
+        self.cfg['datasets'] = result.datasets
+        if result.combinations:
+            self.cfg['model_dataset_combinations'] = result.combinations
+        if result.sweep_cooldown_s:
+            self.cfg['sweep_cooldown_s'] = result.sweep_cooldown_s
+        if result.manifest is not None:
+            self._write_manifest(result.manifest)
+
+    def _write_manifest(self, manifest):
+        """原子写入 `<work_dir>/combos_manifest.json`；失败仅告警不阻断。"""
+        work_dir = self.cfg.get('work_dir')
+        if not work_dir:
+            return
+        path = osp.join(work_dir, MANIFEST_FILENAME)
+        tmp_path = f'{path}.tmp{os.getpid()}'
+        try:
+            os.makedirs(work_dir, exist_ok=True)
+            with open(tmp_path, 'w', encoding='utf-8') as f:
+                json.dump(manifest, f, ensure_ascii=False, indent=2)
+                f.flush()
+                os.fsync(f.fileno())
+            os.replace(tmp_path, path)
+        except (OSError, TypeError) as exc:
+            try:
+                if os.path.exists(tmp_path):
+                    os.remove(tmp_path)
+            except OSError:
+                pass
+            self.logger.warning(f'Failed to write sweep manifest {path}: {exc}')
+            return
+        self.logger.info(f'Sweep manifest written to {path}')
+
     def _init_response_anomaly_config(self):
         """Normalize the optional response anomaly detection configuration."""
         global_cfg, configured_top_logprobs = (
@@ -600,7 +648,9 @@ def _fill_dataset_configs(self):
         for dataset_cfg in self.cfg["datasets"]:
             if dataset_cfg.get("infer_cfg", None) is None:
                 continue
-            fill_test_range_use_num_prompts(self.cfg["cli_args"].get("num_prompts"), dataset_cfg)
+            fill_test_range_use_num_prompts(
+                _normalize_sweep_arg(self.cfg["cli_args"].get("num_prompts")), dataset_cfg
+            )
             fill_model_path_if_datasets_need(self.cfg["models"][0], dataset_cfg)
             retriever_cfg = dataset_cfg["infer_cfg"]["retriever"]
             infer_cfg = dataset_cfg["infer_cfg"]
@@ -829,6 +879,12 @@ def _apply_cli_api_model_overrides(self, models):
                 cli_val = getattr(self.args, field, None)
                 if cli_val is None or field not in model_cfg:
                     continue
+                if field == "batch_size":
+                    # 多档 sweep 时 batch_size 是列表，由组合展开器逐副本写入；
+                    # 单值（含长度 1 列表）在此归一成 int 覆盖既有键。
+                    cli_val = _normalize_sweep_arg(cli_val)
+                    if not isinstance(cli_val, int):
+                        continue
                 model_cfg[field] = cli_val
 
     def _load_summarizers_config(self):
@@ -918,9 +974,12 @@ def _dump_and_reload_config(self):
         recur_convert_config_type(self.cfg)
         self.cfg.dump(output_config_path)
         # eval nums set
-        if (self.args.num_prompts and self.args.num_prompts < 0) or self.args.num_prompts == 0:
+        num_prompts = _normalize_sweep_arg(self.args.num_prompts)
+        if isinstance(num_prompts, int) and num_prompts <= 0:
             raise CommandError(TMAN_CODES.INVALID_ARG_VALUE_IN_CMD, "'--num-prompts' must be a positive integer greater than 0.")
-        self.cfg['num_prompts'] = self.args.num_prompts
+        # 多档 sweep 时 num_prompts 已逐组合写入数据集副本的 test_range，
+        # 顶层字段只承载单值语义，列表形态记 None。
+        self.cfg['num_prompts'] = num_prompts if isinstance(num_prompts, int) else None
         # Config is intentally reloaded here to avoid initialized
         # types cannot be serialized
         try: