From 9074bfcb24e138c6b83efaf4ab3a42d66b1f2284 Mon Sep 17 00:00:00 2001 From: Thomas Montfort Date: Thu, 16 Apr 2026 20:16:16 -0700 Subject: [PATCH 1/5] fix(health): use local endpoint registry for canary health checks Replace the PushRouter/discovery/direct() pipeline in the canary health check with a direct in-process call via LocalEndpointRegistry. The previous approach was fragile in Kubernetes because it depended on discovery settling (EndpointSlice + DynamoWorkerMetadata CR correlation) and required an exact instance_id match via PushRouter::direct() that broke on pod restarts due to stale data from previous pods. The LocalEndpointRegistry is already populated during EndpointConfigBuilder::start() via .register_local_engine(), using the same endpoint.name key as health check target registration. This eliminates the discovery race condition entirely. Removed: - RouterCache type alias and router_cache field - get_or_create_router() method - Discovery/PushRouter/direct() logic in send_health_check_request() - Unused imports (Client, Component, Endpoint, Instance, PushRouter, etc.) Relates to DIS-1185 Signed-off-by: Tyler Montfort Co-Authored-By: Claude Opus 4.6 (1M context) --- lib/runtime/src/health_check.rs | 151 +++++++------------------------- 1 file changed, 32 insertions(+), 119 deletions(-) diff --git a/lib/runtime/src/health_check.rs b/lib/runtime/src/health_check.rs index 72ba4f986af4..df995ac1868c 100644 --- a/lib/runtime/src/health_check.rs +++ b/lib/runtime/src/health_check.rs @@ -1,21 +1,16 @@ // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use crate::component::{Client, Component, Endpoint, Instance}; use crate::config::HealthStatus; -use crate::pipeline::PushRouter; -use crate::pipeline::{AsyncEngine, Context, ManyOut, SingleIn}; -use crate::protocols::annotated::Annotated; -use crate::protocols::maybe_error::MaybeError; -use crate::{DistributedRuntime, SystemHealth}; +use crate::engine::AsyncEngine; +use crate::pipeline::SingleIn; +use crate::DistributedRuntime; use futures::StreamExt; use parking_lot::Mutex; -use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::sync::Arc; -use std::time::{Duration, Instant}; +use std::time::Duration; use tokio::task::JoinHandle; -use tokio::time::{MissedTickBehavior, interval}; use tracing::{debug, error, info, warn}; /// Configuration for health check behavior @@ -37,17 +32,10 @@ impl Default for HealthCheckConfig { } } -// Type alias for the router cache to improve readability -// Maps endpoint subject -> router and payload -type RouterCache = - Arc>>>>>; - /// Health check manager that monitors endpoint health pub struct HealthCheckManager { drt: DistributedRuntime, config: HealthCheckConfig, - /// Cache of PushRouters and payloads for each endpoint - router_cache: RouterCache, /// Track per-endpoint health check tasks /// Maps: endpoint_subject -> task_handle endpoint_tasks: Arc>>>, @@ -58,45 +46,10 @@ impl HealthCheckManager { Self { drt, config, - router_cache: Arc::new(Mutex::new(HashMap::new())), endpoint_tasks: Arc::new(Mutex::new(HashMap::new())), } } - /// Get or create a PushRouter for an endpoint - async fn get_or_create_router( - &self, - cache_key: &str, - endpoint: Endpoint, - ) -> anyhow::Result>>> { - let cache_key = cache_key.to_string(); - - // Check cache first - { - let cache = self.router_cache.lock(); - if let Some(router) = cache.get(&cache_key) { - return Ok(router.clone()); - } - } - - // Create a client that discovers instances dynamically for this endpoint - let client = Client::new(endpoint).await?; - - // Create PushRouter - it will use direct routing when we call direct() - let router: Arc>> = Arc::new( - PushRouter::from_client( - client, - crate::pipeline::RouterMode::RoundRobin, // Default mode, we'll use direct() explicitly - ) - .await?, - ); - - // Cache it - self.router_cache.lock().insert(cache_key, router.clone()); - - Ok(router) - } - /// Start the health check manager by spawning per-endpoint monitoring tasks pub async fn start(self: Arc) -> anyhow::Result<()> { // Get all registered endpoints at startup @@ -237,96 +190,56 @@ impl HealthCheckManager { Ok(()) } - /// Send a health check request through AsyncEngine + /// Send a health check request via the local endpoint registry (in-process). + /// + /// This calls the engine directly without going through discovery, routing, + /// or network transport. The local registry is populated during + /// `EndpointConfigBuilder::start()` via `.register_local_engine()`, using + /// the same `endpoint.name` key as health check target registration. async fn send_health_check_request( &self, endpoint_subject: &str, payload: &serde_json::Value, ) -> anyhow::Result<()> { - let target = self + debug!("Sending health check to {} via local registry", endpoint_subject); + + // Look up the engine in the local endpoint registry (in-process call only). + // If not found, the endpoint hasn't finished registering yet — the canary + // will retry on the next cycle. + let engine = self .drt - .system_health() - .lock() - .get_health_check_target(endpoint_subject) + .local_endpoint_registry() + .get(endpoint_subject) .ok_or_else(|| { - anyhow::anyhow!("No health check target found for {}", endpoint_subject) - })?; - - debug!( - "Sending health check to {} (instance_id: {})", - endpoint_subject, target.instance.instance_id - ); - - // Create the Endpoint directly from the Instance info - let namespace = self.drt.namespace(&target.instance.namespace)?; - let component = namespace.component(&target.instance.component)?; - let endpoint = component.endpoint(&target.instance.endpoint); - - // Get or create router for this endpoint - let router = self - .get_or_create_router(endpoint_subject, endpoint) - .await?; - - // Wait for watch stream to discover instances before checking - // This ensures the router's client has populated its instance list - // from etcd before we attempt to send the health check request. - // Without this, the first health check can fail due to a race condition - // where the watch stream hasn't completed its initial discovery yet. - match tokio::time::timeout( - Duration::from_secs(10), // 10 second timeout for discovery - router.client.wait_for_instances(), - ) - .await - { - Ok(Ok(instances)) => { - debug!( - "Health check for {}: watch stream ready, found {} instance(s)", - endpoint_subject, - instances.len() - ); - } - Ok(Err(e)) => { - return Err(anyhow::anyhow!( - "Failed to discover instances for {} during health check: {}", - endpoint_subject, - e - )); - } - Err(_) => { - return Err(anyhow::anyhow!( - "Timeout waiting for instance discovery for {} during health check", + anyhow::anyhow!( + "Endpoint '{}' not found in local registry, engine may still be initializing", endpoint_subject - )); - } - } - - // Create the request context - let request: SingleIn = Context::new(payload.clone()); + ) + })?; // Clone what we need for the spawned task let system_health = self.drt.system_health().clone(); let endpoint_subject_owned = endpoint_subject.to_string(); - let instance_id = target.instance.instance_id; + let payload = payload.clone(); let timeout = self.config.request_timeout; // Spawn task to send health check and wait for response tokio::spawn(async move { let result = tokio::time::timeout(timeout, async { - // Call direct() on the PushRouter to target specific instance - match router.direct(request, instance_id).await { + let request = SingleIn::new(payload); + match engine.generate(request).await { Ok(mut response_stream) => { // Get the first response to verify endpoint is alive let is_healthy = if let Some(response) = response_stream.next().await { - // Check if response indicates an error - if let Some(error) = response.err() { + if response.data.is_some() { + debug!("Health check successful for {}", endpoint_subject_owned); + true + } else { warn!( - "Health check error response from {}: {:?}", - endpoint_subject_owned, error + "Health check got empty response from {}", + endpoint_subject_owned ); false - } else { - debug!("Health check successful for {}", endpoint_subject_owned); - true } } else { warn!( @@ -336,8 +249,8 @@ impl HealthCheckManager { false }; + // Consume the rest of the stream in the background tokio::spawn(async move { - // We need to consume the rest of the stream to avoid warnings on the frontend. response_stream.for_each(|_| async {}).await; }); From bb0e84dd399013f7bb2ab40681f1aafd6b8ec9d0 Mon Sep 17 00:00:00 2001 From: Thomas Montfort Date: Fri, 17 Apr 2026 11:08:58 -0700 Subject: [PATCH 2/5] style(health): fix cargo fmt formatting Signed-off-by: Thomas Montfort Co-Authored-By: Claude Opus 4.6 (1M context) --- lib/runtime/src/health_check.rs | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/lib/runtime/src/health_check.rs b/lib/runtime/src/health_check.rs index df995ac1868c..ccf72cd55a33 100644 --- a/lib/runtime/src/health_check.rs +++ b/lib/runtime/src/health_check.rs @@ -1,10 +1,10 @@ // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +use crate::DistributedRuntime; use crate::config::HealthStatus; use crate::engine::AsyncEngine; use crate::pipeline::SingleIn; -use crate::DistributedRuntime; use futures::StreamExt; use parking_lot::Mutex; use std::collections::HashMap; @@ -201,7 +201,10 @@ impl HealthCheckManager { endpoint_subject: &str, payload: &serde_json::Value, ) -> anyhow::Result<()> { - debug!("Sending health check to {} via local registry", endpoint_subject); + debug!( + "Sending health check to {} via local registry", + endpoint_subject + ); // Look up the engine in the local endpoint registry (in-process call only). // If not found, the endpoint hasn't finished registering yet — the canary From c3a851f1a0148efc30af96ee4e313c5f2b61c8d6 Mon Sep 17 00:00:00 2001 From: Thomas Montfort Date: Fri, 17 Apr 2026 11:23:31 -0700 Subject: [PATCH 3/5] fix(health): fail fast when canary enabled but local engine missing When health_check_payload is set and canary is enabled (DYN_HEALTH_CHECK_ENABLED=true), but the caller forgot to call .register_local_engine() before .start(), the endpoint now fails to start with a clear error instead of silently staying NotReady forever while the canary retries against an empty local registry. Signed-off-by: Thomas Montfort Co-Authored-By: Claude Opus 4.6 (1M context) --- lib/runtime/src/component/endpoint.rs | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/lib/runtime/src/component/endpoint.rs b/lib/runtime/src/component/endpoint.rs index e9ca3940fbc5..357b7f765efc 100644 --- a/lib/runtime/src/component/endpoint.rs +++ b/lib/runtime/src/component/endpoint.rs @@ -144,6 +144,24 @@ impl EndpointConfigBuilder { // Register health check target in SystemHealth if provided if let Some(health_check_payload) = &health_check_payload { + // When canary health checks are enabled, the canary calls the engine + // via LocalEndpointRegistry. Fail fast if the engine wasn't registered + // (caller forgot .register_local_engine() before .start()). + if system_health.lock().health_check_enabled() + && endpoint + .drt() + .local_endpoint_registry() + .get(&endpoint.name) + .is_none() + { + anyhow::bail!( + "Endpoint '{}' has a health_check_payload and canary is enabled, \ + but no local engine is registered. Call .register_local_engine() \ + before .start() so the canary health check can function.", + endpoint.name + ); + } + // Build transport based on request plane mode let transport = build_transport_type(&endpoint, &endpoint_id, connection_id).await?; From 7ba2146a26f17a0098f48aec9e6ecaf9d04bd404 Mon Sep 17 00:00:00 2001 From: Thomas Montfort Date: Fri, 17 Apr 2026 11:42:36 -0700 Subject: [PATCH 4/5] fix(health): use error-based check instead of data-presence check Match the original semantics: treat the first stream response as healthy unless it contains an explicit error. An Annotated item can have data: None without being an error (e.g., annotation/metadata events), so checking data.is_some() was too strict. Signed-off-by: Thomas Montfort Co-Authored-By: Claude Opus 4.6 (1M context) --- lib/runtime/src/health_check.rs | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/lib/runtime/src/health_check.rs b/lib/runtime/src/health_check.rs index ccf72cd55a33..8c2e36ff2aa8 100644 --- a/lib/runtime/src/health_check.rs +++ b/lib/runtime/src/health_check.rs @@ -5,6 +5,7 @@ use crate::DistributedRuntime; use crate::config::HealthStatus; use crate::engine::AsyncEngine; use crate::pipeline::SingleIn; +use crate::protocols::maybe_error::MaybeError; use futures::StreamExt; use parking_lot::Mutex; use std::collections::HashMap; @@ -232,17 +233,21 @@ impl HealthCheckManager { let request = SingleIn::new(payload); match engine.generate(request).await { Ok(mut response_stream) => { - // Get the first response to verify endpoint is alive + // Get the first response to verify endpoint is alive. + // Check for errors rather than data presence: an Annotated + // item can have data: None without being an error (e.g., + // annotation/metadata events). Only explicit errors indicate + // an unhealthy endpoint. let is_healthy = if let Some(response) = response_stream.next().await { - if response.data.is_some() { - debug!("Health check successful for {}", endpoint_subject_owned); - true - } else { + if let Some(error) = response.err() { warn!( - "Health check got empty response from {}", - endpoint_subject_owned + "Health check error response from {}: {:?}", + endpoint_subject_owned, error ); false + } else { + debug!("Health check successful for {}", endpoint_subject_owned); + true } } else { warn!( From 4200e0dc3293c62c037c54ccfa416f1227553c59 Mon Sep 17 00:00:00 2001 From: tmontfort Date: Fri, 17 Apr 2026 13:56:03 -0700 Subject: [PATCH 5/5] remove verbose comments Signed-off-by: tmontfort --- lib/runtime/src/component/endpoint.rs | 3 --- lib/runtime/src/health_check.rs | 15 ++------------- 2 files changed, 2 insertions(+), 16 deletions(-) diff --git a/lib/runtime/src/component/endpoint.rs b/lib/runtime/src/component/endpoint.rs index 357b7f765efc..5c7277b41a05 100644 --- a/lib/runtime/src/component/endpoint.rs +++ b/lib/runtime/src/component/endpoint.rs @@ -144,9 +144,6 @@ impl EndpointConfigBuilder { // Register health check target in SystemHealth if provided if let Some(health_check_payload) = &health_check_payload { - // When canary health checks are enabled, the canary calls the engine - // via LocalEndpointRegistry. Fail fast if the engine wasn't registered - // (caller forgot .register_local_engine() before .start()). if system_health.lock().health_check_enabled() && endpoint .drt() diff --git a/lib/runtime/src/health_check.rs b/lib/runtime/src/health_check.rs index 8c2e36ff2aa8..d9106e0a5d27 100644 --- a/lib/runtime/src/health_check.rs +++ b/lib/runtime/src/health_check.rs @@ -192,11 +192,6 @@ impl HealthCheckManager { } /// Send a health check request via the local endpoint registry (in-process). - /// - /// This calls the engine directly without going through discovery, routing, - /// or network transport. The local registry is populated during - /// `EndpointConfigBuilder::start()` via `.register_local_engine()`, using - /// the same `endpoint.name` key as health check target registration. async fn send_health_check_request( &self, endpoint_subject: &str, @@ -207,9 +202,6 @@ impl HealthCheckManager { endpoint_subject ); - // Look up the engine in the local endpoint registry (in-process call only). - // If not found, the endpoint hasn't finished registering yet — the canary - // will retry on the next cycle. let engine = self .drt .local_endpoint_registry() @@ -234,10 +226,7 @@ impl HealthCheckManager { match engine.generate(request).await { Ok(mut response_stream) => { // Get the first response to verify endpoint is alive. - // Check for errors rather than data presence: an Annotated - // item can have data: None without being an error (e.g., - // annotation/metadata events). Only explicit errors indicate - // an unhealthy endpoint. + // Check for errors let is_healthy = if let Some(response) = response_stream.next().await { if let Some(error) = response.err() { warn!( @@ -257,8 +246,8 @@ impl HealthCheckManager { false }; - // Consume the rest of the stream in the background tokio::spawn(async move { + // We need to consume the rest of the stream to avoid warnings on the frontend. response_stream.for_each(|_| async {}).await; });