Skip to content

feat(dataset): verbatim payload replay fast path + content-addressed mmap dataset cache#1104

Open
ajcasagrande wants to merge 8 commits into
mainfrom
ajc/payload-bytes-mmap-cache
Open

feat(dataset): verbatim payload replay fast path + content-addressed mmap dataset cache#1104
ajcasagrande wants to merge 8 commits into
mainfrom
ajc/payload-bytes-mmap-cache

Conversation

@ajcasagrande

@ajcasagrande ajcasagrande commented Jul 3, 2026

Copy link
Copy Markdown
Contributor

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_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, 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, InferenceClient now 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.dat is the dominant startup cost of large replays — for 1M-request runs, minutes of work repeated on every byte-identical rerun. The new aiperf.dataset.mmap_cache keys an entry on:

  • sha256 of the input file/directory bytes (or inline records: content),
  • tokenizer identity (name, revision, trust_remote_code, apply_chat_template),
  • every input/prompt/synthesis/seed/timing setting that bakes into the decoded bytes,
  • a manifest version pin (bumped whenever decoded content semantics change).

Entries live 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) 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 skip inputs.json generation (the dump would be huge or just the verbatim source bytes); when raw-payload datasets do emit inputs.json, payloads are exported verbatim instead of re-run through the endpoint formatter.

Carry-alongs

  • MemoryMapFormat enum (conversation | payload_bytes) + format field on MemoryMapClientMetadata
  • adopt_existing_files() on the mmap backing store for the cache HIT path
  • DatasetManager now publishes DatasetConfigurationFailedNotification on configure failure (TimingManager already subscribes to it; previously nothing published it, so a configure failure left TimingManager waiting out the full 300s dataset timeout)
  • filelock added as an explicit dependency (already present transitively via huggingface-hub/datasets)
  • Env knobs: AIPERF_DATASET_MMAP_CACHE_ENABLED, AIPERF_DATASET_MMAP_CACHE_DIR, AIPERF_DATASET_PREFORMAT_PAYLOADS (docs regenerated)

Deferred interactions (land with their features)

  • Cache-key inputs for loader-specific knobs (e.g. Weka trace-reconstruction settings) land with those loaders; the key schema and manifest version pin (v25) are shared so entries stay compatible.
  • The PAYLOAD_BYTES incompatibility guards for body-mutating features (cache-bust marker injection, --use-dynamo-conv-aware-routing / nvext.session_control overlay) land with those flags — neither exists on main yet. Follow-up when they merge: refuse PAYLOAD_BYTES at _select_mmap_format since pre-encoded bytes bypass per-request body mutation.

Testing

  • pytest tests/unit/ -n auto: 13844 passed, 94 skipped, 0 failed
  • New unit tests:
    • tests/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 predicate
    • tests/unit/dataset/test_dataset_manager_cache.py (5 tests): MISS-then-populate, HIT skips tokenizer + composer and still publishes DatasetConfiguredNotification, tokenizer-change invalidation, cache-disabled path, HIT never restores inputs.json
    • tests/unit/dataset/test_payload_mmap.py (4 tests): PAYLOAD_BYTES write/read round-trip, per-turn offsets, format gating of get_conversation/get_payload_bytes
    • tests/unit/dataset/test_payload_formatting.py (6 tests): shared payload-formatting generator
    • tests/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-out
  • ruff format / ruff check, make check-ergonomics, make check-ruff-baselined: clean

Summary by CodeRabbit

  • New Features
    • Added a content-addressed, concurrency-safe disk cache for mmap dataset artifacts, with configurable enable/dir behavior.
    • Added PAYLOAD_BYTES dataset format support, including payload-bytes fast paths across worker and request handling.
  • Bug Fixes
    • Made cache HIT/MISS handling more robust (re-check under lock, stricter format rules) and serving resilient for payload-bytes datasets.
    • Skipped inputs.json for trace/verbatim-style datasets and improved request validation.
  • Documentation
    • Documented new AIPERF_DATASET_* environment variables for mmap caching and payload preformatting.
  • Tests
    • Expanded coverage for cache threading offload, payload formatting, payload mmap behavior, and updated trace expectations.

…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>
@github-actions

github-actions Bot commented Jul 3, 2026

Copy link
Copy Markdown

Try out this PR

Quick install:

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

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@79f6b5b38b4a358e5b0714513897f48546758ade

Last updated for commit: 79f6b5bBrowse code

@github-actions github-actions Bot added the feat label Jul 3, 2026
@github-actions

github-actions Bot commented Jul 3, 2026

Copy link
Copy Markdown

@codecov

codecov Bot commented Jul 3, 2026

Copy link
Copy Markdown

@coderabbitai

coderabbitai Bot commented Jul 3, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 4bdbc3fc-eccc-45a3-8c0a-f9e8edcfc70a

📥 Commits

Reviewing files that changed from the base of the PR and between 3964602 and 79f6b5b.

📒 Files selected for processing (2)
  • src/aiperf/dataset/dataset_manager.py
  • tests/unit/dataset/test_dataset_manager_cache.py
🚧 Files skipped from review as they are similar to previous changes (1)
  • src/aiperf/dataset/dataset_manager.py

Walkthrough

This 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.

Changes

Mmap Cache and Payload-Bytes Fast Path

Layer / File(s) Summary
Enums, config, and metadata
src/aiperf/common/enums/*, src/aiperf/common/environment.py, src/aiperf/common/models/dataset_models.py, pyproject.toml, docs/environment-variables.md
Adds MemoryMapFormat, new dataset mmap settings, MemoryMapClientMetadata.format, the filelock dependency, and environment documentation.
Payload-byte storage and serving
src/aiperf/dataset/memory_map_utils.py, src/aiperf/dataset/payload_formatting.py, tests/unit/dataset/test_payload_mmap.py, tests/unit/dataset/test_payload_formatting.py
Adds per-turn payload-byte storage, payload offset/index support, get_payload_bytes, conversation fallback rejection for PAYLOAD_BYTES, and shared payload formatting helpers.
Cache keying, lookup, and locking
src/aiperf/dataset/mmap_cache.py, src/aiperf/dataset/mmap_cache_lock.py, tests/unit/dataset/test_mmap_cache.py
Adds cache manifests, key computation, lookup/restore/populate, trace/verbatim gating, and async cross-process cache locking.
DatasetManager cache and payload orchestration
src/aiperf/dataset/dataset_manager.py, tests/unit/dataset/test_dataset_manager_cache.py, tests/unit/dataset/test_dataset_manager_failure_propagation.py, tests/unit/dataset/test_dataset_manager_payload_paths.py, tests/integration/*trace*.py, tests/harness/utils.py
Adds cache HIT/MISS control flow, payload preformatting, mmap format selection, cache-hit restoration, failure notifications, serving fallbacks, and trace output expectation updates.
Request bytes path and worker fast path
src/aiperf/workers/worker.py, src/aiperf/workers/inference_client.py, src/aiperf/transports/aiohttp_transport.py, tests/unit/workers/*, tests/unit/transports/*
Adds payload-bytes request handling, worker fast path execution, and bytes-aware transport encoding with corresponding unit tests.

Estimated code review effort: 4 (Complex) | ~75 minutes

Poem

A rabbit hopped through mmap trees,
Caching bytes with practiced ease.
Locks held tight, no double toil,
Payloads streamed without recoil.
Hop, hop, cache HIT — oh what a treat! 🐇✨

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title clearly captures the two main changes: verbatim payload replay fast path and content-addressed mmap dataset caching.
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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 coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 2

🧹 Nitpick comments (1)
src/aiperf/dataset/dataset_manager.py (1)

677-689: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Centralize the mmap path layout
_run_mmap_paths and MemoryMapDatasetBackingStore.__init__ both hardcode the same aiperf_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

📥 Commits

Reviewing files that changed from the base of the PR and between 66f228e and 44fd6dd.

📒 Files selected for processing (20)
  • docs/environment-variables.md
  • pyproject.toml
  • src/aiperf/common/enums/__init__.py
  • src/aiperf/common/enums/enums.py
  • src/aiperf/common/environment.py
  • src/aiperf/common/models/dataset_models.py
  • src/aiperf/dataset/dataset_manager.py
  • src/aiperf/dataset/memory_map_utils.py
  • src/aiperf/dataset/mmap_cache.py
  • src/aiperf/dataset/mmap_cache_lock.py
  • src/aiperf/dataset/payload_formatting.py
  • src/aiperf/transports/aiohttp_transport.py
  • src/aiperf/workers/inference_client.py
  • src/aiperf/workers/worker.py
  • tests/unit/dataset/test_dataset_manager_cache.py
  • tests/unit/dataset/test_mmap_cache.py
  • tests/unit/dataset/test_payload_formatting.py
  • tests/unit/dataset/test_payload_mmap.py
  • tests/unit/workers/test_inference_client.py
  • tests/unit/workers/test_inference_client_payload_bytes_adversarial.py

Comment thread src/aiperf/dataset/dataset_manager.py Outdated
Comment thread src/aiperf/dataset/dataset_manager.py Outdated
@datadog-official

This comment has been minimized.

@ajcasagrande ajcasagrande added the AgentX Feature for AgentX label Jul 3, 2026
…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>

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick comments (1)
tests/unit/dataset/test_dataset_manager_payload_paths.py (1)

379-379: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Use 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

📥 Commits

Reviewing files that changed from the base of the PR and between 44fd6dd and 4011759.

📒 Files selected for processing (15)
  • src/aiperf/dataset/dataset_manager.py
  • tests/harness/utils.py
  • tests/integration/agentic_code_gen/test_profile_with_synthesized_dataset.py
  • tests/integration/test_burst_gpt_trace.py
  • tests/integration/test_mooncake_trace.py
  • tests/integration/test_sagemaker_data_capture.py
  • tests/unit/dataset/test_dataset_manager_cache.py
  • tests/unit/dataset/test_dataset_manager_failure_propagation.py
  • tests/unit/dataset/test_dataset_manager_payload_paths.py
  • tests/unit/dataset/test_mmap_cache.py
  • tests/unit/dataset/test_payload_mmap.py
  • tests/unit/transports/test_aiohttp_transport.py
  • tests/unit/transports/test_aiohttp_transport_video.py
  • tests/unit/workers/test_inference_client_payload_bytes_adversarial.py
  • tests/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

ajcasagrande and others added 3 commits July 6, 2026 12:35
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>
ajcasagrande added a commit that referenced this pull request Jul 7, 2026
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 lkomali left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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()

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

AgentX Feature for AgentX feat

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants