Add num-trials as cli arg (#109)

* add num-trials as cli arg

* add test and format

* fix test

* use global var for default num trials

* implement copilot suggestions

Author: jasonlizhengjian

src/flexible_inference_benchmark/main.py

@@ -38,6 +38,10 @@
 
 logger = logging.getLogger(__name__)
 
+# Default value for num_trials argument
+DEFAULT_NUM_TRIALS = 10
+MAX_TRIALS = 100  # Maximum trials for prompt generation, warn if exceeded
+
 
 def return_random_image_by_size(width: int, height: int, convert_to_base64: bool = False) -> Any:
 
@@ -198,6 +202,8 @@ def generate_prompts(
             "User selected sharegpt dataset. "
             "Ignoring prompt length distribution and following the prompts from the dataset."
         )
+        if args.num_trials != DEFAULT_NUM_TRIALS:  # Check if user specified custom value
+            logger.warning("num_trials parameter is ignored for ShareGPT dataset as prompts are pre-defined")
         prompt_cls = ShareGPT(filename, tokenizer, output_token_dist)
     else:
         logger.info(f"User selected {args.dataset_name} dataset. Generating prompt from distributions.")
@@ -216,7 +222,12 @@ def generate_prompts(
         if args.prefix_len:
             prompt_cls = (
                 Random.with_prefix_len(
-                    args.prefix_len, input_prompt_dist, output_token_dist, tokenizer, args.ignore_input_distribution
+                    args.prefix_len,
+                    input_prompt_dist,
+                    output_token_dist,
+                    tokenizer,
+                    args.ignore_input_distribution,
+                    args.num_trials,
                 )
                 if args.dataset_name == "random"
                 else Textfile.with_prefix_len(
@@ -226,13 +237,19 @@ def generate_prompts(
                     output_token_dist,
                     tokenizer,
                     args.ignore_input_distribution,
+                    args.num_trials,
                 )
             )
         else:
             prefix_text = args.prefix_text or ""
             prompt_cls = (
                 Random.with_prefix_str(
-                    prefix_text, input_prompt_dist, output_token_dist, tokenizer, args.ignore_input_distribution
+                    prefix_text,
+                    input_prompt_dist,
+                    output_token_dist,
+                    tokenizer,
+                    args.ignore_input_distribution,
+                    args.num_trials,
                 )
                 if args.dataset_name == "random"
                 else Textfile.with_prefix_str(
@@ -242,6 +259,7 @@ def generate_prompts(
                     output_token_dist,
                     tokenizer,
                     args.ignore_input_distribution,
+                    args.num_trials,
                 )
             )
 
@@ -492,6 +510,15 @@ def add_benchmark_subparser(subparsers: argparse._SubParsersAction) -> Any:  # t
         help="Number of input tokens to use for validation prompts (default: 128).",
     )
 
+    benchmark_parser.add_argument(
+        "--num-trials",
+        type=int,
+        default=DEFAULT_NUM_TRIALS,
+        help="Number of attempts to achieve exact token count when generating prompts (default: 10). "
+        "Used for 'random' and 'other' datasets. Higher values improve token count precision "
+        "but may slow down prompt generation. Ignored for ShareGPT datasets.",
+    )
+
     return benchmark_parser
 
 
@@ -604,6 +631,12 @@ def fail(msg: str) -> None:
         if args.dataset_path and not args.dataset_name:
             args.dataset_name = "other"
 
+        # Validate num_trials parameter
+        if args.num_trials <= 0:
+            fail("Number of trials must be positive")
+        if args.num_trials > MAX_TRIALS:
+            logger.warning(f"High num_trials value ({args.num_trials}) may slow down prompt generation")
+
     return args
 
 

tests/conftest.py

@@ -69,6 +69,7 @@ def args_configs():
         "temperature": 0.0,
         "top_p": None,
         "top_k": None,
+        "num_trials": 8,
     }
 
     sharegpt_sample_data_path = "tests/data/sharegpt_sample_test_data.json"

tests/test_data.py

@@ -6,6 +6,7 @@
 import json
 import flexible_inference_benchmark.engine.data as data
 import flexible_inference_benchmark.engine.distributions as distributions
+from flexible_inference_benchmark.main import parse_args
 from sharegpt_data import SHAREGPT_DATA
 
 @pytest.mark.parametrize("ignore_input_distribution", [True, False])
@@ -57,3 +58,32 @@ def test_sharegpt():
     if os.path.exists("sharegpt_test.json"):
         os.remove("sharegpt_test.json")
     assert random_data.shape == (10,3)
+
+def test_num_trials_cli_argument():
+    """Test that num_trials CLI argument is properly parsed and validated."""
+    import sys
+    
+    # Test default value
+    original_argv = sys.argv
+    try:
+        sys.argv = ['fib', 'benchmark', '--model', 'test', '--base-url', 'http://test']
+        args = parse_args()
+        assert args.num_trials == 10
+        
+        # Test custom value
+        sys.argv = ['fib', 'benchmark', '--model', 'test', '--base-url', 'http://test', '--num-trials', '5']
+        args = parse_args()
+        assert args.num_trials == 5
+        
+        # Test validation - zero value should fail
+        sys.argv = ['fib', 'benchmark', '--model', 'test', '--base-url', 'http://test', '--num-trials', '0']
+        with pytest.raises(SystemExit):
+            parse_args()
+            
+        # Test validation - negative value should fail
+        sys.argv = ['fib', 'benchmark', '--model', 'test', '--base-url', 'http://test', '--num-trials', '-1']
+        with pytest.raises(SystemExit):
+            parse_args()
+            
+    finally:
+        sys.argv = original_argv