refactor: remove legacy NATS metrics reporting infrastructure

Remove all NATS metrics collection and reporting code from the runtime:

- Remove NATS metric constants (nats_client, nats_service modules)
- Remove ComponentNatsServerPrometheusMetrics struct and impl
- Remove DRTNatsClientPrometheusMetrics struct and impl
- Remove nats_metrics_worker background task
- Remove NATS metrics tests
- Remove NATS metrics helper functions
- Remove stats_handler infrastructure (EndpointStatsHandler, service.rs)
- Update tests to not filter NATS metrics
- Update runtime examples to remove stats_handler usage
- Clarify distributed_runtime.md documentation

This removes ~1200 lines of code while maintaining all passing tests.
The system no longer collects or reports NATS client/service metrics.

Fix assert_eq! argument order in metrics tests

The assert_eq! macro arguments were swapped relative to the error message
template. Fixed by swapping 'Expected' and 'Actual' labels in the error
message to match the actual order of arguments (actual, expected).

Signed-off-by: Keiven Chang <keivenchang@users.noreply.github.com>

Author: keivenchang

docs/design_docs/distributed_runtime.md

@@ -53,7 +53,7 @@ The hierarchy and naming in etcd and NATS may change over time, and this documen
 
   For etcd, it also creates a primary lease and spin up a background task to keep the lease alive. All objects registered under this `DistributedRuntime` use this lease_id to maintain their life cycle. There is also a cancellation token that is tied to the primary lease. When the cancellation token is triggered or the background task failed, the primary lease is revoked or expired and the kv pairs stored with this lease_id is removed.
 - `Namespace`: `Namespace`s are primarily a logical grouping mechanism and is not registered in etcd. It provides the root path for all components under this `Namespace`.
-- `Component`: When a `Component` object is created, similar to `Namespace`, it isn't be registered in etcd. When `create_service` is called, it creates a NATS service group using `{namespace_name}.{service_name}` for metrics and registers a service in the registry of the `Component`, where the registry is an internal data structure that tracks all services and endpoints within the `DistributedRuntime`.
+- `Component`: When a `Component` object is created, similar to `Namespace`, it isn't be registered in etcd. When `create_service` is called, it creates a NATS service group using `{namespace_name}.{service_name}` as the service identifier and registers a service in the registry of the `Component`, where the registry is an internal data structure that tracks all services and endpoints within the `DistributedRuntime`.
 - `Endpoint`: When an Endpoint object is created and started, it performs two key registrations:
   - NATS Registration: The endpoint is registered with the NATS service group created during service creation. The endpoint is assigned a unique subject following the naming: `{namespace_name}.{service_name}.{endpoint_name}-{lease_id_hex}`.
   - etcd Registration: The endpoint information is stored in etcd at a path following the naming: `/services/{namespace}/{component}/{endpoint}-{lease_id}`. Note that the endpoints of different workers of the same type (i.e., two `VllmPrefillWorker`s in one deployment) share the same `Namespace`, `Component`, and `Endpoint` name. They are distinguished by their different primary `lease_id` of their `DistributedRuntime`.

lib/runtime/examples/service_metrics/README.md

@@ -27,13 +27,7 @@ Annotated { data: Some("o"), id: None, event: None, comment: None }
 Annotated { data: Some("r"), id: None, event: None, comment: None }
 Annotated { data: Some("l"), id: None, event: None, comment: None }
 Annotated { data: Some("d"), id: None, event: None, comment: None }
-ServiceSet { services: [ServiceInfo { name: "dynamo_init_backend_720278f8", id: "eOHMc4ndRw8s5flv4WOZx7", version: "0.0.1", started: "2025-02-26T18:54:04.917294605Z", endpoints: [EndpointInfo { name: "dynamo_init_backend_720278f8-generate-694d951a80e06abf", subject: "dynamo_init_backend_720278f8.generate-694d951a80e06abf", data: Some(Metrics(Object {"average_processing_time": Number(53662), "data": Object {"val": Number(10)}, "last_error": String(""), "num_errors": Number(0), "num_requests": Number(2), "processing_time": Number(107325), "queue_group": String("q")})) }] }] }
-```
-
-Note the following stats in the output demonstrate the custom
-`stats_handler` attached to the service in `server.rs` is being invoked:
-```
-data: Some(Metrics(Object {..., "data": Object {"val": Number(10)}, ...)
+ServiceSet { services: [ServiceInfo { name: "dynamo_init_backend_720278f8", id: "eOHMc4ndRw8s5flv4WOZx7", version: "0.0.1", started: "2025-02-26T18:54:04.917294605Z", endpoints: [EndpointInfo { name: "dynamo_init_backend_720278f8-generate-694d951a80e06abf", subject: "dynamo_init_backend_720278f8.generate-694d951a80e06abf", data: Some(Metrics(Object {"average_processing_time": Number(53662), "last_error": String(""), "num_errors": Number(0), "num_requests": Number(2), "processing_time": Number(107325), "queue_group": String("q")})) }] }] }
 ```
 
 If you start two copies of the server, you will see two entries being emitted.

lib/runtime/examples/service_metrics/src/bin/service_server.rs

@@ -1,7 +1,7 @@
 // SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-use service_metrics::{DEFAULT_NAMESPACE, MyStats};
+use service_metrics::DEFAULT_NAMESPACE;
 
 use dynamo_runtime::{
     DistributedRuntime, Runtime, Worker, logging,
@@ -63,11 +63,6 @@ async fn backend(runtime: DistributedRuntime) -> anyhow::Result<()> {
     component
         .endpoint("generate")
         .endpoint_builder()
-        .stats_handler(|stats| {
-            println!("stats: {:?}", stats);
-            let stats = MyStats { val: 10 };
-            serde_json::to_value(stats).unwrap()
-        })
         .handler(ingress)
         .start()
         .await

lib/runtime/examples/service_metrics/src/lib.rs

@@ -1,12 +1,4 @@
 // SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-use serde::{Deserialize, Serialize};
-
 pub const DEFAULT_NAMESPACE: &str = "dynamo";
-
-#[derive(Serialize, Deserialize)]
-// Dummy Stats object to demonstrate how to attach a custom stats handler
-pub struct MyStats {
-    pub val: u32,
-}

lib/runtime/examples/system_metrics/src/lib.rs

@@ -18,13 +18,6 @@ pub const DEFAULT_NAMESPACE: &str = "dyn_example_namespace";
 pub const DEFAULT_COMPONENT: &str = "dyn_example_component";
 pub const DEFAULT_ENDPOINT: &str = "dyn_example_endpoint";
 
-/// Stats structure returned by the endpoint's stats handler
-#[derive(serde::Serialize, serde::Deserialize, Debug, Clone)]
-pub struct MyStats {
-    // Example value for demonstration purposes
-    pub val: i32,
-}
-
 /// Custom metrics for system stats with data bytes tracking
 #[derive(Clone, Debug)]
 pub struct MySystemStatsMetrics {
@@ -103,17 +96,7 @@ pub async fn backend(drt: DistributedRuntime, endpoint_name: Option<&str>) -> an
     // Use the factory pattern - single line factory call with metrics
     let ingress = Ingress::for_engine(RequestHandler::with_metrics(system_metrics))?;
 
-    endpoint
-        .endpoint_builder()
-        .stats_handler(|_stats| {
-            println!("Stats handler called with stats: {:?}", _stats);
-            // TODO(keivenc): return a real stats object
-            let stats = MyStats { val: 10 };
-            serde_json::to_value(stats).unwrap()
-        })
-        .handler(ingress)
-        .start()
-        .await?;
+    endpoint.endpoint_builder().handler(ingress).start().await?;
 
     Ok(())
 }

lib/runtime/src/component.rs

@@ -43,7 +43,6 @@ use super::{DistributedRuntime, Runtime, traits::*, transports::nats::Slug, util
 
 use crate::pipeline::network::{PushWorkHandler, ingress::push_endpoint::PushEndpoint};
 use crate::protocols::EndpointId;
-use crate::service::ComponentNatsServerPrometheusMetrics;
 use async_nats::{
     rustls::quic,
     service::{Service, ServiceExt},
@@ -52,7 +51,6 @@ use derive_builder::Builder;
 use derive_getters::Getters;
 use educe::Educe;
 use serde::{Deserialize, Serialize};
-use service::EndpointStatsHandler;
 use std::{collections::HashMap, hash::Hash, sync::Arc};
 use validator::{Validate, ValidationError};
 
@@ -62,7 +60,6 @@ mod component;
 mod endpoint;
 mod namespace;
 mod registry;
-pub mod service;
 
 pub use client::Client;
 pub use endpoint::build_transport_type;
@@ -79,8 +76,6 @@ pub enum TransportType {
 #[derive(Default)]
 pub struct RegistryInner {
     pub(crate) services: HashMap<String, Service>,
-    pub(crate) stats_handlers:
-        HashMap<String, Arc<parking_lot::Mutex<HashMap<String, EndpointStatsHandler>>>>,
 }
 
 #[derive(Clone)]
@@ -279,11 +274,6 @@ impl ComponentBuilder {
 
     pub fn build(self) -> Result<Component, anyhow::Error> {
         let component = self.build_internal()?;
-        // If this component is using NATS, gather it's metrics
-        let drt = component.drt();
-        if drt.request_plane().is_nats() {
-            drt.start_stats_service(component.clone());
-        }
         Ok(component)
     }
 }

lib/runtime/src/component/endpoint.rs

@@ -4,14 +4,13 @@
 use std::sync::Arc;
 
 use anyhow::Result;
-pub use async_nats::service::endpoint::Stats as EndpointStats;
 use derive_builder::Builder;
 use derive_getters::Dissolve;
 use educe::Educe;
 use tokio_util::sync::CancellationToken;
 
 use crate::{
-    component::{Endpoint, Instance, TransportType, service::EndpointStatsHandler},
+    component::{Endpoint, Instance, TransportType},
     distributed::RequestPlaneMode,
     pipeline::network::{PushWorkHandler, ingress::push_endpoint::PushEndpoint},
     protocols::EndpointId,
@@ -30,11 +29,6 @@ pub struct EndpointConfig {
     #[educe(Debug(ignore))]
     handler: Arc<dyn PushWorkHandler>,
 
-    /// Stats handler
-    #[educe(Debug(ignore))]
-    #[builder(default, private)]
-    _stats_handler: Option<EndpointStatsHandler>,
-
     /// Additional labels for metrics
     #[builder(default, setter(into))]
     metrics_labels: Option<Vec<(String, String)>>,
@@ -56,13 +50,6 @@ impl EndpointConfigBuilder {
         Self::default().endpoint(endpoint)
     }
 
-    pub fn stats_handler<F>(self, handler: F) -> Self
-    where
-        F: FnMut(EndpointStats) -> serde_json::Value + Send + Sync + 'static,
-    {
-        self._stats_handler(Some(Box::new(handler)))
-    }
-
     /// Register an async engine in the local endpoint registry for direct in-process calls
     pub fn register_local_engine(
         self,
@@ -80,46 +67,19 @@ impl EndpointConfigBuilder {
     }
 
     pub async fn start(self) -> Result<()> {
-        let (
-            endpoint,
-            handler,
-            stats_handler,
-            metrics_labels,
-            graceful_shutdown,
-            health_check_payload,
-        ) = self.build_internal()?.dissolve();
+        let (endpoint, handler, metrics_labels, graceful_shutdown, health_check_payload) =
+            self.build_internal()?.dissolve();
         let connection_id = endpoint.drt().connection_id();
         let endpoint_id = endpoint.id();
 
         tracing::debug!("Starting endpoint: {endpoint_id}");
 
-        let service_name = endpoint.component.service_name();
-
         let metrics_labels: Option<Vec<(&str, &str)>> = metrics_labels
             .as_ref()
             .map(|v| v.iter().map(|(k, v)| (k.as_str(), v.as_str())).collect());
         // Add metrics to the handler. The endpoint provides additional information to the handler.
         handler.add_metrics(&endpoint, metrics_labels.as_deref())?;
 
-        // Insert the stats handler. depends on NATS.
-        if let Some(stats_handler) = stats_handler {
-            let registry = endpoint.drt().component_registry().inner.lock().await;
-            let handler_map = registry
-                .stats_handlers
-                .get(&service_name)
-                .cloned()
-                .expect("no stats handler registry; this is unexpected");
-            // There is something wrong with the stats handler map I think.
-            // Here the connection_id is included, but in component/service.rs add_stats_service it uses service_name,
-            // no connection id so it's per-endpoint not per-instance. Doesn't match.
-            // To not block current refactor I am keeping previous behavior, but I think needs
-            // investigation.
-            handler_map.lock().insert(
-                nats::instance_subject(&endpoint_id, connection_id),
-                stats_handler,
-            );
-        }
-
         // This creates a child token of the runtime's endpoint_shutdown_token. That token is
         // cancelled first as part of graceful shutdown. See Runtime::shutdown.
         let endpoint_shutdown_token = endpoint.drt().child_token();