[cleanup] move arg parsing into a separate function (#178)

Pure refactor with no code functionality change.

Author: xy12181

benchmarks/benchmark_serving.py

@@ -706,136 +706,7 @@ def sample_warmup_requests(requests):
         break
 
 
-def main(args: argparse.Namespace):
-  print(args)
-  random.seed(args.seed)
-  np.random.seed(args.seed)
-
-  model_id = args.model
-  tokenizer_id = args.tokenizer
-  use_hf_tokenizer = args.use_hf_tokenizer
-
-  prefill_quota = AsyncCounter(init_value=3)
-  active_req_quota = AsyncCounter(init_value=450)
-
-  api_url = f"{args.server}:{args.port}"
-
-  tokenizer = get_tokenizer(model_id, tokenizer_id, use_hf_tokenizer)
-  if tokenizer == "test" or args.dataset == "test":
-    input_requests = mock_requests(
-        args.total_mock_requests
-    )  # e.g. [("AB", 2, "AB", 3)]
-  else:
-    dataset = []
-    if args.dataset == "openorca":
-      dataset = load_openorca_dataset_pkl(args.dataset_path)
-    elif args.dataset == "sharegpt":
-      dataset = load_sharegpt_dataset(
-          args.dataset_path,
-          args.conversation_starter,
-      )
-
-    # A given args.max_output_length value is the max generation step,
-    # when the args.max_output_length is default to None, the sample's golden
-    # output length will be used to decide the generation step.
-    input_requests = sample_requests(
-        dataset=dataset,
-        tokenizer=tokenizer,
-        num_requests=args.num_prompts,
-        max_output_length=args.max_output_length,
-    )
-
-  warmup_requests = None
-  if args.warmup_mode == "full":
-    warmup_requests = input_requests
-  elif args.warmup_mode == "sampled":
-    warmup_requests = list(sample_warmup_requests(input_requests)) * 2
-
-  if warmup_requests:
-    print(f"Warmup (mode: {args.warmup_mode}) is starting.")
-    _, _ = asyncio.run(
-        benchmark(
-            api_url=api_url,
-            tokenizer=tokenizer,
-            input_requests=warmup_requests,
-            request_rate=args.request_rate,
-            disable_tqdm=args.disable_tqdm,
-            prefill_quota=prefill_quota,
-            active_req_quota=active_req_quota,
-            is_warmup=True,
-        )
-    )
-    print(f"Warmup (mode: {args.warmup_mode}) has completed.")
-
-  # TODO: Replace this with warmup complete signal once supported.
-  # Wait for server completely warmup before running the benchmark.
-  time.sleep(5)
-
-  benchmark_result, request_outputs = asyncio.run(
-      benchmark(
-          api_url=api_url,
-          tokenizer=tokenizer,
-          input_requests=input_requests,
-          request_rate=args.request_rate,
-          disable_tqdm=args.disable_tqdm,
-          prefill_quota=prefill_quota,
-          active_req_quota=active_req_quota,
-      )
-  )
-
-  # Process output
-  output = [output.to_dict() for output in request_outputs]
-  if args.run_eval:
-    eval_json = eval_accuracy(output)
-
-  # Save config and results to json
-  if args.save_result:
-    # dimensions values are strings
-    dimensions_json = {}
-    # metrics values are numerical
-    metrics_json = {}
-
-    # Setup
-    current_dt = datetime.now().strftime("%Y%m%d-%H%M%S")
-    dimensions_json["date"] = current_dt
-    dimensions_json["model_id"] = model_id
-    dimensions_json["tokenizer_id"] = tokenizer_id
-    if args.additional_metadata_metrics_to_save is not None:
-      dimensions_json = {
-          **dimensions_json,
-          **json.loads(args.additional_metadata_metrics_to_save),
-      }
-    metrics_json["num_prompts"] = args.num_prompts
-
-    # Traffic
-    metrics_json["request_rate"] = args.request_rate
-    metrics_json = {**metrics_json, **benchmark_result}
-    if args.run_eval:
-      metrics_json = {**metrics_json, **eval_json}
-
-    final_json = {}
-    final_json["metrics"] = metrics_json
-    final_json["dimensions"] = dimensions_json
-
-    # Save to file
-    base_model_id = model_id.split("/")[-1]
-    file_name = (
-        f"JetStream-{args.request_rate}qps-{base_model_id}-{current_dt}.json"
-    )
-    with open(file_name, "w", encoding="utf-8") as outfile:
-      json.dump(final_json, outfile)
-
-  if args.save_request_outputs:
-    file_path = args.request_outputs_file_path
-    with open(file_path, "w", encoding="utf-8") as output_file:
-      json.dump(
-          output,
-          output_file,
-          indent=4,
-      )
-
-
-if __name__ == "__main__":
+def parse_args() -> argparse.Namespace:
   parser = argparse.ArgumentParser(
       description="Benchmark the online serving throughput."
   )
@@ -909,7 +780,6 @@ def main(args: argparse.Namespace):
       default=150,
       help="The maximum number of mock requests to send for benchmark testing.",
   )
-
   parser.add_argument(
       "--max-output-length",
       type=int,
@@ -926,7 +796,6 @@ def main(args: argparse.Namespace):
           "the output length of the golden dataset would be passed."
       ),
   )
-
   parser.add_argument("--seed", type=int, default=0)
   parser.add_argument(
       "--disable-tqdm",
@@ -977,7 +846,138 @@ def main(args: argparse.Namespace):
       choices=["human", "gpt", "both"],
       help="What entity should be the one starting the conversations.",
   )
+  return parser.parse_args()
 
-  parsed_args = parser.parse_args()
+
+def main(args: argparse.Namespace):
+  print(args)
+  random.seed(args.seed)
+  np.random.seed(args.seed)
+
+  model_id = args.model
+  tokenizer_id = args.tokenizer
+  use_hf_tokenizer = args.use_hf_tokenizer
+
+  prefill_quota = AsyncCounter(init_value=3)
+  active_req_quota = AsyncCounter(init_value=450)
+
+  api_url = f"{args.server}:{args.port}"
+
+  tokenizer = get_tokenizer(model_id, tokenizer_id, use_hf_tokenizer)
+  if tokenizer == "test" or args.dataset == "test":
+    input_requests = mock_requests(
+        args.total_mock_requests
+    )  # e.g. [("AB", 2, "AB", 3)]
+  else:
+    dataset = []
+    if args.dataset == "openorca":
+      dataset = load_openorca_dataset_pkl(args.dataset_path)
+    elif args.dataset == "sharegpt":
+      dataset = load_sharegpt_dataset(
+          args.dataset_path,
+          args.conversation_starter,
+      )
+
+    # A given args.max_output_length value is the max generation step,
+    # when the args.max_output_length is default to None, the sample's golden
+    # output length will be used to decide the generation step.
+    input_requests = sample_requests(
+        dataset=dataset,
+        tokenizer=tokenizer,
+        num_requests=args.num_prompts,
+        max_output_length=args.max_output_length,
+    )
+
+  warmup_requests = None
+  if args.warmup_mode == "full":
+    warmup_requests = input_requests
+  elif args.warmup_mode == "sampled":
+    warmup_requests = list(sample_warmup_requests(input_requests)) * 2
+
+  if warmup_requests:
+    print(f"Warmup (mode: {args.warmup_mode}) is starting.")
+    _, _ = asyncio.run(
+        benchmark(
+            api_url=api_url,
+            tokenizer=tokenizer,
+            input_requests=warmup_requests,
+            request_rate=args.request_rate,
+            disable_tqdm=args.disable_tqdm,
+            prefill_quota=prefill_quota,
+            active_req_quota=active_req_quota,
+            is_warmup=True,
+        )
+    )
+    print(f"Warmup (mode: {args.warmup_mode}) has completed.")
+
+  # TODO: Replace this with warmup complete signal once supported.
+  # Wait for server completely warmup before running the benchmark.
+  time.sleep(5)
+
+  benchmark_result, request_outputs = asyncio.run(
+      benchmark(
+          api_url=api_url,
+          tokenizer=tokenizer,
+          input_requests=input_requests,
+          request_rate=args.request_rate,
+          disable_tqdm=args.disable_tqdm,
+          prefill_quota=prefill_quota,
+          active_req_quota=active_req_quota,
+      )
+  )
+
+  # Process output
+  output = [output.to_dict() for output in request_outputs]
+  if args.run_eval:
+    eval_json = eval_accuracy(output)
+
+  # Save config and results to json
+  if args.save_result:
+    # dimensions values are strings
+    dimensions_json = {}
+    # metrics values are numerical
+    metrics_json = {}
+
+    # Setup
+    current_dt = datetime.now().strftime("%Y%m%d-%H%M%S")
+    dimensions_json["date"] = current_dt
+    dimensions_json["model_id"] = model_id
+    dimensions_json["tokenizer_id"] = tokenizer_id
+    if args.additional_metadata_metrics_to_save is not None:
+      dimensions_json = {
+          **dimensions_json,
+          **json.loads(args.additional_metadata_metrics_to_save),
+      }
+    metrics_json["num_prompts"] = args.num_prompts
+
+    # Traffic
+    metrics_json["request_rate"] = args.request_rate
+    metrics_json = {**metrics_json, **benchmark_result}
+    if args.run_eval:
+      metrics_json = {**metrics_json, **eval_json}
+
+    final_json = {}
+    final_json["metrics"] = metrics_json
+    final_json["dimensions"] = dimensions_json
+
+    # Save to file
+    base_model_id = model_id.split("/")[-1]
+    file_name = (
+        f"JetStream-{args.request_rate}qps-{base_model_id}-{current_dt}.json"
+    )
+    with open(file_name, "w", encoding="utf-8") as outfile:
+      json.dump(final_json, outfile)
+
+  if args.save_request_outputs:
+    file_path = args.request_outputs_file_path
+    with open(file_path, "w", encoding="utf-8") as output_file:
+      json.dump(
+          output,
+          output_file,
+          indent=4,
+      )
+
+
+if __name__ == "__main__":
   gc.disable()
-  main(parsed_args)
+  main(parse_args())