Add --render-endpoint to prepare_data.py for vLLM-delegated tokenization (#652)#665
Add --render-endpoint to prepare_data.py for vLLM-delegated tokenization (#652)#665WindChimeRan wants to merge 15 commits into
Conversation
Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughAdds a vLLM render endpoint preprocessing path to the data generation pipeline. A new ChangesvLLM Render Endpoint Preprocessing
Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
|
The quality checks have failed. Please run |
Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (5)
scripts/prepare_data.py (1)
193-206: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
chat_template_kwargsis not exposed via the CLI.
load_and_preprocess_datasetacceptschat_template_kwargsand forwards it intobuild_dataset_from_render(e.g. forenable_thinking), butprepare_data.pynever passes it, so it is alwaysNonefrom the CLI. If template kwargs are meant to be usable with the render path, add a corresponding argument; otherwise this is fine to defer.🤖 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/prepare_data.py` around lines 193 - 206, The CLI path in prepare_data.py is dropping chat template options, so load_and_preprocess_dataset never receives chat_template_kwargs even though it forwards them to build_dataset_from_render. Expose a matching CLI argument in the prepare_data entrypoint, pass it through the load_and_preprocess_dataset call, and wire it into the existing argument parsing alongside the other dataset/render options so template-specific flags like enable_thinking can be used from the command line.src/speculators/data_generation/preprocessing.py (1)
686-714: 🚀 Performance & Scalability | 🔵 TrivialSerial per-row render requests will bottleneck large datasets.
Each conversation triggers a blocking
render_conversationcall inside a single-threaded loop, so throughput is gated by per-request latency (the PR notes the render path has no concurrency). For multi-million-row corpora this is materially slower than the paralleldataset.map(num_proc=...)default path. Consider batching or a boundedThreadPoolExecutorover rows (the render endpoint is I/O-bound), or document the expected throughput limitation for operators.🤖 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/speculators/data_generation/preprocessing.py` around lines 686 - 714, The per-row rendering in the preprocessing loop is still single-threaded and will bottleneck large datasets. Update the render path around the loop that calls _normalize_conversation, _adapt_conv_for_vllm, and render_conversation to use bounded concurrency (for example, a ThreadPoolExecutor) or batched requests so multiple conversations can render in parallel without overloading the endpoint. Keep the existing per-conversation error handling and dropped_count/log.error behavior intact when refactoring.tests/unit/data_generation/test_render_client.py (1)
84-86: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider asserting retry behavior is not triggered for client errors.
test_error_raisesonly verifies aRenderErroris raised on 500. Given retries wrap this function, a test confirming the request count for a non-retriable error (e.g. 400) would lock in the intended retry policy referenced in therender_client.pyreview.🤖 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/data_generation/test_render_client.py` around lines 84 - 86, The current `test_error_raises` only checks that `_call_render` raises `RenderError`, but it does not verify that client errors are not retried; update this test in `test_render_client.py` to assert the retry policy explicitly by using `_make_transport` with a non-retriable status like 400 and confirming the request is attempted only once. Keep the existing `RenderError` assertion, and add a request-count check around `_call_render` so the behavior of the retry wrapper in the render client stays locked in.tests/integration/datagen/test_render_truncation.py (1)
23-23: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMake the vLLM interpreter path configurable.
VLLM_PYTHONis a hardcoded developer-specific absolute path. The test will silently skip on any other machine via theshutil.whichguard, so it effectively only runs in one environment. The sibling multimodal test already supports aRENDER_ENDPOINTenv override; consider mirroring that here (e.g. readVLLM_PYTHONfrom the environment with this path as a fallback) so the test is portable.♻️ Suggested change
-VLLM_PYTHON = "/home/ran/workspace/vllm/.venv/bin/python" +VLLM_PYTHON = os.environ.get( + "VLLM_PYTHON", "/home/ran/workspace/vllm/.venv/bin/python" +)🤖 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/integration/datagen/test_render_truncation.py` at line 23, The hardcoded VLLM_PYTHON absolute path in test_render_truncation.py makes the test machine-specific and can cause it to skip elsewhere; update the test setup so VLLM_PYTHON is read from an environment override with the current path as a fallback, similar to the sibling multimodal test’s RENDER_ENDPOINT handling, and keep the existing shutil.which guard working against the resolved value.tests/integration/datagen/test_render_multimodal.py (1)
35-38: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAvoid hardcoded developer path and
/tmpconstants.Two portability/hygiene issues here:
VLLM_PYTHONis a developer-specific absolute path (same as the truncation test); prefer reading it from the environment with a fallback.MEDIA_DIR/IMG_PATHhardcode/tmp(Ruff S108) and the generated image is never cleaned up. Since theimagefixture is module-scoped, consider deriving the path fromtmp_path_factoryso the file is isolated and auto-removed.🤖 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/integration/datagen/test_render_multimodal.py` around lines 35 - 38, Replace the hardcoded developer-specific runtime and temp paths in test_render_multimodal by sourcing VLLM_PYTHON from an environment variable with a safe fallback, and stop using fixed /tmp constants for MEDIA_DIR and IMG_PATH. Update the multimodal test setup and any related fixture usage to derive the image location from tmp_path_factory so the generated file is isolated per run and cleaned up automatically, while keeping MODEL and the image fixture wiring intact.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.
Inline comments:
In `@src/speculators/data_generation/render_client.py`:
- Around line 49-53: The response parsing in render_client.py currently lets a
missing token_ids field raise a bare KeyError, which escapes the RenderError
handling used by build_dataset_from_render. Update the logic around resp.json()
and the return in the render client function to validate required fields and
catch malformed payloads, then re-raise any missing/invalid response data as
RenderError so the caller’s per-row exception handling can drop only the bad
row.
- Around line 44-47: The render client currently raises RenderError for every
non-OK response, which causes `@with_retries` to replay deterministic 4xx
failures. Update the request handling in render_client.py’s render flow so that
client errors (4xx) are treated as non-retriable, either by raising a distinct
exception or otherwise bypassing the retry decorator, while keeping transient
transport/5xx cases retriable. Make sure the behavior is applied in the same
place where resp.status_code is checked and RenderError is constructed.
---
Nitpick comments:
In `@scripts/prepare_data.py`:
- Around line 193-206: The CLI path in prepare_data.py is dropping chat template
options, so load_and_preprocess_dataset never receives chat_template_kwargs even
though it forwards them to build_dataset_from_render. Expose a matching CLI
argument in the prepare_data entrypoint, pass it through the
load_and_preprocess_dataset call, and wire it into the existing argument parsing
alongside the other dataset/render options so template-specific flags like
enable_thinking can be used from the command line.
In `@src/speculators/data_generation/preprocessing.py`:
- Around line 686-714: The per-row rendering in the preprocessing loop is still
single-threaded and will bottleneck large datasets. Update the render path
around the loop that calls _normalize_conversation, _adapt_conv_for_vllm, and
render_conversation to use bounded concurrency (for example, a
ThreadPoolExecutor) or batched requests so multiple conversations can render in
parallel without overloading the endpoint. Keep the existing per-conversation
error handling and dropped_count/log.error behavior intact when refactoring.
In `@tests/integration/datagen/test_render_multimodal.py`:
- Around line 35-38: Replace the hardcoded developer-specific runtime and temp
paths in test_render_multimodal by sourcing VLLM_PYTHON from an environment
variable with a safe fallback, and stop using fixed /tmp constants for MEDIA_DIR
and IMG_PATH. Update the multimodal test setup and any related fixture usage to
derive the image location from tmp_path_factory so the generated file is
isolated per run and cleaned up automatically, while keeping MODEL and the image
fixture wiring intact.
In `@tests/integration/datagen/test_render_truncation.py`:
- Line 23: The hardcoded VLLM_PYTHON absolute path in test_render_truncation.py
makes the test machine-specific and can cause it to skip elsewhere; update the
test setup so VLLM_PYTHON is read from an environment override with the current
path as a fallback, similar to the sibling multimodal test’s RENDER_ENDPOINT
handling, and keep the existing shutil.which guard working against the resolved
value.
In `@tests/unit/data_generation/test_render_client.py`:
- Around line 84-86: The current `test_error_raises` only checks that
`_call_render` raises `RenderError`, but it does not verify that client errors
are not retried; update this test in `test_render_client.py` to assert the retry
policy explicitly by using `_make_transport` with a non-retriable status like
400 and confirming the request is attempted only once. Keep the existing
`RenderError` assertion, and add a request-count check around `_call_render` so
the behavior of the retry wrapper in the render client stays locked in.
🪄 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: Pro
Run ID: 20ed4253-978d-49dd-8fcb-31df9476f1f2
📒 Files selected for processing (7)
scripts/prepare_data.pysrc/speculators/data_generation/preprocessing.pysrc/speculators/data_generation/render_client.pytests/integration/datagen/test_render_multimodal.pytests/integration/datagen/test_render_truncation.pytests/unit/data_generation/test_build_dataset_from_render.pytests/unit/data_generation/test_render_client.py
rahul-tuli
left a comment
rahul-tuli
left a comment
There was a problem hiding this comment.
Thank you for working on this @WindChimeRan, the diff looks good to me pending a few small comments, could you respond and update with latest main, and I will approve the workflows to run, great job on this!
|
This pull request has merge conflicts that must be resolved before it can be |
Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
…ssing # Conflicts: # src/speculators/data_generation/preprocessing.py
|
Thanks for the review! All addressed @rahul-tuli @shanjiaz |
Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
|
LGTM, pending CI! |
Fixes the failing quality-checks CI job (ruff format --check). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WTsu9vFfCfPsTTE8PgfEyi Signed-off-by: Ranran Haoran Zhang <ranranhaoranzhang@gmail.com>
|
This pull request has merge conflicts that must be resolved before it can be |
Resolve conflicts from RFC vllm-project#583 (vllm-project#690): - prepare_data.py: drop duplicate token_freq_path (upstream moved it above the overwrite guard); keep chat_template_kwargs. - preprocessing.py: docstring only; keep render-path wording and add upstream's allow_empty_output entry. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JkFcT8oXWTMmSWc93JuDbd
Head branch was pushed to by a user without write access
|
this has design conflicts with #774 |
|
This pull request has merge conflicts that must be resolved before it can be |
4 participants
Purpose
Adds a
--render-endpointflag toprepare_data.pythat delegates tokenization to a running vLLM render server instead of doing it locally with HFapply_chat_template.The --render-endpoint should point at the same vllm serve instance used for hidden-state extraction. The render endpoint runs on the same server with zero GPU overhead. A separate vllm launch render instance is only needed for GPU-less offline preprocessing.
vLLM prerequisite PR has been merged, bun not released yet: vllm-project/vllm#46846
When
--render-endpointis set, raw conversations are sent toPOST /v1/chat/completions/renderand the returnedtoken_ids+loss_maskare used directly, eliminating the need for local chat-template rendering, assistant-pattern detection, and regex-based loss masking.Motivation: online training already requires a vLLM instance, so tokenization should happen once on the vLLM side rather than being duplicated between HF and vLLM. See RFC #652 for discussion.
Tokenization
{% generation %}): token_ids and mask both from /render → 1 tokenization, server-side, no local re-tokenization.Corner cases
The current renderer endpoint + fallback may not fully guarantee the correctness of some corner cases,
e.g., #659 , #645, #690 need to investigate further
Tests
Limitation:
Checklist
I have filled in: