Skip to content

Commit 97ffb52

Browse files
SunnySoldier357cursoragent
SunnySoldier357
and
cursoragent
committed
Extract gateway payload decoders into standalone eval_protocol.tracing
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>
1 parent aee752b commit 97ffb52

15 files changed

Lines changed: 535 additions & 184 deletions

‎eval_protocol/adapters/fireworks_tracing.py‎

Lines changed: 38 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -15,10 +15,8 @@
1515
import os
1616

1717
from eval_protocol.models import EvaluationRow, InputMetadata, ExecutionMetadata, Message
18+
from eval_protocol.tracing import PayloadType, decode_payloads
1819
from .base import BaseAdapter
19-
from .lp_deserializer import decompress_and_parse_lp
20-
from .pti_deserializer import decompress_and_parse_pti
21-
from .r3_deserializer import decompress_and_parse_r3
2220
from .utils import extract_messages_from_data
2321
from ..common_utils import get_user_agent
2422

@@ -103,60 +101,45 @@ def convert_trace_dict_to_evaluation_row(
103101
):
104102
break # Break early if we've found all the metadata we need
105103

106-
# Extract router replay payloads when present
104+
# Decode out-of-band gateway payloads (router replay, logprobs, prompt
105+
# token ids) via the standalone tracing decoder registry, then map the
106+
# decoded values onto the row. Format/decoding lives in
107+
# ``eval_protocol.tracing``; this adapter only does EvaluationRow glue.
107108
payloads = trace.get("payloads")
108109
if isinstance(payloads, dict):
109-
router_replay = payloads.get("router_replay")
110-
if isinstance(router_replay, dict) and router_replay.get("data"):
111-
try:
112-
matrices, r3_meta = decompress_and_parse_r3(router_replay["data"])
113-
if execution_metadata.extra is None:
114-
execution_metadata.extra = {}
115-
execution_metadata.extra["routing_matrices"] = matrices
116-
execution_metadata.extra["routing_metadata"] = r3_meta
117-
except Exception as e:
118-
logger.warning("Failed to decompress R3 payload for trace %s: %s", trace.get("id"), e)
119-
120-
logprobs_payload = payloads.get("logprobs")
121-
if isinstance(logprobs_payload, dict) and logprobs_payload.get("data"):
122-
try:
123-
logprobs, token_ids, lp_meta = decompress_and_parse_lp(logprobs_payload["data"])
124-
if execution_metadata.extra is None:
125-
execution_metadata.extra = {}
126-
execution_metadata.extra["completion_logprobs"] = logprobs
127-
if token_ids is not None:
128-
execution_metadata.extra["completion_token_ids"] = token_ids
129-
execution_metadata.extra["logprobs_metadata"] = lp_meta
130-
131-
for i in range(len(messages) - 1, -1, -1):
132-
if messages[i].role == "assistant":
133-
content_entries = [{"logprob": lp} for lp in logprobs]
134-
if token_ids is not None:
135-
for entry, tid in zip(content_entries, token_ids):
136-
entry["token_id"] = tid
137-
messages[i].logprobs = {"content": content_entries}
138-
break
139-
except Exception as e:
140-
logger.warning(
141-
"Failed to decompress logprobs payload for trace %s: %s",
142-
trace.get("id"),
143-
e,
144-
)
145-
146-
prompt_ids_payload = payloads.get("prompt_token_ids")
147-
if isinstance(prompt_ids_payload, dict) and prompt_ids_payload.get("data"):
148-
try:
149-
prompt_token_ids, pti_meta = decompress_and_parse_pti(prompt_ids_payload["data"])
150-
if execution_metadata.extra is None:
151-
execution_metadata.extra = {}
152-
execution_metadata.extra["prompt_token_ids"] = prompt_token_ids
153-
execution_metadata.extra["prompt_token_ids_metadata"] = pti_meta
154-
except Exception as e:
155-
logger.warning(
156-
"Failed to decompress prompt token IDs payload for trace %s: %s",
157-
trace.get("id"),
158-
e,
159-
)
110+
decoded = decode_payloads(
111+
payloads,
112+
on_error=lambda pt, e: logger.warning(
113+
"Failed to decode %s payload for trace %s: %s", pt.value, trace.get("id"), e
114+
),
115+
)
116+
if decoded and execution_metadata.extra is None:
117+
execution_metadata.extra = {}
118+
119+
if (dp := decoded.get(PayloadType.ROUTER_REPLAY)) is not None:
120+
execution_metadata.extra["routing_matrices"] = dp.value
121+
execution_metadata.extra["routing_metadata"] = dp.metadata
122+
123+
if (dp := decoded.get(PayloadType.LOGPROBS)) is not None:
124+
logprobs = dp.value
125+
token_ids = dp.extras.get("token_ids")
126+
execution_metadata.extra["completion_logprobs"] = logprobs
127+
if token_ids is not None:
128+
execution_metadata.extra["completion_token_ids"] = token_ids
129+
execution_metadata.extra["logprobs_metadata"] = dp.metadata
130+
131+
for i in range(len(messages) - 1, -1, -1):
132+
if messages[i].role == "assistant":
133+
content_entries = [{"logprob": lp} for lp in logprobs]
134+
if token_ids is not None:
135+
for entry, tid in zip(content_entries, token_ids):
136+
entry["token_id"] = tid
137+
messages[i].logprobs = {"content": content_entries}
138+
break
139+
140+
if (dp := decoded.get(PayloadType.PROMPT_TOKEN_IDS)) is not None:
141+
execution_metadata.extra["prompt_token_ids"] = dp.value
142+
execution_metadata.extra["prompt_token_ids_metadata"] = dp.metadata
160143

161144
return EvaluationRow(
162145
messages=messages,

‎eval_protocol/adapters/pti_deserializer.py‎

Lines changed: 0 additions & 98 deletions
This file was deleted.

‎eval_protocol/tracing/README.md‎

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
# `eval_protocol.tracing` — Fireworks tracing-gateway payload decoders
2+
3+
Standalone helpers for decoding the out-of-band **payloads** the Fireworks
4+
tracing gateway stores alongside a trace (prompt token IDs, completion logprobs,
5+
router-replay routing matrices).
6+
7+
This package is intentionally self-contained: it depends only on the stdlib and
8+
`zstandard`. It does **not** import `EvaluationRow`, rollout processors, or any
9+
other Eval Protocol machinery, so you can use it even if you are not using EP for
10+
rollouts — just point at it for extracting gateway payloads.
11+
12+
## What is a "payload"?
13+
14+
When you read a trace with payloads included:
15+
16+
```
17+
GET {gateway}/v1/traces?rollout_id=...&include_payloads=true
18+
```
19+
20+
each trace carries a `payloads` object like:
21+
22+
```json
23+
{
24+
"payloads": {
25+
"prompt_token_ids": {
26+
"manifest": { "PayloadVersion": "pti/v1", "...": "..." },
27+
"data": "<base64 of zstd-compressed bytes>"
28+
},
29+
"logprobs": { "manifest": { "PayloadVersion": "lp/v1" }, "data": "..." },
30+
"router_replay": { "manifest": { "PayloadVersion": "r3/v1" }, "data": "..." }
31+
}
32+
}
33+
```
34+
35+
The `data` field is `base64(zstd(raw_bytes))`. Each payload type has its own
36+
`raw_bytes` encoding (`pti/v1` is a JSON int array; `lp/v1` and `r3/v1` are packed
37+
binary). This package hides all of that.
38+
39+
## Usage
40+
41+
Decode everything at once (the common case):
42+
43+
```python
44+
from eval_protocol.tracing import decode_payloads, PayloadType
45+
46+
decoded = decode_payloads(trace["payloads"])
47+
48+
if PayloadType.PROMPT_TOKEN_IDS in decoded:
49+
token_ids = decoded[PayloadType.PROMPT_TOKEN_IDS].value # List[int]
50+
51+
if PayloadType.LOGPROBS in decoded:
52+
lp = decoded[PayloadType.LOGPROBS]
53+
logprobs = lp.value # List[float]
54+
token_ids = lp.extras.get("token_ids") # Optional[List[int]]
55+
56+
if PayloadType.ROUTER_REPLAY in decoded:
57+
matrices = decoded[PayloadType.ROUTER_REPLAY].value # List[Optional[str]]
58+
```
59+
60+
If you have the whole trace dict, `decode_trace(trace)` reaches into
61+
`trace["payloads"]` for you.
62+
63+
Decode a single payload:
64+
65+
```python
66+
from eval_protocol.tracing import decode_payload, PayloadType
67+
68+
dp = decode_payload(PayloadType.PROMPT_TOKEN_IDS, trace["payloads"]["prompt_token_ids"]["data"])
69+
dp.value # List[int]
70+
```
71+
72+
### Error handling
73+
74+
`decode_payloads` isolates per-payload failures: if one payload fails to decode,
75+
the others are still returned. Pass `on_error=callback(payload_type, exc)` to
76+
control logging (defaults to a warning):
77+
78+
```python
79+
decode_payloads(payloads, on_error=lambda pt, e: print(f"{pt} failed: {e}"))
80+
```
81+
82+
## Return type
83+
84+
`decode_payloads` / `decode_trace` return `Dict[PayloadType, DecodedPayload]`.
85+
86+
`DecodedPayload` fields:
87+
88+
| field | meaning |
89+
|----------------|-------------------------------------------------------------------|
90+
| `payload_type` | `PayloadType` enum member |
91+
| `value` | decoded value (type depends on `payload_type`, see below) |
92+
| `metadata` | decoded header/manifest metadata (token counts, scope, etc.) |
93+
| `extras` | type-specific extras (e.g. logprobs `token_ids`) |
94+
95+
`value` by type:
96+
97+
| `PayloadType` | `value` | notes |
98+
|---------------------|--------------------------|----------------------------------------------|
99+
| `PROMPT_TOKEN_IDS` | `List[int]` | prompt token ids |
100+
| `LOGPROBS` | `List[float]` | per completion token; ids in `extras["token_ids"]` (or `None`) |
101+
| `ROUTER_REPLAY` | `List[Optional[str]]` | per-token base64 routing matrices; `None` where absent |
102+
103+
## Adding a new payload type
104+
105+
1. Add a member to `PayloadType` in `types.py`.
106+
2. Add a `decode_<name>(data_b64) -> DecodedPayload` function in a new module.
107+
3. Register it in `PAYLOAD_DECODERS` in `registry.py`.
108+
109+
`decode_payloads` picks it up automatically.

‎eval_protocol/tracing/__init__.py‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
"""Decode Fireworks tracing-gateway payloads.
2+
3+
Standalone, dependency-light helpers (stdlib + ``zstandard`` only) for turning
4+
the binary/JSON ``payloads`` returned by the Fireworks tracing gateway
5+
(``GET /traces?include_payloads=true``) into Python values. No EvaluationRow or
6+
rollout machinery required -- usable on its own.
7+
8+
Typical use::
9+
10+
from eval_protocol.tracing import decode_payloads, PayloadType
11+
12+
decoded = decode_payloads(trace["payloads"])
13+
decoded[PayloadType.PROMPT_TOKEN_IDS].value # List[int]
14+
decoded[PayloadType.LOGPROBS].value # List[float]
15+
decoded[PayloadType.ROUTER_REPLAY].value # List[Optional[str]]
16+
17+
See ``README.md`` in this package for details.
18+
"""
19+
20+
from __future__ import annotations
21+
22+
from .logprobs import decode_logprobs
23+
from .prompt_token_ids import decode_prompt_token_ids
24+
from .registry import (
25+
PAYLOAD_DECODERS,
26+
decode_payload,
27+
decode_payloads,
28+
decode_trace,
29+
)
30+
from .router_replay import decode_router_replay
31+
from .types import DecodedPayload, PayloadType
32+
33+
__all__ = [
34+
"PayloadType",
35+
"DecodedPayload",
36+
"PAYLOAD_DECODERS",
37+
"decode_payload",
38+
"decode_payloads",
39+
"decode_trace",
40+
"decode_prompt_token_ids",
41+
"decode_logprobs",
42+
"decode_router_replay",
43+
]

‎eval_protocol/tracing/_decompress.py‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
"""Shared helper for the Fireworks tracing-gateway payload decoders."""
2+
3+
from __future__ import annotations
4+
5+
import base64
6+
7+
import zstandard as zstd
8+
9+
10+
def decompress_b64(data_b64: str) -> bytes:
11+
"""Base64-decode then zstd-decompress a gateway ``payloads.*.data`` blob.
12+
13+
The gateway stores every payload as ``base64(zstd(raw_bytes))``; this is the
14+
common first step every decoder shares before interpreting ``raw_bytes``.
15+
"""
16+
compressed = base64.b64decode(data_b64)
17+
return zstd.ZstdDecompressor().decompress(compressed)

0 commit comments

Comments
 (0)