fix(dataset): make MemoryMapDatasetClient.get_conversation thread-safe

[qiaoxj07]

Problem

MemoryMapDatasetClient.get_conversation reads its record through the mmap's shared file position:

# src/aiperf/dataset/memory_map_utils.py
self.data_mmap.seek(offset_info.offset)
conv_bytes = self.data_mmap.read(offset_info.size)

MemoryMapDatasetClientStore.get_conversation dispatches this onto a thread pool (run_in_executor), so several threads share one mmap object and 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 bogus

MemoryMapSerializationError: Failed to decode conversation data:
  Invalid JSON: trailing characters at line 1 column N
  (or: EOF while parsing a string at line 1 column N)

on a random conversation — and, when the foreign slice happens to parse, as the wrong conversation being served with no error at all.

The two sibling readers in the same class (get_payload_bytes, get_payload_turn) already use position-free slicing; only get_conversation did not.

Fix

Use the same position-free slice.

Evidence this is a read-time race, not data corruption

From a 46-run benchmark sweep that all shared one mmap cache entry:

• 4 runs hit the error, each on a different conversation. A bad record would fail identically in every run.
• The on-disk files are perfect. cmp on the 6.4 GB dataset.dat between the cache source and the per-run copy → byte-identical; index.dat md5 identical; sizes identical.
• Each named conversation decodes fine when read standalone from the same file with the same index.
• The error 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; another had index size 18,094,940 but its JSON closed at column 1,154,734, i.e. it landed inside a different, shorter record.

Same bytes, same index, different victim each time.

Reproduction

A standalone stdlib script — one mmap, 8 threads, 16,000 reads of 64 KB–1 MB records, sys.setswitchinterval(1e-6) to widen the GIL handoff window:

seek()+read()  wrong reads: 143 / 16000
slice          wrong reads:   0 / 16000

At the default 5 ms switch interval a short synthetic loop does not hit the window (320,000 fast reads in 0.7 s → 0 hits), so I am not claiming a default-settings repro. The production workload widens the window considerably: records up to 18 MB, page faults against a 6.4 GB file (the method's own docstring notes "mmap reads can block on page faults"), and an asyncio loop competing for the GIL.

Tests

Adds tests/unit/dataset/test_mmap_conversation_concurrency.py:

• test_get_conversation_is_thread_safe — 8 threads × 150 lookups over 24 conversations of varying sizes; every result must match the requested id.
• test_get_conversation_ignores_shared_mmap_position — a stale mmap position must not change what is read.

Note on verification

I could not execute the test suite locally (no environment with the project's dependencies on this machine), so CI should be the gate on these tests. The one-line production change is a direct substitution of an API the same class already uses twice, and the race itself is demonstrated by the standalone script above.

[github-actions]

Try out this PR

Quick install:

pip install --upgrade --force-reinstall git+https://github.com/ai-dynamo/aiperf.git@0d1ef21f988df20fbd8ab52ed8843c479441b07a

Recommended with virtual environment (using uv):

uv venv --python 3.12 && source .venv/bin/activate
uv pip install --upgrade --force-reinstall git+https://github.com/ai-dynamo/aiperf.git@0d1ef21f988df20fbd8ab52ed8843c479441b07a

---

Last updated for commit: 0d1ef21 • Browse code

[qiaoxj07]

Reproduced on the real dataset, with the shipped package

Ran an A/B directly inside a production container against a live 6.4 GB dataset.dat (9,843 conversations), importing the installed aiperf@818c3a5a2 — same class, same index, same _deserialize_conversation, the only variable being the read implementation. Read-only; nothing in site-packages was touched.

dataset: 9843 conversations
probing 32 of them, record sizes 1,266 .. 3,781,642 bytes

=== forced GIL switch interval (1e-06s), 8 threads ===
  seek   reads=3200  wrong_id=0  decode_err=34   (1.5s)
  slice  reads=3200  wrong_id=0  decode_err=0    (1.5s)

=== default GIL switch interval (0.005s), 8 threads ===
  seek   reads=3200  wrong_id=0  decode_err=0    (1.4s)
  slice  reads=3200  wrong_id=0  decode_err=0    (1.4s)

34 of 3,200 reads (1.06%) failed on the shipped seek()+read() path; 0 on the slice path. Every failure was a MemoryMapSerializationError — the exact exception class seen in the field — rather than a wrong-but-valid conversation, which is what you would expect: a cross-read almost always lands mid-record and yields invalid JSON.

The default-interval arm shows 0/3,200 for both. That is not evidence against the race: in production the observed rate is roughly one failure per ~500k conversation reads, so 3,200 reads is three orders of magnitude short of the sample needed. I did not run the ~500k-read version because the only nodes with a live mmap were running benchmarks whose measurements it would perturb.

[ajcasagrande]

I can confirm this is a valid fix.

Its also even more valid to combine this with removing the run_in_executor call in the first place. it was added under false pretenses.