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

[ajcasagrande]

The bug

MemoryMapDatasetClient.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.get_conversation dispatched this 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.

The two sibling readers (get_payload_bytes, get_payload_turn) already used position-free slicing; only get_conversation did not.

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

Reported against a 6.4 GB dataset in SemiAnalysisAI/aiperf#34, from a 46-run sweep sharing one mmap cache entry:

• 4 runs failed, each on a different conversation. A bad record would fail identically every run.
• The on-disk files are perfect. cmp on the 6.4 GB dataset.dat between cache source and per-run copy → byte-identical; index.dat md5 identical.
• Each named conversation decodes fine when read standalone from the same file with the same index.
• The 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.

Changes

1. Position-free slice in get_conversation — the fix.
2. madvise(MADV_WILLNEED) + page walk at open, gated on the new AIPERF_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.
3. Drop run_in_executor from all three client-store readers. The hop existed to absorb blocking major faults, which (2) eliminates.
4. Flatten index.offsets into dict[str, tuple[int, int]] at load time, dropping per-lookup Pydantic attribute access. Wire format unchanged.

Performance proof

The executor hop costs 18x

64 KB read, page-cache warm:

	time
direct on loop	1.08 µs
run_in_executor	19.41 µs

+18.33 µs per lookup, 18.0x. Paid per request/turn on the worker hot path.

It also made tail latency worse

Event-loop tick lateness under continuous reads, warmed pool:

64 KB records	p50	p99	max
idle (baseline)	0.081	0.230	0.338 ms
executor	0.059	0.410	0.639 ms
inline	0.008	0.045	0.128 ms

Inline wins on every percentile at typical record sizes. At 8 MB the executor buys ~2x median at the cost of 2–4x p99 and max — the wrong trade for a tool whose output is latency percentiles.

Why the executor never protected the loop anyway

Neither mmap read path releases the GIL. From CPython Modules/mmapmodule.c:

• mmap_read_method → PyBytes_FromStringAndSize(&self->data[self->pos], num_bytes) at :318 — no Py_BEGIN_ALLOW_THREADS
• mmap_subscript slice → PyBytes_FromStringAndSize(self->data + start, slicelen) at :999 — same
• The only Py_BEGIN_ALLOW_THREADS in the module are dealloc/close (:149/:162/:200/:218) and fstat at construction (:1294)

A major fault inside that memcpy stalls every Python thread regardless of which one runs it.

Prefault cost at open

dataset	cold	warm
16 MB	0.9 ms	0.4 ms
256 MB	12.6 ms	3.4 ms
1024 MB	40.0 ms	13.4 ms

~250 ms one-time for a 6.4 GB dataset. Set AIPERF_DATASET_MMAP_PREFAULT=false to trade tail-latency predictability for faster startup.

seek()+read() vs slice is within noise

record	seek+read	slice	delta
4 KB	0.35 µs	0.31 µs	0.04 µs
64 KB	1.13 µs	1.10 µs	0.03 µs
1 MB	21.25 µs	20.87 µs	0.38 µs
8 MB	278.99 µs	272.66 µs	6.34 µs

Both are the same memcpy. The fix is free — it is a correctness change, not a perf change.

End-to-end: no measurable throughput change

Large-payload replay, 32 concurrency, 640 requests, alternating branch/main to control for drift:

branch:  697, 702, 714, 721, 734 req/s
main:    705, 742, 757 req/s

Overlapping distributions. ~43 µs saved per request against ~42 ms request latency is 0.1% — below noise.

Being explicit: this PR is a correctness fix plus CPU and allocation hygiene. It is not a throughput win, and I could not surface one against a mock server. The savings would matter where the client is CPU-bound (high request rates, many workers), which I have no way to demonstrate here.

Behavioral equivalence was verified separately: same seed, same config, branch vs main → inputs.json byte-identical (111,041 bytes, cmp clean) and every deterministic metric matching exactly, including output_sequence_length.avg to full float precision.

Tests

tests/unit/dataset/test_mmap_conversation_concurrency.py:

• test_get_conversation_ignores_competing_seek_between_seek_and_read — _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 exact reported symptom (asked for conv-0, received conv-2's bytes, Invalid JSON: EOF while parsing a string at line 1 column 947).

  Worth noting for reviewers: merely leaving a stale position behind does not catch this bug — the reader's own seek() overwrites it on the way in and still returns the right record. The competing seek has to land mid-operation.

• test_get_conversation_concurrent_readers_returns_own_bytes — 8 threads × 150 lookups over 24 records of varying size. Weaker: does not reproduce at the default switch interval, only ever suggestive.
• test_get_conversation_matches_written_data_under_prefault_setting — reads agree with MMAP_PREFAULT on and off.

Verification

• pytest tests/unit/ — 17,700 passed
• pytest -m component_integration — 396 passed
• pre-commit — clean
• Functional runs against the mock server: 200-request synthetic, 400-conversation 64-way concurrent stress with ISL 19–2204 (zero decode errors — the bug's signature), and MMAP_PREFAULT=false

Notes for reviewers

• AIPERF_DATASET_MMAP_PREFAULT defaults to True, so every run now pays the startup pass in the table above. It is the safety argument for removing the executor hop; disabling one without the other is not recommended.
• The _offsets name now exists in both MemoryMapDatasetBackingStore (writer, ConversationOffset) and MemoryMapDatasetClient (reader, tuple[int, int]). Different classes, different value types — mildly confusing when grepping.

Fixes the issue reported in SemiAnalysisAI#34.

Summary by CodeRabbit

• New Features

  • Added the AIPERF_DATASET_MMAP_PREFAULT setting, enabled by default, to preload mapped dataset data into the operating system’s page cache.
  • The setting can be disabled when faster startup is preferred over preloading.

• Bug Fixes

  • Improved reliability and consistency of concurrent dataset reads and conversation retrieval.

• Documentation

  • Documented the new setting, its default behavior, and its startup-latency tradeoff.

[github-actions]

Try out this PR

Quick install:

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

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@898d6acb0f9f284bf6cc5f9d086f064fe5f52e31

---

Last updated for commit: 898d6ac • Browse code

[github-actions]

Fern Docs Preview: https://nvidia-preview-01a0361f-be76-70d7-ae8f-91004d90bf7c.docs.buildwithfern.com/aiperf/dev

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

Walkthrough

The change adds the default-enabled AIPERF_DATASET_MMAP_PREFAULT setting. It removes executor-based dataset reads, caches conversation offsets, prefaults mapped pages when enabled, and reads slices without shared mmap position state. Tests cover concurrency, seek interference, and both prefault settings.

Changes

Memory-mapped dataset reads

Layer / File(s)	Summary
Dataset prefault configuration src/aiperf/common/environment.py, docs/environment-variables.md	Adds the default-enabled MMAP_PREFAULT setting and documents its page-cache behavior and startup cost.
Direct client read execution src/aiperf/dataset/memory_map_utils.py	Removes unused event-loop state and executor dispatch from conversation and payload reads.
Prefaulting and concurrent lookup src/aiperf/dataset/memory_map_utils.py, tests/unit/dataset/test_mmap_conversation_concurrency.py	Caches conversation offsets, optionally prefaults mmap pages, reads by slicing, and validates concurrent reads, seek interference, and both prefault settings.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Poem

I’m a rabbit with pages to warm,
Slicing conversations keeps reads from harm.
The cache hops quickly, the mmap stays still,
Prefault or skip it—choose as you will.
Tests thump their paws: every read fits the bill.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 69.57% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly summarizes the main changes: thread-safe mmap reads, page prefaulting, and removal of executor hops.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 2

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/aiperf/dataset/memory_map_utils.py`:
- Around line 422-424: Move mmap client construction, _prefault_data_mmap(), and
all direct mmap slicing performed by the async read methods into a dedicated
executor or worker. Update the affected async methods and initialization paths,
including the methods around the listed ranges, to await executor submissions
while preserving position-free reads and existing results. Ensure no mmap
page-faulting or other blocking file access runs on the event loop, regardless
of MMAP_PREFAULT or later OS eviction.

In `@tests/unit/dataset/test_mmap_conversation_concurrency.py`:
- Around line 43-45: Add type annotations throughout the targeted test helpers
and test functions: annotate tmp_path and monkeypatch in _build_client and the
affected tests, item and the return type in _InterleavingMmap.__getitem__, and
name and the return type in _InterleavingMmap.__getattr__. Ensure every function
parameter and return value in these symbols has an appropriate annotation
without changing behavior.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 5104bd35-511b-47c3-9607-60da5c48a297

📥 Commits

Reviewing files that changed from the base of the PR and between c2f5e9d and 28a7745.

📒 Files selected for processing (4)

• docs/environment-variables.md
• src/aiperf/common/environment.py
• src/aiperf/dataset/memory_map_utils.py
• tests/unit/dataset/test_mmap_conversation_concurrency.py

[codecov]

Codecov Report

❌ Patch coverage is 88.88889% with 2 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/aiperf/dataset/memory_map_utils.py	88.23%	0 Missing and 2 partials ⚠️

📢 Thoughts on this report? Let us know!

[debermudez]

Overall assessment: Fix is correct and minimal — mmap[a:b] is the right call and matches the sibling readers that were already safe. The _InterleavingMmap test is a nice deterministic reproducer, and the evidence chain in the PR description is model-quality bug forensics. Two issues worth discussing before merge.

Fix order:

1. (Recommended) _prefault_data_mmap blocks the async event loop during startup — see inline comment.
2. (Optional) Docstring for get_conversation slightly overstates the no-fault guarantee when MMAP_PREFAULT=False.

What's working well: position-free slice fix, _InterleavingMmap deterministic race test, concurrent stress test with varied sizes, graceful madvise fallback, index flattening.

[ajcasagrande]

Review addressed

Addressed 3 of 5 review threads.

Fixed:

1. _build_client now stops the backing store, so its @on_stop cleanup unlinks the mmap files instead of leaking them (commit f8e24ea)
2. get_conversation docstring now records that major faults are still possible when MMAP_PREFAULT is disabled, and why they're left on the loop thread (commit 562fdfa)
3. Replied to the _prefault_data_mmap() event-loop concern with the measurements — no code change: the thread hop measured ~5x worse on the one genuinely cold rep, because CPython's mmap slice holds the GIL while copying

Skipped (2 threads):

1. coderabbitai: "Keep mmap I/O off the event loop" — already resolved; withdrawn by the reviewer after the GIL/benchmark data
2. coderabbitai: "Add the missing type annotations" — already resolved in 4678074

🤖 Generated with Claude Code

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (2)

tests/unit/dataset/test_mmap_conversation_concurrency.py (2)

161-164: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Assert that the setting controls prefaulting.

Both settings return the same conversations if prefaulting is skipped, always enabled, or inverted. Wrap MemoryMapDatasetClient._prefault_data_mmap() with a counter and assert one call when prefault is True and zero calls when it is False.

Based on PR objectives, this test must validate prefault behavior, not only read output.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unit/dataset/test_mmap_conversation_concurrency.py` around lines 161 -
164, Update the test around _open_client in the prefault parametrization to spy
on or wrap MemoryMapDatasetClient._prefault_data_mmap(), count invocations, and
assert the count is one when prefault is True and zero when prefault is False,
while retaining the existing conversation-read assertions.

---

96-98: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick win

Offload concurrent readers from the event loop.

Replace the blocking list(pool.map(_read_many, range(_THREADS))) call with asyncio.to_thread() tasks awaited by asyncio.gather().

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unit/dataset/test_mmap_conversation_concurrency.py` around lines 96 -
98, Update the concurrent-reader test around the ThreadPoolExecutor usage to run
blocking reader work via asyncio.to_thread() and await the resulting tasks with
asyncio.gather(). Remove the synchronous list(pool.map(_read_many,
range(_THREADS))) execution while preserving all _THREADS reader invocations and
exception propagation.

Source: Coding guidelines

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tests/unit/dataset/test_mmap_conversation_concurrency.py`:
- Around line 66-74: Move MemoryMapDatasetClient construction inside the try
block so store cleanup always reaches the finally block when initialization
fails. Track whether client creation succeeded and call client.close() only for
a successfully constructed client, while always awaiting store.stop().

---

Nitpick comments:
In `@tests/unit/dataset/test_mmap_conversation_concurrency.py`:
- Around line 161-164: Update the test around _open_client in the prefault
parametrization to spy on or wrap MemoryMapDatasetClient._prefault_data_mmap(),
count invocations, and assert the count is one when prefault is True and zero
when prefault is False, while retaining the existing conversation-read
assertions.
- Around line 96-98: Update the concurrent-reader test around the
ThreadPoolExecutor usage to run blocking reader work via asyncio.to_thread() and
await the resulting tasks with asyncio.gather(). Remove the synchronous
list(pool.map(_read_many, range(_THREADS))) execution while preserving all
_THREADS reader invocations and exception propagation.

🪄 Autofix

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: cafb45bf-10c5-4ca6-9810-15550910c17b

📥 Commits

Reviewing files that changed from the base of the PR and between 4678074 and 562fdfa.

📒 Files selected for processing (2)

• src/aiperf/dataset/memory_map_utils.py
• tests/unit/dataset/test_mmap_conversation_concurrency.py

🚧 Files skipped from review as they are similar to previous changes (1)

• src/aiperf/dataset/memory_map_utils.py

[debermudez]

Re-reviewed. Both prior findings resolved:

1. _prefault_data_mmap event-loop concern — verified your rebuttal directly against CPython 3.12 mmapmodule.c: neither mmap_read_method nor mmap_subscript release the GIL, so an executor hop wouldn't have helped. Closing this, no code change needed.
2. Docstring fix (562fdfa) — confirmed present.

Re-validated the fix itself: mutation-tested _InterleavingMmap (reverting to seek()+read() reproduces the exact reported symptom), and ran a real CLI benchmark against the mock server (200 conversations, wide ISL variance, 32-way concurrency) through the actual MemoryMapDatasetClientStore path — 0% errors, no decode errors.

Two CodeRabbit test nitpicks remain open (prefault-call-count assertion, ThreadPoolExecutor vs asyncio.to_thread in test) — both optional, non-blocking.

LGTM.