Allow direct payloads to be passed into GenAI-Perf

genai-perf/genai_perf/inputs/inputs.py

@@ -18,6 +18,7 @@
 from pathlib import Path
 from typing import Any, Dict, List, Tuple
 
+import genai_perf.logging as logging
 from genai_perf.exceptions import GenAIPerfException
 from genai_perf.inputs.converters.output_format_converter_factory import (
     OutputFormatConverterFactory,
@@ -31,6 +32,9 @@
 )
 from genai_perf.inputs.inputs_config import InputsConfig
 from genai_perf.inputs.retrievers.input_retriever_factory import InputRetrieverFactory
+from genai_perf.utils import load_json_str
+
+logger = logging.getLogger(__name__)
 
 
 class Inputs:
@@ -51,13 +55,24 @@ def create_inputs(self) -> None:
         (in a JSON dictionary) to a file.
         """
         self._check_for_valid_args()
-        input_retriever = InputRetrieverFactory.create(self.config)
-        generic_dataset = input_retriever.retrieve_data()
 
-        json_in_pa_format = self._convert_generic_dataset_to_output_format(
-            generic_dataset,
-        )
-        self._write_json_to_file(json_in_pa_format)
+        if self.config.payload_mode:
+            logger.debug(
+                f"Using direct payload mode with payload: {self.config.payload_path}"
+            )
+            if self.config.payload_path is None:
+                raise GenAIPerfException(
+                    "Payload path not detected when using direct payload mode."
+                )
+            self._use_direct_payload(self.config.payload_path)
+        else:
+            input_retriever = InputRetrieverFactory.create(self.config)
+            generic_dataset = input_retriever.retrieve_data()
+
+            json_in_pa_format = self._convert_generic_dataset_to_output_format(
+                generic_dataset,
+            )
+            self._write_json_to_file(json_in_pa_format)
 
     def _check_for_valid_args(self) -> None:
         self._check_for_tokenzier_if_input_type_is_synthetic()
@@ -102,3 +117,29 @@ def _check_for_valid_length(self) -> None:
             raise GenAIPerfException(
                 f"starting_index: {self.config.length} must be larger than {MINIMUM_LENGTH}."
             )
+
+    def _use_direct_payload(self, payload_path: str) -> None:
+        """
+        Reads the direct payload file and writes it to the output directory.
+        """
+        try:
+            with open(payload_path, "r") as file:
+                raw_lines = file.readlines()
+
+            formatted_payload: Dict[str, List] = {"data": []}
+
+            for line in raw_lines:
+                if line.strip():
+                    entry = load_json_str(line.strip())
+                    formatted_payload["data"].append({"payload": [entry]})
+
+            filename = self.config.output_dir / DEFAULT_INPUT_DATA_JSON
+            with open(filename, "w") as output_file:
+                output_file.write(json.dumps(formatted_payload, indent=2))
+
+        except FileNotFoundError:
+            raise GenAIPerfException(f"Payload file '{payload_path}' not found.")
+        except json.JSONDecodeError as e:
+            raise GenAIPerfException(f"Invalid JSONL format in payload file: {e}")
+        except Exception as e:
+            raise GenAIPerfException(f"Error processing payload file: {e}")

genai-perf/genai_perf/inputs/inputs_config.py

@@ -129,6 +129,12 @@ class InputsConfig:
     # The directory where all arifacts are saved
     output_dir: Path = Path("")
 
+    # If true, the input data is a direct payload
+    payload_mode: bool = False
+
+    # Path to the direct payload file
+    payload_path: Optional[str] = None
+
     ########################################
     # Synthetic Prompt Generation Parameters
     ########################################

genai-perf/genai_perf/parser.py

@@ -496,7 +496,7 @@ def _convert_str_to_enum_entry(args, option, enum):
 
 
 def file_or_directory(value: str) -> Path:
-    if value.startswith("synthetic:"):
+    if value.startswith("synthetic:") or value.startswith("payload:"):
         return Path(value)
     else:
         path = Path(value)
@@ -761,6 +761,9 @@ def _add_input_args(parser):
         "You can repeat this flag for multiple headers.",
     )
 
+    ## TODO: Update this to have an option for --input-file
+    ## payload:<path_to_payload_file>. Not doing this now to
+    ## avoid making it harder to rebase the delay feature branch.
     input_group.add_argument(
         "--input-file",
         type=file_or_directory,

genai-perf/genai_perf/subcommand/common.py

@@ -154,6 +154,14 @@ def create_config_options(args: Namespace) -> InputsConfig:
     except ValueError as e:
         raise GenAIPerfException(e)
 
+    input_file = str(args.input_file)
+    if input_file.startswith("payload:"):
+        payload_mode = True
+        payload_path = input_file.split("payload:")[1]
+    else:
+        payload_mode = False
+        payload_path = None
+
     return InputsConfig(
         input_type=args.prompt_source,
         output_format=args.output_format,
@@ -185,6 +193,8 @@ def create_config_options(args: Namespace) -> InputsConfig:
         output_dir=args.artifact_dir,
         num_prefix_prompts=args.num_prefix_prompts,
         prefix_prompt_length=args.prefix_prompt_length,
+        payload_mode=payload_mode,
+        payload_path=payload_path,
     )