Add timing metrics to offline data generation pipeline

[orestis-z]

Summary

• Instruments the offline data generation script with performance timing using time.perf_counter()
• Times each generate_hidden_states_async() call (vLLM request latency) and each shutil.move() call (file write latency) independently
• Adds live throughput to tqdm postfix (rps = samples/s, vllm = avg request ms, write = avg write ms)
• Logs a final timing summary with aggregate stats after pipeline completion
• Uses plain logging + tqdm (not speculators.metrics), per the approach agreed in #717

Companion to the response regeneration timing PR (#803). Both address INFERENG-8704.

Test plan

• Existing unit tests pass (tests/unit/data_generation/test_offline.py — 18 tests)
• ruff check and ruff format pass
• Manual run against a vLLM endpoint to verify tqdm postfix and final summary output

🤖 Generated with Claude Code

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 99fc64bd-12a5-4339-8102-44666525c256

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 Walkthrough

Walkthrough

Changes

The offline hidden-state generator now tracks shared success, error, request, and timing statistics across workers, displays derived progress metrics, and logs aggregate elapsed time and average latencies at completion.

Generation metrics

Layer / File(s)	Summary
Shared statistics setup scripts/data_generation_offline.py	Initializes shared run statistics and passes them into worker tasks.
Worker timing and progress reporting scripts/data_generation_offline.py	Measures vLLM and write durations, updates outcome counters, and reports rates and average latencies in the progress bar.
Final timing summary scripts/data_generation_offline.py	Logs total elapsed time and average vLLM and write durations after processing completes.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly matches the main change: adding timing metrics to the offline data generation pipeline.
Description check	✅ Passed	The description directly explains the timing instrumentation, live metrics, and final summary added by this PR.
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.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

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

[coderabbitai]

🧹 Nitpick comments (1)

scripts/data_generation_offline.py (1)

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

Remove redundant stats["requests"] counter.

The stats["requests"] counter is incremented only alongside stats["ok"] upon successful processing, meaning both keys will always hold the exact same value. You can simplify the dictionary and its usage by removing stats["requests"] entirely and just relying on stats["ok"].

• scripts/data_generation_offline.py#L276-L300: Remove the stats["requests"] += 1 increment and use stats["ok"] for the postfix average calculations (e.g., stats['ok'] / elapsed).
• scripts/data_generation_offline.py#L375-L382: Remove the "requests": 0 initialization from the stats dictionary.
• scripts/data_generation_offline.py#L428-L438: Update the final timing logs to check against and use stats["ok"] as the denominator instead of stats["requests"].

🤖 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 `@scripts/data_generation_offline.py` around lines 276 - 300, Remove the
redundant stats["requests"] counter and rely on stats["ok"] throughout: in
scripts/data_generation_offline.py lines 276-300, delete its increment and use
stats["ok"] for postfix rate and timing averages; in lines 375-382, remove its
initialization; in lines 428-438, use stats["ok"] for the denominator and
zero-check in final timing logs.

🤖 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 `@scripts/data_generation_offline.py`:
- Around line 276-300: Remove the redundant stats["requests"] counter and rely
on stats["ok"] throughout: in scripts/data_generation_offline.py lines 276-300,
delete its increment and use stats["ok"] for postfix rate and timing averages;
in lines 375-382, remove its initialization; in lines 428-438, use stats["ok"]
for the denominator and zero-check in final timing logs.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: ce1a8185-c4b1-4afb-9700-b36a622022c1

📥 Commits

Reviewing files that changed from the base of the PR and between c2d7a6c and 803f53a.

📒 Files selected for processing (1)

• scripts/data_generation_offline.py

[WindChimeRan]

the deletion of "requests": 0, break the code (this is from the fix to the coderabbit). : KeyError: 'requests' hangs the pipeline on the first successful sample

[speculatorsbot]

Clean, well-scoped instrumentation that aligns with the approach agreed in #717. Agree with @WindChimeRan's finding — the three stats['requests'] references in the finally block need to be replaced with stats['ok'] to fix the KeyError introduced by the second commit. Recommend approving once that's fixed.

🤖 Generated with Claude Code using the /pr-review skill

[speculatorsbot]

The stats['requests'] KeyError flagged by @WindChimeRan is fixed in 7e608eb — all three f-string references now correctly use stats['ok']. No other issues. Recommend approving.

🤖 Generated with Claude Code using the /pr-review skill

[WindChimeRan]

One semaphore issues. then LGTM

---

I found this script overcomplicated. e.g., some of the semaphore usages are over-engineering. need to clean it up later.

[WindChimeRan]

write_s excludes validation, which still holds the write semaphore.

Suggest change:

Suggested change

	write_s = time.perf_counter() - t_write
	if validate_outputs:
	def _load_and_check(
	path=target_hidden_states_path,
	tokens=item["input_ids"],
	):
	loaded = load_file(path)
	check_hidden_states(loaded, tokens)
	await asyncio.to_thread(_load_and_check)
	write_s = time.perf_counter() - t_write

[orestis-z]

Good catch on the semaphore scope, but I think the current placement is intentional — write_s is meant to measure shutil.move (the actual disk write), not the optional validation read-back.

The tqdm postfix and summary log label this as "write" / "avg file write", so it should reflect write latency specifically. Validation (load_file + check_hidden_states) is a read + integrity check — bundling it in would inflate the metric and make it misleading, especially since it's gated behind --validate-outputs.

The write semaphore limits concurrent disk I/O (both writes and validation reads), but it's a concurrency limiter, not a semantic boundary for timing. If validation timing is useful it could be its own metric in a follow-up.

[WindChimeRan]

LGTM. Recommend accept.

[mergify]

Merge Protections

🟢 Merge protection satisfied — ready to merge.

Show 1 satisfied protection

🟢 Require approval from approved reviewers list

All pull requests must have at least one approving review from a member of the approved reviewers list before merging.

• any of:
  • approved-reviews-by = shanjiaz
  • approved-reviews-by = dsikka
  • approved-reviews-by = fynnsu
  • approved-reviews-by = orestis-z
  • approved-reviews-by = rahul-tuli

[shanjiaz]

Thanks for adding this!

[rahul-tuli]

Q: Won't this just measure the rename time? Since the server has written the hidden_states_already?