diff --git a/components/src/dynamo/common/export_forward_pass_metrics.py b/components/src/dynamo/common/export_forward_pass_metrics.py new file mode 100644 index 000000000000..b93d299af7eb --- /dev/null +++ b/components/src/dynamo/common/export_forward_pass_metrics.py @@ -0,0 +1,218 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Export Dynamo ForwardPassMetrics to a local newline-delimited JSON socket. + +This process remains on the Dynamo side of the integration boundary: it uses +``FpmEventSubscriber`` for discovery and event-plane transport, decodes the +versioned Dynamo payload, and forwards a stable record to a local collector +such as Tachometer. +""" + +from __future__ import annotations + +import argparse +import asyncio +import logging +import os +import time +from collections.abc import Sequence + +import msgspec + +from dynamo.common.forward_pass_metrics import ForwardPassMetrics, decode + +logger = logging.getLogger(__name__) + + +class FpmExportRecord(msgspec.Struct, frozen=True, gc=False): # type: ignore[call-arg] + """One event-plane FPM record sent to the local collection socket.""" + + namespace: str + component: str + publisher_id: int + sequence: int + published_at_ms: int + metrics: ForwardPassMetrics + + +_json_encoder = msgspec.json.Encoder() + + +def encode_export_record(record: FpmExportRecord) -> bytes: + """Encode one record as an NDJSON line.""" + return _json_encoder.encode(record) + b"\n" + + +def _socket_path(value: str) -> str: + if value.startswith("unix://"): + return value.removeprefix("unix://") + return value + + +async def _connect_sink(path: str, timeout: float) -> asyncio.StreamWriter: + deadline = time.monotonic() + timeout + last_error: OSError | None = None + while time.monotonic() < deadline: + try: + _reader, writer = await asyncio.open_unix_connection(path) + return writer + except OSError as error: + last_error = error + await asyncio.sleep(0.2) + raise TimeoutError(f"Timed out connecting to FPM sink {path}: {last_error}") + + +async def _receive_component( + *, + subscriber, + namespace: str, + component: str, + records: asyncio.Queue[FpmExportRecord], +) -> None: + while True: + envelope = await asyncio.to_thread(subscriber.recv_envelope) + if envelope is None: + logger.info("FPM stream closed for component=%s", component) + return + + publisher_id, sequence, published_at_ms, payload = envelope + metrics = decode(payload) + if metrics is None: + continue + + await records.put( + FpmExportRecord( + namespace=namespace, + component=component, + publisher_id=publisher_id, + sequence=sequence, + published_at_ms=published_at_ms, + metrics=metrics, + ) + ) + + +async def _write_records( + writer: asyncio.StreamWriter, records: asyncio.Queue[FpmExportRecord] +) -> None: + while True: + record = await records.get() + writer.write(encode_export_record(record)) + await writer.drain() + + +async def export( + *, + namespace: str, + components: Sequence[str], + endpoint_name: str, + discovery_backend: str, + request_plane: str, + socket_path: str, + connect_timeout: float, +) -> None: + from dynamo.llm import FpmEventSubscriber + from dynamo.runtime import DistributedRuntime + + writer = await _connect_sink(socket_path, connect_timeout) + logger.info("Connected to FPM collection sink at %s", socket_path) + + loop = asyncio.get_running_loop() + runtime = DistributedRuntime(loop, discovery_backend, request_plane) + subscribers = [] + records: asyncio.Queue[FpmExportRecord] = asyncio.Queue(maxsize=100_000) + + for component in components: + endpoint = runtime.endpoint(f"{namespace}.{component}.{endpoint_name}") + subscriber = FpmEventSubscriber(endpoint) + subscribers.append((component, subscriber)) + + writer_task = asyncio.create_task(_write_records(writer, records)) + receiver_tasks = [ + asyncio.create_task( + _receive_component( + subscriber=subscriber, + namespace=namespace, + component=component, + records=records, + ) + ) + for component, subscriber in subscribers + ] + + logger.info( + "Exporting forward-pass-metrics components=%s via %s event plane", + ",".join(components), + os.environ.get("DYN_EVENT_PLANE", "zmq"), + ) + + try: + tasks = [writer_task, *receiver_tasks] + done, _pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) + for task in done: + task.result() + raise RuntimeError("An FPM export task stopped unexpectedly") + finally: + for _component, subscriber in subscribers: + subscriber.shutdown() + for task in receiver_tasks: + task.cancel() + writer_task.cancel() + await asyncio.gather(*receiver_tasks, writer_task, return_exceptions=True) + writer.close() + await writer.wait_closed() + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Export Dynamo ForwardPassMetrics to a local NDJSON socket" + ) + parser.add_argument("--namespace", default="dynamo") + parser.add_argument( + "--component", + action="append", + dest="components", + help="Component to subscribe to; repeat for multiple components", + ) + parser.add_argument("--endpoint", default="generate") + parser.add_argument( + "--discovery-backend", + default=os.environ.get("DYN_DISCOVERY_BACKEND", "etcd"), + ) + parser.add_argument( + "--request-plane", default=os.environ.get("DYN_REQUEST_PLANE", "tcp") + ) + parser.add_argument( + "--socket", + required=True, + help="Unix socket path, optionally prefixed with unix://", + ) + parser.add_argument("--connect-timeout", type=float, default=60.0) + return parser.parse_args() + + +def main() -> None: + from dynamo.runtime.logging import configure_dynamo_logging + + configure_dynamo_logging() + args = _parse_args() + components = args.components or ["backend"] + try: + asyncio.run( + export( + namespace=args.namespace, + components=components, + endpoint_name=args.endpoint, + discovery_backend=args.discovery_backend, + request_plane=args.request_plane, + socket_path=_socket_path(args.socket), + connect_timeout=args.connect_timeout, + ) + ) + except KeyboardInterrupt: + logger.info("Stopped FPM exporter") + + +if __name__ == "__main__": + main() diff --git a/components/src/dynamo/common/tests/test_export_forward_pass_metrics.py b/components/src/dynamo/common/tests/test_export_forward_pass_metrics.py new file mode 100644 index 000000000000..b2536e90387e --- /dev/null +++ b/components/src/dynamo/common/tests/test_export_forward_pass_metrics.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import msgspec +import pytest + +from dynamo.common.export_forward_pass_metrics import ( + FpmExportRecord, + encode_export_record, +) +from dynamo.common.forward_pass_metrics import ( + ForwardPassMetrics, + ScheduledRequestMetrics, +) + +pytestmark = [pytest.mark.pre_merge, pytest.mark.unit, pytest.mark.gpu_0] + + +def test_encode_export_record_is_ndjson_and_preserves_envelope_metadata(): + record = FpmExportRecord( + namespace="dynamo", + component="backend", + publisher_id=17, + sequence=23, + published_at_ms=1_700_000_000_123, + metrics=ForwardPassMetrics( + worker_id="worker-1", + dp_rank=2, + counter_id=99, + wall_time=0.025, + scheduled_requests=ScheduledRequestMetrics( + num_decode_requests=4, + sum_decode_kv_tokens=4096, + ), + ), + ) + + encoded = encode_export_record(record) + + assert encoded.endswith(b"\n") + decoded = msgspec.json.decode(encoded) + assert decoded["namespace"] == "dynamo" + assert decoded["component"] == "backend" + assert decoded["publisher_id"] == 17 + assert decoded["sequence"] == 23 + assert decoded["published_at_ms"] == 1_700_000_000_123 + assert decoded["metrics"]["worker_id"] == "worker-1" + assert decoded["metrics"]["dp_rank"] == 2 + assert decoded["metrics"]["counter_id"] == 99 + assert decoded["metrics"]["scheduled_requests"]["num_decode_requests"] == 4 diff --git a/lib/bindings/python/rust/llm/fpm.rs b/lib/bindings/python/rust/llm/fpm.rs index 47d57572c934..4698522a5e5e 100644 --- a/lib/bindings/python/rust/llm/fpm.rs +++ b/lib/bindings/python/rust/llm/fpm.rs @@ -422,7 +422,7 @@ pub(crate) struct FpmEventSubscriber { // recv mode state (lazily initialised on first recv() call) recv_started: Arc, - rx: Arc>>>>, + rx: Arc>>>, // tracking mode state tracking_started: Arc, @@ -438,6 +438,102 @@ pub(crate) struct FpmEventSubscriber { worker_model_cards: Arc>, } +struct ReceivedFpmEvent { + publisher_id: u64, + sequence: u64, + published_at: u64, + payload: Vec, +} + +type FpmEnvelope = (u64, u64, u64, Vec); + +impl FpmEventSubscriber { + fn ensure_recv_started(&self) -> PyResult<()> { + if self.tracking_started.load(Ordering::SeqCst) { + return Err(PyRuntimeError::new_err( + "Cannot receive events after start_tracking()", + )); + } + + if self.recv_started.swap(true, Ordering::SeqCst) { + return Ok(()); + } + + let component = self.component.clone(); + let cancel = self.cancel.clone(); + let (tx, rx_new) = tokio::sync::mpsc::unbounded_channel::(); + + { + let mut guard = self.rx.lock().map_err(|e| to_pyerr(format!("{e}")))?; + *guard = Some(rx_new); + } + + let rt = component.drt().runtime().secondary(); + rt.spawn(async move { + let mut subscriber = match EventSubscriber::for_component(&component, FPM_TOPIC).await { + Ok(s) => s, + Err(e) => { + tracing::error!("FPM subscriber (recv): failed to create: {e}"); + return; + } + }; + + tracing::info!("FPM subscriber (recv): listening for forward-pass-metrics events"); + + loop { + tokio::select! { + biased; + _ = cancel.cancelled() => { + tracing::info!("FPM subscriber (recv): shutting down"); + break; + } + event = subscriber.next() => { + match event { + Some(Ok(envelope)) => { + let received = ReceivedFpmEvent { + publisher_id: envelope.publisher_id, + sequence: envelope.sequence, + published_at: envelope.published_at, + payload: envelope.payload.to_vec(), + }; + if tx.send(received).is_err() { + tracing::info!( + "FPM subscriber (recv): receiver dropped, exiting" + ); + break; + } + } + Some(Err(e)) => { + tracing::warn!("FPM subscriber (recv): event error: {e}"); + } + None => { + tracing::info!("FPM subscriber (recv): stream ended"); + break; + } + } + } + } + } + }); + + Ok(()) + } + + fn blocking_recv_event(&self, py: Python) -> PyResult> { + self.ensure_recv_started()?; + let rx = self.rx.clone(); + py.allow_threads(move || { + let mut guard = rx + .lock() + .map_err(|e| to_pyerr(format!("lock poisoned: {e}")))?; + match guard.as_mut() { + Some(rx) => Ok(rx.blocking_recv()), + None => Ok(None), + } + }) + } +} + #[pymethods] impl FpmEventSubscriber { /// Create a subscriber that auto-discovers FPM publishers. @@ -469,77 +565,23 @@ impl FpmEventSubscriber { /// /// Returns the raw msgspec payload, or None if the stream is closed. fn recv(&self, py: Python) -> PyResult>> { - if self.tracking_started.load(Ordering::SeqCst) { - return Err(PyRuntimeError::new_err( - "Cannot call recv() after start_tracking()", - )); - } - - // Lazily start the recv-mode subscriber task on the first call. - if !self.recv_started.swap(true, Ordering::SeqCst) { - let component = self.component.clone(); - let cancel = self.cancel.clone(); - let (tx, rx_new) = tokio::sync::mpsc::unbounded_channel::>(); - - { - let mut guard = self.rx.lock().map_err(|e| to_pyerr(format!("{e}")))?; - *guard = Some(rx_new); - } - - let rt = component.drt().runtime().secondary(); - rt.spawn(async move { - let mut subscriber = - match EventSubscriber::for_component(&component, FPM_TOPIC).await { - Ok(s) => s, - Err(e) => { - tracing::error!("FPM subscriber (recv): failed to create: {e}"); - return; - } - }; - - tracing::info!("FPM subscriber (recv): listening for forward-pass-metrics events"); - - loop { - tokio::select! { - biased; - _ = cancel.cancelled() => { - tracing::info!("FPM subscriber (recv): shutting down"); - break; - } - event = subscriber.next() => { - match event { - Some(Ok(envelope)) => { - if tx.send(envelope.payload.to_vec()).is_err() { - tracing::info!( - "FPM subscriber (recv): receiver dropped, exiting" - ); - break; - } - } - Some(Err(e)) => { - tracing::warn!("FPM subscriber (recv): event error: {e}"); - } - None => { - tracing::info!("FPM subscriber (recv): stream ended"); - break; - } - } - } - } - } - }); - } + Ok(self.blocking_recv_event(py)?.map(|event| event.payload)) + } - let rx = self.rx.clone(); - py.allow_threads(move || { - let mut guard = rx - .lock() - .map_err(|e| to_pyerr(format!("lock poisoned: {e}")))?; - match guard.as_mut() { - Some(rx) => Ok(rx.blocking_recv()), - None => Ok(None), - } - }) + /// Blocking receive of the next message with its event-plane envelope metadata. + /// + /// Returns `(publisher_id, sequence, published_at_ms, payload)` or `None` if + /// the stream is closed. This shares recv mode with `recv()`; callers may use + /// either method, but neither can be used after `start_tracking()`. + fn recv_envelope(&self, py: Python) -> PyResult> { + Ok(self.blocking_recv_event(py)?.map(|event| { + ( + event.publisher_id, + event.sequence, + event.published_at, + event.payload, + ) + })) } /// Start background tracking of the latest FPM per `(worker_id, dp_rank)`. diff --git a/lib/bindings/python/src/dynamo/_core.pyi b/lib/bindings/python/src/dynamo/_core.pyi index 5a4010ef9451..c0888e8c6ecc 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -1264,6 +1264,19 @@ class FpmEventSubscriber: """ ... + def recv_envelope(self) -> Optional[tuple[int, int, int, bytes]]: + """ + Blocking receive of the next message with event-plane metadata. + + This shares recv mode with ``recv()`` and releases the GIL while + waiting. It cannot be used after ``start_tracking()``. + + Returns: + ``(publisher_id, sequence, published_at_ms, payload)`` or None + if the stream is closed. + """ + ... + def start_tracking(self) -> None: """ Start background tracking of the latest FPM per (worker_id, dp_rank).