Add output token replay annotations#1063
Conversation
Try out this PRQuick install: pip install --upgrade --force-reinstall git+https://github.com/ai-dynamo/aiperf.git@2d2accb348cac368623ccb0c417e88e5c2a5c6d6Recommended 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@2d2accb348cac368623ccb0c417e88e5c2a5c6d6Last updated for commit: |
WalkthroughAdds output token replay support to the trace-based dataset loader. A new ChangesOutput Token Replay Feature
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. Comment |
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 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/loader/output_replay.py`:
- Around line 13-23: The pre-commit check is failing due to missing docstrings
(D103) on three public functions: effective_replay_key,
output_replay_id_annotation, and merge_output_replay_annotation. Add docstrings
to each of these functions immediately after their function definition lines,
describing what each function does, its parameters, and what it returns. The
docstrings should follow the standard Python docstring format (e.g.,
triple-quoted strings) to satisfy the ruff-baselined pre-commit requirement.
In `@tests/unit/dataset/loader/test_mooncake_payload_mode.py`:
- Around line 239-244: Add complete type annotations to the test methods. For
the test_replay_request_id_injects_output_replay_annotation method and the other
three test methods mentioned in the comment, add the return type annotation ->
None to each method signature and provide a proper type hint for the
mock_prompt_generator parameter (which should be annotated with the appropriate
fixture or Mock type). Ensure all four test methods have complete type hints for
both parameters and return types per coding guidelines.
- Around line 42-53: The pytest.raises match parameters contain regex
metacharacters (.*) that should be in raw strings to avoid RUF043 lint warnings.
In both test functions test_mooncake_trace_rejects_output_length_mismatch and
test_mooncake_trace_rejects_replay_tokens_without_output_length, add the r
prefix to the match string arguments so that the regex patterns are treated as
raw strings. Change match="output_length.*len" to match=r"output_length.*len"
and match="output_length.*required" to match=r"output_length.*required"
respectively.
🪄 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: b91c741e-da48-4094-9df2-27ef0cc26cda
📒 Files selected for processing (5)
docs/benchmark-modes/trace-replay.mdsrc/aiperf/dataset/loader/models.pysrc/aiperf/dataset/loader/mooncake_trace.pysrc/aiperf/dataset/loader/output_replay.pytests/unit/dataset/loader/test_mooncake_payload_mode.py
| def effective_replay_key( | ||
| request_id: str | None, | ||
| session_id: str | None, | ||
| turn_index: int, | ||
| line_index: int, | ||
| ) -> str: | ||
| if request_id is not None and request_id.strip(): | ||
| return request_id.strip() | ||
| if session_id is not None and session_id.strip(): | ||
| return f"{session_id.strip()}:{turn_index}" | ||
| return f"line:{line_index}" |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Add docstrings to exported helpers to unblock pre-commit.
ruff-baselined is currently failing on D103 for effective_replay_key (Line 13), output_replay_id_annotation (Line 26), and merge_output_replay_annotation (Line 30), so this change cannot merge until those docstrings are added.
🛠️ Proposed fix
def effective_replay_key(
request_id: str | None,
session_id: str | None,
turn_index: int,
line_index: int,
) -> str:
+ """Derive a stable replay key from request/session/line context.
+
+ Priority: request_id, then session_id:turn_index, then line:<line_index>.
+ """
if request_id is not None and request_id.strip():
return request_id.strip()
if session_id is not None and session_id.strip():
return f"{session_id.strip()}:{turn_index}"
return f"line:{line_index}"
def output_replay_id_annotation(replay_key: str) -> str:
+ """Format the replay-key annotation value for nvext.annotations."""
return f"{OUTPUT_REPLAY_ID_ANNOTATION_PREFIX}{replay_key}"
def merge_output_replay_annotation(
body: dict[str, Any] | None, replay_key: str | None
) -> dict[str, Any] | None:
+ """Merge output replay annotation into body['nvext']['annotations'].
+
+ Returns the original body unchanged when replay_key is None.
+ """
if replay_key is None:
return bodyBased on learnings from pipeline failures, pre-commit is failing specifically on D103 for these three public functions.
Also applies to: 26-27, 30-57
🧰 Tools
🪛 GitHub Actions: Pre-commit / 0_pre-commit.txt
[error] 13-13: check-ruff-baselined failed: ruff-baselined reported new violation [D103] in effective_replay_key — Missing docstring in public function. Hook exited with code 2.
🪛 GitHub Actions: Pre-commit / pre-commit
[error] 13-13: check-ruff-baselined failed (exit code 2). ruff violation [D103]: in effective_replay_key — Missing docstring in public function. Step: . .venv/bin/activate && python tools/ruff_baselined.py
🤖 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/loader/output_replay.py` around lines 13 - 23, The
pre-commit check is failing due to missing docstrings (D103) on three public
functions: effective_replay_key, output_replay_id_annotation, and
merge_output_replay_annotation. Add docstrings to each of these functions
immediately after their function definition lines, describing what each function
does, its parameters, and what it returns. The docstrings should follow the
standard Python docstring format (e.g., triple-quoted strings) to satisfy the
ruff-baselined pre-commit requirement.
Source: Pipeline failures
| with pytest.raises(ValidationError, match="output_length.*len"): | ||
| MooncakeTrace( | ||
| text_input="Hello", | ||
| output_length=2, | ||
| output_token_ids=[10, 11, 12], | ||
| ) | ||
|
|
||
|
|
||
| def test_mooncake_trace_rejects_replay_tokens_without_output_length() -> None: | ||
| with pytest.raises(ValidationError, match="output_length.*required"): | ||
| MooncakeTrace(text_input="Hello", output_token_ids=[10]) | ||
|
|
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor
🧩 Analysis chain
🏁 Script executed:
# First, locate the test file
find . -type f -name "test_mooncake_payload_mode.py" | head -5Repository: ai-dynamo/aiperf
Length of output: 212
🏁 Script executed:
# If found, read the relevant lines from the file
if [ -f tests/unit/dataset/loader/test_mooncake_payload_mode.py ]; then
sed -n '40,55p' tests/unit/dataset/loader/test_mooncake_payload_mode.py
fiRepository: ai-dynamo/aiperf
Length of output: 718
🏁 Script executed:
# Also check the full context around those lines
if [ -f tests/unit/dataset/loader/test_mooncake_payload_mode.py ]; then
wc -l tests/unit/dataset/loader/test_mooncake_payload_mode.py
fiRepository: ai-dynamo/aiperf
Length of output: 214
🌐 Web query:
Ruff RUF043 rule regex pytest raises match raw string
💡 Result:
The Ruff rule RUF043 (pytest-raises-ambiguous-pattern) identifies cases where a non-raw string is passed to the match parameter of pytest.raises while containing unescaped regex metacharacters [1]. Because pytest interprets the match argument as a regular expression, using a standard string can create ambiguity about whether the pattern is intended to be a literal string or a regex [1]. To resolve this rule, you should clarify your intent by choosing one of the following approaches: 1. If you intended the string to be a regex: Use a raw string literal (e.g., match=r"pattern") to explicitly signal that regex metacharacters are intentional [1]. 2. If you intended the string to be a literal match: Use re.escape to ensure the string is treated as a literal sequence (e.g., match=re.escape("pattern")) [1]. 3. If you prefer to escape manually: Use backslashes to escape any regex metacharacters within the standard string [1]. Currently, RUF043 does not provide an automated fix [2][3]. This is intentional because Ruff cannot programmatically determine whether the developer's original intent was to use a regex or a literal string, and choosing the wrong fix would alter the test's behavior [4][3]. Therefore, developers must manually choose the approach that matches their specific testing requirement [1].
Citations:
- 1: https://docs.astral.sh/ruff/rules/pytest-raises-ambiguous-pattern/
- 2: [
ruff] Ambiguous pattern passed topytest.raises()(RUF043) astral-sh/ruff#14966 - 3: Consider adding a fix for
pytest-raises-ambiguous-pattern(RUF043) astral-sh/ruff#16690 - 4: [
ruff] Stabilizepytest-raises-ambiguous-pattern(RUF043) astral-sh/ruff#16657
🏁 Script executed:
# Search for all pytest.raises with match parameter in this file
rg "pytest\.raises.*match=" tests/unit/dataset/loader/test_mooncake_payload_mode.py -A 1Repository: ai-dynamo/aiperf
Length of output: 691
Use raw regex literals in pytest.raises(match=...) assertions.
Ruff flags these patterns with RUF043 (pytest-raises-ambiguous-pattern), which can fail lint checks in CI. Both assertions at lines 42 and 49 need raw string prefixes to clarify that the regex metacharacters (.*) are intentional.
Suggested fix
- with pytest.raises(ValidationError, match="output_length.*len"):
+ with pytest.raises(ValidationError, match=r"output_length.*len"):
MooncakeTrace(
text_input="Hello",
output_length=2,
output_token_ids=[10, 11, 12],
)
- with pytest.raises(ValidationError, match="output_length.*required"):
+ with pytest.raises(ValidationError, match=r"output_length.*required"):
MooncakeTrace(text_input="Hello", output_token_ids=[10])🧰 Tools
🪛 Ruff (0.15.18)
[warning] 42-42: Pattern passed to match= contains metacharacters but is neither escaped nor raw
(RUF043)
[warning] 51-51: Pattern passed to match= contains metacharacters but is neither escaped nor raw
(RUF043)
🤖 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/loader/test_mooncake_payload_mode.py` around lines 42 -
53, The pytest.raises match parameters contain regex metacharacters (.*) that
should be in raw strings to avoid RUF043 lint warnings. In both test functions
test_mooncake_trace_rejects_output_length_mismatch and
test_mooncake_trace_rejects_replay_tokens_without_output_length, add the r
prefix to the match string arguments so that the regex patterns are treated as
raw strings. Change match="output_length.*len" to match=r"output_length.*len"
and match="output_length.*required" to match=r"output_length.*required"
respectively.
Source: Linters/SAST tools
| def test_replay_request_id_injects_output_replay_annotation( | ||
| self, | ||
| tmp_path: Path, | ||
| default_cfg: CLIConfig, | ||
| mock_prompt_generator, | ||
| ): |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor
🧩 Analysis chain
🏁 Script executed:
# First, locate and examine the test file
fd test_mooncake_payload_mode.pyRepository: ai-dynamo/aiperf
Length of output: 210
🏁 Script executed:
# Get an outline of the file to see the structure
ast-grep outline tests/unit/dataset/loader/test_mooncake_payload_mode.pyRepository: ai-dynamo/aiperf
Length of output: 1477
🏁 Script executed:
# Read the specific lines mentioned in the review comment (239-244 and others)
sed -n '235,250p' tests/unit/dataset/loader/test_mooncake_payload_mode.py
sed -n '260,275p' tests/unit/dataset/loader/test_mooncake_payload_mode.py
sed -n '296,310p' tests/unit/dataset/loader/test_mooncake_payload_mode.py
sed -n '320,335p' tests/unit/dataset/loader/test_mooncake_payload_mode.pyRepository: ai-dynamo/aiperf
Length of output: 2009
🏁 Script executed:
# Check the imports at the top of the file to understand fixture types
head -50 tests/unit/dataset/loader/test_mooncake_payload_mode.pyRepository: ai-dynamo/aiperf
Length of output: 1637
Add full type annotations to test method signatures.
The four test methods are missing -> None return type annotations and have an untyped mock_prompt_generator parameter. Per coding guidelines, all functions must have complete type hints for parameters and return types.
Suggested fix
def test_replay_request_id_injects_output_replay_annotation(
self,
tmp_path: Path,
default_cfg: CLIConfig,
- mock_prompt_generator,
-):
+ mock_prompt_generator: Mock,
+) -> None: def test_replay_key_uses_session_turn_index_without_request_id(
self,
tmp_path: Path,
default_cfg: CLIConfig,
- mock_prompt_generator,
-):
+ mock_prompt_generator: Mock,
+) -> None: def test_replay_key_uses_physical_line_index_without_request_or_session(
self,
tmp_path: Path,
default_cfg: CLIConfig,
- mock_prompt_generator,
-):
+ mock_prompt_generator: Mock,
+) -> None: def test_replay_annotation_is_injected_into_payload_mode_raw_payload(
self,
tmp_path: Path,
default_cfg: CLIConfig,
- mock_prompt_generator,
-):
+ mock_prompt_generator: Mock,
+) -> None:🤖 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/loader/test_mooncake_payload_mode.py` around lines 239 -
244, Add complete type annotations to the test methods. For the
test_replay_request_id_injects_output_replay_annotation method and the other
three test methods mentioned in the comment, add the return type annotation ->
None to each method signature and provide a proper type hint for the
mock_prompt_generator parameter (which should be annotated with the appropriate
fixture or Mock type). Ensure all four test methods have complete type hints for
both parameters and return types per coding guidelines.
Source: Coding guidelines
Summary
Validation
Summary by CodeRabbit
New Features
request_idfield for stable request identification.Documentation