fix(health): use local endpoint registry for canary health checks

[tmonty12]

Overview:

The canary health check used PushRouter/discovery/direct() to send a test request back to itself through the full distributed pipeline. This was fragile in Kubernetes because it depended on discovery settling (EndpointSlice + DynamoWorkerMetadata CR correlation), and PushRouter::direct(instance_id) required an exact instance_id match that broke on pod restarts due to stale data from previous pods. The wait_for_instances() call only checked that any instance existed — not that the correct instance was routable — so the canary would discover a stale pod, then fail with "instance_id not found" when trying to direct-route to itself.

This PR replaces the entire discovery/routing path with a direct in-process call via LocalEndpointRegistry, which is already populated during EndpointConfigBuilder::start() using the same endpoint.name key. This eliminates all discovery timing dependencies and instance_id matching issues.

Details:

• lib/runtime/src/health_check.rs:
  • Replaced send_health_check_request() to use drt.local_endpoint_registry().get(endpoint_subject) → engine.generate() (same pattern as call_lora_endpoint in system_status_server.rs)
  • Removed RouterCache type alias, router_cache field, and get_or_create_router() method
  • Removed unused imports: Client, Component, Endpoint, Instance, PushRouter, Context, ManyOut, Annotated, MaybeError, Serialize, Deserialize, Instant, MissedTickBehavior, interval
  • No changes to the timer/notifier/per-endpoint task lifecycle (start(), spawn_endpoint_health_check_task(), spawn_new_endpoint_monitor())

Where should the reviewer start?

lib/runtime/src/health_check.rs — specifically the new send_health_check_request() method (line ~240). The old method was ~140 lines of discovery/routing logic; the new one is ~80 lines with a single local_endpoint_registry.get() → engine.generate() call.

Related Issues:

• Relates to DIS-1185

Summary by CodeRabbit

• Refactor
  • Simplified internal health check mechanism by removing router-discovery machinery and updating endpoint resolution to use local registry. No public API changes.

[nnshah1]

Nice simplification — eliminating the discovery/routing path for canary is a big improvement.

One concern: register_local_engine() is called on the builder (optional), not inside start(). Today all backends go through the Python binding which always calls it (lib.rs:881). But if a future Rust-native backend uses EndpointConfigBuilder with a health_check_payload and forgets register_local_engine(), the canary silently skips it — just logs "endpoint not found in local registry" every cycle with no obvious error.

Suggestion: could start() automatically register the local engine when health_check_payload is provided? Or at minimum, emit a warn! in start() if a payload is set but no local engine is registered — that way it's not a silent failure.

Also heads up — this PR includes the health_check_enabled gate from #8165 (now merged to main). You'll want to rebase to pick up the latest version which renames set_ready() → set_endpoint_registered() and trims the comments.

[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.

[devin-ai-integration]

Devin Review found 2 potential issues.

View 2 additional findings in Devin Review.

[coderabbitai]

Walkthrough

Refactored HealthCheckManager to remove router-discovery and routing machinery. Reworked health check request flow from routing via discovered router to querying the local endpoint registry directly. Updated request/response handling to use engine invocation instead of router operations.

Changes

Cohort / File(s)

Summary

Health Check Manager Refactor lib/runtime/src/health_check.rs

Removed RouterCache type alias, router_cache field, and get_or_create_router() async helper. Reworked send_health_check_request() to use drt.local_endpoint_registry().get() instead of router discovery. Replaced router.direct(request, instance_id) with engine.generate(SingleIn::new(payload)) and changed error detection from response.err() to response.data.is_some(). Updated control-flow and logging for in-process local registry path.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~50 minutes

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly describes the main change: replacing PushRouter/discovery machinery with LocalEndpointRegistry in the health check module.
Description check	✅ Passed	The description covers all required template sections: Overview (problem and solution), Details (specific code changes with file references), Where should the reviewer start (specific file and line guidance), and Related Issues (with the DIS-1185 linear issue link).
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

---

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

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (2)

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

232-253: ⚠️ Potential issue | 🟠 Major

Use response.is_ok() instead of response.data.is_some() to check health.

The current check incorrectly treats any frame with data: None as unhealthy. The Annotated type defines is_ok() as event.as_deref() != Some("error")—the event field signals error status, not data presence. Metadata and event frames (SSE keep-alives, comments, ids) commonly have data: None and a non-error event, and would incorrectly flip the endpoint to NotReady with the current logic. Other in-process callers like lib/llm/src/http/service/openai.rs (lines 1708–1709) and lib/llm/src/http/service/anthropic.rs (lines 395–396) validate both data.is_none() AND event != "error" when inspecting stream frames. Replace with response.is_ok() at line 237, and update the data.is_none() branch (lines 244–250) to check response.is_err() instead.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@lib/runtime/src/health_check.rs` around lines 232 - 253, The health check
currently treats any frame with response.data == None as failure; instead use
the Annotated success helpers: when reading the first frame from the stream
produced by engine.generate(SingleIn::new(payload)) check response.is_ok() to
mark healthy, and change the branch that currently checks data.is_none() to
check response.is_err() to log warnings/failure; update usage around
response_stream.next().await, response.is_ok(), response.is_err(), and keep
endpoint_subject_owned in the log messages.

---

107-109: ⚠️ Potential issue | 🟡 Minor

Downgrade log level for expected startup race condition.

The PR description correctly notes that during startup, missing engines mean "the endpoint hasn't finished registering yet—the canary will retry on the next cycle." However, the error at lines 107–109 logs this via error!, which will spam error-level telemetry on every canary cycle during normal startup until register_local_engine() fires. Consider returning a typed error (or Ok(()) with a debug log) for the "not yet registered" case so the caller can log it at debug!/warn! instead of error!.

Note: The concern about deregistration causing "stuck-at-Ready" is unfounded. LocalEndpointRegistry contains no remove(), unregister(), or clear() methods—once an engine is registered, it cannot be deregistered. The registry's design prevents this scenario entirely.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@lib/runtime/src/health_check.rs` around lines 107 - 109, The current error
log in the health-check loop indiscriminately logs any failure from
manager.send_health_check_request as error!, causing noisy telemetry for an
expected startup race; change send_health_check_request to return a typed Result
with a specific variant for "endpoint not registered" (e.g.,
HealthCheckError::EndpointNotRegistered or MissingEngine) and then update the
caller in health_check.rs to match the returned error: if it is the
NotRegistered/MissingEngine variant, log at debug! (or Ok(()) silently) with a
brief message referencing that registration is pending (register_local_engine),
otherwise log true failures at error!. Ensure you reference and update the
send_health_check_request function and the match in the health-check loop that
currently logs the Err(e).

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Outside diff comments:
In `@lib/runtime/src/health_check.rs`:
- Around line 232-253: The health check currently treats any frame with
response.data == None as failure; instead use the Annotated success helpers:
when reading the first frame from the stream produced by
engine.generate(SingleIn::new(payload)) check response.is_ok() to mark healthy,
and change the branch that currently checks data.is_none() to check
response.is_err() to log warnings/failure; update usage around
response_stream.next().await, response.is_ok(), response.is_err(), and keep
endpoint_subject_owned in the log messages.
- Around line 107-109: The current error log in the health-check loop
indiscriminately logs any failure from manager.send_health_check_request as
error!, causing noisy telemetry for an expected startup race; change
send_health_check_request to return a typed Result with a specific variant for
"endpoint not registered" (e.g., HealthCheckError::EndpointNotRegistered or
MissingEngine) and then update the caller in health_check.rs to match the
returned error: if it is the NotRegistered/MissingEngine variant, log at debug!
(or Ok(()) silently) with a brief message referencing that registration is
pending (register_local_engine), otherwise log true failures at error!. Ensure
you reference and update the send_health_check_request function and the match in
the health-check loop that currently logs the Err(e).

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 1b4a2f25-978a-45f0-9dfc-ee842b467251

📥 Commits

Reviewing files that changed from the base of the PR and between 94ee2aa and bb0e84d.

📒 Files selected for processing (1)

• lib/runtime/src/health_check.rs

[tmonty12]

Nice simplification — eliminating the discovery/routing path for canary is a big improvement.

One concern: register_local_engine() is called on the builder (optional), not inside start(). Today all backends go through the Python binding which always calls it (lib.rs:881). But if a future Rust-native backend uses EndpointConfigBuilder with a health_check_payload and forgets register_local_engine(), the canary silently skips it — just logs "endpoint not found in local registry" every cycle with no obvious error.

Suggestion: could start() automatically register the local engine when health_check_payload is provided? Or at minimum, emit a warn! in start() if a payload is set but no local engine is registered — that way it's not a silent failure.

Also heads up — this PR includes the health_check_enabled gate from #8165 (now merged to main). You'll want to rebase to pick up the latest version which renames set_ready() → set_endpoint_registered() and trims the comments.

Good point. The python binding serve_endpoint will always call register_local_engine but not the rust native API for EndpointConfigBuilder. I've made it actually error instead of warn if the canary is enabled with a registered health check payload but no locally registered engine. The user is intending to have canary enabled but it functionally wouldn't work - this fails fast instead of undesired behavior.

[nnshah1]

LGTM - but would do a simplify comments pass - comments seem a bit verbose and not meaningful without context.

[devin-ai-integration]

Devin Review found 1 new potential issue.

View 6 additional findings in Devin Review.