feat(metrics): add NATS client metrics to prometheus_metrics_fmt

[keivenchang]

Overview:

This PR adds comprehensive NATS metrics collection and Prometheus metrics management to the Dynamo runtime, including endpoint-grouped metrics and callback-based metric registration.

Details:

Simple Flow:

• Client hits /metrics
• Component callback pulls from SRV.STATS.
• NATS callback pulls client metrics
• Both get copied into Prometheus structures
• Return formatted metrics to client

Where should the reviewer start?

• lib/runtime/src/metrics.rs - Core metrics registry implementation and callback system
• lib/runtime/src/transports/nats.rs - NATS metrics collection and registration
• lib/runtime/src/component.rs - Component metrics integration and automatic registration
• lib/runtime/src/distributed.rs - DistributedRuntime metrics management methods

Related Issues: (use one of the action words Closes / Fixes / Resolves / Relates to)

Relates to DIS-298

Summary by CodeRabbit

• New Features

  • Introduced a structured metrics registry system, allowing for registration and execution of custom metrics callbacks grouped by prefix.
  • Added integration of NATS client metrics, exposing internal NATS statistics as Prometheus metrics.
  • Enhanced metrics retrieval to execute registered callbacks before collecting Prometheus metrics.

• Bug Fixes

  • Standardized naming and logging for the system server (HTTP) across the runtime.

• Tests

  • Added tests to verify metrics registry callback functionality and NATS client metrics integration.

• Refactor

  • Improved organization of metrics registration and callback management for extensibility and clarity.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

Walkthrough

This change refactors and extends the metrics system in the distributed runtime. It introduces a MetricsRegistryEntry struct to manage Prometheus registries and runtime callbacks by prefix, updates the registration and execution of metrics and callbacks, integrates NATS client metrics, and standardizes naming from "HTTP server" to "system server (HTTP)" throughout the codebase. New tests verify the enhanced metrics system.

Changes

Cohort / File(s)	Change Summary
DistributedRuntime Metrics System Refactor lib/runtime/src/distributed.rs, lib/runtime/src/lib.rs, lib/runtime/src/metrics.rs	Introduced MetricsRegistryEntry to encapsulate Prometheus registries and runtime callbacks by prefix. Updated DistributedRuntime to use this structure, renamed related fields, and added methods for registering metrics and callbacks, executing callbacks, and retrieving prefixes. Refactored metrics registration and formatting to use new helper methods and callback execution. Added and updated tests for metrics and callbacks.
System Server Naming Standardization lib/runtime/src/distributed.rs, lib/runtime/src/http_server.rs	Renamed all references from "HTTP server" to "system server (HTTP)" in documentation, logs, and tests for consistency. Improved logging and comments accordingly.
NATS Client Metrics Integration lib/runtime/src/transports/nats.rs, lib/runtime/src/distributed.rs, lib/runtime/src/metrics.rs	Added NatsClientMetrics struct to expose NATS client statistics as Prometheus metrics. Registered metrics and callbacks with the distributed runtime, updating metrics from NATS client stats. Updated tests to verify presence of NATS metrics.
Test Enhancements and Adjustments lib/runtime/src/metrics.rs, lib/runtime/src/http_server.rs	Added and updated tests for metrics registry, callbacks, and Prometheus output. Adjusted tests to filter out NATS metrics where necessary and verify their presence separately. Improved test coverage for new features.

Sequence Diagram(s)

sequenceDiagram
    participant User
    participant DistributedRuntime
    participant MetricsRegistryEntry
    participant NatsClient
    participant Prometheus

    User->>DistributedRuntime: add_prometheus_metric(prefix, collector)
    DistributedRuntime->>MetricsRegistryEntry: add collector to registry

    User->>DistributedRuntime: add_metrics_callback(prefix, callback)
    DistributedRuntime->>MetricsRegistryEntry: add callback to callbacks list

    User->>DistributedRuntime: execute_metrics_callbacks(prefix)
    DistributedRuntime->>MetricsRegistryEntry: execute all callbacks
    MetricsRegistryEntry-->>DistributedRuntime: return results

    NatsClient->>DistributedRuntime: add_metrics(drt)
    DistributedRuntime->>MetricsRegistryEntry: register NATS metrics & callback

    Prometheus->>DistributedRuntime: scrape metrics endpoint
    DistributedRuntime->>MetricsRegistryEntry: execute callbacks
    DistributedRuntime->>Prometheus: return formatted metrics

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~40 minutes

Possibly related PRs

• chore: metrics endpoint variables renamed from HTTP_SERVER->SYSTEM #1934: Refactors and standardizes naming from "HTTP server" to "system server" in distributed runtime and configuration, overlapping directly with terminology and code changes in this PR.

Poem

A hop and a skip through metrics anew,
Prefixes sorted, callbacks in queue.
NATS now counts its bytes and its state,
System server logs look first-rate!
With every gauge and registry set,
This bunny’s metrics are the best yet.
🐇📊

Note

⚡️ Unit Test Generation is now available in beta!

Learn more here, or try it out under "Finishing Touches" below.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

• Review comments: Directly reply to a review comment made by CodeRabbit. Example:
  • I pushed a fix in commit <commit_id>, please review it.
  • Explain this complex logic.
  • Open a follow-up GitHub issue for this discussion.
• Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
  • @coderabbitai explain this code block.
• PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
  • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
  • @coderabbitai read src/utils.ts and explain its main purpose.
  • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.

Support

Need help? Create a ticket on our support page for assistance with any issues or questions.

CodeRabbit Commands (Invoked using PR comments)

• @coderabbitai pause to pause the reviews on a PR.
• @coderabbitai resume to resume the paused reviews.
• @coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
• @coderabbitai full review to do a full review from scratch and review all the files again.
• @coderabbitai summary to regenerate the summary of the PR.
• @coderabbitai generate sequence diagram to generate a sequence diagram of the changes in this PR.
• @coderabbitai resolve resolve all the CodeRabbit review comments.
• @coderabbitai configuration to show the current CodeRabbit configuration for the repository.
• @coderabbitai help to get help.

Other keywords and placeholders

• Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
• Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
• Add @coderabbitai anywhere in the PR title to generate the title automatically.

Documentation and Community

• Visit our Documentation for detailed information on how to use CodeRabbit.
• Join our Discord Community to get help, request features, and share feedback.
• Follow us on X/Twitter for updates and announcements.

[coderabbitai]

Actionable comments posted: 4

🧹 Nitpick comments (3)

lib/runtime/src/lib.rs (1)

191-198: Document the partial clone behavior more prominently

The Clone implementation only clones the Prometheus registry and resets callbacks to empty. While this is documented in the inline comment, this surprising behavior should be more prominently documented at the struct level to prevent confusion.

Add documentation to the struct definition:

 /// Structure to hold Prometheus registries and associated callbacks for a given prefix
+/// 
+/// Note: When cloning this struct, only the Prometheus registry is cloned. 
+/// Callbacks cannot be cloned and will be reset to an empty list in the cloned instance.
 pub struct MetricsRegistryEntry {

lib/runtime/src/transports/nats.rs (1)

76-76: Remove or explain commented code

The commented line //nats_metrics.copy_from_nats_client_stats(&self.client); should either be removed or explained with a comment if it's intentionally disabled.

-        //nats_metrics.copy_from_nats_client_stats(&self.client);
-

lib/runtime/src/http_server.rs (1)

351-357: Make test assertion order-agnostic

Prometheus doesn’t guarantee metric ordering. Building a single string and asserting exact equality can cause flaky tests when the registry order changes. Prefer comparing sets, sorting both strings, or asserting that expected lines are contained in the output.

📜 Review details

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 78ffe3d and b8d1fc4.

📒 Files selected for processing (5)

• lib/runtime/src/distributed.rs (5 hunks)
• lib/runtime/src/http_server.rs (12 hunks)
• lib/runtime/src/lib.rs (2 hunks)
• lib/runtime/src/metrics.rs (6 hunks)
• lib/runtime/src/transports/nats.rs (3 hunks)

🧰 Additional context used

🧠 Learnings (6)

📓 Common learnings

Learnt from: kthui
PR: ai-dynamo/dynamo#1424
File: lib/runtime/src/pipeline/network/egress/push_router.rs:204-209
Timestamp: 2025-06-13T22:07:24.843Z
Learning: The codebase uses async-nats version 0.40, not the older nats crate. Error handling should use async_nats::error::Error variants, not nats::Error variants.

📚 Learning: the asyncenginecontextprovider trait in lib/runtime/src/engine.rs was intentionally changed from `se...

Learnt from: ryanolson
PR: ai-dynamo/dynamo#1919
File: lib/runtime/src/engine.rs:168-168
Timestamp: 2025-07-14T21:25:56.930Z
Learning: The AsyncEngineContextProvider trait in lib/runtime/src/engine.rs was intentionally changed from `Send + Sync + Debug` to `Send + Debug` because the Sync bound was overly constraining. The trait should only require Send + Debug as designed.

Applied to files:

• lib/runtime/src/distributed.rs
• lib/runtime/src/http_server.rs
• lib/runtime/src/lib.rs
• lib/runtime/src/metrics.rs

📚 Learning: in lib/runtime/src/component/client.rs, the current mutex usage in get_or_create_dynamic_instance_so...

Learnt from: grahamking
PR: ai-dynamo/dynamo#1962
File: lib/runtime/src/component/client.rs:270-273
Timestamp: 2025-07-16T12:41:12.543Z
Learning: In lib/runtime/src/component/client.rs, the current mutex usage in get_or_create_dynamic_instance_source is temporary while evaluating whether the mutex can be dropped entirely. The code currently has a race condition between try_lock and lock().await, but this is acknowledged as an interim state during the performance optimization process.

Applied to files:

• lib/runtime/src/distributed.rs
• lib/runtime/src/http_server.rs
• lib/runtime/src/lib.rs
• lib/runtime/src/metrics.rs

📚 Learning: in rust/tokio applications, when background tasks use channels for communication, dropping the sende...

Learnt from: oandreeva-nv
PR: ai-dynamo/dynamo#1195
File: lib/llm/tests/block_manager.rs:150-152
Timestamp: 2025-06-02T19:37:27.666Z
Learning: In Rust/Tokio applications, when background tasks use channels for communication, dropping the sender automatically signals task termination when the receiver gets `None`. The `start_batching_publisher` function in `lib/llm/tests/block_manager.rs` demonstrates this pattern: when the `KVBMDynamoRuntimeComponent` is dropped, its `batch_tx` sender is dropped, causing `rx.recv()` to return `None`, which triggers cleanup and task termination.

Applied to files:

• lib/runtime/src/http_server.rs

📚 Learning: the codebase uses async-nats version 0.40, not the older nats crate. error handling should use async...

Learnt from: kthui
PR: ai-dynamo/dynamo#1424
File: lib/runtime/src/pipeline/network/egress/push_router.rs:204-209
Timestamp: 2025-06-13T22:07:24.843Z
Learning: The codebase uses async-nats version 0.40, not the older nats crate. Error handling should use async_nats::error::Error variants, not nats::Error variants.

Applied to files:

• lib/runtime/src/transports/nats.rs

📚 Learning: in async-nats, the "no responders" error is represented as async_nats::client::requesterrorkind::nor...

Learnt from: kthui
PR: ai-dynamo/dynamo#1424
File: lib/runtime/src/pipeline/network/egress/push_router.rs:204-209
Timestamp: 2025-06-13T22:32:05.022Z
Learning: In async-nats, the "no responders" error is represented as async_nats::client::RequestErrorKind::NoResponders, not async_nats::Error::NoResponders. Use err.downcast_ref::<async_nats::client::RequestError>() and then check request_err.kind() against RequestErrorKind::NoResponders.

Applied to files:

• lib/runtime/src/transports/nats.rs

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)

• GitHub Check: pre-merge-rust (lib/runtime/examples)
• GitHub Check: pre-merge-rust (lib/bindings/python)
• GitHub Check: pre-merge-rust (.)
• GitHub Check: Build and Test - vllm

🔇 Additional comments (4)

lib/runtime/src/distributed.rs (2)

113-118: LGTM! Good error handling for metrics registration

The approach of logging warnings instead of failing initialization is appropriate for metrics registration. This ensures the runtime remains functional even if metrics setup fails.

---

297-338: Well-structured test coverage

The test module provides comprehensive coverage of the metrics registry functionality, including callback execution, prefix handling, and Prometheus metric registration.

lib/runtime/src/metrics.rs (1)

1064-1157: Excellent test strategy for dynamic metrics

The approach of filtering NATS metrics for deterministic comparisons while separately verifying their presence is well-designed. This pattern effectively handles testing systems with dynamic components.

lib/runtime/src/http_server.rs (1)

183-185: Terminology update looks consistent

Renaming the log prefix to “System server (HTTP)” keeps the wording aligned across the file.