Conversation
|
No actionable comments were generated in the recent review. 🎉 📝 WalkthroughWalkthroughConsolidates DP-aware behavior into Changes
Sequence Diagram(s)mermaid Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
⚔️ Resolve merge conflicts (beta)
Comment |
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 significantly refactors the worker management system by consolidating data-parallel (DP) awareness directly into the 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
|
|
Hi @slin1237, the branch Please use one of the following formats:
Allowed types:
|
|
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:
|
Delete DPAwareWorker and DPAwareWorkerBuilder, moving DP-awareness into BasicWorker as three optional fields (dp_rank, dp_size, dp_base_url). What changed: - worker.rs: Added optional dp_rank/dp_size/dp_base_url fields to BasicWorker. Overrode 6 Worker trait methods (is_dp_aware, base_url, dp_rank, dp_size, prepare_request, endpoint_url) inline. Deleted DPAwareWorker struct and its full Worker impl (~147 lines). Removed normalised_url() in favor of base_url(). Updated Debug impl and all DP tests. - worker_builder.rs: Added dp_config(rank, size) to BasicWorkerBuilder which captures the current URL as base, then formats url@rank. Deleted DPAwareWorkerBuilder and its 16 passthrough setter methods (~112 lines). Updated builder tests. - mod.rs: Removed DPAwareWorkerBuilder from public exports. - create_worker.rs: Switched from DPAwareWorkerBuilder::new() to BasicWorkerBuilder::new().dp_config(). - update_worker_properties.rs: Eliminated separate DP/non-DP builder branches into a single BasicWorkerBuilder path with conditional dp_config() call. - router.rs: Removed dp_aware field, extract_dp_rank(), and worker_base_url() helper. Refactored send_typed_request to accept &dyn Worker and use prepare_request()/endpoint_url() unconditionally, making the router fully DP-agnostic. route_simple_request now uses worker.base_url() directly. Why: DPAwareWorker was pure boilerplate — ~250 lines of delegation that added no behavior beyond what three optional fields provide. The router also duplicated DP logic (URL parsing, body injection) that the Worker trait already abstracts, marked with a TODO to fix. This consolidation removes all duplication and makes the router DP-agnostic. How: DP-awareness is now opt-in via BasicWorkerBuilder::dp_config(). Non-DP workers have None for all three fields, so all trait defaults (is_dp_aware=false, base_url=url, prepare_request=passthrough) apply with zero overhead. The router uses a single code path through prepare_request() and endpoint_url() for all workers. Signed-off-by: Simo Lin <simo.lin@oracle.com>
![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
This is an excellent refactoring that significantly simplifies the codebase by consolidating DPAwareWorker into BasicWorker. The removal of boilerplate and duplicated logic in the HTTP router is a great improvement for maintainability. The changes are well-structured and the new dp_config builder pattern is clean. I have a minor suggestion to further improve code quality and safety, but overall this is a very solid pull request.
| async fn prepare_request(&self, mut req: serde_json::Value) -> WorkerResult<serde_json::Value> { | ||
| if let Some(rank) = self.dp_rank { | ||
| if let Some(map) = req.as_object_mut() { | ||
| map.insert("data_parallel_rank".to_string(), serde_json::json!(rank)); |
There was a problem hiding this comment.
The key "data_parallel_rank" is hardcoded here. It was previously defined as a const in router.rs. It would be good practice to define it as a const at the module level (e.g., const DP_RANK_KEY: &str = "data_parallel_rank";) and reuse it here. This avoids magic strings and improves maintainability.
![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
🤖 Fix all issues with AI agents
In `@model_gateway/src/core/steps/worker/local/update_worker_properties.rs`:
- Around line 88-92: The DP-aware branch silently masks broken invariants by
using unwrap_or for dp_rank()/dp_size(); replace the fallbacks with expect()
calls so a missing value fails fast (e.g., call worker.dp_rank().expect("dp_rank
must be Some when is_dp_aware() is true") and similarly for dp_size()) before
passing to builder.dp_config, or if the fallback behavior is intentional, add a
short comment next to the worker.is_dp_aware() check explaining why default
values 0 and 1 are safe and desired.
When is_dp_aware() is true, dp_rank and dp_size are guaranteed Some by construction. Replace silent fallbacks (0, 1) with expect() to fail fast on invariant violations instead of masking bugs. Signed-off-by: Simo Lin <simo.lin@oracle.com>
…_type Move the three data-parallel identity fields (dp_rank, dp_size, dp_base_url) from BasicWorker to WorkerSpec in the protocols crate. These are config/identity fields set at construction and never mutated, so they belong alongside other worker identity in WorkerSpec. What changed: - protocols/src/worker.rs: added dp_base_url, dp_rank, dp_size to WorkerSpec with serde(default, skip_serializing_if) for backwards compatibility - model_gateway/src/core/worker.rs: replaced 6 hardcoded DP trait defaults with implementations reading from self.metadata().spec; removed 3 DP fields and 6 method overrides from BasicWorker; simplified Debug impl; removed default_model_type from WorkerMetadata and inlined ModelType::LLM in supports_endpoint() fallback; updated 2 test struct literals - model_gateway/src/core/worker_builder.rs: removed 3 DP fields from builder struct and all 3 constructors; dp_config() now writes to self.spec.*; removed default_model_type from build(); removed unused ModelType import - bindings/golang/src/policy.rs: removed default_model_type from WorkerMetadata construction and ModelType from imports Why: after consolidating DPAwareWorker into BasicWorker (#434), the DP fields were still on BasicWorker alongside runtime state. Moving them to WorkerSpec makes DP info visible in GET /workers API responses (via WorkerInfo's serde(flatten) on spec), gives any future Worker implementor DP support for free via trait defaults, and removes ~30 lines of duplicated overrides. How: DP fields use Option<T> with serde(default) so old clients that don't send these fields get None (backwards compatible). The Worker trait defaults now read from metadata().spec instead of returning hardcoded None/false. default_model_type was always ModelType::LLM at all 4 construction sites so it was inlined.
…_type Move the three data-parallel identity fields (dp_rank, dp_size, dp_base_url) from BasicWorker to WorkerSpec in the protocols crate. These are config/identity fields set at construction and never mutated, so they belong alongside other worker identity in WorkerSpec. What changed: - protocols/src/worker.rs: added dp_base_url, dp_rank, dp_size to WorkerSpec with serde(default, skip_serializing_if) for backwards compatibility - model_gateway/src/core/worker.rs: replaced 6 hardcoded DP trait defaults with implementations reading from self.metadata().spec; removed 3 DP fields and 6 method overrides from BasicWorker; simplified Debug impl; removed default_model_type from WorkerMetadata and inlined ModelType::LLM in supports_endpoint() fallback; updated 2 test struct literals - model_gateway/src/core/worker_builder.rs: removed 3 DP fields from builder struct and all 3 constructors; dp_config() now writes to self.spec.*; removed default_model_type from build(); removed unused ModelType import - bindings/golang/src/policy.rs: removed default_model_type from WorkerMetadata construction and ModelType from imports Why: after consolidating DPAwareWorker into BasicWorker (#434), the DP fields were still on BasicWorker alongside runtime state. Moving them to WorkerSpec makes DP info visible in GET /workers API responses (via WorkerInfo's serde(flatten) on spec), gives any future Worker implementor DP support for free via trait defaults, and removes ~30 lines of duplicated overrides. How: DP fields use Option<T> with serde(default) so old clients that don't send these fields get None (backwards compatible). The Worker trait defaults now read from metadata().spec instead of returning hardcoded None/false. default_model_type was always ModelType::LLM at all 4 construction sites so it was inlined. Signed-off-by: Simo Lin <simo.lin@oracle.com>
…_type Move the three data-parallel identity fields (dp_rank, dp_size, dp_base_url) from BasicWorker to WorkerSpec in the protocols crate. These are config/identity fields set at construction and never mutated, so they belong alongside other worker identity in WorkerSpec. What changed: - protocols/src/worker.rs: added dp_base_url, dp_rank, dp_size to WorkerSpec with serde(default, skip_serializing_if) for backwards compatibility - model_gateway/src/core/worker.rs: replaced 6 hardcoded DP trait defaults with implementations reading from self.metadata().spec; removed 3 DP fields and 6 method overrides from BasicWorker; simplified Debug impl; removed default_model_type from WorkerMetadata and inlined ModelType::LLM in supports_endpoint() fallback; updated 2 test struct literals - model_gateway/src/core/worker_builder.rs: removed 3 DP fields from builder struct and all 3 constructors; dp_config() now writes to self.spec.*; removed default_model_type from build(); removed unused ModelType import - bindings/golang/src/policy.rs: removed default_model_type from WorkerMetadata construction and ModelType from imports Why: after consolidating DPAwareWorker into BasicWorker (#434), the DP fields were still on BasicWorker alongside runtime state. Moving them to WorkerSpec makes DP info visible in GET /workers API responses (via WorkerInfo's serde(flatten) on spec), gives any future Worker implementor DP support for free via trait defaults, and removes ~30 lines of duplicated overrides. How: DP fields use Option<T> with serde(default) so old clients that don't send these fields get None (backwards compatible). The Worker trait defaults now read from metadata().spec instead of returning hardcoded None/false. default_model_type was always ModelType::LLM at all 4 construction sites so it was inlined. Signed-off-by: Simo Lin <simo.lin@oracle.com>
Signed-off-by: Simo Lin <simo.lin@oracle.com> Signed-off-by: ppraneth <pranethparuchuri@gmail.com>
|
Hi @slin1237, the branch Please use one of the following formats:
Allowed types:
|
1 participant
Summary
Consolidates
DPAwareWorkerintoBasicWorkerby making DP-awareness optional fields instead of a separate wrapper struct, and removes duplicated DP logic from the HTTP router. Net -280 lines.What changed
worker.rsdp_rank,dp_size,dp_base_url) toBasicWorker. Overrode 6 Worker trait methods inline. DeletedDPAwareWorkerstruct + full Worker impl (~147 lines). Removednormalised_url()(superseded bybase_url()).worker_builder.rsdp_config(rank, size)toBasicWorkerBuilder— captures current URL as base, formatsurl@rank. DeletedDPAwareWorkerBuilder+ 16 passthrough setters (~112 lines).mod.rsDPAwareWorkerBuilderfrom exports.create_worker.rsDPAwareWorkerBuilder::new(url, rank, size)→BasicWorkerBuilder::new(url).dp_config(rank, size)update_worker_properties.rsBasicWorkerBuilderpath with conditional.dp_config().router.rsdp_awarefield,extract_dp_rank(),worker_base_url().send_typed_requestnow takes&dyn Workerand usesprepare_request()/endpoint_url()unconditionally — router is fully DP-agnostic.Why
DPAwareWorkerwas pure boilerplate: ~250 lines delegating every Worker trait method to an innerBasicWorker, plus a separate builder duplicating 16 setter methods. The HTTP router also re-implemented DP logic (URL parsing, body injection) that the Worker trait already provides — with an existingTODO (rui): Better accommodate to the Worker abstraction.Three optional fields on
BasicWorkerprovide the same behavior with zero delegation overhead.How
DP-awareness is opt-in via
BasicWorkerBuilder::dp_config(rank, size). Non-DP workers haveNonefor all three fields, so trait defaults apply (is_dp_aware()=false,base_url()=url(),prepare_request()=passthrough) with zero overhead. The router uses a single code path throughprepare_request()andendpoint_url()for all workers.What stays unchanged
RouterConfig.dp_awareconfig flag — still used for worker creation/discovery decisionsfind_workers_by_url()prefix matching — still needed for worker lookup/removalWorkerRemovalRequest.dp_aware,WorkerUpdateWorkflowData.dp_aware— lifecycle plumbingDiscoverDPInfoStep— checks config to decide whether to query DP infoTest plan
cargo check --all-targets --all-features— passcargo test -p smg -- worker— all 92 worker tests pass (12 DP tests rewritten)cargo clippy -p smg --all-targets --all-features -- -D warnings— zero warningsgrep -r "DPAwareWorker"— zero remaining referencesSummary by CodeRabbit
Refactor
Breaking Changes