diff --git a/src/aiperf/dataset/memory_map_utils.py b/src/aiperf/dataset/memory_map_utils.py index 649d3010ac..de4612fb63 100644 --- a/src/aiperf/dataset/memory_map_utils.py +++ b/src/aiperf/dataset/memory_map_utils.py @@ -780,8 +780,11 @@ def get_conversation(self, conversation_id: str) -> Conversation: offset_info = self.index.offsets[conversation_id] try: - self.data_mmap.seek(offset_info.offset) - conv_bytes = self.data_mmap.read(offset_info.size) + conv_bytes = bytes( + self.data_mmap[ + offset_info.offset : offset_info.offset + offset_info.size + ] + ) _logger.debug( lambda: f"Loading conversation '{conversation_id}': offset={offset_info.offset}, size={offset_info.size} bytes" diff --git a/tests/unit/dataset/test_payload_mmap.py b/tests/unit/dataset/test_payload_mmap.py index 2520db2674..a38bc39719 100644 --- a/tests/unit/dataset/test_payload_mmap.py +++ b/tests/unit/dataset/test_payload_mmap.py @@ -1,6 +1,9 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from concurrent.futures import ThreadPoolExecutor +from threading import Barrier + import orjson import pytest @@ -22,6 +25,27 @@ def _make_raw_conversation( return Conversation(session_id=session_id, turns=turns) +class _RacingCursorMmap: + """Expose deterministic shared-cursor races while preserving slice reads.""" + + def __init__(self, data: bytes) -> None: + self._data = data + self._position = 0 + self._seek_barrier = Barrier(2) + + def __getitem__(self, key: slice) -> bytes: + return self._data[key] + + def seek(self, offset: int) -> None: + self._position = offset + self._seek_barrier.wait(timeout=5) + + def read(self, size: int) -> bytes: + start = self._position + self._position += size + return self._data[start : start + size] + + @pytest.mark.asyncio async def test_payload_mmap_round_trip(tmp_path, monkeypatch): """Test writing and reading payload bytes through the mmap backing store.""" @@ -94,6 +118,55 @@ async def test_conversation_format_returns_none_for_payload_bytes( await store.stop() +@pytest.mark.asyncio +async def test_conversation_mmap_parallel_reads_do_not_share_cursor( + tmp_path, monkeypatch +): + """Parallel conversation reads must use independent mmap byte ranges.""" + monkeypatch.setenv("AIPERF_DATASET_MMAP_BASE_PATH", str(tmp_path)) + + store = MemoryMapDatasetBackingStore(benchmark_id="test_parallel_reads") + await store.initialize() + + conversations = { + "conv-1": _make_raw_conversation( + "conv-1", [{"messages": [{"role": "user", "content": "a" * 1024}]}] + ), + "conv-2": _make_raw_conversation( + "conv-2", [{"messages": [{"role": "user", "content": "b" * 2048}]}] + ), + } + await store.add_conversations(conversations) + await store.finalize() + + metadata = store.get_client_metadata() + client = MemoryMapDatasetClient( + metadata.data_file_path, + metadata.index_file_path, + ) + original_mmap = client.data_mmap + client.data_mmap = _RacingCursorMmap(bytes(original_mmap[:])) + + try: + with ThreadPoolExecutor(max_workers=2) as executor: + futures = { + conversation_id: executor.submit( + client.get_conversation, conversation_id + ) + for conversation_id in conversations + } + loaded = { + conversation_id: future.result(timeout=5) + for conversation_id, future in futures.items() + } + + assert loaded == conversations + finally: + client.data_mmap = original_mmap + client.close() + await store.stop() + + @pytest.mark.asyncio async def test_client_store_get_payload_bytes(tmp_path, monkeypatch): """Test MemoryMapDatasetClientStore.get_payload_bytes async wrapper."""