diff --git a/genai-perf/genai_perf/inputs/inputs.py b/genai-perf/genai_perf/inputs/inputs.py index 72c6b15e..051e6b91 100644 --- a/genai-perf/genai_perf/inputs/inputs.py +++ b/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}") diff --git a/genai-perf/genai_perf/inputs/inputs_config.py b/genai-perf/genai_perf/inputs/inputs_config.py index 5b1429f5..a355304a 100644 --- a/genai-perf/genai_perf/inputs/inputs_config.py +++ b/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 ######################################## diff --git a/genai-perf/genai_perf/parser.py b/genai-perf/genai_perf/parser.py index 58a63573..0349ee49 100644 --- a/genai-perf/genai_perf/parser.py +++ b/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:. 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, diff --git a/genai-perf/genai_perf/subcommand/common.py b/genai-perf/genai_perf/subcommand/common.py index 21a2e8db..8cba075c 100644 --- a/genai-perf/genai_perf/subcommand/common.py +++ b/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, )