From c5ae3272b9889d9b01eb15b7b5515992235b4076 Mon Sep 17 00:00:00 2001
From: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
Date: Wed, 25 Jun 2025 23:13:56 +0800
Subject: [PATCH] feat: Make benchmark_serving part of the library (#5428)

Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
---
 tensorrt_llm/serve/scripts/__init__.py        |  0
 .../serve/scripts/benchmark_dataset.py        |  3 ++-
 .../serve/scripts/benchmark_serving.py        | 19 ++++++++++---------
 3 files changed, 12 insertions(+), 10 deletions(-)
 create mode 100644 tensorrt_llm/serve/scripts/__init__.py

diff --git a/tensorrt_llm/serve/scripts/__init__.py b/tensorrt_llm/serve/scripts/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tensorrt_llm/serve/scripts/benchmark_dataset.py b/tensorrt_llm/serve/scripts/benchmark_dataset.py
index ab7f45e893..59ce5d045b 100644
--- a/tensorrt_llm/serve/scripts/benchmark_dataset.py
+++ b/tensorrt_llm/serve/scripts/benchmark_dataset.py
@@ -26,10 +26,11 @@ from typing import Any, Callable, Optional, Union
 
 import numpy as np
 import pandas as pd
-from benchmark_utils import download_and_cache_file
 from datasets import load_dataset
 from transformers import PreTrainedTokenizerBase
 
+from .benchmark_utils import download_and_cache_file
+
 logger = logging.getLogger(__name__)
 
 # -----------------------------------------------------------------------------
diff --git a/tensorrt_llm/serve/scripts/benchmark_serving.py b/tensorrt_llm/serve/scripts/benchmark_serving.py
index fb41378977..13df3eeab1 100644
--- a/tensorrt_llm/serve/scripts/benchmark_serving.py
+++ b/tensorrt_llm/serve/scripts/benchmark_serving.py
@@ -31,18 +31,19 @@ from datetime import datetime
 from typing import Any, Optional
 
 import numpy as np
-from backend_request_func import (ASYNC_REQUEST_FUNCS,
-                                  OPENAI_COMPATIBLE_BACKENDS, RequestFuncInput,
-                                  RequestFuncOutput, get_tokenizer)
-from benchmark_dataset import (AIMODataset, BurstGPTDataset,
-                               ConversationDataset, HuggingFaceDataset,
-                               InstructCoderDataset, RandomDataset,
-                               SampleRequest, ShareGPTDataset, SonnetDataset,
-                               VisionArenaDataset)
-from benchmark_utils import convert_to_pytorch_benchmark_format, write_to_json
 from tqdm.asyncio import tqdm
 from transformers import PreTrainedTokenizerBase
 
+from .backend_request_func import (ASYNC_REQUEST_FUNCS,
+                                   OPENAI_COMPATIBLE_BACKENDS, RequestFuncInput,
+                                   RequestFuncOutput, get_tokenizer)
+from .benchmark_dataset import (AIMODataset, BurstGPTDataset,
+                                ConversationDataset, HuggingFaceDataset,
+                                InstructCoderDataset, RandomDataset,
+                                SampleRequest, ShareGPTDataset, SonnetDataset,
+                                VisionArenaDataset)
+from .benchmark_utils import convert_to_pytorch_benchmark_format, write_to_json
+
 MILLISECONDS_TO_SECONDS_CONVERSION = 1000