Repository navigation
fix: pre-existing issues in worker health config, bootstrap parsing, and model card cloning - #415
Merged
Merged
Conversation
slin1237
requested review from
CatherineSue,
key4ng and
whybeyoung
as code owners
February 12, 2026 20:53
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughWorker health is split: specs carry optional Changes
Sequence Diagram(s)sequenceDiagram
participant Registry as WorkerRegistry
participant Checker as HealthChecker
participant Worker as Worker (metadata + endpoint)
participant Notify as tokio::Notify
Registry->>Registry: compute next-deadlines per worker from health_config
Registry->>Checker: wait for earliest deadline or Notify
Note over Checker,Notify: Checker awaits Notify or sleep-until(deadline)
Checker->>Worker: perform health probe (timeout from health_config)
Worker-->>Checker: probe result (success/failure)
Checker->>Registry: report result -> update next-deadline/metrics
Registry->>Notify: notify Checker if earlier deadline requires wakeup
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 3 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches🧪 Generate unit tests (beta)
Comment |
Contributor
Summary of ChangesHello @slin1237, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request resolves several identified issues to enhance the robustness and correctness of worker management. It improves how per-worker health check configurations are applied, ensures accurate parsing of bootstrap host URLs, guarantees that worker connection attempts are always made, and prevents the loss of user-defined model card metadata during worker initialization. These changes contribute to more reliable worker registration and operation within the system. Highlights
Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces several important fixes for pre-existing issues, enhancing the robustness and correctness of the worker configuration and management. The changes to handle per-worker health config overrides are well-implemented using HealthCheckUpdate. The fix for parsing DP-aware bootstrap URLs is a great improvement, and the correction for max_connection_attempts prevents potential stalls. Additionally, the refactoring of build_model_card to correctly clone model card data is a clean and effective solution. Overall, these are excellent and much-needed fixes. I have one minor suggestion to simplify the implementation of parse_bootstrap_host for better maintainability.
Comment on lines
+244
to
269
| // Strip DP rank suffix (e.g., "http://host:8080@3" -> "http://host:8080") | ||
| let clean_url = match url.rfind('@') { | ||
| Some(at_pos) if url[at_pos + 1..].chars().all(|c| c.is_ascii_digit()) => &url[..at_pos], | ||
| _ => url, | ||
| }; | ||
|
|
||
| // Try parsing as-is first. If the URL lacks a scheme (e.g., "worker1:8080"), | ||
| // Url::parse may treat the host as a scheme — detect this via missing host_str() | ||
| // and fall back to prefixing "http://". | ||
| let try_parse = |u: &str| -> Option<String> { | ||
| url::Url::parse(u) | ||
| .ok() | ||
| .and_then(|p| p.host_str().map(|h| h.to_string())) | ||
| }; | ||
|
|
||
| if let Some(host) = try_parse(clean_url) { | ||
| host | ||
| } else if !clean_url.contains("://") { | ||
| try_parse(&format!("http://{}", clean_url)).unwrap_or_else(|| { | ||
| tracing::warn!("Failed to parse URL '{}', defaulting to localhost", url); | ||
| "localhost".to_string() | ||
| } | ||
| }) | ||
| } else { | ||
| tracing::warn!("Failed to parse URL '{}', defaulting to localhost", url); | ||
| "localhost".to_string() | ||
| } |
Contributor
There was a problem hiding this comment.
The logic for parsing the bootstrap host is correct and handles the described edge cases well. However, the implementation can be slightly simplified to reduce code duplication, specifically the tracing::warn! and default to "localhost" part.
This refactoring consolidates the error handling path, making the function's flow a bit more linear and easier to follow.
// Strip DP rank suffix (e.g., "http://host:8080@3" -> "http://host:8080")
let clean_url = match url.rfind('@') {
Some(at_pos) if url[at_pos + 1..].chars().all(|c| c.is_ascii_digit()) => &url[..at_pos],
_ => url,
};
// Try parsing as-is first. If the URL lacks a scheme (e.g., "worker1:8080"),
// Url::parse may treat the host as a scheme — detect this via missing host_str()
// and fall back to prefixing "http://".
let try_parse = |u: &str| -> Option<String> {
url::Url::parse(u)
.ok()
.and_then(|p| p.host_str().map(|h| h.to_string()))
};
if let Some(host) = try_parse(clean_url) {
return host;
}
if !clean_url.contains("://") {
if let Some(host) = try_parse(&format!("http://{}", clean_url)) {
return host;
}
}
tracing::warn!("Failed to parse URL '{}', defaulting to localhost", url);
"localhost".to_string()![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
model_gateway/src/core/steps/worker/local/create_worker.rs (1)
191-223:⚠️ Potential issue | 🟡 MinorFallback to
num_labelsis skipped whenid2label_jsonis empty/invalid.
Ifid2label_jsonis present but empty or malformed, theelse ifprevents the fallback mapping from being created, despite the comment saying “if id2label wasn’t parsed.”🔧 Proposed fix: track whether id2label was successfully parsed
- if let Some(id2label_json) = labels.get("id2label_json") { - if !id2label_json.is_empty() { - // Parse JSON: keys are string indices, values are label names - if let Ok(string_map) = serde_json::from_str::<HashMap<String, String>>(id2label_json) { - // Convert string keys ("0", "1") to u32 keys (0, 1) - let id2label: HashMap<u32, String> = string_map - .into_iter() - .filter_map(|(k, v)| k.parse::<u32>().ok().map(|idx| (idx, v))) - .collect(); - - if !id2label.is_empty() { - card = card.with_id2label(id2label); - debug!("Parsed id2label with {} classes", card.num_labels); - } - } - } - } - // Fallback: if num_labels is set but id2label wasn't parsed, create default labels - // Match logic in serving_classify.py::_get_id2label_mapping - else if let Some(num_labels_str) = labels.get("num_labels") { + let mut parsed_id2label = false; + if let Some(id2label_json) = labels.get("id2label_json") { + if !id2label_json.is_empty() { + // Parse JSON: keys are string indices, values are label names + if let Ok(string_map) = serde_json::from_str::<HashMap<String, String>>(id2label_json) { + // Convert string keys ("0", "1") to u32 keys (0, 1) + let id2label: HashMap<u32, String> = string_map + .into_iter() + .filter_map(|(k, v)| k.parse::<u32>().ok().map(|idx| (idx, v))) + .collect(); + + if !id2label.is_empty() { + card = card.with_id2label(id2label); + debug!("Parsed id2label with {} classes", card.num_labels); + parsed_id2label = true; + } + } + } + } + // Fallback: if num_labels is set but id2label wasn't parsed, create default labels + // Match logic in serving_classify.py::_get_id2label_mapping + if !parsed_id2label { + if let Some(num_labels_str) = labels.get("num_labels") { if let Ok(num_labels) = num_labels_str.parse::<u32>() { if num_labels > 0 { // Create default mapping: {0: "LABEL_0", 1: "LABEL_1", ...} let id2label: HashMap<u32, String> = (0..num_labels) .map(|i| (i, format!("LABEL_{}", i))) .collect(); card = card.with_id2label(id2label); debug!("Created default id2label with {} classes", num_labels); } } - } + } + }
🤖 Fix all issues with AI agents
In `@model_gateway/src/core/steps/worker/local/detect_connection.rs`:
- Around line 117-129: The debug log currently prints the optional override
config.health.timeout_secs which can be None; compute the resolved timeout first
(the `timeout` value derived from
`config.health.timeout_secs.unwrap_or(app_context.router_config.health_check.timeout_secs)`)
and update the `debug!` call in detect_connection.rs to log that resolved
`timeout` (and optionally which source was used) instead of the raw optional so
the message reflects the actual timeout used by the health check.
slin1237
force-pushed
the
open-fix-n/2
branch
from
February 12, 2026 21:12
24cb230 to
3685097
Compare
github-actions
Bot
added
protocols
3 tasks done
Member
|
@Mergifyio refresh |
Contributor
✅ Pull request refreshed |
Member
|
@Mergifyio update |
Contributor
|
Member
|
@Mergifyio rebase |
Contributor
|
Member
|
@Mergifyio rebase |
Contributor
☑️ Nothing to do, the required conditions are not metDetails
|
Member
|
@Mergifyio update |
Contributor
☑️ Nothing to do, the required conditions are not metDetails
|
Contributor
|
Hi @slin1237, this PR has merge conflicts that must be resolved before it can be merged. Please rebase your branch: git fetch origin main
git rebase origin/main
# resolve any conflicts, then:
git push --force-with-lease |
mergify
Bot
added
the
needs-rebase
slin1237
force-pushed
the
open-fix-n/2
branch
from
February 16, 2026 02:02
3685097 to
5813709
Compare
Contributor
|
Hi @slin1237, the DCO sign-off check has failed. All commits must include a To fix existing commits: # Sign off the last N commits (replace N with the number of unsigned commits)
git rebase HEAD~N --signoff
git push --force-with-leaseTo sign off future commits automatically:
|
mergify
Bot
removed
the
needs-rebase
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Fix all issues with AI agents
In `@model_gateway/src/core/steps/worker/external/create_workers.rs`:
- Around line 53-64: The code unconditionally forces merged.disable_health_check
= true after applying per-worker config (via config.health.apply_to), which
prevents user overrides; change the logic to only set
merged.disable_health_check = true when the user has not provided an explicit
value (i.e., treat true as the default). Locate the block building
(health_config, health_endpoint), check the structure that holds the user's
disable_health_check (the result of config.health.apply_to and/or the config
type backing it), and set merged.disable_health_check only when that field is
unset/None/has the zero/default value; if the config type has no “unset” state,
instead adjust apply_to or the precedence so user-provided values win, or if
override must be forbidden, update the comment to remove the “Users can still
override” line.
In `@model_gateway/src/core/worker_builder.rs`:
- Around line 242-246: The current DP-rank stripping uses url.rfind('@') and
then checks url[at_pos + 1..].chars().all(|c| c.is_ascii_digit()), which returns
true for an empty suffix and implicitly strips a trailing '@'; update the
condition used when computing clean_url to ensure the suffix after '@' is
non-empty (e.g., check at_pos + 1 < url.len() or that the suffix is not empty)
in addition to all characters being ASCII digits, so only valid numeric suffixes
are stripped; update the logic around clean_url (and any code that relies on it)
and add a unit test for the edge case URL ending with '@' (e.g.,
"http://host:8080@") to assert the trailing '@' handling is explicit and
correct.
slin1237
force-pushed
the
open-fix-n/2
branch
from
February 16, 2026 02:25
5813709 to
a7f5bf4
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@model_gateway/src/core/worker_registry.rs`:
- Around line 665-673: The scheduled-check map next_check currently retains
entries for workers that later set disable_health_check=true; build a set of
disabled worker URLs (e.g., iterate workers and collect worker.url().to_string()
where worker.metadata().health_config.disable_health_check is true) and change
the retain call to drop entries that are either not in active_urls or are in
that disabled set (e.g., next_check.retain(|url,_| active_urls.contains(url) &&
!disabled_urls.contains(url))). Keep the existing loop that inserts
next_check.entry(...) only for workers that are not disabled.
slin1237
force-pushed
the
open-fix-n/2
branch
from
February 16, 2026 06:19
a7f5bf4 to
e4996d8
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Fix all issues with AI agents
In `@model_gateway/src/core/worker_builder.rs`:
- Around line 564-599: Add tests covering userinfo and trailing-@ edge cases for
parse_bootstrap_host and, if failing, adjust parse_bootstrap_host logic: add a
test `test_parse_bootstrap_host_edge_cases` asserting
parse_bootstrap_host("http://user@host:8080") == "host" and
parse_bootstrap_host("http://host:8080@") == "host"; ensure the function treats
an '@' inside userinfo (before hostname) correctly by extracting the hostname
after userinfo, and do not strip a trailing '@' with an empty suffix as a DP
rank removal step unless a numeric rank follows; run/update DPAwareWorkerBuilder
tests (e.g., test_dp_aware_worker_bootstrap_host) to confirm behavior remains
unchanged.
In `@model_gateway/src/service_discovery.rs`:
- Around line 455-462: Add a clarifying comment above the assignment to
spec.max_connection_attempts explaining why we multiply success_threshold by 20
here (instead of the 10x used in job_queue.rs): note that service discovery
(location: service_discovery.rs around spec.max_connection_attempts and
app_context.router_config.health_check.success_threshold) handles workers with
per-worker health overrides and so needs extra retry headroom, hence the 20x
multiplier.
…and model card cloning What changed: - protocols/src/worker.rs: Changed WorkerSpec.health from HealthCheckConfig (all required fields) to HealthCheckUpdate (all Option<T>) for proper PATCH semantics. Added Default derive and is_empty() helper on HealthCheckUpdate. - model_gateway/src/core/worker.rs: Added health_config: HealthCheckConfig to WorkerMetadata for resolved runtime config. Updated all runtime access sites (check_health_async, grpc/http health checks, worker_registry health filter) from spec.health.X to health_config.X. - model_gateway/src/core/worker_builder.rs: Builder now stores resolved health_config on WorkerMetadata via apply_to(). Fixed parse_bootstrap_host to strip @rank suffix from DP-aware URLs and handle bare host:port inputs. - model_gateway/src/core/job_queue.rs: Removed incorrect assignment of full HealthCheckConfig to spec.health. Clamped max_connection_attempts with .max(1) to prevent 0 attempts when success_threshold is 0. - model_gateway/src/service_discovery.rs: Same max_connection_attempts clamp and removed spec.health assignment. - model_gateway/src/core/steps/worker/local/create_worker.rs: Changed build_health_config to use config.health.apply_to(&base). Changed build_model_card to clone full proto_card preserving all user-supplied fields. - model_gateway/src/core/steps/worker/external/create_workers.rs: Same apply_to() pattern for health config merge. - model_gateway/src/core/steps/worker/local/detect_connection.rs: Fall back to router timeout when per-worker override not set. - model_gateway/src/core/steps/worker/local/update_worker_properties.rs: Merge health updates against resolved health_config instead of spec.health. - bindings/golang/src/policy.rs: Added health_config field to WorkerMetadata construction. Why: - WorkerSpec.health as HealthCheckConfig silently reset unspecified fields to defaults, making per-worker health overrides impossible. Changing to HealthCheckUpdate gives proper partial-update semantics where only Some fields override router defaults. - parse_bootstrap_host misinterpreted DP-aware URLs (e.g. http://host:8080@3) because @ is an RFC 3986 userinfo delimiter, extracting "3" as host instead of "host". Also failed on bare host:port inputs. - max_connection_attempts was 0 when success_threshold was 0 (0 * N = 0), preventing any connection attempts. - build_model_card selectively copied only 5 fields from proto_card, silently dropping user-supplied fields like display_name, provider, and context_length. How: - WorkerSpec.health is now HealthCheckUpdate (partial overrides). The resolved config lives on WorkerMetadata.health_config, populated at build time via HealthCheckUpdate::apply_to(&base). Runtime code reads health_config directly. - parse_bootstrap_host strips trailing @digits before parsing and uses a try_parse helper that detects missing host_str() to handle schemeless URLs. - .max(1) clamp on all 3 success_threshold multiplication sites. - build_model_card now does config.models.find(id).cloned() to get the full card. Signed-off-by: Simo Lin <simo.lin@oracle.com>
slin1237
force-pushed
the
open-fix-n/2
branch
from
February 16, 2026 06:53
e4996d8 to
c2a8d5c
Compare
This was referenced Feb 16, 2026
ppraneth
pushed a commit
that referenced
this pull request
Feb 18, 2026
…and model card cloning (#415) Signed-off-by: Simo Lin <linsimo.mark@gmail.com> Signed-off-by: ppraneth <pranethparuchuri@gmail.com>
This was referenced Feb 24, 2026
3 tasks done
This was referenced Mar 18, 2026
8 of 9 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Fixes several pre-existing issues identified during PR #412 review:
WorkerSpec.healthwasHealthCheckConfig(all required fields), so any unspecified fields reset to defaults instead of inheriting router config. Changed toHealthCheckUpdate(allOption<T>) with proper merge viaapply_to().parse_bootstrap_hostbroke on DP-aware URLs — URLs likehttp://host:8080@3had@misinterpreted as RFC 3986 userinfo delimiter, extracting"3"as host instead of"host".max_connection_attemptscould be 0 — whensuccess_thresholdis 0,0 * N = 0prevented any connection attempts.build_model_carddropped user-supplied fields — selectively copied only 5 fields from proto_card, silently losingdisplay_name,provider,context_length, etc.What changed
protocols/src/worker.rsWorkerSpec.health:HealthCheckConfig→HealthCheckUpdate; addedDefault+is_empty()onHealthCheckUpdatemodel_gateway/src/core/worker.rshealth_config: HealthCheckConfigtoWorkerMetadata; updated all runtime access sites (~8)model_gateway/src/core/worker_builder.rsapply_to(); fixedparse_bootstrap_hostfor DP URLs and bare hosts; added 4 testsmodel_gateway/src/core/job_queue.rsspec.healthassignment; clampedmax_connection_attemptswith.max(1)model_gateway/src/service_discovery.rsspec.healthassignmentmodel_gateway/src/core/steps/worker/local/create_worker.rsbuild_health_configusesapply_to();build_model_cardclones full proto_cardmodel_gateway/src/core/steps/worker/external/create_workers.rsapply_to()pattern for health mergemodel_gateway/src/core/steps/worker/local/detect_connection.rsmodel_gateway/src/core/steps/worker/local/update_worker_properties.rshealth_configmodel_gateway/src/core/worker_registry.rshealth_configbindings/golang/src/policy.rshealth_configfield toWorkerMetadataconstructionTest plan
cargo build— compiles cleanlycargo clippy --all-targets— no warningscargo test— 344+ tests pass, 0 failurestest_parse_bootstrap_host_normal_url,test_parse_bootstrap_host_dp_aware_url,test_parse_bootstrap_host_bare_host,test_dp_aware_worker_bootstrap_hosthealth_config(resolved) instead ofspec.health(partial)HealthCheckUpdateserializes withskip_serializing_none,WorkerSpec.healthskips when emptySummary by CodeRabbit
New Features
Bug Fixes
Refactor