feat(dataset): verbatim payload replay fast path + content-addressed mmap dataset cache

docs/environment-variables.md

@@ -86,6 +86,9 @@ Dataset loading and configuration. Controls timeouts and behavior for dataset lo
 |----------------------|---------|-------------|-------------|
 | `AIPERF_DATASET_CONFIGURATION_TIMEOUT` | `300.0` | ≥ 1.0, ≤ 100000.0 | Timeout in seconds for dataset configuration operations |
 | `AIPERF_DATASET_MMAP_BASE_PATH` | `None` | — | Base path for memory-mapped dataset files. If None, uses system temp directory. Set to a shared filesystem path for Kubernetes mounted volumes. Example: AIPERF_DATASET_MMAP_BASE_PATH=/mnt/shared-pvc creates files at /mnt/shared-pvc/aiperf_mmap_{benchmark_id}/ |
+| `AIPERF_DATASET_MMAP_CACHE_ENABLED` | `True` | — | If True, AIPerf reuses memory-mapped dataset files across runs whose input bytes, tokenizer identity, and prompt/input settings are byte-identical. Set to False to force every run to re-tokenize and re-write its mmap files. Cache misses still produce byte-identical mmap files to a non-cached run. |
+| `AIPERF_DATASET_MMAP_CACHE_DIR` | `None` | — | Directory holding the content-addressed mmap cache. If None, defaults to ~/.cache/aiperf/dataset_mmap. Each cache entry lives under a `dir/key` subpath and contains dataset.dat, index.dat, and manifest.json. No automatic eviction is implemented yet -- delete the directory to reclaim disk. |
+| `AIPERF_DATASET_PREFORMAT_PAYLOADS` | `False` | — | If True, pre-encode single-turn / self-contained synthetic conversations to the PAYLOAD_BYTES mmap fast path at dataset-build time so workers stream the bytes verbatim and skip per-request encoding. This is a throughput optimization that DROPS input-tokenization metrics (input_sequence_length, image counts) because the structured prompt is discarded. Default False keeps the structured-turns (CONVERSATION) path so those metrics are computed. Datasets that natively ship raw payloads (raw_payload / inputs_json / mooncake-with-payload) always use PAYLOAD_BYTES regardless of this flag. |
 | `AIPERF_DATASET_PUBLIC_DATASET_TIMEOUT` | `300.0` | ≥ 1.0, ≤ 100000.0 | Timeout in seconds for public dataset loading operations |
 | `AIPERF_DATASET_MEDIA_DOWNLOAD_TIMEOUT` | `60.0` | ≥ 1.0, ≤ 100000.0 | Timeout in seconds per media URL download when inline encoding is required |
 | `AIPERF_DATASET_MEDIA_DOWNLOAD_MAX_CONCURRENCY` | `10` | ≥ 1, ≤ 100 | Maximum number of concurrent media URL downloads |

pyproject.toml

@@ -26,6 +26,7 @@ dependencies = [
   "cyclopts>=4,<5",
   "fastapi>=0.115,<1",
   "ffmpeg-python~=0.2.0",
+  "filelock>=3.13", # Cross-process populate lock for the mmap dataset cache
   # datasets pulls pyarrow (its Arrow core), which has no Windows-on-ARM wheel
   # (apache/arrow#47195). Excluded on win_arm64 so the base install works there;
   # HF-backed --public-dataset is gated on that platform, but synthetic /

src/aiperf/common/enums/__init__.py

@@ -35,6 +35,7 @@
     IPVersion,
     LifecycleState,
     MediaType,
+    MemoryMapFormat,
     MessageType,
     ModelSelectionStrategy,
     OptimizationDirection,
@@ -122,6 +123,7 @@
     "ImageSourceSamplingStrategy",
     "LifecycleState",
     "MediaType",
+    "MemoryMapFormat",
     "MessageType",
     "MetricDictValueTypeT",
     "MetricFlags",

src/aiperf/common/enums/enums.py

@@ -334,6 +334,16 @@ class MediaType(CaseInsensitiveStrEnum):
     VIDEO = "video"
 
 
+class MemoryMapFormat(CaseInsensitiveStrEnum):
+    """Storage format for memory-mapped dataset files."""
+
+    CONVERSATION = "conversation"
+    """Each entry is a JSON-serialized Conversation object."""
+
+    PAYLOAD_BYTES = "payload_bytes"
+    """Each entry is pre-encoded payload bytes for verbatim API replay."""
+
+
 class MessageType(CaseInsensitiveStrEnum):
     """The various types of messages that can be sent between services.
 

src/aiperf/common/environment.py

@@ -240,6 +240,32 @@ class _DatasetSettings(BaseSettings):
         "Example: AIPERF_DATASET_MMAP_BASE_PATH=/mnt/shared-pvc "
         "creates files at /mnt/shared-pvc/aiperf_mmap_{benchmark_id}/",
     )
+    MMAP_CACHE_ENABLED: bool = Field(
+        default=True,
+        description="If True, AIPerf reuses memory-mapped dataset files across runs whose "
+        "input bytes, tokenizer identity, and prompt/input settings are byte-identical. "
+        "Set to False to force every run to re-tokenize and re-write its mmap files. "
+        "Cache misses still produce byte-identical mmap files to a non-cached run.",
+    )
+    MMAP_CACHE_DIR: Path | None = Field(
+        default=None,
+        description="Directory holding the content-addressed mmap cache. If None, defaults to "
+        "~/.cache/aiperf/dataset_mmap. Each cache entry lives under a `dir/key` subpath and contains "
+        "dataset.dat, index.dat, and manifest.json. "
+        "No automatic eviction is implemented yet -- delete the directory to reclaim disk.",
+    )
+    PREFORMAT_PAYLOADS: bool = Field(
+        default=False,
+        description="If True, pre-encode single-turn / self-contained synthetic "
+        "conversations to the PAYLOAD_BYTES mmap fast path at dataset-build time "
+        "so workers stream the bytes verbatim and skip per-request encoding. This "
+        "is a throughput optimization that DROPS input-tokenization metrics "
+        "(input_sequence_length, image counts) because the structured prompt is "
+        "discarded. Default False keeps the structured-turns (CONVERSATION) path "
+        "so those metrics are computed. Datasets that natively ship raw payloads "
+        "(raw_payload / inputs_json / mooncake-with-payload) always use "
+        "PAYLOAD_BYTES regardless of this flag.",
+    )
     PUBLIC_DATASET_TIMEOUT: float = Field(
         ge=1.0,
         le=100000.0,

src/aiperf/common/models/dataset_models.py

@@ -11,6 +11,7 @@
     ConversationBranchMode,
     ConversationContextMode,
     MediaType,
+    MemoryMapFormat,
 )
 from aiperf.common.models.base_models import AIPerfBaseModel
 from aiperf.common.models.branch import ConversationBranchInfo
@@ -43,6 +44,11 @@ class MemoryMapClientMetadata(DatasetClientMetadata):
 
     client_type: DatasetClientStoreType = DatasetClientStoreType.MEMORY_MAP
 
+    format: MemoryMapFormat = Field(
+        default=MemoryMapFormat.CONVERSATION,
+        description="Storage format of the memory-mapped dataset files "
+        "(serialized Conversations vs pre-encoded per-turn payload bytes).",
+    )
     data_file_path: Path = Field(
         ...,
         description="Path to the memory-mapped data file containing serialized conversations.",