feat(core): add per-model retry config to WorkerRegistry - #821
Conversation
|
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:
📝 WalkthroughWalkthroughAdds per-model retry override storage and APIs to WorkerRegistry and updates the worker registration step to apply per-worker retry overrides to each model a worker serves (sequential last-write-wins when multiple workers target the same model IDs). Changes
Sequence Diagram(s)sequenceDiagram
participant Registrar as Register Step
participant Worker as Worker (metadata/resilience)
participant Registry as WorkerRegistry
Registrar->>Worker: iterate registered workers
Worker-->>Registrar: metadata(), resilience(), worker_model_ids()
Registrar->>Registry: set_model_retry_config(model_id, retry_config, enabled)
Registry-->>Registrar: store/overwrite per-model retry state
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~22 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Summary of ChangesHello, 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 enhances the system's resilience by introducing granular, per-model retry configuration capabilities. Previously, retry settings were applied globally, which was insufficient for diverse worker groups. The changes enable the system to store and manage distinct retry policies for individual models, ensuring that each model's specific requirements are met. This is a foundational step towards a more robust and adaptable worker resilience framework, allowing for tailored retry behavior based on the characteristics of different models. Highlights
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. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
The pull request effectively implements per-model retry configuration within the WorkerRegistry, addressing the need for differentiated retry characteristics across various worker groups. The solution follows a clear last-write-wins semantic and includes proper cleanup when workers are removed. The changes are well-tested with new unit tests covering the core logic. The code is clean, readable, and integrates smoothly with the existing architecture.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@model_gateway/src/core/steps/worker/shared/register.rs`:
- Around line 56-74: The loop only iterates worker.models() so workers that only
expose a model via worker.model_id() (labels) never get their retry config
saved; replace the iteration with the same fallback used by
WorkerRegistry::worker_model_ids (or call WorkerRegistry::worker_model_ids) to
obtain model IDs for the worker, then call
app_context.worker_registry.set_model_retry_config(model_id,
resolved.retry.clone(), resolved.retry_enabled) for each returned id; use
worker.resilience() to get resolved.retry and resolved.retry_enabled as shown.
- Around line 66-72: The loop repeatedly calls resolved.retry.clone() for each
model; compute let retry_cfg = resolved.retry.clone() once before iterating
(using worker.resilience() result already in resolved) and pass retry_cfg (and
resolved.retry_enabled) into app_context.worker_registry.set_model_retry_config
inside the for loop so you avoid redundant clones while preserving behavior of
worker.resilience(), worker.models(), and set_model_retry_config.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: bc6bdb75-1fce-4805-8c43-eed6fbb3128e
📒 Files selected for processing (2)
model_gateway/src/core/steps/worker/shared/register.rsmodel_gateway/src/core/worker_registry.rs
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 73fcff4e8e
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| pub fn set_model_retry_config(&self, model_id: &str, config: RetryConfig, enabled: bool) { | ||
| self.model_retry_configs | ||
| .insert(model_id.to_string(), config); | ||
| self.model_retry_enabled | ||
| .insert(model_id.to_string(), enabled); |
There was a problem hiding this comment.
If the same model is exposed by more than one backend group, this key is too coarse. WorkerSelection::get_candidates() still partitions candidates by worker_type, connection_mode, runtime_type, and optional provider, so a local SGLang gpt-4o worker and an external OpenAI gpt-4o worker can coexist. With only model_id here, whichever worker registers last overwrites the other's retry policy, so a later lookup cannot return the retry settings for the group that was actually selected.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Valid concern — matches how PolicyRegistry works today (also keyed by model_id only). This is part of the broader WorkerGroup consolidation planned as future work. Keeping consistent with the existing pattern for now.
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@model_gateway/src/core/worker_registry.rs`:
- Around line 296-303: set_model_retry_config currently does two separate
inserts into model_retry_configs and model_retry_enabled which can interleave
under concurrent writes; replace the two DashMaps with a single DashMap keyed by
model_id that stores a combined value (e.g., a small struct or tuple containing
RetryConfig and enabled flag) and update set_model_retry_config to perform one
atomic insert into that single map (references: set_model_retry_config,
model_retry_configs, model_retry_enabled, RetryConfig); also update any read
sites to access the combined value and adjust types accordingly so
config+enabled are always updated together.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 1c8c7323-878d-4890-8862-9c77d3534042
📒 Files selected for processing (2)
model_gateway/src/core/steps/worker/shared/register.rsmodel_gateway/src/core/worker_registry.rs
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d51eb2468b
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
model_gateway/src/core/worker_registry.rs (1)
296-303:⚠️ Potential issue | 🟠 MajorConfig and enabled flag updates are non-atomic across two maps.
At Line 299-302,
set_model_retry_configwritesRetryConfigandenabledin separate operations. Concurrent updates for the samemodel_idcan leave a mixed pair (config from writer A, enabled from writer B), which breaks strict per-write last-write-wins semantics.💡 Proposed refactor: store a single combined value per model
+#[derive(Debug, Clone)] +struct ModelRetrySettings { + config: RetryConfig, + enabled: bool, +} @@ - model_retry_configs: Arc<DashMap<String, RetryConfig>>, - model_retry_enabled: Arc<DashMap<String, bool>>, + model_retry_settings: Arc<DashMap<String, ModelRetrySettings>>, @@ - model_retry_configs: Arc::new(DashMap::new()), - model_retry_enabled: Arc::new(DashMap::new()), + model_retry_settings: Arc::new(DashMap::new()), @@ - model_retry_configs: self.model_retry_configs.clone(), - model_retry_enabled: self.model_retry_enabled.clone(), + model_retry_settings: self.model_retry_settings.clone(), @@ pub fn get_retry_config(&self, model_id: &str) -> Option<RetryConfig> { - self.model_retry_configs - .get(model_id) - .map(|entry| entry.value().clone()) + self.model_retry_settings + .get(model_id) + .map(|entry| entry.value().config.clone()) } @@ pub fn get_retry_enabled(&self, model_id: &str) -> Option<bool> { - self.model_retry_enabled - .get(model_id) - .map(|entry| *entry.value()) + self.model_retry_settings + .get(model_id) + .map(|entry| entry.value().enabled) } @@ pub fn set_model_retry_config(&self, model_id: &str, config: RetryConfig, enabled: bool) { - self.model_retry_configs - .insert(model_id.to_string(), config); - self.model_retry_enabled - .insert(model_id.to_string(), enabled); + self.model_retry_settings.insert( + model_id.to_string(), + ModelRetrySettings { config, enabled }, + ); }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/core/worker_registry.rs` around lines 296 - 303, set_model_retry_config currently updates model_retry_configs and model_retry_enabled in two separate operations which can interleave; change to a single atomic update by storing a combined value instead of two maps or by performing both writes under the same lock. Concretely, introduce a small struct (e.g., ModelRetry { config: RetryConfig, enabled: bool }) and replace model_retry_configs/model_retry_enabled with a single map keyed by model_id, then update set_model_retry_config to insert that combined struct in one operation (or wrap the two inserts in the same mutex/critical section if retaining two maps).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@model_gateway/src/core/worker_registry.rs`:
- Around line 453-462: The cleanup of model_retry_configs/model_retry_enabled
races with concurrent registrations because you check model_index emptiness then
remove retry state without synchronizing; wrap the model-index mutation and
retry-state updates for a given model in a model-scoped critical section (e.g.,
a per-model Mutex or RwLock) so that the emptiness check and subsequent removal
of entries are atomic with respect to concurrent register/remove operations;
apply this locking around the code that mutates self.model_index and the block
that touches self.model_retry_configs and self.model_retry_enabled (the section
using self.model_index.get(&model_id) and
self.model_retry_configs.remove(&model_id)/self.model_retry_enabled.remove(&model_id)).
---
Duplicate comments:
In `@model_gateway/src/core/worker_registry.rs`:
- Around line 296-303: set_model_retry_config currently updates
model_retry_configs and model_retry_enabled in two separate operations which can
interleave; change to a single atomic update by storing a combined value instead
of two maps or by performing both writes under the same lock. Concretely,
introduce a small struct (e.g., ModelRetry { config: RetryConfig, enabled: bool
}) and replace model_retry_configs/model_retry_enabled with a single map keyed
by model_id, then update set_model_retry_config to insert that combined struct
in one operation (or wrap the two inserts in the same mutex/critical section if
retaining two maps).
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: ae5be2df-1548-42c4-9e95-86bf0d388dbd
📒 Files selected for processing (1)
model_gateway/src/core/worker_registry.rs
There was a problem hiding this comment.
💡 Codex Review
https://github.com/lightseekorg/smg/blob/1ce6e2fa33b58fe616ac9780464de592484474fd/model_gateway/src/core/worker_registry.rs#L375-L376
register() removes the old worker from each previous model when the same URL is re-registered, but the new retry-state cleanup only exists in remove(), so this path never clears model_retry_*. Fresh evidence: test_re_register_same_url_refreshes_all_model_indexes() now documents same-URL model swaps as a supported path. If a gpt-4o worker with retry overrides is replaced by the same URL serving o3/o4-mini, the stale gpt-4o override survives; a later gpt-4o worker with no overrides will then inherit that stale policy because RegisterWorkersStep intentionally skips default-config workers.
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
Store retry config per model in WorkerRegistry (last write wins). When a worker registers with non-empty retry overrides, its resolved RetryConfig is stored for the model group. Routers will use this instead of a single global retry config. - Add model_retry_configs and model_retry_enabled DashMaps - Add get/set methods for per-model retry config - Store retry config during worker registration (RegisterWorkersStep) - Clean up retry config when last worker for a model is removed - 2 new tests: last-write-wins semantics and cleanup on removal Signed-off-by: Chang Su <chang.s.su@oracle.com>
- Fall back to worker.model_id() when worker.models() is empty, matching WorkerRegistry::worker_model_ids() behavior. Fixes retry config not being stored for wildcard/external workers. - Clone retry config once outside the inner model loop. Signed-off-by: Chang Su <chang.s.su@oracle.com>
Make worker_model_ids() public and use it in RegisterWorkersStep instead of duplicating the model ID resolution logic. Signed-off-by: Chang Su <chang.s.su@oracle.com>
… race Move per-model retry config cleanup from remove_worker_from_model_index (called during both removal and re-registration) to remove() (called only during actual worker removal). This prevents re-registering the same URL from clearing retry config when the worker is temporarily the last one for a model. Signed-off-by: Chang Su <chang.s.su@oracle.com>
1ce6e2f to
a3d6517
Compare
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@model_gateway/src/core/worker_registry.rs`:
- Around line 603-612: The condition does two separate lookups on
self.model_index for model_id; replace them with a single lookup to avoid
duplicated access and possible inconsistency by using a match/if-let on
self.model_index.get(&model_id) (e.g., match self.model_index.get(&model_id) {
None | Some(v) if v.is_empty() => { self.model_retry_configs.remove(&model_id);
self.model_retry_enabled.remove(&model_id); }, _ => {} }) so the cleanup of
model_retry_configs and model_retry_enabled happens only when the single-lookup
result indicates no entry or an empty entry for model_id.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 9c9e230a-cbbf-4ce5-87dd-5dbe82cf6417
📒 Files selected for processing (2)
model_gateway/src/core/steps/worker/shared/register.rsmodel_gateway/src/core/worker_registry.rs
…del_retry_enabled Instead of a separate enabled flag, set max_retries=1 when retries are disabled (matching RouterConfig::effective_retry_config() pattern). This removes model_retry_enabled DashMap entirely and simplifies the retry config storage to a single DashMap. Also simplify double model_index lookup in remove() cleanup. Signed-off-by: Chang Su <chang.s.su@oracle.com>
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
Restore comments that were lost during PR #836 refactor: - "clone needed for DashMap key ownership" on type/connection index updates - "no-op if mesh is not enabled" on mesh sync blocks Signed-off-by: Chang Su <chang.s.su@oracle.com>
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
2 participants
Description
Part of the per-worker resilience refactor series: #799 → #803 →
#811→ this PR → router migration → cleanup.Problem
Retry config is global — all routers store a single
retry_configfromRouterConfig. Different worker groups (e.g., local SGLang vs external OpenAI) have different retry characteristics but share the same config.Solution
Store per-model retry config in
WorkerRegistryusing last-write-wins semantics. When a worker registers with non-empty retry overrides inWorkerSpec.resilience, its resolvedRetryConfigis stored for the model group. Subsequent workers with retry overrides overwrite it. Workers with no overrides don't change the stored config.Retry config is cleaned up when the last worker for a model is removed.
This follows the same pattern as
PolicyRegistry(per-model policy) but lives inWorkerRegistryto avoid a separate registry. Routers will look up retry config from the registry at request time in a subsequent PR.Design decision: retry at router level, not per-worker
Retries re-select workers on each attempt (like Envoy/NGINX), so retry config belongs to the worker group, not individual workers. See updated design doc at
.claude/docs/plans/2026-03-14-per-worker-resilience-design.md.Changes
model_retry_configsandmodel_retry_enabledDashMaps toWorkerRegistryget_retry_config(),get_retry_enabled(),set_model_retry_config()methodsRegisterWorkersStep— stores resolved config when worker has non-empty retry overridesremove_worker_from_model_index()when last worker is removedTest Plan
cargo test -p smg --lib core::worker_registry— all 8 tests pass (6 existing + 2 new)cargo test -p smg --lib— all 437 tests pass, 0 failuresChecklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspassesSummary by CodeRabbit
Summary by CodeRabbit