From ce2dec952bb19ab44ef1657c598d54711c6821ac Mon Sep 17 00:00:00 2001 From: Anthony Casagrande Date: Mon, 3 Aug 2026 17:39:50 -0700 Subject: [PATCH 1/5] fix(dataset/mmap): make get_conversation thread-safe, prefault pages, drop executor hop ## The bug `get_conversation` read its record through the mmap's *shared file position*: self.data_mmap.seek(offset_info.offset) conv_bytes = self.data_mmap.read(offset_info.size) `MemoryMapDatasetClientStore` dispatched that onto the default executor, so N threads shared one mmap object and could interleave: thread A: seek(offset_a) thread B: seek(offset_b) thread A: read(size_a) -> bytes from offset_b, length size_a Thread A silently gets another conversation's bytes. It surfaces as a spurious MemoryMapSerializationError: Failed to decode conversation data: Invalid JSON: trailing characters at line 1 column N on a *random* conversation, or -- when the foreign slice happens to parse -- as the wrong conversation served with no error at all. Reported against a 6.4GB dataset (SemiAnalysisAI/aiperf#34): four runs sharing one mmap cache entry each failed on a *different* conversation while `cmp` showed the on-disk files byte-identical, and each named conversation decoded fine when read standalone. The reported offsets are consistent with reading the right *length* from the wrong *place* -- one failure reported "EOF while parsing at column 1801628" where the index size for that record is exactly 1,801,628. The two sibling readers (`get_payload_bytes`, `get_payload_turn`) already used position-free slicing; only `get_conversation` did not. Now it does. ## Three further changes in the same path 1. `madvise(MADV_WILLNEED)` plus a page walk at open, gated on the new `Environment.DATASET.MMAP_PREFAULT` (default True). No read then takes a major fault mid-benchmark, where the fault latency would land inside a measured request. Workers share the kernel page cache, so N workers cost one disk read rather than N serialized ones. 2. Drop `loop.run_in_executor` from all three client-store readers. The hop existed to absorb blocking major faults, which (1) eliminates. 3. Flatten `index.offsets` into `dict[str, tuple[int, int]]` at load time, dropping per-lookup Pydantic attribute access. Wire format unchanged; `self.index` stays for callers iterating `index.conversation_ids`. ## Measurements Executor hop, 64KB read, page-cache warm: direct on loop 1.08us run_in_executor 19.41us 18.0x, +18.33us per lookup Event-loop tick lateness under continuous reads, warmed pool: 64KB records p50 p99 max idle 0.081 0.230 0.338 ms executor 0.059 0.410 0.639 ms inline 0.008 0.045 0.128 ms The executor was worse on every percentile at typical record sizes. At 8MB it bought ~2x median at the cost of 2-4x p99 and max -- the wrong trade for a tool whose output is latency percentiles. Prefault cost at open, by dataset size: 16 MB 0.9 ms cold 0.4 ms warm 256 MB 12.6 ms cold 3.4 ms warm 1024 MB 40.0 ms cold 13.4 ms warm seek()+read() vs slice is within noise (0.03us at 64KB, 6.3us at 8MB) -- both are the same memcpy under the GIL. Neither mmap read path releases the GIL (CPython Modules/mmapmodule.c:318 and :999 both call PyBytes_FromStringAndSize with no Py_BEGIN_ALLOW_THREADS), which is why the executor could not protect the loop from a fault in the first place. End-to-end throughput is unchanged: ~43us saved per request against ~42ms request latency is below run-to-run noise. This is a correctness fix plus CPU and allocation hygiene, not a throughput win. ## Tests `tests/unit/dataset/test_mmap_conversation_concurrency.py`: - `_InterleavingMmap` injects a competing seek between the reader's seek and its read, making the race deterministic. Mutation-verified: reintroducing seek()+read() fails it with the reported production symptom (asked for conv-0, received conv-2's bytes, "Invalid JSON: EOF while parsing a string"). Note that merely leaving a stale position behind does NOT catch this -- the reader's own seek() overwrites it on the way in. - 8 threads x 150 lookups over 24 records of varying size. Weaker: it does not reproduce at the default switch interval and can only ever be suggestive. - Reads agree with MMAP_PREFAULT on and off. Signed-off-by: Anthony Casagrande Co-authored-by: Xianjie Qiao <5410381+qiaoxj07@users.noreply.github.com> --- docs/environment-variables.md | 1 + src/aiperf/common/environment.py | 10 ++ src/aiperf/dataset/memory_map_utils.py | 58 ++++--- .../test_mmap_conversation_concurrency.py | 161 ++++++++++++++++++ 4 files changed, 210 insertions(+), 20 deletions(-) create mode 100644 tests/unit/dataset/test_mmap_conversation_concurrency.py diff --git a/docs/environment-variables.md b/docs/environment-variables.md index babdfddb27..2fc940a13c 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -99,6 +99,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 | diff --git a/src/aiperf/common/environment.py b/src/aiperf/common/environment.py index a736ea15ac..64de1bd7d6 100644 --- a/src/aiperf/common/environment.py +++ b/src/aiperf/common/environment.py @@ -323,6 +323,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 " diff --git a/src/aiperf/dataset/memory_map_utils.py b/src/aiperf/dataset/memory_map_utils.py index 649d3010ac..c9dee7e8cf 100644 --- a/src/aiperf/dataset/memory_map_utils.py +++ b/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,9 @@ 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. Args: conversation_id: Session ID of the conversation @@ -432,11 +432,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 +449,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 +466,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 +667,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 +738,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 +790,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) diff --git a/tests/unit/dataset/test_mmap_conversation_concurrency.py b/tests/unit/dataset/test_mmap_conversation_concurrency.py new file mode 100644 index 0000000000..5ccfc1fd7f --- /dev/null +++ b/tests/unit/dataset/test_mmap_conversation_concurrency.py @@ -0,0 +1,161 @@ +# 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 concurrent.futures import ThreadPoolExecutor + +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])])], + ) + + +async def _build_client( + tmp_path, benchmark_id: str, count: int +) -> MemoryMapDatasetClient: + """Write ``count`` conversations and open a client over them.""" + 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() + return MemoryMapDatasetClient( + metadata.data_file_path, + metadata.index_file_path, + ) + + +@pytest.mark.asyncio +async def test_get_conversation_concurrent_readers_returns_own_bytes( + tmp_path, monkeypatch +) -> None: + """Concurrent get_conversation() calls must not read each other's bytes.""" + monkeypatch.setenv("AIPERF_DATASET_MMAP_BASE_PATH", str(tmp_path)) + client = await _build_client(tmp_path, "test_conv_concurrency", _CONVERSATION_COUNT) + + 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}" + + try: + with ThreadPoolExecutor(max_workers=_THREADS) as pool: + # list() resolves every future so exceptions propagate. + list(pool.map(_read_many, range(_THREADS))) + finally: + client.close() + + +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): + return self._real[item] + + def __len__(self) -> int: + return len(self._real) + + def __getattr__(self, name): + return getattr(self._real, name) + + +@pytest.mark.asyncio +async def test_get_conversation_ignores_competing_seek_between_seek_and_read( + tmp_path, 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)) + client = await _build_client(tmp_path, "test_conv_position", 3) + + try: + 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:") + finally: + client.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("prefault", [True, False]) # fmt: skip +async def test_get_conversation_matches_written_data_under_prefault_setting( + tmp_path, 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 + ) + client = await _build_client(tmp_path, f"test_conv_prefault_{prefault}", 5) + + try: + 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}:") + finally: + client.close() From 4678074a089c9f58a838cc1e282fb850712518e3 Mon Sep 17 00:00:00 2001 From: Anthony Casagrande Date: Mon, 3 Aug 2026 17:50:28 -0700 Subject: [PATCH 2/5] test(dataset/mmap): annotate test helper signatures Adds the missing type hints on _InterleavingMmap.__getitem__ / __getattr__ and on the tmp_path / monkeypatch fixture parameters, per the repo standard that every function parameter and return value is annotated. No behaviour change. Signed-off-by: Anthony Casagrande --- .../dataset/test_mmap_conversation_concurrency.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/tests/unit/dataset/test_mmap_conversation_concurrency.py b/tests/unit/dataset/test_mmap_conversation_concurrency.py index 5ccfc1fd7f..3f4f0e3d28 100644 --- a/tests/unit/dataset/test_mmap_conversation_concurrency.py +++ b/tests/unit/dataset/test_mmap_conversation_concurrency.py @@ -14,6 +14,8 @@ import mmap from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import Any import pytest @@ -41,7 +43,7 @@ def _make_conversation(index: int) -> Conversation: async def _build_client( - tmp_path, benchmark_id: str, count: int + tmp_path: Path, benchmark_id: str, count: int ) -> MemoryMapDatasetClient: """Write ``count`` conversations and open a client over them.""" store = MemoryMapDatasetBackingStore( @@ -61,7 +63,7 @@ async def _build_client( @pytest.mark.asyncio async def test_get_conversation_concurrent_readers_returns_own_bytes( - tmp_path, monkeypatch + 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)) @@ -106,19 +108,19 @@ def read(self, size: int) -> bytes: self._real.seek(self._foreign_offset) return self._real.read(size) - def __getitem__(self, item): + def __getitem__(self, item: slice | int) -> bytes | int: return self._real[item] def __len__(self) -> int: return len(self._real) - def __getattr__(self, name): + 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, monkeypatch + tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: """An interleaved foreign seek must not change what get_conversation reads. @@ -143,7 +145,7 @@ async def test_get_conversation_ignores_competing_seek_between_seek_and_read( @pytest.mark.asyncio @pytest.mark.parametrize("prefault", [True, False]) # fmt: skip async def test_get_conversation_matches_written_data_under_prefault_setting( - tmp_path, monkeypatch, prefault: bool + 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)) From f8e24ea490666aa7cca288300782ce076a55a0cb Mon Sep 17 00:00:00 2001 From: Anthony Casagrande Date: Tue, 4 Aug 2026 22:38:34 +0000 Subject: [PATCH 3/5] test(dataset/mmap): stop backing store so mmap files are unlinked _build_client dropped the MemoryMapDatasetBackingStore without stopping it, so its @on_stop cleanup never ran and the dataset.dat/index.dat files it created were left on disk after each test. Replace it with an _open_client asynccontextmanager that closes the client and then awaits store.stop(), which also removes the duplicated try/finally in each test. Co-Authored-By: Claude Opus 5 --- .../test_mmap_conversation_concurrency.py | 57 ++++++++++--------- 1 file changed, 31 insertions(+), 26 deletions(-) diff --git a/tests/unit/dataset/test_mmap_conversation_concurrency.py b/tests/unit/dataset/test_mmap_conversation_concurrency.py index 3f4f0e3d28..6a3ab1e9de 100644 --- a/tests/unit/dataset/test_mmap_conversation_concurrency.py +++ b/tests/unit/dataset/test_mmap_conversation_concurrency.py @@ -13,7 +13,9 @@ """ import mmap +from collections.abc import AsyncIterator from concurrent.futures import ThreadPoolExecutor +from contextlib import asynccontextmanager from pathlib import Path from typing import Any @@ -42,10 +44,16 @@ def _make_conversation(index: int) -> Conversation: ) -async def _build_client( +@asynccontextmanager +async def _open_client( tmp_path: Path, benchmark_id: str, count: int -) -> MemoryMapDatasetClient: - """Write ``count`` conversations and open a client over them.""" +) -> 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 ) @@ -55,10 +63,15 @@ async def _build_client( await store.finalize() metadata = store.get_client_metadata() - return MemoryMapDatasetClient( + client = MemoryMapDatasetClient( metadata.data_file_path, metadata.index_file_path, ) + try: + yield client + finally: + client.close() + await store.stop() @pytest.mark.asyncio @@ -67,22 +80,22 @@ async def test_get_conversation_concurrent_readers_returns_own_bytes( ) -> None: """Concurrent get_conversation() calls must not read each other's bytes.""" monkeypatch.setenv("AIPERF_DATASET_MMAP_BASE_PATH", str(tmp_path)) - client = await _build_client(tmp_path, "test_conv_concurrency", _CONVERSATION_COUNT) + 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}" - 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}" - - try: with ThreadPoolExecutor(max_workers=_THREADS) as pool: # list() resolves every future so exceptions propagate. list(pool.map(_read_many, range(_THREADS))) - finally: - client.close() class _InterleavingMmap: @@ -128,9 +141,7 @@ async def test_get_conversation_ignores_competing_seek_between_seek_and_read( foreign record's bytes. """ monkeypatch.setenv("AIPERF_DATASET_MMAP_BASE_PATH", str(tmp_path)) - client = await _build_client(tmp_path, "test_conv_position", 3) - - try: + 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) @@ -138,8 +149,6 @@ async def test_get_conversation_ignores_competing_seek_between_seek_and_read( assert conversation.session_id == "conv-0" assert conversation.turns[0].texts[0].contents[0].startswith("conv-0:") - finally: - client.close() @pytest.mark.asyncio @@ -152,12 +161,8 @@ async def test_get_conversation_matches_written_data_under_prefault_setting( monkeypatch.setattr( "aiperf.common.environment.Environment.DATASET.MMAP_PREFAULT", prefault ) - client = await _build_client(tmp_path, f"test_conv_prefault_{prefault}", 5) - - try: + 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}:") - finally: - client.close() From 562fdfa3411a85f88c9c5bc65d15d3ee13ee12b4 Mon Sep 17 00:00:00 2001 From: Anthony Casagrande Date: Tue, 4 Aug 2026 22:39:05 +0000 Subject: [PATCH 4/5] docs(dataset/mmap): note major faults still possible when MMAP_PREFAULT is off The get_conversation docstring stated reads don't trigger major faults, but that only holds with MMAP_PREFAULT enabled. Record the disabled case and why those faults are left on the loop thread rather than moved to an executor. Co-Authored-By: Claude Opus 5 --- src/aiperf/dataset/memory_map_utils.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/aiperf/dataset/memory_map_utils.py b/src/aiperf/dataset/memory_map_utils.py index c9dee7e8cf..53b667980b 100644 --- a/src/aiperf/dataset/memory_map_utils.py +++ b/src/aiperf/dataset/memory_map_utils.py @@ -423,6 +423,11 @@ async def get_conversation(self, conversation_id: str) -> Conversation: 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 From e4661823a46be7a0660b4b8b776884d3b53d4780 Mon Sep 17 00:00:00 2001 From: Anthony Casagrande Date: Tue, 4 Aug 2026 22:46:47 +0000 Subject: [PATCH 5/5] test(dataset/mmap): stop store even if client construction fails MemoryMapDatasetClient(...) ran before the try block, so a failure opening the mmap skipped the finally and left store.stop() uncalled -- the same leak the guard was added to prevent. Construct inside the guard and close only a client that was actually built. Co-Authored-By: Claude Opus 5 --- .../dataset/test_mmap_conversation_concurrency.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/tests/unit/dataset/test_mmap_conversation_concurrency.py b/tests/unit/dataset/test_mmap_conversation_concurrency.py index 6a3ab1e9de..7339b962ae 100644 --- a/tests/unit/dataset/test_mmap_conversation_concurrency.py +++ b/tests/unit/dataset/test_mmap_conversation_concurrency.py @@ -63,14 +63,18 @@ async def _open_client( await store.finalize() metadata = store.get_client_metadata() - client = MemoryMapDatasetClient( - metadata.data_file_path, - metadata.index_file_path, - ) + 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: - client.close() + if client is not None: + client.close() await store.stop()