perf(router): stream the data path instead of buffering it

miles/dashboard/collector.py

@@ -7,8 +7,8 @@
 from dataclasses import dataclass, field
 from typing import Any, ClassVar
 
-from miles.dashboard.events import PhaseEvent, TrajectoryEvent
-from miles.dashboard.store import DashboardStore, Record, Stream
+from miles.dashboard.events import PhaseEvent, RequestEvent, TrajectoryEvent
+from miles.dashboard.store import DashboardStore, Record, stream_for
 
 logger = logging.getLogger(__name__)
 
@@ -59,8 +59,12 @@ def push_trajectories(self, batch: list[TrajectoryEvent]) -> None:
         for event in batch:
             self._append(event)
 
+    def push_requests(self, batch: list[RequestEvent]) -> None:
+        for event in batch:
+            self._append(event)
+
     def _append(self, record: Record) -> None:
-        stream = Stream.PHASES if isinstance(record, PhaseEvent) else Stream.TRAJECTORIES
+        stream = stream_for(record)
         with self._lock:
             if self._store.buffered_count(stream) >= self.MAX_BUFFERED_PER_STREAM:
                 self._dropped_since_flush += self._store.drop_oldest_buffered(stream)

miles/dashboard/events.py

@@ -53,6 +53,28 @@ class TrajectoryEventKind(StrEnum):
 }
 
 
+@dataclass
+class RequestEvent:
+    """One rollout request's wall-clock marks, keyed by miles.utils.request_timing.
+
+    The marks, not the leg durations derived from them: the reader places every
+    leg at the time it happened, and durations are a subtraction away.
+    ``engine_stages`` is the engine's own breakdown of the one leg it owns.
+    """
+
+    rollout_id: int
+    request_id: str
+    group_index: int
+    sample_indices: list[int]
+    marks: dict[str, float]
+    engine_stages: dict[str, float]
+    worker: str
+    resp_bytes: int
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+
 @dataclass
 class TrajectoryEvent:
     ts: float

miles/dashboard/hooks.py

@@ -10,7 +10,18 @@
 from contextlib import contextmanager
 from pathlib import Path
 
-from miles.dashboard.events import STAGE_KINDS, PhaseEvent, TrajectoryEvent
+from miles.dashboard.events import STAGE_KINDS, PhaseEvent, RequestEvent, TrajectoryEvent
+from miles.utils.request_timing import (
+    ROUTER_TIMING_HEADER,
+    ROUTER_WIRE_KEYS,
+    ROUTER_WORKER_HEADER,
+    SGLD_STAGES_HEADER,
+    SGLD_TIMING_HEADER,
+    SGLD_WIRE_KEYS,
+    Marks,
+    parse_header,
+    parse_stages,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -119,8 +130,82 @@ def flush(self) -> None:
             logger.warning("dashboard trajectory sink flush failed; dropping events", exc_info=True)
 
 
+class RequestSink:
+    def __init__(self, handle) -> None:
+        self.handle = handle
+        self._buffer: list[RequestEvent] = []
+        self._lock = threading.Lock()
+        self._last_flush = time.monotonic()
+
+    def record(self, tracer: RequestTracer, samples) -> None:
+        try:
+            sample = samples[0]
+            event = RequestEvent(
+                rollout_id=tracer.rollout_id,
+                request_id=sample.request_id or "",
+                group_index=sample.group_index if sample.group_index is not None else -1,
+                sample_indices=[s.index if s.index is not None else -1 for s in samples],
+                marks={name: round(ts, 6) for name, ts in tracer.marks.marks.items()},
+                engine_stages=tracer.engine_stages,
+                worker=tracer.worker,
+                resp_bytes=tracer.resp_bytes,
+            )
+            with self._lock:
+                self._buffer.append(event)
+                batch = self._take_batch_if_due()
+            if batch:
+                self.handle.push_requests.remote(batch)
+        except Exception:  # noqa: BLE001
+            logger.warning("dashboard request sink failed; dropping events", exc_info=True)
+
+    def _take_batch_if_due(self) -> list[RequestEvent] | None:
+        if len(self._buffer) < BATCH_MAX_EVENTS and time.monotonic() - self._last_flush < BATCH_MAX_SECONDS:
+            return None
+        batch, self._buffer = self._buffer, []
+        self._last_flush = time.monotonic()
+        return batch
+
+    def flush(self) -> None:
+        try:
+            with self._lock:
+                batch, self._buffer = self._buffer, []
+            if batch:
+                _ray_get(self.handle.push_requests.remote(batch))
+        except Exception:  # noqa: BLE001
+            logger.warning("dashboard request sink flush failed; dropping events", exc_info=True)
+
+
+class RequestTracer:
+    """Collects one rollout request's marks; ``done`` is a no-op with the
+    dashboard off, so the rollout path can carry a tracer unconditionally."""
+
+    def __init__(self, rollout_id: int) -> None:
+        self.rollout_id = rollout_id
+        self.marks = Marks()
+        self.engine_stages: dict[str, float] = {}
+        self.worker = ""
+        self.resp_bytes = 0
+
+    def mark(self, name: str) -> None:
+        self.marks.mark(name)
+
+    def absorb_response(self, result) -> None:
+        """Take the sender's own marks plus the ones the reply carries back."""
+        self.marks.absorb(result.marks)
+        headers = {key.lower(): value for key, value in result.headers.items()}
+        self.marks.absorb(parse_header(headers.get(SGLD_TIMING_HEADER), SGLD_WIRE_KEYS))
+        self.marks.absorb(parse_header(headers.get(ROUTER_TIMING_HEADER), ROUTER_WIRE_KEYS))
+        self.engine_stages = parse_stages(headers.get(SGLD_STAGES_HEADER))
+        self.worker = headers.get(ROUTER_WORKER_HEADER, "")
+
+    def done(self, samples) -> None:
+        if _request_sink is not None:
+            _request_sink.record(self, samples)
+
+
 _phase_sink: PhaseSink | None = None
 _trajectory_sink: TrajectorySink | None = None
+_request_sink: RequestSink | None = None
 _rollout_id = -1
 _GPU_SAMPLER: GpuUtilSampler | None = None
 
@@ -152,13 +237,18 @@ def register_rollout_manager(args) -> None:
         return
     attach_phase_sink(handle, "rollout")
     attach_trajectory_sink(handle)
+    attach_request_sink(handle)
 
 
 def set_rollout_id(rollout_id: int) -> None:
     global _rollout_id
     _rollout_id = rollout_id
 
 
+def current_rollout_id() -> int:
+    return _rollout_id
+
+
 def record_trajectory(sample) -> None:
     if _trajectory_sink is not None:
         _trajectory_sink.record(sample, _rollout_id)
@@ -201,8 +291,14 @@ def attach_trajectory_sink(handle) -> None:
         _trajectory_sink = TrajectorySink(handle)
 
 
+def attach_request_sink(handle) -> None:
+    global _request_sink
+    if _request_sink is None:
+        _request_sink = RequestSink(handle)
+
+
 def detach_and_flush() -> None:
-    global _phase_sink, _trajectory_sink, _GPU_SAMPLER
+    global _phase_sink, _trajectory_sink, _request_sink, _GPU_SAMPLER
     from miles.utils.timer import Timer
 
     if _phase_sink is not None:
@@ -212,6 +308,9 @@ def detach_and_flush() -> None:
     if _trajectory_sink is not None:
         _trajectory_sink.flush()
         _trajectory_sink = None
+    if _request_sink is not None:
+        _request_sink.flush()
+        _request_sink = None
     if _GPU_SAMPLER is not None:
         _GPU_SAMPLER.stop()
         _GPU_SAMPLER = None

miles/dashboard/store.py

@@ -8,7 +8,7 @@
 from enum import StrEnum
 from pathlib import Path
 
-from miles.dashboard.events import PhaseEvent, TrajectoryEvent
+from miles.dashboard.events import PhaseEvent, RequestEvent, TrajectoryEvent
 
 
 RUN_DIR_PREFIX = "run_"
@@ -17,20 +17,25 @@
 class Stream(StrEnum):
     PHASES = "phases"
     TRAJECTORIES = "trajectories"
+    REQUESTS = "requests"
 
 
-Record = PhaseEvent | TrajectoryEvent
+Record = PhaseEvent | TrajectoryEvent | RequestEvent
 
 
-def _stream(record: Record) -> Stream:
+def stream_for(record: Record) -> Stream:
     if isinstance(record, PhaseEvent):
         return Stream.PHASES
+    if isinstance(record, RequestEvent):
+        return Stream.REQUESTS
     return Stream.TRAJECTORIES
 
 
 def _timestamp(record: Record) -> float:
     if isinstance(record, PhaseEvent):
         return record.t1
+    if isinstance(record, RequestEvent):
+        return max(record.marks.values(), default=0.0)
     return record.ts
 
 
@@ -63,7 +68,7 @@ def write_meta(self, *, run_name: str, start_ts: float, args: dict) -> None:
         )
 
     def append(self, record: Record) -> None:
-        self._buffers[_stream(record)].append(record)
+        self._buffers[stream_for(record)].append(record)
 
     def buffered_count(self, stream: Stream) -> int:
         return len(self._buffers[stream])