feat(dataset): verbatim payload replay fast path + content-addressed mmap dataset cache#1104
feat(dataset): verbatim payload replay fast path + content-addressed mmap dataset cache#1104ajcasagrande wants to merge 8 commits into
Conversation
…mmap dataset cache Two coupled throughput features for large trace replays: PAYLOAD_BYTES mmap fast path. When every turn of every conversation carries a verbatim Turn.raw_payload (raw_payload / inputs_json / mooncake_trace with a payload field, or the new opt-in AIPERF_DATASET_PREFORMAT_PAYLOADS pre-formatting pass), DatasetManager writes the mmap dataset in a new PAYLOAD_BYTES format: each turn's request body is pre-encoded to JSON bytes at dataset-configure time. Workers detect the format from the dataset-configured notification and serve credits straight from the mmap (get_payload_bytes), skipping session/conversation deserialization, endpoint payload formatting, and re-encoding per request; the transport sends the bytes verbatim. InferenceClient now also canonicalises dict payloads to bytes once before transport dispatch, so the transport no longer re-dumps the payload it was handed (multipart endpoints keep the structured dict so FormData still works). DAG descendants (agent_depth > 0) keep the session path. Content-addressed cross-run mmap dataset cache. Tokenizing and reconstructing a multi-GB trace into dataset.dat/index.dat is the dominant startup cost of large replays (for 1M-request runs, minutes of work repeated on every rerun). The new aiperf.dataset.mmap_cache keys an entry on sha256 of the input file/directory bytes + tokenizer identity (name, revision, trust_remote_code, apply_chat_template) + every input/prompt/synthesis/seed setting that bakes into the decoded bytes, and stores dataset.dat, index.dat, and a versioned manifest under ~/.cache/aiperf/dataset_mmap (AIPERF_DATASET_MMAP_CACHE_DIR to relocate, AIPERF_DATASET_MMAP_CACHE_ENABLED=0 to disable). A HIT restores files by hardlink (copy fallback), skips tokenizer + composer entirely, and adopts the stored format; a MISS runs the normal pipeline under a cross-process flock (double-checked, with a cache-complete bypass for stale locks) and populates the entry atomically on the way out. Version-mismatched, corrupt, or partial entries are treated as MISSes. Cache-served dataset types (explicit config type only): mooncake_trace, bailian_trace, burst_gpt_trace, sagemaker_data_capture, raw_payload, inputs_json. These trace/verbatim types also skip inputs.json generation (the dump would be huge or just the verbatim source bytes); raw-payload datasets that do emit inputs.json export the payloads verbatim instead of re-running the endpoint formatter. Carry-alongs: MemoryMapFormat enum (conversation | payload_bytes), format field on MemoryMapClientMetadata, per-turn payload offsets in the mmap index, adopt_existing_files() on the backing store for the cache HIT path, DatasetConfigurationFailedNotification publish on configure failure (TimingManager already subscribes), and a filelock dependency for the populate lock. Environment knobs: AIPERF_DATASET_MMAP_CACHE_ENABLED, AIPERF_DATASET_MMAP_CACHE_DIR, AIPERF_DATASET_PREFORMAT_PAYLOADS. Co-authored-by: Cam Quilici <cjquilici@gmail.com> Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>
Try out this PRQuick install: pip install --upgrade --force-reinstall git+https://github.com/ai-dynamo/aiperf.git@79f6b5b38b4a358e5b0714513897f48546758adeRecommended 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@79f6b5b38b4a358e5b0714513897f48546758adeLast updated for commit: |
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (1)
WalkthroughThis PR adds a content-addressed mmap dataset cache, a PAYLOAD_BYTES mmap storage format, DatasetManager cache HIT/MISS orchestration with payload preformatting, and verbatim payload-bytes fast paths across worker, inference client, and transport layers. ChangesMmap Cache and Payload-Bytes Fast Path
Estimated code review effort: 4 (Complex) | ~75 minutes Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
src/aiperf/dataset/dataset_manager.py (1)
677-689: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCentralize the mmap path layout
_run_mmap_pathsandMemoryMapDatasetBackingStore.__init__both hardcode the sameaiperf_mmap_{benchmark_id}/dataset.dat(.zst)scheme. Extracting that construction into a shared helper would keep restore/populate and the store in sync if the layout ever changes.🤖 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 `@src/aiperf/dataset/dataset_manager.py` around lines 677 - 689, Centralize the mmap path construction currently duplicated in _run_mmap_paths and MemoryMapDatasetBackingStore.__init__. Extract the shared aiperf_mmap_{benchmark_id}/dataset.dat(.zst) and index.dat(.zst) layout into a common helper, then have _run_mmap_paths and the backing store call that helper so restore/populate and storage stay in sync if the scheme changes.
🤖 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/dataset_manager.py`:
- Around line 247-248: The post-run cache population in DatasetManager is doing
blocking filesystem work inline, which can stall the event loop. Update the code
around the _populate_cache_after_run() call in the run-completion path to
offload it with asyncio.to_thread, so the cache copy/link work for dataset.dat
and index.dat runs outside the loop. Keep the existing _cache_key_for_run guard
and apply the change at the call site rather than inside the helper.
- Around line 713-722: The synchronous cache file I/O in dataset configuration
is blocking the event loop, so update the cache-hit and cache-miss paths in
dataset_manager’s _configure_from_cache_hit and the populate call site to run
mmap_cache.restore_to_run_dir and mmap_cache.populate off-thread. Wrap these
helpers with asyncio.to_thread(...) or convert the cache helpers to async so
hardlink/copy work, including the shutil.copyfile fallback, does not execute
inside the async path.
---
Nitpick comments:
In `@src/aiperf/dataset/dataset_manager.py`:
- Around line 677-689: Centralize the mmap path construction currently
duplicated in _run_mmap_paths and MemoryMapDatasetBackingStore.__init__. Extract
the shared aiperf_mmap_{benchmark_id}/dataset.dat(.zst) and index.dat(.zst)
layout into a common helper, then have _run_mmap_paths and the backing store
call that helper so restore/populate and storage stay in sync if the scheme
changes.
🪄 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: 7574f68e-f431-4b5b-80da-dd9861c535b3
📒 Files selected for processing (20)
docs/environment-variables.mdpyproject.tomlsrc/aiperf/common/enums/__init__.pysrc/aiperf/common/enums/enums.pysrc/aiperf/common/environment.pysrc/aiperf/common/models/dataset_models.pysrc/aiperf/dataset/dataset_manager.pysrc/aiperf/dataset/memory_map_utils.pysrc/aiperf/dataset/mmap_cache.pysrc/aiperf/dataset/mmap_cache_lock.pysrc/aiperf/dataset/payload_formatting.pysrc/aiperf/transports/aiohttp_transport.pysrc/aiperf/workers/inference_client.pysrc/aiperf/workers/worker.pytests/unit/dataset/test_dataset_manager_cache.pytests/unit/dataset/test_mmap_cache.pytests/unit/dataset/test_payload_formatting.pytests/unit/dataset/test_payload_mmap.pytests/unit/workers/test_inference_client.pytests/unit/workers/test_inference_client_payload_bytes_adversarial.py
This comment has been minimized.
This comment has been minimized.
…sets Trace / verbatim datasets (mooncake, sagemaker, burst_gpt, raw_payload, inputs_json) now intentionally skip inputs.json generation because their per-turn payloads are verbatim source bytes and the dump would be huge. Add has_all_outputs_except_inputs to the harness and update the affected integration tests to assert inputs.json is absent instead of present. The auto-detected burst_gpt test keeps has_all_outputs: the skip predicate keys off the CONFIG dataset type, so auto-detected traces still emit inputs.json. Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>
…path A worker that subscribes after the DatasetConfiguredNotification broadcast (much more likely now that a cache HIT makes dataset configure nearly instant) keeps _dataset_client=None and falls back to the DatasetManager conversation request. On PAYLOAD_BYTES stores that request crashed with MemoryMapSerializationError (get_conversation is unsupported), failing the whole run with 'All N inference request(s) failed'. Reconstruct a minimal Conversation from the per-turn payload bytes instead: each turn carries the verbatim raw_payload, so the worker's session path still replays the authored bytes. Applies to both the conversation and the turn fallback handlers via a shared _get_conversation_for_serving helper. Also add unit coverage for the preformat/select-format/inputs-json payload pathways in DatasetManager. Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>
Targeted unit tests for the patch regions codecov flagged: - mmap_cache: directory-content hashing, populate races (winner-stays, leftover tmp dir, os.replace loser, missing sources), partial entries (missing dataset/index files, legacy inputs.json flags), accuracy-mode and synthetic-run cache gates, public-dataset source identity, prompt cache_bust exclusion. - mmap_cache_lock: manifest-presence lock bypass (pre-acquire and mid-retry), SoftFileLock fallback on flock-less filesystems, unrelated NotImplementedError propagation, release() error swallowing. - dataset_manager: hit-under-lock after a stale pre-lock miss, corrupt manifest metadata falling back to full configure, lookup/key-computation crash degradation, populate-after-run guards and failure tolerance, failure-notification fan-out (publish success and publish crash). - memory_map_utils: get_conversation rejection on PAYLOAD_BYTES stores, uninitialized client store guard, adopt_existing_files preconditions. - worker: payload-bytes fast path gating (format/DAG-depth/missing-payload) and error recording; real _make_first_token_callback factory. - transports: multipart dict payloads build FormData while pre-encoded bytes pass verbatim (send_request and video submit). - inference_client: trace logging skips the per-turn dump when the payload-bytes path has no turns. Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>
The accuracy-mode gate test passed vacuously: without an input file the synthetic-run gate also returns None, so the accuracy branch was never executed. Pair the accuracy benchmark with a mooncake trace file and assert cfg.accuracy.enabled so the None can only come from the accuracy gate. Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/unit/dataset/test_dataset_manager_payload_paths.py (1)
379-379: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse a raw string for the regex
match=pattern.
"mixed[\\s\\S]*raw_payload"relies on double-escaping to express\s/\S; a raw string (r"mixed[\s\S]*raw_payload") is clearer and avoids the Ruff RUF043 warning about ambiguous regex metacharacters.♻️ Proposed fix
- with pytest.raises(ValueError, match="mixed[\\s\\S]*raw_payload"): + with pytest.raises(ValueError, match=r"mixed[\s\S]*raw_payload"):🤖 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_dataset_manager_payload_paths.py` at line 379, Update the pytest assertion in test_dataset_manager_payload_paths to use a raw string for the ValueError match pattern instead of a double-escaped normal string. Replace the match argument in the pytest.raises call so the regex is written as a raw pattern, keeping the same text but avoiding the Ruff RUF043 ambiguous regex warning.Source: Linters/SAST tools
🤖 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.
Nitpick comments:
In `@tests/unit/dataset/test_dataset_manager_payload_paths.py`:
- Line 379: Update the pytest assertion in test_dataset_manager_payload_paths to
use a raw string for the ValueError match pattern instead of a double-escaped
normal string. Replace the match argument in the pytest.raises call so the regex
is written as a raw pattern, keeping the same text but avoiding the Ruff RUF043
ambiguous regex warning.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 19b1de19-032b-475d-b5c6-124e34858e66
📒 Files selected for processing (15)
src/aiperf/dataset/dataset_manager.pytests/harness/utils.pytests/integration/agentic_code_gen/test_profile_with_synthesized_dataset.pytests/integration/test_burst_gpt_trace.pytests/integration/test_mooncake_trace.pytests/integration/test_sagemaker_data_capture.pytests/unit/dataset/test_dataset_manager_cache.pytests/unit/dataset/test_dataset_manager_failure_propagation.pytests/unit/dataset/test_dataset_manager_payload_paths.pytests/unit/dataset/test_mmap_cache.pytests/unit/dataset/test_payload_mmap.pytests/unit/transports/test_aiohttp_transport.pytests/unit/transports/test_aiohttp_transport_video.pytests/unit/workers/test_inference_client_payload_bytes_adversarial.pytests/unit/workers/test_worker.py
✅ Files skipped from review due to trivial changes (1)
- tests/integration/test_burst_gpt_trace.py
🚧 Files skipped from review as they are similar to previous changes (2)
- tests/unit/workers/test_inference_client_payload_bytes_adversarial.py
- src/aiperf/dataset/dataset_manager.py
test_populate_cache_after_run_missing_files_skips unlinked the run mmap files while the dataset client store still had them memory-mapped, which Windows rejects with PermissionError (WinError 32). Stop the manager (and thus close the mmap handles) before deleting the files. Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>
Address review feedback: _populate_cache_after_run() and mmap_cache.restore_to_run_dir() perform multi-GiB shutil.copyfile operations, which blocked the DatasetManager event loop when called inline from async configure paths. Wrap both call sites in asyncio.to_thread, matching the branch's existing offload pattern (flock acquire, mmap writes, mkdir). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>
…ncio.to_thread) Pins b457db3: a recording asyncio.to_thread wrapper asserts both _populate_cache_after_run and mmap_cache.restore_to_run_dir are dispatched through to_thread (and still execute for real). Reverting either call site to a direct sync call fails the corresponding test. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>
Residue port from carve commit 44fd6dd (PR #1104, ajc/payload-bytes-mmap-cache). src/aiperf/dataset/mmap_cache_lock.py imports filelock directly, but agentx only had it as a transitive dependency (huggingface-hub/torch), which could silently break under dependency upgrades. The rest of the carve (verbatim payload replay fast path + content-addressed mmap cache) originated on agentx and is fully present, extended with weka dataset support. (cherry picked from commit 44fd6dd, in part) Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com> Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
lkomali
left a comment
There was a problem hiding this comment.
A few things from a closer pass on the cache correctness/robustness. Nice PR overall — the once-encoding refactor and the DatasetConfigurationFailedNotification fan-out both check out, and I verified the PAYLOAD_BYTES offsets are exact. Three comments below.
| # flips the stored mmap FORMAT (CONVERSATION vs PAYLOAD_BYTES) and the | ||
| # decoded bytes, and a cache HIT adopts the stored format verbatim, so a | ||
| # warm entry built with the other setting serves the wrong format. | ||
| "preformat_payloads": Environment.DATASET.PREFORMAT_PAYLOADS, |
There was a problem hiding this comment.
When PREFORMAT_PAYLOADS is on, the stored bytes are produced by endpoint.format_payload() (payload_formatting.py), which bakes in stream/stream_options, endpoint.extra, and the max_tokens-vs-max_completion_tokens field name. But the key only captures endpoint_type and model_name — none of those three. So two runs with preformat on, same trace/model/endpoint-type but differing only in --streaming / --extra-inputs / --use-legacy-max-tokens, land on the same key, and the second gets a HIT serving the first run's bytes (e.g. "stream": true it never asked for). That breaks the "a MISS is byte-identical to a HIT" guarantee. Could we fold endpoint.streaming, endpoint.extra, and use_legacy_max_tokens into the key when preformat_payloads is True? They only change the stored bytes in that mode, so gating them on the flag keeps the key stable for the common (non-preformat) path.
| wakes after the winner populates uses the cached entry instead of | ||
| repeating the work. | ||
| """ | ||
| cache_hit = self._try_cache_lookup() |
There was a problem hiding this comment.
This lookup runs synchronously on the event loop, and computing the key hashes the entire input file (or walks the whole input dir hashing every file) via hash_file_bytes/hash_dir_contents. For the multi-GB traces this cache targets, that blocks the DatasetManager loop — and its heartbeat/command handlers — for the full hash during configure. The populate/restore paths were already moved to asyncio.to_thread; this lookup/hash path is the one that wasn't. Could we offload it the same way, e.g. await asyncio.to_thread(self._try_cache_lookup) (and the _lookup_under_lock call)?
| acquired = await asyncio.to_thread( | ||
| _blocking_acquire, lock, timeout, lock_path, _cache_complete | ||
| ) | ||
| except NotImplementedError as e: |
There was a problem hiding this comment.
This only catches NotImplementedError (the SoftFileLock fallback). If _blocking_acquire hits its timeout it raises filelock.Timeout, which isn't caught here and propagates out of _do_profile_configure, failing the whole run. The cache-complete bypass covers a populator killed after it wrote the manifest — but one killed before completing (on the SoftFileLock/NFS path, where the lock tombstone persists) leaves the entry incomplete, so the bypass never fires: every waiter blocks the full timeout and then the run dies instead of degrading to an uncached local populate. Since populate is idempotent/atomic, could we catch Timeout here and proceed unlocked? Narrow (NFS-only), lower priority than the other two.
Summary
Two coupled startup/throughput features for large verbatim trace replays:
PAYLOAD_BYTES mmap fast path
When every turn of every conversation carries a verbatim
Turn.raw_payload(raw_payload/inputs_json/mooncake_tracewith a payload field, or the new opt-inAIPERF_DATASET_PREFORMAT_PAYLOADSpre-formatting pass), DatasetManager writes the mmap dataset in a newPAYLOAD_BYTESformat: each turn's request body is pre-encoded to JSON bytes at dataset-configure time, with per-turn offsets in the mmap index.Workers detect the format from the dataset-configured notification and serve credits straight from the mmap (
get_payload_bytes), skipping session/conversation deserialization, endpoint payload formatting, and per-request encoding; the transport sends the bytes verbatim. DAG descendants (agent_depth > 0) keep the session path. As part of this,InferenceClientnow canonicalises dict payloads to bytes once before transport dispatch, so the transport no longer re-encodes the payload it was handed (multipart endpoints keep the structured dict so FormData still works).Content-addressed cross-run mmap dataset cache
Tokenizing and reconstructing a multi-GB trace into
dataset.dat/index.datis the dominant startup cost of large replays — for 1M-request runs, minutes of work repeated on every byte-identical rerun. The newaiperf.dataset.mmap_cachekeys an entry on:records:content),trust_remote_code,apply_chat_template),Entries live under
~/.cache/aiperf/dataset_mmap(AIPERF_DATASET_MMAP_CACHE_DIRto relocate,AIPERF_DATASET_MMAP_CACHE_ENABLED=0to disable). A HIT restores files by hardlink (copy fallback) and skips tokenizer + composer entirely; a MISS runs the normal pipeline under a cross-process flock (double-checked locking, with a cache-complete bypass so a SIGKILLed populator's stale lock can't wedge waiters) and atomically populates the entry on the way out. Version-mismatched, corrupt, or partial entries are treated as MISSes. Cache misses still produce byte-identical mmap files to a non-cached run.Cache-served dataset types (explicit config type only):
mooncake_trace,bailian_trace,burst_gpt_trace,sagemaker_data_capture,raw_payload,inputs_json. These trace/verbatim types also skipinputs.jsongeneration (the dump would be huge or just the verbatim source bytes); when raw-payload datasets do emitinputs.json, payloads are exported verbatim instead of re-run through the endpoint formatter.Carry-alongs
MemoryMapFormatenum (conversation|payload_bytes) +formatfield onMemoryMapClientMetadataadopt_existing_files()on the mmap backing store for the cache HIT pathDatasetConfigurationFailedNotificationon configure failure (TimingManager already subscribes to it; previously nothing published it, so a configure failure left TimingManager waiting out the full 300s dataset timeout)filelockadded as an explicit dependency (already present transitively via huggingface-hub/datasets)AIPERF_DATASET_MMAP_CACHE_ENABLED,AIPERF_DATASET_MMAP_CACHE_DIR,AIPERF_DATASET_PREFORMAT_PAYLOADS(docs regenerated)Deferred interactions (land with their features)
PAYLOAD_BYTESincompatibility guards for body-mutating features (cache-bust marker injection,--use-dynamo-conv-aware-routing/nvext.session_controloverlay) land with those flags — neither exists onmainyet. Follow-up when they merge: refusePAYLOAD_BYTESat_select_mmap_formatsince pre-encoded bytes bypass per-request body mutation.Testing
pytest tests/unit/ -n auto: 13844 passed, 94 skipped, 0 failedtests/unit/dataset/test_mmap_cache.py(43 tests): cache-key stability/collision sensitivity (input bytes, tokenizer identity, settings, seeds, synthesis multipliers,PREFORMAT_PAYLOADS), populate/lookup round-trip, manifest version pin (v25), corrupt/partial/compressed-mismatch MISSes, hardlink restore + copy fallback, run-dir cleanup preserving the cache entry, populate-lock serialization/timeout, trace-verbatim gate predicatetests/unit/dataset/test_dataset_manager_cache.py(5 tests): MISS-then-populate, HIT skips tokenizer + composer and still publishesDatasetConfiguredNotification, tokenizer-change invalidation, cache-disabled path, HIT never restoresinputs.jsontests/unit/dataset/test_payload_mmap.py(4 tests): PAYLOAD_BYTES write/read round-trip, per-turn offsets, format gating ofget_conversation/get_payload_bytestests/unit/dataset/test_payload_formatting.py(6 tests): shared payload-formatting generatortests/unit/workers/test_inference_client_payload_bytes_adversarial.py(12 tests): payload_bytes -> raw_payload -> format_payload priority chain, empty-turns guard, verbatim byte forwarding (unicode/invalid JSON), multipart dict carve-outruff format/ruff check,make check-ergonomics,make check-ruff-baselined: cleanSummary by CodeRabbit
PAYLOAD_BYTESdataset format support, including payload-bytes fast paths across worker and request handling.inputs.jsonfor trace/verbatim-style datasets and improved request validation.AIPERF_DATASET_*environment variables for mmap caching and payload preformatting.