feat(libsy): add pluggable routing affinity#50
Conversation
85d291e to
24264aa
Compare
1df67dc to
aa72aa0
Compare
fcdc0ad to
eab29f6
Compare
55e2f0f to
36d53c3
Compare
ayushag-nv
changed the title
| /// Id of the task the request belongs to. | ||
| pub task_id: Option<String>, | ||
| /// Agent-specific lineage and semantic routing signals. | ||
| pub agent_context: Option<Box<AgentContext>>, |
|
|
||
| /// Namespaced stable identity used to retain a model assignment. | ||
| #[derive(Clone, Debug, Hash, PartialEq, Eq)] | ||
| pub struct AffinityKey { |
There was a problem hiding this comment.
i believe there's another rust piece for affinity, keep either one and delete the other
| // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | ||
| // SPDX-License-Identifier: Apache-2.0 | ||
|
|
||
| //! Demo proxy combining Switchyard HTTP/translation with libsy random routing and affinity. |
There was a problem hiding this comment.
all of this can be under libsy/bin or libsy/example
| } | ||
|
|
||
| #[async_trait] | ||
| impl Profile for LibsyAffinityProfile { |
There was a problem hiding this comment.
lets forget anything Profile*
Signed-off-by: ayushag <ayushag@nvidia.com>
Signed-off-by: ayushag <ayushag@nvidia.com>
Signed-off-by: ayushag <ayushag@nvidia.com>
Signed-off-by: ayushag <ayushag@nvidia.com>
Signed-off-by: Greg Clark <grclark@nvidia.com>
924763d to
adfc7bc
Compare
WalkthroughChangesAffinity Routing
CI Toolchain Updates
Estimated code review effort: 4 (Complex) | ~60 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Comment |
![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
🧹 Nitpick comments (1)
crates/libsy/src/algorithms/rand.rs (1)
427-448: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winForce both requests to miss before either retains.
This test allows the first future to complete retention before the second is polled, so it can pass even if first-writer atomicity regresses. Use a barrier with different proposals, then assert both callers receive one canonical assignment.
As per coding guidelines, “Add concise Rust comments or
///documentation for important behavior-encoding tests and complex concurrency logic.”🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/libsy/src/algorithms/rand.rs` around lines 427 - 448, Strengthen concurrent_first_subagent_turns_use_one_assignment by synchronizing both requests with a barrier so they miss before either retains its proposal, using different proposed models. Assert both completed requests receive the same canonical selected_model, and add a concise Rust comment documenting the first-writer atomicity behavior being tested.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@crates/libsy/README.md`:
- Around line 50-57: Add the missing Arc and Algorithm imports to the README
Rust example so the snippet compiles as shown, while preserving the existing
RandomAlgo and SubAgentAffinity imports and usage.
In `@demo/libsy-proxy/src/main.rs`:
- Around line 123-141: Update libsy_request to validate and reject requests with
stream: true before decoding or forwarding them upstream. Return the existing
invalid-request error type with a clear streaming-unsupported message, while
preserving the current non-streaming request construction path.
---
Nitpick comments:
In `@crates/libsy/src/algorithms/rand.rs`:
- Around line 427-448: Strengthen
concurrent_first_subagent_turns_use_one_assignment by synchronizing both
requests with a barrier so they miss before either retains its proposal, using
different proposed models. Assert both completed requests receive the same
canonical selected_model, and add a concise Rust comment documenting the
first-writer atomicity behavior being tested.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 9157671b-adad-47d6-877d-ecb42423c775
⛔ Files ignored due to path filters (1)
Cargo.lockis excluded by!**/*.lock,!Cargo.lock
📒 Files selected for processing (13)
.github/workflows/ci.yml.github/workflows/perf.ymlCargo.tomlcrates/libsy-protocol/src/envelope.rscrates/libsy/Cargo.tomlcrates/libsy/README.mdcrates/libsy/src/affinity.rscrates/libsy/src/algorithms/rand.rscrates/libsy/src/lib.rscrates/libsy/tests/affinity.rsdemo/libsy-proxy/Cargo.tomldemo/libsy-proxy/README.mddemo/libsy-proxy/src/main.rs
| ```rust | ||
| use libsy::affinity::SubAgentAffinity; | ||
| use libsy::RandomAlgo; | ||
|
|
||
| let algo: Arc<dyn Algorithm> = Arc::new( | ||
| RandomAlgo::new(targets) | ||
| .with_affinity(Arc::new(SubAgentAffinity::new())), | ||
| ); |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
FILE="crates/libsy/README.md"
# Show the relevant section with line numbers
sed -n '1,120p' "$FILE" | nl -ba | sed -n '1,120p'Repository: NVIDIA-NeMo/Switchyard
Length of output: 200
🏁 Script executed:
#!/bin/bash
set -euo pipefail
FILE="crates/libsy/README.md"
awk 'NR>=1 && NR<=120 { printf "%4d %s\n", NR, $0 }' "$FILE"Repository: NVIDIA-NeMo/Switchyard
Length of output: 6391
Add the missing imports to the example
Arc and Algorithm are used in the snippet but not imported, so the example does not compile as shown.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@crates/libsy/README.md` around lines 50 - 57, Add the missing Arc and
Algorithm imports to the README Rust example so the snippet compiles as shown,
while preserving the existing RandomAlgo and SubAgentAffinity imports and usage.
| fn libsy_request(input: &ProfileInput, translation: &TranslationEngine) -> Result<Request> { | ||
| let format = wire_format(input.request.request_type()); | ||
| let decoded = translation | ||
| .decode_request(format, input.request.body(), &TranslationPolicy::default()) | ||
| .map_err(|error| SwitchyardError::InvalidRequest(error.to_string()))?; | ||
| let mut metadata = metadata_from_headers(&input.metadata.headers); | ||
| metadata | ||
| .extra_metadata | ||
| .get_or_insert_with(Default::default) | ||
| .insert( | ||
| "inbound_format".to_string(), | ||
| inbound_format(input.request.request_type()).to_string(), | ||
| ); | ||
|
|
||
| Ok(Request { | ||
| llm_request: decoded.request, | ||
| raw_request: Some(input.request.body().clone()), | ||
| metadata: Some(metadata), | ||
| }) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Reject streaming before making the upstream call.
A request with stream: true is forwarded, but Line 71 later requires a buffered body, producing an avoidable server error after external work has begun.
Proposed validation
let decoded = translation
.decode_request(format, input.request.body(), &TranslationPolicy::default())
.map_err(|error| SwitchyardError::InvalidRequest(error.to_string()))?;
+ if decoded.request.stream {
+ return Err(SwitchyardError::InvalidRequest(
+ "libsy-proxy currently supports buffered requests only".to_string(),
+ ));
+ }📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| fn libsy_request(input: &ProfileInput, translation: &TranslationEngine) -> Result<Request> { | |
| let format = wire_format(input.request.request_type()); | |
| let decoded = translation | |
| .decode_request(format, input.request.body(), &TranslationPolicy::default()) | |
| .map_err(|error| SwitchyardError::InvalidRequest(error.to_string()))?; | |
| let mut metadata = metadata_from_headers(&input.metadata.headers); | |
| metadata | |
| .extra_metadata | |
| .get_or_insert_with(Default::default) | |
| .insert( | |
| "inbound_format".to_string(), | |
| inbound_format(input.request.request_type()).to_string(), | |
| ); | |
| Ok(Request { | |
| llm_request: decoded.request, | |
| raw_request: Some(input.request.body().clone()), | |
| metadata: Some(metadata), | |
| }) | |
| fn libsy_request(input: &ProfileInput, translation: &TranslationEngine) -> Result<Request> { | |
| let format = wire_format(input.request.request_type()); | |
| let decoded = translation | |
| .decode_request(format, input.request.body(), &TranslationPolicy::default()) | |
| .map_err(|error| SwitchyardError::InvalidRequest(error.to_string()))?; | |
| if decoded.request.stream { | |
| return Err(SwitchyardError::InvalidRequest( | |
| "libsy-proxy currently supports buffered requests only".to_string(), | |
| )); | |
| } | |
| let mut metadata = metadata_from_headers(&input.metadata.headers); | |
| metadata | |
| .extra_metadata | |
| .get_or_insert_with(Default::default) | |
| .insert( | |
| "inbound_format".to_string(), | |
| inbound_format(input.request.request_type()).to_string(), | |
| ); | |
| Ok(Request { | |
| llm_request: decoded.request, | |
| raw_request: Some(input.request.body().clone()), | |
| metadata: Some(metadata), | |
| }) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@demo/libsy-proxy/src/main.rs` around lines 123 - 141, Update libsy_request to
validate and reject requests with stream: true before decoding or forwarding
them upstream. Return the existing invalid-request error type with a clear
streaming-unsupported message, while preserving the current non-streaming
request construction path.
There was a problem hiding this comment.
Let's remove this [demo/libsy-proxy]
| } | ||
|
|
||
| #[derive(Default, Deserialize)] | ||
| struct CodexTurnMetadata { |
There was a problem hiding this comment.
Can these fields be combined with protocol::Metadata?
| /// Explicit `x-switchyard-*` headers win. NeMo Relay and Dynamo correlation | ||
| /// headers are accepted without linking either runtime. Codex's structured turn | ||
| /// metadata is preferred over its compatibility projections. | ||
| pub fn metadata_from_headers(headers: &BTreeMap<String, Vec<String>>) -> Metadata { |
There was a problem hiding this comment.
this is fine for now. I think we'll wan to consolidate header -> metadata processing in the future
|
Feedback from the Codex agent: Would it be okay to move this normalization into shared ingress/request-metadata infrastructure in this PR? If that would expand or conflict with the affinity work already in flight, Lin Jia can follow up with a separate PR instead. |
4 participants
Why
Sub-agent routing should be an optional persistence policy, not a separate routing algorithm. A normal algorithm should select a target on the first eligible request, then reuse that selection for the same stable identity. Root-agent traffic and algorithms without affinity must preserve their existing behavior.
What
AgentAwareOrchAlgoPOC with a public, reusableAffinitytraitSessionAffinityandSubAgentAffinitypolicies with bounded first-writer-wins stateAgentContext::is_subagentmetadatax-switchyard-*lineage headers without adding a runtime dependencyRandomAlgooptionally retain its final selection through anyAffinityimplementationSubAgentAffinityis enabledNoopAlgounchanged because it does not select a routable modelThe affinity state is bounded to 4,096 process-local assignments.
SubAgentAffinitykeys by(session_id, agent_id)and intentionally excludestask_id, so one child remains on the same model as its task/turn metadata changes.How
Affinityowns only request eligibility, stable key construction, and retained state. Algorithms remain responsible for model selection. On an affinity miss,RandomAlgoruns its normal selection, atomically retains the first winner, and routes to the canonical retained target. On a hit, it bypasses reselection. Requests the policy cannot key continue through ordinary random routing.AffinityKeyis namespaced and public, whileAffinityStateprovides the reusable bounded storage. The integration contract test implements a custom task-based affinity policy outside thelibsycrate to verify the extension point is usable by downstream libraries.Header normalization remains a separate adapter. Explicit
x-switchyard-*headers take precedence; Codex, Relay, and Dynamo headers are converted into neutral libsy metadata. No direct NeMo Relay dependency is required.Try it
Start the buffered demo against NVIDIA Inference Hub:
Send the first turn for a child agent:
The first response reports the provider model and a rationale like:
Repeat the request with the same
session-idandthread-id, but changex-switchyard-task-idtotask-2and the turn metadata toturn-2. The response should use the same model and report:The demo currently requires
"stream": false; adapting an unmodified streaming Codex CLI remains outside this focused change.Where to Start Review
crates/libsy/src/affinity.rs— public trait, state, built-in policies, and header normalizationcrates/libsy/src/algorithms/rand.rs— optional affinity at the final selection pointcrates/libsy/tests/affinity.rs— downstream custom-policy contractdemo/libsy-proxy/src/main.rs— live HTTP/provider wiringcrates/libsy-protocol/src/envelope.rs— explicit child-agent signalTest Plan
cargo fmt --all --check— passcargo clippy --workspace --all-targets -- -D warnings— passcargo test --workspace— passuv run ruff check .— passlibsy-proxy— pass:random routing selected and retainedaffinity reusedtests/e2e/run — incomplete: 17 passed, then 5 setup errors reached--maxfail=5; 21 tests did not run because the existing passthrough fixture invokes the removedswitchyard passthroughCLI commandThe single-target behavior is covered by
single_target_subagent_is_retained_without_reselection; the live demo currently uses its two configured provider targets.Summary by CodeRabbit
New Features
Documentation
Tests