feat: token-level capture for RL training (Switchyard ↔ NeMo Gym)#63
feat: token-level capture for RL training (Switchyard ↔ NeMo Gym)#63drifold wants to merge 6 commits into
Conversation
…capture Signed-off-by: drifold <drifold@nvidia.com>
Signed-off-by: drifold <drifold@nvidia.com>
Signed-off-by: drifold <drifold@nvidia.com>
WalkthroughAdds configurable token capture for routed LLM requests, including session propagation, token-record persistence and retrieval, launcher integration, target-specific parameters, synthetic streaming, and extensive tests. ChangesToken capture contracts and configuration
Capture processing and retrieval
CLI and launcher integration
End-to-end capture validation
Estimated code review effort: 4 (Complex) | ~60 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
switchyard/cli/switchyard_cli.py (1)
905-916: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winCLI help text: grammar and scope inaccuracy.
"an
token_capture_engine" should be "atoken_capture_engine"; also, the flag applies toserveas well (per_cmd_serve'sresolve_rl_log_dirusage), not justlaunchsessions as stated.✏️ Proposed fix
- "file per request/response pair) for `launch` sessions. When the " - "route bundle declares an `token_capture_engine` target, traces are " + "file per request/response pair) for `launch` and `serve` sessions. " + "When the route bundle declares a `token_capture_engine` target, traces are "🤖 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 `@switchyard/cli/switchyard_cli.py` around lines 905 - 916, Update the help text for the --enable-rl-logging argument to use “a `token_capture_engine`” and describe the flag as applying to both launch and serve sessions, consistent with _cmd_serve’s resolve_rl_log_dir usage; preserve the existing trace and placement guidance.switchyard/lib/processors/rl_logging_response_processor.py (1)
192-199: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick winNewly public
format_trace_tool_choiceis missing a docstring.Renamed from private
_format_tool_choiceto a public helper but no docstring was added, unlike its siblingformat_trace_tools/build_trace_messages/build_trace_token_count. As per coding guidelines,switchyard/**/*.pycode requires "docstrings for public APIs."📝 Proposed fix
def format_trace_tool_choice(tool_choice: object) -> str: + """Extract the tool-choice ``type`` string, defaulting to ``"auto"``.""" if isinstance(tool_choice, str): return tool_choice🤖 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 `@switchyard/lib/processors/rl_logging_response_processor.py` around lines 192 - 199, Add a concise docstring to the public format_trace_tool_choice function describing its accepted tool-choice input and returned choice string, including the "auto" fallback behavior. Keep the existing logic unchanged and match the documentation style of format_trace_tools, build_trace_messages, and build_trace_token_count.Source: Coding guidelines
🧹 Nitpick comments (3)
tests/test_token_capture.py (1)
158-196: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse
pytest-mockinstead ofmonkeypatch.This is general mocking and should use the repository-standard
mockerfixture.Proposed change
+from pytest_mock import MockerFixture + def test_serve_installs_capture_pair_for_declaring_bundle( - monkeypatch: pytest.MonkeyPatch, + mocker: MockerFixture, tmp_path: Path, ) -> None: ... - monkeypatch.setattr(cli, "load_route_bundle_table", _fake_load) - monkeypatch.setattr(cli, "build_and_serve", lambda *a, **k: None) + mocker.patch.object(cli, "load_route_bundle_table", side_effect=_fake_load) + mocker.patch.object(cli, "build_and_serve")As per coding guidelines,
tests/**/*.pymust “usepytest-mockfor general mocking.”🤖 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/test_token_capture.py` around lines 158 - 196, Replace the monkeypatch fixture and its setattr calls in test_serve_installs_capture_pair_for_declaring_bundle with the repository-standard pytest-mock mocker fixture, using mocker.patch for cli.load_route_bundle_table and cli.build_and_serve while preserving the existing fake behavior and assertions.Source: Coding guidelines
switchyard/lib/endpoints/token_capture_endpoint.py (1)
66-75: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winBlocking file I/O inside an async route handler.
Path.is_dir(),session_dir.glob(...), andrecord_path.read_text()are synchronous disk calls executed directly on the event loop insideasync def get_session_completions. Under concurrent requests or a session with many records, this can stall other in-flight requests on the same worker.♻️ Proposed direction
- completions: list[dict[str, Any]] = [] - for record_path in sorted(session_dir.glob("*.json")): - try: - completions.append(json.loads(record_path.read_text())) - except (OSError, json.JSONDecodeError) as exc: - ... + def _read_all() -> list[dict[str, Any]]: + out: list[dict[str, Any]] = [] + for record_path in sorted(session_dir.glob("*.json")): + try: + out.append(json.loads(record_path.read_text())) + except (OSError, json.JSONDecodeError) as exc: + logger.warning(...) + return out + completions = await asyncio.to_thread(_read_all)🤖 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 `@switchyard/lib/endpoints/token_capture_endpoint.py` around lines 66 - 75, Update the async get_session_completions handler so all synchronous filesystem operations, including Path.is_dir(), session_dir.glob(), and record_path.read_text(), run off the event loop via the project’s established async file-I/O or thread-offloading mechanism. Preserve sorted record processing, JSON parsing, and warning behavior for unreadable records.switchyard/lib/processors/token_capture_response_processor.py (1)
149-169: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winBlocking file I/O runs inside the async request path.
_write_recordperforms synchronousopen/os.replace/os.chmodon every captured turn from withinprocess(), blocking the event loop. Mirrors the existingRlLoggingResponseProcessor._write_entrypattern, so not a regression, but worth moving off the loop (e.g.,asyncio.to_thread) given the "async-only code" guideline. As per coding guidelines,switchyard/**/*.pycode should be "async-only code."🤖 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 `@switchyard/lib/processors/token_capture_response_processor.py` around lines 149 - 169, The synchronous file operations in _write_record block the async request path. Update the token capture processor’s process flow to invoke _write_record through asyncio.to_thread (or an equivalent async offloading mechanism), preserving its existing arguments and record-writing behavior while keeping filesystem I/O off the event loop.Source: Coding guidelines
🤖 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 `@switchyard/lib/backends/llm_target.py`:
- Line 125: Reformat the ValueError message in the token-capture validation near
token_capture_engine so every physical line stays within the 100-character
limit, while preserving the existing error text and supported-values details.
- Around line 105-134: The llm_target_with_token_capture function drops backend
connection settings when rebuilding the LlmTarget. Preserve target.base_url,
target.api_key, and target.timeout_secs by forwarding each field into the new
LlmTarget alongside the existing properties.
In `@switchyard/lib/processors/token_capture_request_processor.py`:
- Around line 28-33: The session-ID trust check rejects valid custom re-run IDs,
causing token capture to silently drop them. In
switchyard/lib/processors/token_capture_request_processor.py:28-33 and :80-91,
and switchyard/cli/launchers/launch_intake_config.py:205-222, change the
launcher-to-capture flow so intake.session_id is explicitly validated or marked
as launcher-generated before capture_session_id_for_api_key() and
_resolve_session_id() use it; preserve rejection of untrusted credential-like
values without silently discarding valid custom session IDs.
In `@switchyard/lib/processors/token_capture_response_processor.py`:
- Around line 246-329: Update _synthesize_stream so
CompletionUsage.model_validate is guarded against partial or malformed usage
dictionaries; when validation fails, set final_usage to None and continue
yielding the synthesized stream normally.
In `@tests/test_token_capture.py`:
- Around line 667-668: Update the traversal test around the client request to
use an encoded parent segment (%2E%2E) so the path reaches the handler, and
create a sentinel record outside the sessions directory under tmp_path. Assert
the request still returns 404, ensuring a vulnerable path join cannot read the
sentinel and return 200.
---
Outside diff comments:
In `@switchyard/cli/switchyard_cli.py`:
- Around line 905-916: Update the help text for the --enable-rl-logging argument
to use “a `token_capture_engine`” and describe the flag as applying to both
launch and serve sessions, consistent with _cmd_serve’s resolve_rl_log_dir
usage; preserve the existing trace and placement guidance.
In `@switchyard/lib/processors/rl_logging_response_processor.py`:
- Around line 192-199: Add a concise docstring to the public
format_trace_tool_choice function describing its accepted tool-choice input and
returned choice string, including the "auto" fallback behavior. Keep the
existing logic unchanged and match the documentation style of
format_trace_tools, build_trace_messages, and build_trace_token_count.
---
Nitpick comments:
In `@switchyard/lib/endpoints/token_capture_endpoint.py`:
- Around line 66-75: Update the async get_session_completions handler so all
synchronous filesystem operations, including Path.is_dir(), session_dir.glob(),
and record_path.read_text(), run off the event loop via the project’s
established async file-I/O or thread-offloading mechanism. Preserve sorted
record processing, JSON parsing, and warning behavior for unreadable records.
In `@switchyard/lib/processors/token_capture_response_processor.py`:
- Around line 149-169: The synchronous file operations in _write_record block
the async request path. Update the token capture processor’s process flow to
invoke _write_record through asyncio.to_thread (or an equivalent async
offloading mechanism), preserving its existing arguments and record-writing
behavior while keeping filesystem I/O off the event loop.
In `@tests/test_token_capture.py`:
- Around line 158-196: Replace the monkeypatch fixture and its setattr calls in
test_serve_installs_capture_pair_for_declaring_bundle with the
repository-standard pytest-mock mocker fixture, using mocker.patch for
cli.load_route_bundle_table and cli.build_and_serve while preserving the
existing fake behavior and assertions.
🪄 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: b88bd523-4e03-4b97-bcf5-ad10301f09b9
📒 Files selected for processing (15)
switchyard/cli/launch_command.pyswitchyard/cli/launchers/claude_code_launcher.pyswitchyard/cli/launchers/codex_cli_launcher.pyswitchyard/cli/launchers/launch_intake_config.pyswitchyard/cli/launchers/openclaw_launcher.pyswitchyard/cli/route_bundle.pyswitchyard/cli/switchyard_cli.pyswitchyard/lib/backends/llm_target.pyswitchyard/lib/endpoints/token_capture_endpoint.pyswitchyard/lib/processors/rl_logging_response_processor.pyswitchyard/lib/processors/token_capture_request_processor.pyswitchyard/lib/processors/token_capture_response_processor.pytests/test_token_capture.pytests/test_token_capture_target.pytests/test_token_capture_translation.py
…t, traversal test Signed-off-by: drifold <drifold@nvidia.com>
|
Dafna: Could you take a look at this to see if they should be fixed? Feel free to leverage the possible fix or use anything you find appropriate.
|
|
@drifold can you please make sure CodeRabbit comments are addressed and resolve them if so? |
…y field Signed-off-by: Lin Jia <linj@nvidia.com>
…ovenance Signed-off-by: drifold <drifold@nvidia.com>
|
@linj-glitch - All fixed as suggested |
|
@ryan-lempka - done |
3 participants
What
Adds opt-in token-level capture to Switchyard's RL-logging path. When a route target declares token_capture_engine: vllm, Switchyard requests exact token data from vLLM, records one unified trace per model call (message text + prompt_token_ids / generation_token_ids / generation_log_probs), groups records by the proxy_x_session_id header, and exposes them via GET /v1/sessions/{session_id}/completions. Streaming-required harnesses (Claude Code) are supported by capturing a non-streaming upstream response and synthesizing the client stream. Implements the agreed Switchyard Token Capture design.
Why
Enables RL training of unmodified coding harnesses (Claude Code, Codex) through Switchyard. RL needs the training signal attached to the exact tokens the policy sampled — re-tokenizing text after the fact drifts the token boundaries and log-probs, making training silently off-policy. Capturing vLLM's exact prompt_token_ids / generation_token_ids / generation_log_probs at the proxy boundary preserves that behavior-policy fidelity. NeMo Gym maps its rollout ID to the session, retrieves the token-faithful records, and builds trajectories. Switchyard stays generic — no Gym dependency.
How tested
Checklist
Notes for reviewers
Summary by CodeRabbit
/v1/sessions/{session_id}/completionsendpoint.