Decode prompt token trace payloads#455
Open
SunnySoldier357 wants to merge 6 commits into
Open
Conversation
Hydrate prompt_token_ids from Fireworks tracing payloads so RemoteRolloutProcessor can pass token-native prompt IDs through assistant turn metadata. Co-authored-by: Cursor <cursoragent@cursor.com>
|
You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard. |
Move the per-payload binary deserializers out of adapters/ into a dependency-light eval_protocol/tracing package: a PayloadType StrEnum, DecodedPayload, and a decode_payloads registry (master decode) usable without EvaluationRow/rollout machinery. Refactor FireworksTracingAdapter to use the registry instead of three copy-pasted decode blocks, and decode pti/v1 as zstd(JSON int array) to match the gateway. Adds tests/tracing and a README. Co-authored-by: Cursor <cursoragent@cursor.com>
Replace the 4-line narrating comment over the payload-decode block with a single intent line; the code already shows what it does. Co-authored-by: Cursor <cursoragent@cursor.com>
Export only PayloadType, DecodedPayload, and decode_payloads/decode_payload/ decode_trace from the package __init__. The per-type decoders and PAYLOAD_DECODERS are internal building blocks, still reachable via submodules. Co-authored-by: Cursor <cursoragent@cursor.com>
… module The shared helper had a single caller (prompt_token_ids); logprobs/router_replay already inline the same base64+zstd step. Inline it there too for consistency and remove the dedicated _decompress.py file. Co-authored-by: Cursor <cursoragent@cursor.com>
Replace the Dict[str, Any] `extras` field (only ever holding logprobs token_ids) with a typed `token_ids: Optional[List[int]]` field. Callers now get a real type (dp.token_ids) instead of Any from extras.get(...). Update adapter, tests, README. Co-authored-by: Cursor <cursoragent@cursor.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
eval_protocol/tracing/package for decoding Fireworks tracing-gateway payloads, independent ofEvaluationRow/rollout machinery (stdlib +zstandardonly). It is the "folder you can point at" for extracting gateway payloads even without using EP for rollouts.PayloadTypeStrEnum+DecodedPayload+ a masterdecode_payloads(payloads)registry that runs every decoder over a trace'spayloadsin one call (plusdecode_payload()anddecode_trace()). Adding a payload type is one decoder + onePAYLOAD_DECODERSentry.pti/v1aszstd(JSON int array)to match the gateway (fw-ai/fireworks#29923). Thelp/v1andr3/v1decoders are moved verbatim.FireworksTracingAdapterto usedecode_payloadsinstead of three copy-pasted decode blocks;execution_metadata.extrakeys are unchanged. Removedadapters/{pti,lp,r3}_deserializer.py.tests/tracing/test_registry.pyand a packageREADME.md.Test plan
conda run -n pilot python -m pytest tests/tracing/test_registry.py tests/adapters/test_fireworks_tracing_prompt_token_ids.py tests/adapters/test_fireworks_tracing_logprobs.py tests/adapters/test_lp_deserializer.py tests/adapters/test_r3_deserializer.py(pass; one pre-existing skip that needs the mono serializer).accounts/pyroworks-dev/deployments/malaysia2-intended-butterfly:scripts/test_remote_rollout_prompt_token_ids.pyhydratesassistant_turn_payloads[0].prompt_token_idswith 13 ids.prompt_token_idsandlogprobsdecoded live through the new registry;router_replaycovered by unit tests (deployment has MoE stats disabled).Requires fw-ai/fireworks#29923 (the gateway must emit
pti/v1aszstd(JSON)); land together.Note
Medium Risk
Refactors trace payload decoding used by rollouts and eval rows; behavior is intended to be backward-compatible for existing extra keys, but incorrect decoding or merge logic could affect training/eval consumers of logprobs, routing, and prompt token IDs.
Overview
Introduces
eval_protocol.tracing, a small stdlib +zstandardpackage that decodes Fireworks tracing-gateway out-of-band payloads (prompt_token_ids/pti/v1,logprobs/lp/v1,router_replay/r3/v1) viadecode_payloads,decode_payload, anddecode_trace, with per-payload error isolation and a registry for new types.FireworksTracingAdapterno longer inlines LP/R3 decode logic; it callsdecode_payloadsand maps results onto the sameexecution_metadata.extrakeys as before, plus newprompt_token_idsandprompt_token_ids_metadata. LP/R3 parsing moves fromadapters/*_deserializerintoeval_protocol.tracing(those adapter deserializers are removed).Rollout path:
_merge_payloads_into_longest_rownow carriesprompt_token_ids(and metadata) intoassistant_turn_payloadsper turn. Tests and a live E2E script (scripts/test_remote_rollout_prompt_token_ids.py) cover registry behavior and remote rollout hydration. Depends on gateway emittingpti/v1as zstd(JSON).Reviewed by Cursor Bugbot for commit 38fe102. Bugbot is set up for automated code reviews on this repo. Configure here.