chore: Make nats_client private at crate level, various tidy up (#4513)

Signed-off-by: Graham King <grahamk@nvidia.com>

Author: grahamking

lib/llm/src/audit/bus.rs

@@ -2,22 +2,22 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use super::handle::AuditRecord;
-use std::sync::{Arc, OnceLock};
+use std::sync::OnceLock;
 use tokio::sync::broadcast;
 
-static BUS: OnceLock<broadcast::Sender<Arc<AuditRecord>>> = OnceLock::new();
+static BUS: OnceLock<broadcast::Sender<AuditRecord>> = OnceLock::new();
 
 pub fn init(capacity: usize) {
-    let (tx, _rx) = broadcast::channel::<Arc<AuditRecord>>(capacity);
+    let (tx, _rx) = broadcast::channel::<AuditRecord>(capacity);
     let _ = BUS.set(tx);
 }
 
-pub fn subscribe() -> broadcast::Receiver<Arc<AuditRecord>> {
+pub fn subscribe() -> broadcast::Receiver<AuditRecord> {
     BUS.get().expect("audit bus not initialized").subscribe()
 }
 
 pub fn publish(rec: AuditRecord) {
     if let Some(tx) = BUS.get() {
-        let _ = tx.send(Arc::new(rec));
+        let _ = tx.send(rec);
     }
 }

lib/llm/src/audit/config.rs

@@ -10,11 +10,11 @@ pub struct AuditPolicy {
 
 static POLICY: OnceLock<AuditPolicy> = OnceLock::new();
 
+/// Audit is enabled if we have at least one sink
 pub fn init_from_env() -> AuditPolicy {
-    let enabled = std::env::var("DYN_AUDIT_ENABLED")
-        .map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
-        .unwrap_or(false);
-    AuditPolicy { enabled }
+    AuditPolicy {
+        enabled: std::env::var("DYN_AUDIT_SINKS").is_ok(),
+    }
 }
 
 pub fn policy() -> AuditPolicy {

lib/llm/src/audit/sink.rs

@@ -1,8 +1,10 @@
 // SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 // SPDX-License-Identifier: Apache-2.0
 
+use anyhow::Context as _;
 use async_nats::jetstream;
 use async_trait::async_trait;
+use dynamo_runtime::transports::nats;
 use std::sync::Arc;
 use tokio::sync::broadcast;
 
@@ -36,7 +38,7 @@ pub struct NatsSink {
 }
 
 impl NatsSink {
-    pub fn new(nats_client: &dynamo_runtime::transports::nats::Client) -> Self {
+    pub fn new(nats_client: dynamo_runtime::transports::nats::Client) -> Self {
         let subject = std::env::var("DYN_AUDIT_NATS_SUBJECT")
             .unwrap_or_else(|_| "dynamo.audit.v1".to_string());
         Self {
@@ -64,37 +66,32 @@ impl AuditSink for NatsSink {
     }
 }
 
-fn parse_sinks_from_env(
-    nats_client: Option<&dynamo_runtime::transports::nats::Client>,
-) -> Vec<Arc<dyn AuditSink>> {
+async fn parse_sinks_from_env() -> anyhow::Result<Vec<Arc<dyn AuditSink>>> {
     let cfg = std::env::var("DYN_AUDIT_SINKS").unwrap_or_else(|_| "stderr".into());
     let mut out: Vec<Arc<dyn AuditSink>> = Vec::new();
     for name in cfg.split(',').map(|s| s.trim().to_lowercase()) {
         match name.as_str() {
             "stderr" | "" => out.push(Arc::new(StderrSink)),
             "nats" => {
-                if let Some(client) = nats_client {
-                    out.push(Arc::new(NatsSink::new(client)));
-                } else {
-                    tracing::warn!(
-                        "NATS sink requested but no DistributedRuntime NATS client available; skipping"
-                    );
-                }
+                let nats_client = nats::ClientOptions::default()
+                    .connect()
+                    .await
+                    .context("Attempting to connect NATS sink from env var DYN_AUDIT_SINKS")?;
+                out.push(Arc::new(NatsSink::new(nats_client)));
             }
             // "pg"   => out.push(Arc::new(PostgresSink::from_env())),
             other => tracing::warn!(%other, "audit: unknown sink ignored"),
         }
     }
-    out
+    Ok(out)
 }
 
 /// spawn one worker per sink; each subscribes to the bus (off hot path)
-pub fn spawn_workers_from_env(drt: &dynamo_runtime::DistributedRuntime) {
-    let nats_client = drt.nats_client();
-    let sinks = parse_sinks_from_env(nats_client);
+pub async fn spawn_workers_from_env() -> anyhow::Result<()> {
+    let sinks = parse_sinks_from_env().await?;
     for sink in sinks {
         let name = sink.name();
-        let mut rx: broadcast::Receiver<Arc<AuditRecord>> = bus::subscribe();
+        let mut rx: broadcast::Receiver<AuditRecord> = bus::subscribe();
         tokio::spawn(async move {
             loop {
                 match rx.recv().await {
@@ -110,4 +107,5 @@ pub fn spawn_workers_from_env(drt: &dynamo_runtime::DistributedRuntime) {
         });
     }
     tracing::info!("Audit sinks ready.");
+    Ok(())
 }

lib/llm/src/block_manager/controller.rs

@@ -39,10 +39,8 @@ pub struct Controller<Locality: LocalityProvider, Metadata: BlockMetadata> {
 impl<Locality: LocalityProvider, Metadata: BlockMetadata> Controller<Locality, Metadata> {
     pub async fn new(
         block_manager: KvBlockManager<Locality, Metadata>,
-        mut component: dynamo_runtime::component::Component,
+        component: dynamo_runtime::component::Component,
     ) -> anyhow::Result<Self> {
-        component.add_stats_service().await?;
-
         let handler = ControllerHandler::new(block_manager.clone());
         let engine = Ingress::for_engine(handler.clone())?;
 

lib/llm/src/entrypoint/input.rs

@@ -117,7 +117,7 @@ pub async fn run_input(
             .and_then(|v| v.parse().ok())
             .unwrap_or(1024);
         crate::audit::bus::init(cap);
-        crate::audit::sink::spawn_workers_from_env(&drt);
+        crate::audit::sink::spawn_workers_from_env().await?;
         tracing::info!(cap, "Audit initialized");
     }
 

lib/llm/src/entrypoint/input/endpoint.rs

@@ -32,14 +32,9 @@ pub async fn run(
     let cancel_token = distributed_runtime.primary_token().clone();
     let endpoint_id: EndpointId = path.parse()?;
 
-    let mut component = distributed_runtime
+    let component = distributed_runtime
         .namespace(&endpoint_id.namespace)?
         .component(&endpoint_id.component)?;
-
-    // We can only make the NATS service if we have NATS
-    if distributed_runtime.nats_client().is_some() {
-        component.add_stats_service().await?;
-    }
     let endpoint = component.endpoint(&endpoint_id.name);
 
     let rt_fut: Pin<Box<dyn Future<Output = _> + Send + 'static>> = match engine_config {

lib/llm/src/kv_router/sequence.rs

@@ -1002,8 +1002,7 @@ mod tests {
 
         // Create namespace and shared component for both seq_managers
         let namespace = distributed.namespace("test_cross_instance_sync")?;
-        let mut component = namespace.component("sequences")?;
-        component.add_stats_service().await?;
+        let component = namespace.component("sequences")?;
 
         // Create multi-worker sequence managers with:
         // - Worker 0 with dp_size=2 (dp_ranks 0 and 1)
@@ -1168,8 +1167,7 @@ mod tests {
 
         // Create namespace and shared component for both seq_managers
         let namespace = distributed.namespace("test_no_token_seq_sync")?;
-        let mut component = namespace.component("sequences")?;
-        component.add_stats_service().await?;
+        let component = namespace.component("sequences")?;
 
         // Create multi-worker sequence managers with ALL workers [0, 1, 2]
         // Both use the same component to ensure event synchronization works

lib/llm/src/mocker/engine.rs

@@ -463,8 +463,7 @@ mod integration_tests {
         tracing::info!("✓ Runtime and distributed runtime created");
 
         // Create component for MockVllmEngine (needed for publishers)
-        let mut test_component = distributed.namespace("test")?.component(MOCKER_COMPONENT)?;
-        test_component.add_stats_service().await?;
+        let test_component = distributed.namespace("test")?.component(MOCKER_COMPONENT)?;
         tracing::info!("✓ Test component created");
 
         // Create MockVllmEngine WITH component (enables publishers)

lib/llm/src/preprocessor.rs

@@ -856,7 +856,7 @@ impl
         let request_id = context.id().to_string();
         let original_stream_flag = request.inner.stream.unwrap_or(false);
 
-        // Build audit handle (None if DYN_AUDIT_ENABLED=0)
+        // Build audit handle (None if no DYN_AUDIT_SINKS)
         let mut audit_handle = crate::audit::handle::create_handle(&request, &request_id);
 
         if let Some(ref mut h) = audit_handle {

lib/llm/tests/audit_nats_integration.rs

@@ -27,7 +27,6 @@ mod tests {
     use dynamo_llm::protocols::openai::chat_completions::{
         NvCreateChatCompletionRequest, NvCreateChatCompletionResponse,
     };
-    use dynamo_runtime::Runtime;
     use dynamo_runtime::transports::nats;
     use futures::StreamExt;
     use serde_json::Value;
@@ -48,15 +47,6 @@ mod tests {
             .expect("Failed to connect to NATS server")
     }
 
-    /// Helper to create a test DistributedRuntime with NATS
-    async fn create_test_drt() -> dynamo_runtime::DistributedRuntime {
-        let rt = Runtime::from_current().unwrap();
-        let config = dynamo_runtime::distributed::DistributedConfig::from_settings();
-        dynamo_runtime::DistributedRuntime::new(rt, config)
-            .await
-            .expect("Failed to create DistributedRuntime")
-    }
-
     /// Helper to create a minimal test request
     fn create_test_request(model: &str, store: bool) -> NvCreateChatCompletionRequest {
         let json = serde_json::json!({
@@ -155,7 +145,6 @@ mod tests {
         // Core test: audit records are published to NATS with correct structure
         async_with_vars(
             [
-                ("DYN_AUDIT_ENABLED", Some("1")),
                 ("DYN_AUDIT_SINKS", Some("nats")),
                 ("DYN_AUDIT_NATS_SUBJECT", Some(TEST_SUBJECT)),
             ],
@@ -166,8 +155,7 @@ mod tests {
                 setup_test_stream(&client, &stream_name, TEST_SUBJECT).await;
 
                 bus::init(100);
-                let drt = create_test_drt().await;
-                sink::spawn_workers_from_env(&drt);
+                sink::spawn_workers_from_env().await.unwrap();
                 time::sleep(Duration::from_millis(100)).await;
 
                 // Emit audit record
@@ -212,7 +200,6 @@ mod tests {
 
         async_with_vars(
             [
-                ("DYN_AUDIT_ENABLED", Some("1")),
                 ("DYN_AUDIT_SINKS", Some("nats")),
                 ("DYN_AUDIT_NATS_SUBJECT", Some(TEST_SUBJECT)),
             ],
@@ -223,8 +210,7 @@ mod tests {
                 setup_test_stream(&client, &stream_name, TEST_SUBJECT).await;
 
                 bus::init(100);
-                let drt = create_test_drt().await;
-                sink::spawn_workers_from_env(&drt);
+                sink::spawn_workers_from_env().await.unwrap();
                 time::sleep(Duration::from_millis(100)).await;
 
                 // Request with store=true (should be audited)