Allow extra-inputs to support keys without values (#262)

triton-inference-server · Jan 30, 2025 · 2157c97 · 2157c97
1 parent b93d73c
commit 2157c97
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 19 deletions.
diff --git a/genai-perf/genai_perf/subcommand/common.py b/genai-perf/genai_perf/subcommand/common.py
@@ -83,24 +83,22 @@ def calculate_metrics(args: Namespace, tokenizer: Tokenizer) -> ProfileDataParse
 
 
 def get_extra_inputs_as_dict(args: Namespace) -> Dict[str, Any]:
-    request_inputs = {}
+    request_inputs: Dict[str, Any] = {}
     if args.extra_inputs:
         for input_str in args.extra_inputs:
+            semicolon_count = input_str.count(":")
             if input_str.startswith("{") and input_str.endswith("}"):
                 request_inputs.update(load_json_str(input_str))
-            else:
-                semicolon_count = input_str.count(":")
-                if semicolon_count != 1:
-                    raise ValueError(
-                        f"Invalid input format for --extra-inputs: {input_str}\n"
-                        "Expected input format: 'input_name:value'"
-                    )
+            elif semicolon_count == 0:  # extra input as a flag
+                request_inputs[input_str] = None
+            elif semicolon_count == 1:
                 input_name, value = input_str.split(":", 1)
 
                 if not input_name or not value:
                     raise ValueError(
-                        f"Input name or value is empty in --extra-inputs: {input_str}\n"
-                        "Expected input format: 'input_name:value'"
+                        f"Input name or value is empty in --extra-inputs: "
+                        f"{input_str}\nExpected input format: 'input_name' or "
+                        "'input_name:value'"
                     )
 
                 is_bool = value.lower() in ["true", "false"]
@@ -118,9 +116,15 @@ def get_extra_inputs_as_dict(args: Namespace) -> Dict[str, Any]:
 
                 if input_name in request_inputs:
                     raise ValueError(
-                        f"Input name already exists in request_inputs dictionary: {input_name}"
+                        f"Input name already exists in request_inputs "
+                        f"dictionary: {input_name}"
                     )
                 request_inputs[input_name] = value
+            else:
+                raise ValueError(
+                    f"Invalid input format for --extra-inputs: {input_str}\n"
+                    "Expected input format: 'input_name' or 'input_name:value'"
+                )
 
     return request_inputs
 

diff --git a/genai-perf/tests/test_cli.py b/genai-perf/tests/test_cli.py
@@ -843,27 +843,26 @@ def test_inferred_output_format(self, monkeypatch, args, expected_format):
         [
             (
                 ["--extra-inputs", "hi:"],
-                "Input name or value is empty in --extra-inputs: hi:\nExpected input format: 'input_name:value'",
+                "Input name or value is empty in --extra-inputs: hi:\n"
+                "Expected input format: 'input_name' or 'input_name:value'",
             ),
             (
                 ["--extra-inputs", ":a"],
-                "Input name or value is empty in --extra-inputs: :a\nExpected input format: 'input_name:value'",
+                "Input name or value is empty in --extra-inputs: :a\n"
+                "Expected input format: 'input_name' or 'input_name:value'",
             ),
             (
                 ["--extra-inputs", ":a:"],
-                "Invalid input format for --extra-inputs: :a:\nExpected input format: 'input_name:value'",
-            ),
-            (
-                ["--extra-inputs", "unknown"],
-                "Invalid input format for --extra-inputs: unknown\nExpected input format: 'input_name:value'",
+                "Invalid input format for --extra-inputs: :a:\n"
+                "Expected input format: 'input_name' or 'input_name:value'",
             ),
             (
                 ["--extra-inputs", "test_key:5", "--extra-inputs", "test_key:6"],
                 "Input name already exists in request_inputs dictionary: test_key",
             ),
         ],
     )
-    def test_repeated_extra_arg_warning(self, monkeypatch, args, expected_error):
+    def test_get_extra_inputs_as_dict_warning(self, monkeypatch, args, expected_error):
         combined_args = ["genai-perf", "profile", "-m", "test_model"] + args
         monkeypatch.setattr("sys.argv", combined_args)
 
@@ -991,6 +990,7 @@ def test_compare_not_provided(self, monkeypatch, capsys):
         "extra_inputs_list, expected_dict",
         [
             (["test_key:test_value"], {"test_key": "test_value"}),
+            (["test_key"], {"test_key": None}),
             (
                 ["test_key:1", "another_test_key:2"],
                 {"test_key": 1, "another_test_key": 2},