feat(endpoints): add image UUID cache reuse

[furionw]

Why

AIPerf currently resends repeated images, preventing cache-aware sliding-window benchmarks. The endpoint—not AIPerf—is authoritative for UUID cache state: a UUID-only reference may be valid because another run, process, or prewarm populated it. Dataset-authored UUIDs therefore pass through unchanged, while this opt-in path strips only content AIPerf previously observed. Operators still need adequate cache capacity and session-to-replica affinity.

Effect

Full image: {"image_url":{"url":"image"},"uuid":"u1"}

Cache-only reference: {"image_url":{"url":""},"uuid":"u1"}

What Change

• Normalize UUID-only images while preserving structural and record-copy semantics.
• Always pass authored UUIDs through on the Chat endpoint.
• Make --uuid-and-strip control only AIPerf-managed deduplication.
• Count cache-only references as logical images in image metrics.
• Keep automatic stripping limited to grouped single-turn sessions.

Test Plan

• 1,769 loader and endpoint tests passed; four skipped.
• 60 property and model tests passed.
• uv run pre-commit run --all-files

[github-actions]

Try out this PR

Quick install:

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

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@4cdc4240de02abbbd375c9a89fe327d84b84909a

---

Last updated for commit: 4cdc424 • Browse code

[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

Adds a vLLM-oriented --uuid-and-strip CLI flag and wiring: per-image UUIDs propagate through dataset loader and models, optional in endpoint config, dataset conversion can blank repeated-image contents per conversation, and ChatEndpoint message construction emits uuid fields and optionally strips repeated image URLs.

UUID-Keyed Multimodal Cache Support

Layer / File(s)	Summary
Configuration & Defaults src/aiperf/common/config/config_defaults.py, src/aiperf/common/config/endpoint_config.py	Adds EndpointDefaults.UUID_AND_STRIP = False and registers EndpointConfig.uuid_and_strip with CLI flag --uuid-and-strip (defaults from the new default).
Endpoint Runtime Info src/aiperf/common/models/model_endpoint_info.py	EndpointInfo gains uuid_and_strip: bool, populated from user config.
Data Shape src/aiperf/common/models/dataset_models.py, src/aiperf/dataset/loader/models.py	Image gains uuids: list[str] with a post-validator enforcing length matches contents; SingleTurn gains optional `image_uuids: list[str]
Plural-field Conversion src/aiperf/dataset/loader/mixins.py	When converting plural image fields, per-image UUID lists from loader input are threaded into constructed Image(..., uuids=...) instances.
Dataset Conversion / Deduplication src/aiperf/dataset/loader/single_turn.py	convert_to_conversations reads self.user_config.endpoint.uuid_and_strip; when enabled it calls _dedup_repeated_images_inplace which, per conversation, blanks image.contents[i] for later occurrences of a seen UUID while preserving Image.uuids.
Endpoint Message Construction src/aiperf/endpoints/openai_chat.py	format_payload now computes uuid_and_strip and delegates to _build_messages; message construction refactored (_create_messages, _set_message_content) and a new _append_image_parts emits image_url parts. Behavior: when uuid_and_strip is false or an image has no UUID, legacy non-empty-URL-only parts are emitted; when true and a UUID exists, emitted image_url includes uuid and image_url.url may be empty (for stripped repeats).
Fixtures / Test Helpers tests/unit/endpoints/conftest.py	create_request_info() gains is_final_turn and url_index parameters and forwards them into RequestInfo.
Tests tests/unit/dataset/loader/test_single_turn.py, tests/unit/endpoints/test_chat_endpoint_multimodal_cache_mode.py, tests/unit/endpoints/test_openai_chat_completions.py	Adds tests verifying image_uuids propagation and validation, dataset deduping/strip behavior, ChatEndpoint uuid_and_strip off/on wire shapes (including empty-URL-on-repeat and raw_messages bypass), and updates a mock image in an existing test to include an empty uuids list.
Documentation docs/cli-options.md	Documents aiperf profile --uuid-and-strip semantics, vLLM-only note, and operational note about sizing --mm-processor-cache-gb to cover the working set.

🎯 4 (Complex) | ⏱️ ~45 minutes

🐇 I nibble bytes then hop away,
First send the picture, then I play—
UUIDs stay snug, URLs go light,
Cached hugs in the quiet night.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 64.58% 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
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.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly matches the main change: adding image UUID-based cache reuse support for endpoints.

---

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

[codecov]

Codecov Report

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

Files with missing lines	Patch %	Lines
src/aiperf/dataset/loader/single_turn.py	88.23%	1 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

[github-actions]

Fern Docs Preview: https://nvidia-preview-019f8275-0974-7504-be74-51d541efd15a.docs.buildwithfern.com/aiperf/dev

[coderabbitai]

♻️ Duplicate comments (2)

src/aiperf/endpoints/openai_chat.py (2)

110-112: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Preserve strip key on raw_messages path so final-turn cleanup still executes.

Line 111 returns strip_key=None; then Lines 91-92 cannot evict session state on final turns, leaving stale UUID tracker entries behind.

Proposed fix

         if turns[-1].raw_messages is not None:
-            return turns[-1].raw_messages, None
+            strip_key: str | None = (
+                request_info.x_correlation_id
+                if uuid_and_strip and request_info.x_correlation_id
+                else None
+            )
+            return turns[-1].raw_messages, strip_key

Also applies to: 91-92

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/aiperf/endpoints/openai_chat.py` around lines 110 - 112, The early return
that yields raw_messages currently returns None for strip_key which prevents the
final-turn cleanup from running; update the return(s) that return
turns[-1].raw_messages so they return the existing strip_key value instead of
None (e.g., return turns[-1].raw_messages, strip_key), and make the same change
for the other similar return path near lines 91-92 to ensure session
eviction/UUID tracker cleanup still executes.

---

134-138: ⚠️ Potential issue | 🟠 Major | 🏗️ Heavy lift

Commit UUID strip-state only after a successful send boundary.

Line 137 updates _mm_uuid_sessions while building payload. If the request fails/cancels before ingestion, a retry may strip URLs for UUIDs the backend never primed.

Suggested direction

-        if strip_key is not None and newly_emitted:
-            self._mm_uuid_sessions.setdefault(strip_key, set()).update(newly_emitted)
-        return messages, strip_key
+        # Defer commit until transport reports successful request completion.
+        # Return newly_emitted as commit-delta for caller-managed success hook.
+        return messages, strip_key, newly_emitted

+# On successful completion (outside payload formatting), then:
+self._mm_uuid_sessions.setdefault(strip_key, set()).update(newly_emitted)

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/aiperf/endpoints/openai_chat.py` around lines 134 - 138, The code
currently mutates _mm_uuid_sessions during payload building (using strip_key and
newly_emitted), which can mark UUIDs as primed even if the request fails; change
the flow so that the update to self._mm_uuid_sessions.setdefault(strip_key,
set()).update(newly_emitted) happens only after a confirmed successful
send/ingest boundary. Concretely, have _build_payload (or the function returning
messages, strip_key) return the newly_emitted set (or a success callback token)
rather than applying the update, and then perform the update in the send/ingest
success path (the caller that actually posts the request), using the same
strip_key and newly_emitted identifiers so abandoned builds don’t commit strip
state.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Duplicate comments:
In `@src/aiperf/endpoints/openai_chat.py`:
- Around line 110-112: The early return that yields raw_messages currently
returns None for strip_key which prevents the final-turn cleanup from running;
update the return(s) that return turns[-1].raw_messages so they return the
existing strip_key value instead of None (e.g., return turns[-1].raw_messages,
strip_key), and make the same change for the other similar return path near
lines 91-92 to ensure session eviction/UUID tracker cleanup still executes.
- Around line 134-138: The code currently mutates _mm_uuid_sessions during
payload building (using strip_key and newly_emitted), which can mark UUIDs as
primed even if the request fails; change the flow so that the update to
self._mm_uuid_sessions.setdefault(strip_key, set()).update(newly_emitted)
happens only after a confirmed successful send/ingest boundary. Concretely, have
_build_payload (or the function returning messages, strip_key) return the
newly_emitted set (or a success callback token) rather than applying the update,
and then perform the update in the send/ingest success path (the caller that
actually posts the request), using the same strip_key and newly_emitted
identifiers so abandoned builds don’t commit strip state.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 7f37a432-b16b-44ec-9cc8-78adedb3fb3b

📥 Commits

Reviewing files that changed from the base of the PR and between 6b7e012 and b48a696.

📒 Files selected for processing (6)

• src/aiperf/common/models/dataset_models.py
• src/aiperf/dataset/loader/single_turn.py
• src/aiperf/endpoints/openai_chat.py
• tests/unit/dataset/loader/test_single_turn.py
• tests/unit/endpoints/test_chat_endpoint_multimodal_cache_mode.py
• tests/unit/endpoints/test_openai_chat_completions.py

✅ Files skipped from review due to trivial changes (1)

• src/aiperf/dataset/loader/single_turn.py

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

• src/aiperf/common/models/dataset_models.py
• tests/unit/dataset/loader/test_single_turn.py

[matthewkotila]

Reviewed against origin/main, validated against the actual files, and reproduced wire output against the in-repo mock server. Also validated end-to-end against real vLLM (Qwen/Qwen3-VL-2B-Instruct, --mm-processor-cache-gb 4, --enforce-eager, sliding-window dataset, 4 turns, 1 repeat): TTFT avg 397.80 → 61.36 ms (−85%), TTFT max 1,404.44 → 76.58 ms (−95%), end-to-end duration 12.55 → 9.03 s (−28%). The feature delivers the win the PR description claims. The dedup-at-load-time architecture is clean and the default-off path is byte-identical to main. One coverage gap stands out — the equivalent multi-turn JSONL row shape is unhooked.

Suggested fix order (with rough complexity):

1. 🟡 F1 — multi_turn JSONL loader doesn't apply load-time dedup; wire emits uuid keys but no url strip
2. 🟢 F2 — consider treating "" as an opt-out sentinel on the wire (matches dedup; lets parallel-array form mix UUIDed and non-UUIDed images)
3. 🟢 F3 — first-occurrence empty-content treated as cache-served (dedup signal collision)
4. 🟢 F4 — --uuid-and-strip silently no-ops on non-chat endpoints
5. 🟢 F5 — unused is_final_turn / url_index conftest params
6. 🟡 F6 — flag name describes implementation, not user intent (rename consideration)

Working well: stateless endpoint, strict 1:1 alignment validation at both layers, default-off backward compatibility verified. Pre-deduping at load time means the hot path does no work and there's no concurrency bookkeeping to worry about.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[devin-ai-integration]

Devin Review found 0 new potential issues.

[ajcasagrande]

Overall: Clean, narrowly-scoped, well-guarded opt-in feature. Pydantic validators + the config-level chat-endpoint gate + the multi_turn NotImplementedError are all appropriate, and the 95 targeted unit tests pass. Verified end-to-end with the real aiperf CLI against a request-logging server: dedup and cache-only-reference emission behave exactly as documented — repeated-within-turn retained, cross-turn stripped to url:"", explicit cache-only references passed through.

Two Low-severity findings below, both confirmed at runtime and independently adversarially re-verified. Neither is a correctness bug or a merge blocker.

Fix order (optional polish): F1 (metrics semantics) → F2 (silent degradation). Both are documentation/UX, not code-correctness.

Working well: strict length validation on Image.uuids, the narrow _extend_image_parts override that leaves audio/video untouched, and the hard gates preventing use outside chat + single_turn.

[ajcasagrande]

Approving. Re-verified on ed29c63f (current head) — both findings from my earlier review are resolved, and the one outstanding bot flag is obsolete.

Findings resolved:

• F1 (cache-only refs counted as images): documented as intended in docs/metrics-reference.md — num_images is now defined as logical image references, explicitly noting cache-only URLs count. Good by-design call.
• F2 (flag-off UUID passthrough): fixed. --uuid-and-strip now gates only AIPerf-side dedup; authored UUIDs (including cache-only refs) always pass through on the chat endpoint. Verified Image(uuids=[...]) with strip off renders {"url": "", "uuid": ...}, locked in by test_authored_uuids_pass_through_when_strip_disabled.

CodeRabbit "Major" lifecycle concern: obsolete — _mm_uuid_sessions per-request state is gone; stripping is a deterministic load-time pass in the loader, so retry-re-strip and unbounded-growth no longer apply.

Verification: 97/97 pass across the touched test files (chat cache-mode, stripped-media, endpoint validator, single_turn, multi_turn).

Working well: strict UUID length/empty-string validation, the narrow _extend_image_parts override that leaves audio/video untouched, and the hard gates keeping this to chat + single_turn.