feat: add inputs.json file for dataset traceability (#310)

Author: ajcasagrande

aiperf/common/config/config_defaults.py

@@ -114,6 +114,11 @@ class TurnDelayDefaults:
 class OutputDefaults:
     ARTIFACT_DIRECTORY = Path("./artifacts")
     PROFILE_EXPORT_FILE = Path("profile_export.json")
+    LOG_FOLDER = Path("logs")
+    LOG_FILE = Path("aiperf.log")
+    INPUTS_JSON_FILE = Path("inputs.json")
+    PROFILE_EXPORT_AIPERF_CSV_FILE = Path("profile_export_aiperf.csv")
+    PROFILE_EXPORT_AIPERF_JSON_FILE = Path("profile_export_aiperf.json")
 
 
 @dataclass(frozen=True)

aiperf/common/constants.py

@@ -8,6 +8,7 @@
 MILLIS_PER_SECOND = 1000
 BYTES_PER_MIB = 1024 * 1024
 
+
 GRACEFUL_SHUTDOWN_TIMEOUT = 5.0
 """Default timeout for shutting down services in seconds."""
 

aiperf/common/logging.py

@@ -11,6 +11,7 @@
 
 from aiperf.common.aiperf_logger import _DEBUG, _TRACE, AIPerfLogger
 from aiperf.common.config import ServiceConfig, ServiceDefaults, UserConfig
+from aiperf.common.config.config_defaults import OutputDefaults
 from aiperf.common.enums import ServiceType
 from aiperf.common.enums.ui_enums import AIPerfUIType
 from aiperf.common.factories import ServiceFactory
@@ -112,7 +113,7 @@ def setup_child_process_logging(
 
     if user_config and user_config.output.artifact_directory:
         file_handler = create_file_handler(
-            user_config.output.artifact_directory / "logs", level
+            user_config.output.artifact_directory / OutputDefaults.LOG_FOLDER, level
         )
         root_logger.addHandler(file_handler)
 
@@ -138,9 +139,9 @@ def setup_rich_logging(user_config: UserConfig, service_config: ServiceConfig) -
 
     # Enable file logging for services
     # TODO: Use config to determine if file logging is enabled and the folder path.
-    log_folder = user_config.output.artifact_directory / "logs"
+    log_folder = user_config.output.artifact_directory / OutputDefaults.LOG_FOLDER
     log_folder.mkdir(parents=True, exist_ok=True)
-    file_handler = logging.FileHandler(log_folder / "aiperf.log")
+    file_handler = logging.FileHandler(log_folder / OutputDefaults.LOG_FILE)
     file_handler.setLevel(level)
     file_handler.formatter = logging.Formatter(
         "%(asctime)s - %(name)s - %(levelname)s - %(message)s",
@@ -158,7 +159,7 @@ def create_file_handler(
     """Configure a file handler for logging."""
 
     log_folder.mkdir(parents=True, exist_ok=True)
-    log_file_path = log_folder / "aiperf.log"
+    log_file_path = log_folder / OutputDefaults.LOG_FILE
 
     file_handler = logging.FileHandler(log_file_path, encoding="utf-8")
     file_handler.setLevel(level)

aiperf/common/models/__init__.py

@@ -21,7 +21,9 @@
     Audio,
     Conversation,
     Image,
+    InputsFile,
     Media,
+    SessionPayloads,
     Text,
     Turn,
 )
@@ -86,6 +88,7 @@
     "IOCounters",
     "Image",
     "InferenceServerResponse",
+    "InputsFile",
     "Media",
     "MetricResult",
     "ParsedResponse",
@@ -102,6 +105,7 @@
     "SSEField",
     "SSEMessage",
     "ServiceRunInfo",
+    "SessionPayloads",
     "StatsProtocol",
     "Text",
     "TextResponse",

aiperf/common/models/dataset_models.py

@@ -1,7 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import ClassVar
+from typing import Any, ClassVar
 
 from pydantic import Field
 
@@ -81,3 +81,25 @@ class Conversation(AIPerfBaseModel):
         default=[], description="List of turns in the conversation."
     )
     session_id: str = Field(default="", description="Session ID of the conversation.")
+
+
+class SessionPayloads(AIPerfBaseModel):
+    """A single session, with its session ID and a list of formatted payloads (one per turn)."""
+
+    session_id: str | None = Field(
+        default=None, description="Session ID of the conversation."
+    )
+    payloads: list[dict[str, Any]] = Field(
+        default=[],
+        description="List of formatted payloads in the session (one per turn). These have been formatted for the model and endpoint.",
+    )
+
+
+class InputsFile(AIPerfBaseModel):
+    """A list of all dataset sessions. Each session contains a list of formatted payloads (one per turn).
+    This is similar to the format used by GenAI-Perf for the inputs.json file.
+    """
+
+    data: list[SessionPayloads] = Field(
+        default=[], description="List of all dataset sessions."
+    )

aiperf/controller/system_controller.py

@@ -10,6 +10,7 @@
 
 from aiperf.common.base_service import BaseService
 from aiperf.common.config import ServiceConfig, UserConfig
+from aiperf.common.config.config_defaults import OutputDefaults
 from aiperf.common.config.dev_config import print_developer_mode_warning
 from aiperf.common.constants import (
     AIPERF_DEV_MODE,
@@ -513,7 +514,11 @@ async def _print_post_benchmark_info_and_metrics(self) -> None:
 
     def _print_log_file_info(self, console: Console) -> None:
         """Print the log file info."""
-        log_file = self.user_config.output.artifact_directory / "logs" / "aiperf.log"
+        log_file = (
+            self.user_config.output.artifact_directory
+            / OutputDefaults.LOG_FOLDER
+            / OutputDefaults.LOG_FILE
+        )
         console.print(
             f"[bold green]Log File:[/bold green] [cyan]{log_file.resolve()}[/cyan]"
         )

aiperf/dataset/dataset_manager.py

@@ -4,9 +4,13 @@
 import random
 import time
 
+import aiofiles
+
+from aiperf.clients.model_endpoint_info import ModelEndpointInfo
 from aiperf.common.aiperf_logger import AIPerfLogger
 from aiperf.common.base_component_service import BaseComponentService
 from aiperf.common.config import ServiceConfig, UserConfig
+from aiperf.common.config.config_defaults import OutputDefaults
 from aiperf.common.decorators import implements_protocol
 from aiperf.common.enums import (
     CommAddress,
@@ -16,7 +20,11 @@
     ServiceType,
 )
 from aiperf.common.enums.dataset_enums import CustomDatasetType
-from aiperf.common.factories import ComposerFactory, ServiceFactory
+from aiperf.common.factories import (
+    ComposerFactory,
+    RequestConverterFactory,
+    ServiceFactory,
+)
 from aiperf.common.hooks import on_command, on_request
 from aiperf.common.messages import (
     ConversationRequestMessage,
@@ -29,8 +37,9 @@
     ProfileConfigureCommand,
 )
 from aiperf.common.mixins import ReplyClientMixin
-from aiperf.common.models import Conversation
-from aiperf.common.protocols import ServiceProtocol
+from aiperf.common.models import Conversation, InputsFile
+from aiperf.common.models.dataset_models import SessionPayloads
+from aiperf.common.protocols import RequestConverterProtocol, ServiceProtocol
 from aiperf.common.tokenizer import Tokenizer
 from aiperf.dataset.loader import ShareGPTLoader
 
@@ -87,6 +96,7 @@ async def _profile_configure_command(
         self.info(lambda: f"Configuring dataset for {self.service_id}")
         begin = time.perf_counter()
         await self._configure_dataset()
+        await self._generate_inputs_json_file()
         duration = time.perf_counter() - begin
         self.info(lambda: f"Dataset configured in {duration:.2f} seconds")
 
@@ -104,6 +114,57 @@ async def _configure_tokenizer(self) -> None:
             revision=self.user_config.tokenizer.revision,
         )
 
+    async def _generate_input_payloads(
+        self,
+        model_endpoint: ModelEndpointInfo,
+        request_converter: RequestConverterProtocol,
+    ) -> InputsFile:
+        """Generate input payloads from the dataset for use in the inputs.json file."""
+        inputs = InputsFile()
+        for conversation in self.dataset.values():
+            payloads = await asyncio.gather(
+                *[
+                    request_converter.format_payload(model_endpoint, turn)
+                    for turn in conversation.turns
+                ]
+            )
+            inputs.data.append(
+                SessionPayloads(session_id=conversation.session_id, payloads=payloads)
+            )
+        return inputs
+
+    async def _generate_inputs_json_file(self) -> None:
+        """Generate inputs.json file in the artifact directory."""
+        file_path = (
+            self.user_config.output.artifact_directory / OutputDefaults.INPUTS_JSON_FILE
+        )
+        self.info(f"Generating inputs.json file at {file_path.resolve()}")
+
+        try:
+            start_time = time.perf_counter()
+            file_path.parent.mkdir(parents=True, exist_ok=True)
+
+            model_endpoint = ModelEndpointInfo.from_user_config(self.user_config)
+            request_converter = RequestConverterFactory.create_instance(
+                model_endpoint.endpoint.type,
+            )
+
+            inputs = await self._generate_input_payloads(
+                model_endpoint, request_converter
+            )
+
+            async with aiofiles.open(file_path, "w") as f:
+                await f.write(inputs.model_dump_json(indent=2, exclude_unset=True))
+
+            duration = time.perf_counter() - start_time
+            self.info(f"inputs.json file generated in {duration:.2f} seconds")
+
+        except Exception as e:
+            # Log as warning, but continue to run the benchmark
+            self.warning(
+                f"Error generating inputs.json file at {file_path.resolve()}: {e}"
+            )
+
     async def _configure_dataset(self) -> None:
         if self.user_config is None:
             raise self._service_error("User config is required for dataset manager")

aiperf/exporters/csv_exporter.py

@@ -9,6 +9,7 @@
 
 import aiofiles
 
+from aiperf.common.config.config_defaults import OutputDefaults
 from aiperf.common.decorators import implements_protocol
 from aiperf.common.enums import DataExporterType
 from aiperf.common.enums.metric_enums import MetricFlags
@@ -40,7 +41,9 @@ def __init__(self, exporter_config: ExporterConfig, **kwargs) -> None:
         self._results = exporter_config.results
         self._output_directory = exporter_config.user_config.output.artifact_directory
         self._metric_registry = MetricRegistry
-        self._file_path = self._output_directory / "profile_export_aiperf.csv"
+        self._file_path = (
+            self._output_directory / OutputDefaults.PROFILE_EXPORT_AIPERF_CSV_FILE
+        )
         self._percentile_keys = _percentile_keys_from(STAT_KEYS)
 
     def get_export_info(self) -> FileExportInfo:

aiperf/exporters/json_exporter.py

@@ -7,6 +7,7 @@
 from pydantic import BaseModel
 
 from aiperf.common.config import UserConfig
+from aiperf.common.config.config_defaults import OutputDefaults
 from aiperf.common.constants import NANOS_PER_SECOND
 from aiperf.common.decorators import implements_protocol
 from aiperf.common.enums import DataExporterType, MetricFlags
@@ -45,7 +46,9 @@ def __init__(self, exporter_config: ExporterConfig, **kwargs) -> None:
         self._output_directory = exporter_config.user_config.output.artifact_directory
         self._input_config = exporter_config.user_config
         self._metric_registry = MetricRegistry
-        self._file_path = self._output_directory / "profile_export_aiperf.json"
+        self._file_path = (
+            self._output_directory / OutputDefaults.PROFILE_EXPORT_AIPERF_JSON_FILE
+        )
 
     def get_export_info(self) -> FileExportInfo:
         return FileExportInfo(

tests/data_exporters/test_csv_exporter.py

@@ -8,6 +8,7 @@
 import pytest
 
 from aiperf.common.config import EndpointConfig, ServiceConfig, UserConfig
+from aiperf.common.config.config_defaults import OutputDefaults
 from aiperf.common.enums import EndpointType
 from aiperf.common.models import MetricResult
 from aiperf.exporters.csv_exporter import CsvExporter
@@ -146,7 +147,7 @@ async def test_csv_exporter_writes_two_sections_and_values(
         exporter = CsvExporter(cfg)
         await exporter.export()
 
-        expected = outdir / "profile_export_aiperf.csv"
+        expected = outdir / OutputDefaults.PROFILE_EXPORT_AIPERF_CSV_FILE
         assert expected.exists()
 
         text = _read(expected)
@@ -195,7 +196,7 @@ async def test_csv_exporter_empty_records_creates_empty_file(
         exporter = CsvExporter(cfg)
         await exporter.export()
 
-        expected = outdir / "profile_export_aiperf.csv"
+        expected = outdir / OutputDefaults.PROFILE_EXPORT_AIPERF_CSV_FILE
         assert expected.exists()
         content = _read(expected)
         assert content.strip() == ""
@@ -233,7 +234,7 @@ async def test_csv_exporter_deterministic_sort_order(
         exporter = CsvExporter(cfg)
         await exporter.export()
 
-        text = _read(outdir / "profile_export_aiperf.csv")
+        text = _read(outdir / OutputDefaults.PROFILE_EXPORT_AIPERF_CSV_FILE)
 
         # Request section should list aaa_latency then zzz_latency in order
         # Pull only the request rows region (before the blank line separator).
@@ -288,7 +289,7 @@ async def test_csv_exporter_unit_aware_number_formatting(
         exporter = CsvExporter(cfg)
         await exporter.export()
 
-        text = _read(outdir / "profile_export_aiperf.csv")
+        text = _read(outdir / OutputDefaults.PROFILE_EXPORT_AIPERF_CSV_FILE)
 
         # counts: integer
         assert re.search(r"Input Sequence Length \(tokens\),\s*4096\b", text)