Fix comments

triton-inference-server · Aug 14, 2024 · 8251542 · 8251542
1 parent 407aba5
commit 8251542
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 7 deletions.
diff --git a/genai-perf/genai_perf/parser.py b/genai-perf/genai_perf/parser.py
@@ -778,7 +778,11 @@ def profile_handler(args, extra_args):
             server_metrics_url=DEFAULT_TRITON_METRICS_URL
         )
 
-    Profiler.run(telemetry_data_collector, args=args, extra_args=extra_args)
+    Profiler.run(
+        args=args,
+        extra_args=extra_args,
+        telemetry_data_collector=telemetry_data_collector,
+    )
 
 
 ### Parser Initialization ###

diff --git a/genai-perf/genai_perf/wrapper.py b/genai-perf/genai_perf/wrapper.py
@@ -33,7 +33,7 @@
 from genai_perf.constants import DEFAULT_GRPC_URL, DEFAULT_INPUT_DATA_JSON
 from genai_perf.llm_inputs.llm_inputs import OutputFormat
 from genai_perf.telemetry_data.triton_telemetry_data_collector import (
-    TritonTelemetryDataCollector,
+    TelemetryDataCollector,
 )
 
 logger = logging.getLogger(__name__)
@@ -146,9 +146,9 @@ def build_cmd(args: Namespace, extra_args: Optional[List[str]] = None) -> List[s
 
     @staticmethod
     def run(
-        telemetry_data_collector: TritonTelemetryDataCollector,
         args: Namespace,
         extra_args: Optional[List[str]],
+        telemetry_data_collector: Optional[TelemetryDataCollector] = None,
     ) -> None:
         try:
             if telemetry_data_collector is not None:

diff --git a/genai-perf/tests/test_wrapper.py b/genai-perf/tests/test_wrapper.py
@@ -147,16 +147,16 @@ def test_service_openai(self, monkeypatch, arg):
         assert cmd_string.count(" -i http") == 1
 
     @patch("genai_perf.wrapper.subprocess.run")
-    @patch("genai_perf.wrapper.TritonTelemetryDataCollector")
+    @patch("genai_perf.wrapper.TelemetryDataCollector")
     def test_stdout_verbose(self, mock_telemetry_collector, mock_subprocess_run):
         args = MagicMock()
         args.model = "test_model"
         args.verbose = True
         telemetry_data_collector = mock_telemetry_collector.return_value
         Profiler.run(
-            telemetry_data_collector=telemetry_data_collector,
             args=args,
             extra_args=None,
+            telemetry_data_collector=telemetry_data_collector,
         )
 
         # Check that standard output was not redirected.
@@ -167,16 +167,16 @@ def test_stdout_verbose(self, mock_telemetry_collector, mock_subprocess_run):
             ), "With the verbose flag, stdout should not be redirected."
 
     @patch("genai_perf.wrapper.subprocess.run")
-    @patch("genai_perf.wrapper.TritonTelemetryDataCollector")
+    @patch("genai_perf.wrapper.TelemetryDataCollector")
     def test_stdout_not_verbose(self, mock_telemetry_collector, mock_subprocess_run):
         args = MagicMock()
         args.model = "test_model"
         args.verbose = False
         telemetry_data_collector = mock_telemetry_collector.return_value
         Profiler.run(
-            telemetry_data_collector=telemetry_data_collector,
             args=args,
             extra_args=None,
+            telemetry_data_collector=telemetry_data_collector,
         )
 
         # Check that standard output was redirected.