fix(mmap): make get_conversation thread-safe, prefault pages, drop executor hop

docs/environment-variables.md

@@ -100,6 +100,7 @@ Dataset loading and configuration. Controls timeouts and behavior for dataset lo
 | `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, manifest.json, and (when produced) inputs.json. No automatic eviction is implemented yet -- delete the directory to reclaim disk. |
+| `AIPERF_DATASET_MMAP_PREFAULT` | `True` | — | If True, each memory-mapped dataset client walks every page of the data file at open time (after madvise(MADV_WILLNEED)) to force-populate the OS page cache. Reads afterwards are served warm, so no request pays a major page fault mid-benchmark -- which would otherwise land in the measured latency. Workers share the kernel page cache, so the disk read happens once regardless of worker count. Costs a one-time startup pass proportional to dataset size; set to False to trade predictable tail latency for faster startup on very large datasets. |
 | `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; cache-bust runs always use CONVERSATION regardless. |
 | `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 |

src/aiperf/common/environment.py

@@ -332,6 +332,16 @@ class _DatasetSettings(BaseSettings):
         "dataset.dat, index.dat, manifest.json, and (when produced) inputs.json. "
         "No automatic eviction is implemented yet -- delete the directory to reclaim disk.",
     )
+    MMAP_PREFAULT: bool = Field(
+        default=True,
+        description="If True, each memory-mapped dataset client walks every page of the "
+        "data file at open time (after madvise(MADV_WILLNEED)) to force-populate the OS "
+        "page cache. Reads afterwards are served warm, so no request pays a major page "
+        "fault mid-benchmark -- which would otherwise land in the measured latency. "
+        "Workers share the kernel page cache, so the disk read happens once regardless of "
+        "worker count. Costs a one-time startup pass proportional to dataset size; set to "
+        "False to trade predictable tail latency for faster startup on very large datasets.",
+    )
     PREFORMAT_PAYLOADS: bool = Field(
         default=False,
         description="If True, pre-encode single-turn / self-contained synthetic "

src/aiperf/dataset/memory_map_utils.py

@@ -403,12 +403,10 @@ def __init__(self, client_metadata: MemoryMapClientMetadata, **kwargs) -> None:
         self._data_path: Path = client_metadata.data_file_path
         self._index_path: Path = client_metadata.index_file_path
         self._client: MemoryMapDatasetClient | None = None
-        self._loop: asyncio.AbstractEventLoop | None = None
 
     @on_init
     async def _setup(self) -> None:
         """Open memory-mapped files (read-only)."""
-        self._loop = asyncio.get_running_loop()
         self.debug(
             lambda: f"Opening memory-mapped files: data={self._data_path}, index={self._index_path}"
         )
@@ -421,7 +419,14 @@ async def _setup(self) -> None:
     async def get_conversation(self, conversation_id: str) -> Conversation:
         """Retrieve conversation from memory-mapped file.
 
-        Runs in executor since mmap reads can block on page faults.
+        No executor hop: pages are prefaulted at open (when MMAP_PREFAULT is
+        enabled) so reads don't trigger major faults. Remains async-callable
+        for API parity with other dataset clients.
+
+        When MMAP_PREFAULT is disabled, reads may still take major faults, and
+        those now land on the event loop thread rather than an executor thread.
+        That is deliberate: CPython's mmap slice holds the GIL while copying,
+        so an executor relocates the stall instead of removing it.
 
         Args:
             conversation_id: Session ID of the conversation
@@ -432,11 +437,9 @@ async def get_conversation(self, conversation_id: str) -> Conversation:
         Raises:
             KeyError: If conversation_id not found
         """
-        if self._client is None or self._loop is None:
+        if self._client is None:
             raise RuntimeError("Client store not initialized. Call initialize() first.")
-        return await self._loop.run_in_executor(
-            None, self._client.get_conversation, conversation_id
-        )
+        return self._client.get_conversation(conversation_id)
 
     async def get_payload_bytes(
         self, conversation_id: str, turn_index: int
@@ -451,11 +454,9 @@ async def get_payload_bytes(
             Pre-encoded JSON bytes, or None when the dataset is not in
             PAYLOAD_BYTES format or the turn has no payload.
         """
-        if self._client is None or self._loop is None:
+        if self._client is None:
             raise RuntimeError("Client store not initialized. Call initialize() first.")
-        return await self._loop.run_in_executor(
-            None, self._client.get_payload_bytes, conversation_id, turn_index
-        )
+        return self._client.get_payload_bytes(conversation_id, turn_index)
 
     async def get_payload_turn(
         self, conversation_id: str, turn_index: int
@@ -470,11 +471,9 @@ async def get_payload_turn(
             ``PayloadTurnData`` or None when the dataset is not in
             PAYLOAD_BYTES format or the turn has no payload.
         """
-        if self._client is None or self._loop is None:
+        if self._client is None:
             raise RuntimeError("Client store not initialized. Call initialize() first.")
-        return await self._loop.run_in_executor(
-            None, self._client.get_payload_turn, conversation_id, turn_index
-        )
+        return self._client.get_payload_turn(conversation_id, turn_index)
 
     @on_stop
     async def _cleanup(self) -> None:
@@ -673,6 +672,15 @@ def __init__(self, data_file_path: Path | str, index_file_path: Path | str) -> N
 
             index_data = self.index_mmap.read()
             self.index = MemoryMapDatasetIndex.model_validate_json(index_data)
+            # Flatten to a plain dict[str, tuple[int, int]] to avoid per-lookup
+            # Pydantic attribute overhead on the hot path. Wire format unchanged;
+            # self.index stays for callers that iterate index.conversation_ids.
+            self._offsets: dict[str, tuple[int, int]] = {
+                cid: (o.offset, o.size) for cid, o in self.index.offsets.items()
+            }
+
+            if Environment.DATASET.MMAP_PREFAULT:
+                self._prefault_data_mmap()
 
         except OSError as e:
             self._cleanup_resources()
@@ -735,6 +743,19 @@ def _cleanup_resources(self) -> None:
                 with suppress(Exception):
                     obj.close()
 
+    def _prefault_data_mmap(self) -> None:
+        """Advise WILLNEED and touch every page to populate the page cache.
+
+        Without this, the first read of each page takes a major fault whose
+        latency lands inside a measured request. Workers share the kernel page
+        cache, so N workers cost one disk read rather than N serialized ones.
+        """
+        # madvise is unavailable on some platforms; the page walk still works.
+        with suppress(OSError, AttributeError):
+            self.data_mmap.madvise(mmap.MADV_WILLNEED)
+        for offset in range(0, len(self.data_mmap), mmap.PAGESIZE):
+            _ = self.data_mmap[offset]
+
     def _deserialize_conversation(self, data: bytes) -> Conversation:
         """Deserialize a single conversation from bytes.
 
@@ -774,17 +795,19 @@ def get_conversation(self, conversation_id: str) -> Conversation:
                 "format. Use get_payload_bytes() instead."
             )
 
-        if conversation_id not in self.index.offsets:
+        if conversation_id not in self._offsets:
             raise KeyError(f"Conversation '{conversation_id}' not found in dataset")
 
-        offset_info = self.index.offsets[conversation_id]
+        offset, size = self._offsets[conversation_id]
 
         try:
-            self.data_mmap.seek(offset_info.offset)
-            conv_bytes = self.data_mmap.read(offset_info.size)
+            # Slice rather than seek()+read(): the mmap file position is shared
+            # state, so two concurrent readers can interleave a seek between
+            # another's seek and read and serve foreign bytes.
+            conv_bytes = self.data_mmap[offset : offset + size]
 
             _logger.debug(
-                lambda: f"Loading conversation '{conversation_id}': offset={offset_info.offset}, size={offset_info.size} bytes"
+                lambda: f"Loading conversation '{conversation_id}': offset={offset}, size={size} bytes"
             )
 
             return self._deserialize_conversation(conv_bytes)

tests/unit/dataset/test_mmap_conversation_concurrency.py

@@ -0,0 +1,172 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Concurrency and resource-lifetime tests for MemoryMapDatasetClient.
+
+``get_conversation`` used to read through the mmap's shared file position
+(``seek()`` then ``read()``). Any two readers sharing one client could
+interleave those calls and serve each other's bytes, surfacing either as a
+spurious ``MemoryMapSerializationError`` on a random conversation or, when the
+foreign slice happened to parse, as silently wrong data.
+
+The reads are now position-free slices, so no reader depends on the shared
+position.
+"""
+
+import mmap
+from collections.abc import AsyncIterator
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import asynccontextmanager
+from pathlib import Path
+from typing import Any
+
+import pytest
+
+from aiperf.common.enums import MemoryMapFormat
+from aiperf.common.models import Conversation, Text, Turn
+from aiperf.dataset.memory_map_utils import (
+    MemoryMapDatasetBackingStore,
+    MemoryMapDatasetClient,
+)
+
+# Wildly varying sizes so a cross-read lands mid-record and is caught either as
+# a decode error or as a mismatched session_id.
+_CONVERSATION_COUNT = 24
+_THREADS = 8
+_LOOKUPS_PER_THREAD = 150
+
+
+def _make_conversation(index: int) -> Conversation:
+    """Conversation whose size grows with ``index`` and whose text encodes it."""
+    text = f"conv-{index}:" + ("x" * (64 * (index + 1)))
+    return Conversation(
+        session_id=f"conv-{index}",
+        turns=[Turn(role="user", texts=[Text(contents=[text])])],
+    )
+
+
+@asynccontextmanager
+async def _open_client(
+    tmp_path: Path, benchmark_id: str, count: int
+) -> AsyncIterator[MemoryMapDatasetClient]:
+    """Write ``count`` conversations and open a client over them.
+
+    The backing store is stopped on exit, so its ``@on_stop`` cleanup unlinks
+    the mmap files instead of leaking them for the run's lifetime. The client
+    is closed first: the store's cleanup deletes the files it maps.
+    """
+    store = MemoryMapDatasetBackingStore(
+        benchmark_id=benchmark_id, format=MemoryMapFormat.CONVERSATION
+    )
+    await store.initialize()
+    for i in range(count):
+        await store.add_conversation(f"conv-{i}", _make_conversation(i))
+    await store.finalize()
+
+    metadata = store.get_client_metadata()
+    client: MemoryMapDatasetClient | None = None
+    try:
+        # Constructed inside the guard: if opening the mmap raises, the store
+        # still gets stopped and its files still get unlinked.
+        client = MemoryMapDatasetClient(
+            metadata.data_file_path,
+            metadata.index_file_path,
+        )
+        yield client
+    finally:
+        if client is not None:
+            client.close()
+        await store.stop()
+
+
+@pytest.mark.asyncio
+async def test_get_conversation_concurrent_readers_returns_own_bytes(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """Concurrent get_conversation() calls must not read each other's bytes."""
+    monkeypatch.setenv("AIPERF_DATASET_MMAP_BASE_PATH", str(tmp_path))
+    async with _open_client(
+        tmp_path, "test_conv_concurrency", _CONVERSATION_COUNT
+    ) as client:
+
+        def _read_many(thread_index: int) -> None:
+            for n in range(_LOOKUPS_PER_THREAD):
+                # Stagger the starting id per thread so the interleaving covers
+                # many (offset, size) pairs rather than repeatedly hitting one
+                # record.
+                i = (thread_index + n) % _CONVERSATION_COUNT
+                conv = client.get_conversation(f"conv-{i}")
+                assert conv.session_id == f"conv-{i}"
+
+        with ThreadPoolExecutor(max_workers=_THREADS) as pool:
+            # list() resolves every future so exceptions propagate.
+            list(pool.map(_read_many, range(_THREADS)))
+
+
+class _InterleavingMmap:
+    """mmap wrapper that moves the shared position between seek() and read().
+
+    This is the race, made deterministic. Merely leaving a stale position
+    behind proves nothing: a ``seek()`` + ``read()`` reader overwrites it on
+    the way in and still returns the right record. The bug needs a competing
+    seek to land *after* this reader seeks and *before* it reads.
+
+    A position-free reader goes through ``__getitem__`` and is unaffected.
+    """
+
+    def __init__(self, real: mmap.mmap, foreign_offset: int) -> None:
+        self._real = real
+        self._foreign_offset = foreign_offset
+
+    def seek(self, pos: int) -> None:
+        self._real.seek(pos)
+
+    def read(self, size: int) -> bytes:
+        # The competing reader lands here, between our seek and our read.
+        self._real.seek(self._foreign_offset)
+        return self._real.read(size)
+
+    def __getitem__(self, item: slice | int) -> bytes | int:
+        return self._real[item]
+
+    def __len__(self) -> int:
+        return len(self._real)
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._real, name)
+
+
+@pytest.mark.asyncio
+async def test_get_conversation_ignores_competing_seek_between_seek_and_read(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """An interleaved foreign seek must not change what get_conversation reads.
+
+    Fails against a ``seek()`` + ``read()`` implementation, which serves the
+    foreign record's bytes.
+    """
+    monkeypatch.setenv("AIPERF_DATASET_MMAP_BASE_PATH", str(tmp_path))
+    async with _open_client(tmp_path, "test_conv_position", 3) as client:
+        foreign = client.index.offsets["conv-2"].offset
+        client.data_mmap = _InterleavingMmap(client.data_mmap, foreign)
+
+        conversation = client.get_conversation("conv-0")
+
+        assert conversation.session_id == "conv-0"
+        assert conversation.turns[0].texts[0].contents[0].startswith("conv-0:")
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("prefault", [True, False])  # fmt: skip
+async def test_get_conversation_matches_written_data_under_prefault_setting(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, prefault: bool
+) -> None:
+    """Reads are identical whether or not pages were prefaulted at open."""
+    monkeypatch.setenv("AIPERF_DATASET_MMAP_BASE_PATH", str(tmp_path))
+    monkeypatch.setattr(
+        "aiperf.common.environment.Environment.DATASET.MMAP_PREFAULT", prefault
+    )
+    async with _open_client(tmp_path, f"test_conv_prefault_{prefault}", 5) as client:
+        for i in range(5):
+            conv = client.get_conversation(f"conv-{i}")
+            assert conv.session_id == f"conv-{i}"
+            assert conv.turns[0].texts[0].contents[0].startswith(f"conv-{i}:")