diff --git a/src/aiperf/dataset/memory_map_utils.py b/src/aiperf/dataset/memory_map_utils.py index 649d3010ac..eaddcbc99f 100644 --- a/src/aiperf/dataset/memory_map_utils.py +++ b/src/aiperf/dataset/memory_map_utils.py @@ -780,8 +780,13 @@ 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) + # Slice, not seek()+read(): the mmap file position is shared and + # this runs on a thread pool, so the two calls can interleave. + 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_mmap_conversation_concurrency.py b/tests/unit/dataset/test_mmap_conversation_concurrency.py new file mode 100644 index 0000000000..1fca8f35d2 --- /dev/null +++ b/tests/unit/dataset/test_mmap_conversation_concurrency.py @@ -0,0 +1,91 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Concurrency regression tests for MemoryMapDatasetClient.get_conversation(). + +The store dispatches it via ``run_in_executor``, so several threads read one +shared ``mmap``. Any reader that depends on the mmap's file position can then +serve another conversation's bytes. +""" + +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, +) + +# Distinct, wildly varying sizes so a cross-read lands mid-record and is +# detected 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])])], + ) + + +@pytest.mark.asyncio +async def test_get_conversation_is_thread_safe(tmp_path, monkeypatch): + """Concurrent get_conversation() calls must not read each other's bytes.""" + monkeypatch.setenv("AIPERF_DATASET_MMAP_BASE_PATH", str(tmp_path)) + + store = MemoryMapDatasetBackingStore( + benchmark_id="test_conv_concurrency", format=MemoryMapFormat.CONVERSATION + ) + await store.initialize() + for i in range(_CONVERSATION_COUNT): + await store.add_conversation(f"conv-{i}", _make_conversation(i)) + await store.finalize() + + metadata = store.get_client_metadata() + client = MemoryMapDatasetClient( + metadata.data_file_path, + metadata.index_file_path, + ) + + 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 one hot 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() forces every future to be resolved so exceptions propagate. + list(pool.map(_read_many, range(_THREADS))) + + +@pytest.mark.asyncio +async def test_get_conversation_ignores_shared_mmap_position(tmp_path, monkeypatch): + """A stale mmap file position must not affect what get_conversation reads.""" + monkeypatch.setenv("AIPERF_DATASET_MMAP_BASE_PATH", str(tmp_path)) + + store = MemoryMapDatasetBackingStore( + benchmark_id="test_conv_position", format=MemoryMapFormat.CONVERSATION + ) + await store.initialize() + for i in range(3): + await store.add_conversation(f"conv-{i}", _make_conversation(i)) + await store.finalize() + + metadata = store.get_client_metadata() + client = MemoryMapDatasetClient( + metadata.data_file_path, + metadata.index_file_path, + ) + + # Simulate another thread having just moved the shared position. + client.data_mmap.seek(client.index.offsets["conv-2"].offset) + assert client.get_conversation("conv-0").session_id == "conv-0"