initial commit

Author: SunnySoldier357

eval_protocol/adapters/fireworks_tracing.py

@@ -100,13 +100,29 @@ def convert_trace_dict_to_evaluation_row(
                 ):
                     break  # Break early if we've found all the metadata we need
 
+        # Extract router replay payloads when present
+        payloads = trace.get("payloads")
+        if isinstance(payloads, dict):
+            router_replay = payloads.get("router_replay")
+            if isinstance(router_replay, dict) and router_replay.get("data"):
+                try:
+                    from .r3_deserializer import decompress_and_parse_r3
+
+                    matrices, r3_meta = decompress_and_parse_r3(router_replay["data"])
+                    if execution_metadata.extra is None:
+                        execution_metadata.extra = {}
+                    execution_metadata.extra["routing_matrices"] = matrices
+                    execution_metadata.extra["routing_metadata"] = r3_meta
+                except Exception as e:
+                    logger.warning("Failed to decompress R3 payload for trace %s: %s", trace.get("id"), e)
+
         return EvaluationRow(
             messages=messages,
             tools=tools,
             input_metadata=InputMetadata(
                 row_id=row_id,
                 session_data={
-                    "langfuse_trace_id": trace.get("id"),  # Store the trace ID here
+                    "langfuse_trace_id": trace.get("id"),
                 },
             ),
             execution_metadata=execution_metadata,
@@ -426,6 +442,7 @@ def get_evaluation_rows(
         max_retries: int = 3,
         span_name: Optional[str] = None,
         converter: Optional[TraceDictConverter] = None,
+        include_payloads: bool = False,
     ) -> List[EvaluationRow]:
         """Pull traces from Langfuse via proxy and convert to EvaluationRow format.
 
@@ -449,6 +466,8 @@ def get_evaluation_rows(
             max_retries: Max retry attempts used by proxy (default: 3)
             converter: Optional custom converter implementing TraceDictConverter protocol.
                 If provided, this will be used instead of the default conversion logic.
+            include_payloads: If True, request payload data (e.g., router replay)
+                from the gateway and decompress it into the returned EvaluationRows.
 
         Returns:
             List[EvaluationRow]: Converted evaluation rows
@@ -479,6 +498,7 @@ def get_evaluation_rows(
             "to_timestamp": to_timestamp.isoformat() if to_timestamp else None,
             "sleep_between_gets": sleep_between_gets,
             "max_retries": max_retries,
+            "include_payloads": include_payloads if include_payloads else None,
         }
 
         # Remove None values

eval_protocol/adapters/r3_deserializer.py

@@ -0,0 +1,189 @@
+"""R3/v1 binary deserializer for router-replay payloads.
+
+Implements the inverse of the packed binary format produced by the tracing
+gateway's ``r3_serializer.serialize_r3``.  See that module for the full
+header specification.
+
+The main entry point is :func:`decompress_and_parse_r3`, which accepts the
+base64-encoded compressed blob returned by the gateway's
+``/v1/traces/pointwise?include_payloads=true`` endpoint and produces
+per-token routing matrices in the same ``List[Optional[str]]`` format used
+by the direct inference path (``DeploymentSampler.sample_with_tokens()``).
+"""
+
+from __future__ import annotations
+
+import base64
+import logging
+import math
+import struct
+from enum import IntEnum
+from typing import Any, Dict, List, Optional, Tuple
+
+import zstandard as zstd
+
+logger = logging.getLogger(__name__)
+
+MAGIC = b"R3V1"
+HEADER_FORMAT = "<4sBBBBIIHHIIQ"
+HEADER_SIZE = struct.calcsize(HEADER_FORMAT)  # 36 bytes
+
+
+class _SelectorMode(IntEnum):
+    ALL = 0
+    SUFFIX = 1
+    BITMAP = 2
+
+
+class _RoutingDtype(IntEnum):
+    UINT8 = 1
+    UINT16 = 2
+
+    @property
+    def byte_width(self) -> int:
+        return self.value
+
+
+_SELECTOR_MODE_NAMES = {v: v.name.lower() for v in _SelectorMode}
+_ROUTING_DTYPE_NAMES = {v: v.name.lower() for v in _RoutingDtype}
+
+
+def _parse_header(raw: bytes) -> Dict[str, Any]:
+    if len(raw) < HEADER_SIZE:
+        raise ValueError(
+            f"Payload too short for r3/v1 header: {len(raw)} < {HEADER_SIZE}"
+        )
+
+    (
+        magic,
+        version,
+        selector_mode,
+        routing_dtype,
+        flags,
+        total_token_count,
+        replayed_token_count,
+        num_moe_layers,
+        top_k,
+        replay_start_token,
+        selector_byte_length,
+        matrix_byte_length,
+    ) = struct.unpack(HEADER_FORMAT, raw[:HEADER_SIZE])
+
+    if magic != MAGIC:
+        raise ValueError(f"Bad R3 magic: {magic!r}")
+    if version != 1:
+        raise ValueError(f"Unsupported R3 header version: {version}")
+
+    return {
+        "selector_mode": selector_mode,
+        "routing_dtype": routing_dtype,
+        "flags": flags,
+        "total_token_count": total_token_count,
+        "replayed_token_count": replayed_token_count,
+        "num_moe_layers": num_moe_layers,
+        "top_k": top_k,
+        "replay_start_token": replay_start_token,
+        "selector_byte_length": selector_byte_length,
+        "matrix_byte_length": matrix_byte_length,
+    }
+
+
+def _read_bitmap_positions(
+    selector_bytes: bytes, total_token_count: int
+) -> List[int]:
+    """Return sorted token indices where the bitmap bit is set."""
+    positions: List[int] = []
+    for i in range(total_token_count):
+        byte_idx = i >> 3
+        bit_idx = i & 7
+        if byte_idx < len(selector_bytes) and (selector_bytes[byte_idx] >> bit_idx) & 1:
+            positions.append(i)
+    return positions
+
+
+def decompress_and_parse_r3(
+    data_b64: str,
+) -> Tuple[List[Optional[str]], Dict[str, Any]]:
+    """Decompress and unpack an R3/v1 payload into per-token routing matrices.
+
+    Args:
+        data_b64: Base64-encoded zstd-compressed R3 binary blob, as returned
+            by the tracing gateway in ``payloads.router_replay.data``.
+
+    Returns:
+        A tuple of ``(routing_matrices, metadata)`` where:
+
+        - ``routing_matrices`` is a ``List[Optional[str]]`` of length
+          ``total_token_count``.  Each present position contains a
+          base64-encoded routing matrix (matching the format returned by
+          the direct inference path); absent positions are ``None``.
+        - ``metadata`` is a dict with keys ``num_moe_layers``, ``top_k``,
+          ``routing_dtype``, ``selector_mode``, ``total_token_count``,
+          ``replayed_token_count``, ``replay_start_token``.
+    """
+    compressed = base64.b64decode(data_b64)
+
+    decompressor = zstd.ZstdDecompressor()
+    raw = decompressor.decompress(compressed, max_output_size=len(compressed) * 20)
+
+    header = _parse_header(raw)
+
+    selector_mode = header["selector_mode"]
+    routing_dtype = header["routing_dtype"]
+    total_token_count = header["total_token_count"]
+    replayed_token_count = header["replayed_token_count"]
+    num_moe_layers = header["num_moe_layers"]
+    top_k = header["top_k"]
+    replay_start_token = header["replay_start_token"]
+    selector_byte_length = header["selector_byte_length"]
+    matrix_byte_length = header["matrix_byte_length"]
+
+    dtype_byte_width = _RoutingDtype(routing_dtype).byte_width
+    matrix_elem_size = num_moe_layers * top_k * dtype_byte_width
+
+    body = raw[HEADER_SIZE:]
+    selector_bytes = body[:selector_byte_length]
+    matrix_bytes = body[selector_byte_length : selector_byte_length + matrix_byte_length]
+
+    if matrix_elem_size == 0:
+        replayed_positions: List[int] = []
+    elif selector_mode == _SelectorMode.ALL:
+        replayed_positions = list(range(total_token_count))
+    elif selector_mode == _SelectorMode.SUFFIX:
+        replayed_positions = list(
+            range(replay_start_token, replay_start_token + replayed_token_count)
+        )
+    elif selector_mode == _SelectorMode.BITMAP:
+        replayed_positions = _read_bitmap_positions(selector_bytes, total_token_count)
+    else:
+        raise ValueError(f"Unknown selector_mode: {selector_mode}")
+
+    # Split matrix bytes into per-token chunks and base64-encode each one
+    matrices: List[Optional[str]] = [None] * total_token_count
+    for idx, pos in enumerate(replayed_positions):
+        start = idx * matrix_elem_size
+        end = start + matrix_elem_size
+        if end > len(matrix_bytes):
+            logger.warning(
+                "R3 matrix data truncated at token %d (position %d): "
+                "expected %d bytes but only %d remaining",
+                idx, pos, matrix_elem_size, len(matrix_bytes) - start,
+            )
+            break
+        matrices[pos] = base64.b64encode(matrix_bytes[start:end]).decode("ascii")
+
+    metadata: Dict[str, Any] = {
+        "num_moe_layers": num_moe_layers,
+        "top_k": top_k,
+        "routing_dtype": _ROUTING_DTYPE_NAMES.get(
+            _RoutingDtype(routing_dtype), str(routing_dtype)
+        ),
+        "selector_mode": _SELECTOR_MODE_NAMES.get(
+            _SelectorMode(selector_mode), str(selector_mode)
+        ),
+        "total_token_count": total_token_count,
+        "replayed_token_count": replayed_token_count,
+        "replay_start_token": replay_start_token,
+    }
+
+    return matrices, metadata

eval_protocol/pytest/remote_rollout_processor.py

@@ -35,11 +35,13 @@ def __init__(
         model_base_url: str = "https://tracing.fireworks.ai",
         poll_interval: float = 1.0,
         timeout_seconds: float = 120.0,
+        include_payloads: bool = False,
     ):
         # Prefer constructor-provided configuration. These can be overridden via
         # config.kwargs at call time for backward compatibility.
         self._remote_base_url = remote_base_url
         self._model_base_url = model_base_url
+        self._include_payloads = include_payloads
         if os.getenv("EP_REMOTE_ROLLOUT_PROCESSOR_BASE_URL"):
             self._remote_base_url = os.getenv("EP_REMOTE_ROLLOUT_PROCESSOR_BASE_URL")
         _ep_model_base_url = os.getenv("EP_MODEL_BASE_URL")
@@ -175,7 +177,10 @@ async def _process_row(row: EvaluationRow) -> EvaluationRow:
             row.execution_metadata.rollout_duration_seconds = time.perf_counter() - start_time
 
             def _update_with_trace() -> None:
-                return update_row_with_remote_trace(row, default_fireworks_output_data_loader, model_base_url)
+                return update_row_with_remote_trace(
+                    row, default_fireworks_output_data_loader, model_base_url,
+                    include_payloads=self._include_payloads,
+                )
 
             await asyncio.to_thread(_update_with_trace)  # Update row with remote trace in-place
             return row

eval_protocol/pytest/tracing_utils.py

@@ -22,7 +22,11 @@ def fetch_traces() -> List[EvaluationRow]:
         # Use EP_REMOTE_API_KEY for fetching remote traces, falling back to FIREWORKS_API_KEY
         api_key = os.environ.get("EP_REMOTE_API_KEY") or os.environ.get("FIREWORKS_API_KEY")
         adapter = FireworksTracingAdapter(base_url=base_url, api_key=api_key)
-        return adapter.get_evaluation_rows(tags=[f"rollout_id:{config.rollout_id}"], max_retries=5)
+        return adapter.get_evaluation_rows(
+            tags=[f"rollout_id:{config.rollout_id}"],
+            max_retries=5,
+            include_payloads=config.include_payloads,
+        )
 
     return DynamicDataLoader(generators=[fetch_traces], preprocess_fn=filter_longest_conversation)
 
@@ -148,13 +152,20 @@ def build_init_request(
 
 
 def update_row_with_remote_trace(
-    row: EvaluationRow, output_data_loader: Callable[[DataLoaderConfig], DynamicDataLoader], model_base_url: str
+    row: EvaluationRow,
+    output_data_loader: Callable[[DataLoaderConfig], DynamicDataLoader],
+    model_base_url: str,
+    include_payloads: bool = False,
 ) -> None:
     """Update row with remote trace data using output_data_loader (shared logic)."""
     if not row.execution_metadata.rollout_id:
         return None
 
-    loader_config = DataLoaderConfig(rollout_id=row.execution_metadata.rollout_id, model_base_url=model_base_url)
+    loader_config = DataLoaderConfig(
+        rollout_id=row.execution_metadata.rollout_id,
+        model_base_url=model_base_url,
+        include_payloads=include_payloads,
+    )
     data_loader = output_data_loader(loader_config)
     results = data_loader.load()
     output_rows: List[EvaluationRow] = [r for result in results for r in result.rows]

eval_protocol/types/remote_rollout_processor.py

@@ -39,6 +39,7 @@ class DataLoaderConfig(BaseModel):
 
     rollout_id: str
     model_base_url: Optional[str] = None
+    include_payloads: bool = False
 
 
 class InitRequest(BaseModel):

pyproject.toml

@@ -48,6 +48,7 @@ dependencies = [
     "deepdiff>=6.0.0",
     "websockets>=15.0.1",
     "fastapi>=0.116.1",
+    "zstandard>=0.19.0",
 ]
 
 [project.urls]