feat(prefill-router): add learned prefill complexity routing#140
feat(prefill-router): add learned prefill complexity routing#140mich-xu wants to merge 5 commits into
Conversation
Signed-off-by: michxu <michxu@nvidia.com>
Signed-off-by: michxu <michxu@nvidia.com>
Signed-off-by: michxu <michxu@nvidia.com>
Signed-off-by: michxu <michxu@nvidia.com>
Signed-off-by: michxu <michxu@nvidia.com>
|
WalkthroughChangesPrefill-probe routing
Estimated code review effort: 5 (Critical) | ~120 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: 4
🧹 Nitpick comments (6)
crates/switchyard-components/src/prefill_probe/policy.rs (1)
17-56: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument that only the cost ordering affects routing, plus the tie rule.
Min-max normalizing two values collapses them to
{0.0, 1.0}(or{0.0, 0.0}), soweak_cost: 0.0 / strong_cost: 1.0and0.0 / 1000.0produce identical decisions — the magnitudes operators write in YAML are inert.scorealso silently breaks ties toward weak (margin >= 0.0). Both are non-obvious invariants worth stating on the public items.📝 Proposed doc-comment clarification
- /// Validates the policy and min-max normalizes costs across weak and strong. + /// Validates the policy and min-max normalizes costs across weak and strong. + /// + /// With only two targets the normalization always yields `{0.0, 1.0}` (or + /// `{0.0, 0.0}` when the costs are equal), so only the relative ordering of + /// `weak_cost` and `strong_cost` affects routing — not their magnitudes. pub(crate) fn new(lambda: f64, weak_cost: f64, strong_cost: f64) -> Result<Self> { @@ - /// Returns `1.0` for weak and `0.0` for strong from two mapped probabilities. + /// Returns `1.0` for weak and `0.0` for strong from two mapped probabilities. + /// + /// Equal utilities resolve to weak, so a zero margin routes to the cheap tier. pub(crate) fn score(&self, weak_probability: f64, strong_probability: f64) -> Result<f64> {🤖 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/switchyard-components/src/prefill_probe/policy.rs` around lines 17 - 56, Update the public documentation for the policy constructor and score behavior around `Policy::new` and `Policy::score`: state that min-max normalization makes routing depend only on weak-versus-strong cost ordering, not cost magnitude, and document that equal utility scores are resolved in favor of weak. Do not change the implementation.crates/switchyard-components/src/prefill_probe/scorer.rs (2)
181-195: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winArtifacts are never cleaned up on the failure paths, so
hidden_states_dirgrows without bound.Cleanup only happens after a successful parse (Line 489). When
wait_for_artifacttimes out (~950 ms), when validation rejects the path, or when parsing fails — a casefailed_locked_read_preserves_artifactasserts as intentional — the.safetensorsfile stays on disk forever, one per failed request. On a long-running proxy with a misbehaving probe this fills the volume. Consider a bounded retention/reaper for the directory, or document an external cleanup requirement indocs/vllm-serve-hidden-state.md.🤖 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/switchyard-components/src/prefill_probe/scorer.rs` around lines 181 - 195, Ensure hidden-state artifacts are cleaned up or bounded on every failure path, including wait_for_artifact timeouts, path validation failures, and parse failures while preserving the intentional failed_locked_read_preserves_artifact behavior. Implement a bounded retention/reaper for hidden_states_dir, or document the required external cleanup in docs/vllm-serve-hidden-state.md; do not limit cleanup to the successful parse path.
296-302: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUnreachable duplicate dtype arm.
Lines 254-262 already reject any dtype other than
F32/BF16, so this arm can never execute. The coding guidelines call out avoiding impossible-case handling;matchon the two supported variants with_ => unreachable_via_earlier_checkremoved keeps the reduction loop tight.As per coding guidelines: "Prefer the minimum code that solves the request; do not add speculative features, abstractions, configurability, or impossible-case handling."
♻️ Proposed simplification
- other => { - return Err(hidden_state_error(format!( - "unsupported hidden_states dtype: {other:?}" - ))) - } - } + }Restructure so
element_sizeand the decode closure are selected once, e.g. by matching on the dtype a single time and dispatching to a shared pooling helper parameterized by afn(&[u8]) -> f32decoder.🤖 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/switchyard-components/src/prefill_probe/scorer.rs` around lines 296 - 302, Remove the unreachable duplicate dtype error arm in the scorer’s hidden-state reduction logic. Restructure the flow around the existing dtype validation so F32 and BF16 select element_size and their decode function once, then share the pooling/reduction loop through the existing scorer helper or a minimal function parameterized by the decoder; do not retain impossible-case handling for unsupported dtypes.Source: Coding guidelines
crates/switchyard-py/src/component_bindings/prefill_probe_request_processor.rs (1)
15-19: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd
///docs for the binding type and its lifecycle methods.Sibling Rust items in this PR carry doc comments (
artifact.rs,policy.rs,scorer.rs); this pyclass and itsstartup/shutdown/processmethods have none, and the no-op lifecycle methods in particular need a line explaining why they do nothing (i.e. where checkpoint loading actually happens).As per coding guidelines: "Add documentation for public APIs: ...
///doc comments for public Rust items. Describe behavior, important invariants, and relevant errors."Also applies to: 74-80
🤖 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/switchyard-py/src/component_bindings/prefill_probe_request_processor.rs` around lines 15 - 19, Add /// documentation to PyPrefillProbeRequestProcessor and its startup, shutdown, and process methods, describing their behavior and lifecycle contract. Explicitly explain in the no-op startup and shutdown documentation that checkpoint loading occurs elsewhere, and document relevant processing behavior or errors consistently with sibling binding types such as artifact.rs, policy.rs, and scorer.rs.Source: Coding guidelines
tests/test_route_bundle.py (1)
84-96: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
raising=Falselets a renamed export silently no-op the patch.If
PrefillProbeRequestProcessoris ever renamed inswitchyard_rust/components.py, this fixture would still "succeed" and the test would exercise the real (checkpoint-loading) processor or fail obscurely. Since the lazy__getattr__is whyraising=Falseis needed, assert membership in the export set instead.♻️ Suggested guard
"""Replace checkpoint loading while retaining route construction.""" _FakePrefillProbeRequestProcessor.instances.clear() + assert "PrefillProbeRequestProcessor" in rust_components.__all__ monkeypatch.setattr(🤖 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 `@tests/test_route_bundle.py` around lines 84 - 96, Update the fake_prefill_probe_processor fixture to validate that PrefillProbeRequestProcessor is present in rust_components’ exported names before applying the monkeypatch, while retaining raising=False for lazy __getattr__ support. Keep the existing replacement and instance-reset behavior unchanged.crates/switchyard-components/src/request_processors/prefill_probe_request_processor.rs (1)
18-18: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider a typed decision instead of
f64+0.5threshold.
CostAwareRoutingPolicy::scorealready returns a binary1.0/0.0, soLEARNED_SCORE_THRESHOLDre-derives a decision from a value that is not a probability. HavingProbeScorerreturn an explicit tier enum would remove the magic constant and prevent a future probability-valued scorer from silently bypassing the cost-aware policy.Also applies to: 130-138
🤖 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/switchyard-components/src/request_processors/prefill_probe_request_processor.rs` at line 18, Replace the f64 score-and-threshold decision flow in ProbeScorer with an explicit tier enum returned by CostAwareRoutingPolicy::score, and remove LEARNED_SCORE_THRESHOLD. Update the ProbeScorer call sites around the affected decision logic to branch on the enum, preserving the existing cost-aware routing behavior while preventing probability values from being treated as binary decisions.
🤖 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/switchyard-components/src/prefill_probe/scorer.rs`:
- Around line 170-177: Update the async scoring flow around
read_and_cleanup_hidden_states so the entire synchronous artifact read, lock,
and cleanup operation runs inside tokio::task::spawn_blocking rather than on the
async worker. Preserve the existing hidden_states_dir, hidden_states_path, and
HiddenStateLayout::from_artifact arguments, propagate the blocking task’s join
error and I/O result, then pass the returned features to self.routing.score.
- Around line 95-119: Update HiddenStateProbeScorer::new and its callers to
configure the reqwest client with a finite probe timeout, preferably by
threading a probe_timeout_s value from PrefillProbeConfig. Ensure score() uses
this client so stalled probe requests terminate and the existing strong-tier
fallback on probe failure is preserved.
In
`@crates/switchyard-components/src/request_processors/prefill_probe_request_processor.rs`:
- Around line 125-149: Update the decision_cache used by select_for_input to
prevent unbounded retention of raw probe text: introduce a capacity-bounded
eviction strategy such as an LRU or capped cache, and use hashed keys if needed
to avoid storing full prompts. Preserve cache lookup behavior and scoring-based
strong/weak target selection while ensuring insertions cannot grow the cache
indefinitely.
In `@switchyard/cli/route_bundle.py`:
- Around line 929-950: Update _prefill_probe_switchyard to accept model_id, wrap
PrefillProbeConfig.model_validate(...) in a ValidationError handler, and
re-raise RouteBundleConfigError that includes the model_id and original
validation details. Update its callers to pass the route’s model_id so malformed
prefill-probe configuration identifies the offending route.
---
Nitpick comments:
In `@crates/switchyard-components/src/prefill_probe/policy.rs`:
- Around line 17-56: Update the public documentation for the policy constructor
and score behavior around `Policy::new` and `Policy::score`: state that min-max
normalization makes routing depend only on weak-versus-strong cost ordering, not
cost magnitude, and document that equal utility scores are resolved in favor of
weak. Do not change the implementation.
In `@crates/switchyard-components/src/prefill_probe/scorer.rs`:
- Around line 181-195: Ensure hidden-state artifacts are cleaned up or bounded
on every failure path, including wait_for_artifact timeouts, path validation
failures, and parse failures while preserving the intentional
failed_locked_read_preserves_artifact behavior. Implement a bounded
retention/reaper for hidden_states_dir, or document the required external
cleanup in docs/vllm-serve-hidden-state.md; do not limit cleanup to the
successful parse path.
- Around line 296-302: Remove the unreachable duplicate dtype error arm in the
scorer’s hidden-state reduction logic. Restructure the flow around the existing
dtype validation so F32 and BF16 select element_size and their decode function
once, then share the pooling/reduction loop through the existing scorer helper
or a minimal function parameterized by the decoder; do not retain
impossible-case handling for unsupported dtypes.
In
`@crates/switchyard-components/src/request_processors/prefill_probe_request_processor.rs`:
- Line 18: Replace the f64 score-and-threshold decision flow in ProbeScorer with
an explicit tier enum returned by CostAwareRoutingPolicy::score, and remove
LEARNED_SCORE_THRESHOLD. Update the ProbeScorer call sites around the affected
decision logic to branch on the enum, preserving the existing cost-aware routing
behavior while preventing probability values from being treated as binary
decisions.
In
`@crates/switchyard-py/src/component_bindings/prefill_probe_request_processor.rs`:
- Around line 15-19: Add /// documentation to PyPrefillProbeRequestProcessor and
its startup, shutdown, and process methods, describing their behavior and
lifecycle contract. Explicitly explain in the no-op startup and shutdown
documentation that checkpoint loading occurs elsewhere, and document relevant
processing behavior or errors consistently with sibling binding types such as
artifact.rs, policy.rs, and scorer.rs.
In `@tests/test_route_bundle.py`:
- Around line 84-96: Update the fake_prefill_probe_processor fixture to validate
that PrefillProbeRequestProcessor is present in rust_components’ exported names
before applying the monkeypatch, while retaining raising=False for lazy
__getattr__ support. Keep the existing replacement and instance-reset behavior
unchanged.
🪄 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: 4cccd0b1-8a8f-4ea9-9116-1739e1e788d3
⛔ Files ignored due to path filters (1)
Cargo.lockis excluded by!**/*.lock,!Cargo.lock
📒 Files selected for processing (24)
.agents/skills/switchyard-coding-agent-launchers/SKILL.md.agents/skills/switchyard-lib-core/SKILL.mdbenchmark/routing-profiles/prefill-probe-local.yamlcrates/switchyard-components/Cargo.tomlcrates/switchyard-components/src/lib.rscrates/switchyard-components/src/prefill_probe/artifact.rscrates/switchyard-components/src/prefill_probe/mod.rscrates/switchyard-components/src/prefill_probe/policy.rscrates/switchyard-components/src/prefill_probe/scorer.rscrates/switchyard-components/src/request_processors/mod.rscrates/switchyard-components/src/request_processors/prefill_probe_request_processor.rscrates/switchyard-py/src/component_bindings.rscrates/switchyard-py/src/component_bindings/prefill_probe_request_processor.rsdocs/vllm-serve-hidden-state.mdswitchyard/__init__.pyswitchyard/cli/route_bundle.pyswitchyard/lib/profiles/__init__.pyswitchyard/lib/profiles/prefill_probe_config.pyswitchyard/lib/profiles/prefill_probe_profile_config.pyswitchyard_rust/components.pyswitchyard_rust/components.pyitests/test_configure_model_picker_inputs.pytests/test_prefill_probe_profile.pytests/test_route_bundle.py
| impl HiddenStateProbeScorer { | ||
| /// Builds a scorer from startup-validated artifact and policy resources. | ||
| pub(crate) fn new( | ||
| base_url: impl Into<String>, | ||
| model: impl Into<String>, | ||
| hidden_states_dir: impl Into<PathBuf>, | ||
| artifact: Arc<InferenceArtifact>, | ||
| weak_head_index: usize, | ||
| strong_head_index: usize, | ||
| policy: CostAwareRoutingPolicy, | ||
| ) -> Self { | ||
| let base_url = base_url.into(); | ||
| Self { | ||
| completions_url: format!("{}/chat/completions", base_url.trim_end_matches('/')), | ||
| model: model.into(), | ||
| hidden_states_dir: hidden_states_dir.into(), | ||
| client: reqwest::Client::new(), | ||
| routing: LearnedRouting { | ||
| artifact, | ||
| weak_head_index, | ||
| strong_head_index, | ||
| policy, | ||
| }, | ||
| } | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
Probe HTTP client has no timeout — a hung probe stalls requests indefinitely.
reqwest::Client::new() applies no request timeout, and score() is invoked on the inbound request path for every routing decision. If the vLLM probe accepts the connection but never responds (loaded/wedged server, dropped packets), send().await never returns, so the "fall back to strong on probe failure" behavior never engages and the caller hangs until the client gives up. Note PrefillProbeConfig exposes tier_timeout_s for completion tiers but nothing bounds the probe call.
⏱️ Proposed fix: bound the probe call
+const PROBE_REQUEST_TIMEOUT: Duration = Duration::from_secs(10);
+
impl HiddenStateProbeScorer {
/// Builds a scorer from startup-validated artifact and policy resources.
pub(crate) fn new(
@@
- ) -> Self {
+ ) -> Result<Self> {
let base_url = base_url.into();
- Self {
+ let client = reqwest::Client::builder()
+ .timeout(PROBE_REQUEST_TIMEOUT)
+ .build()
+ .map_err(|error| probe_error(format!("probe client build failed: {error}")))?;
+ Ok(Self {
completions_url: format!("{}/chat/completions", base_url.trim_end_matches('/')),
model: model.into(),
hidden_states_dir: hidden_states_dir.into(),
- client: reqwest::Client::new(),
+ client,
routing: LearnedRouting {
artifact,
weak_head_index,
strong_head_index,
policy,
},
- }
+ })
}
}A configurable probe_timeout_s threaded from PrefillProbeConfig would be preferable if you want operators to tune it.
📝 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.
| impl HiddenStateProbeScorer { | |
| /// Builds a scorer from startup-validated artifact and policy resources. | |
| pub(crate) fn new( | |
| base_url: impl Into<String>, | |
| model: impl Into<String>, | |
| hidden_states_dir: impl Into<PathBuf>, | |
| artifact: Arc<InferenceArtifact>, | |
| weak_head_index: usize, | |
| strong_head_index: usize, | |
| policy: CostAwareRoutingPolicy, | |
| ) -> Self { | |
| let base_url = base_url.into(); | |
| Self { | |
| completions_url: format!("{}/chat/completions", base_url.trim_end_matches('/')), | |
| model: model.into(), | |
| hidden_states_dir: hidden_states_dir.into(), | |
| client: reqwest::Client::new(), | |
| routing: LearnedRouting { | |
| artifact, | |
| weak_head_index, | |
| strong_head_index, | |
| policy, | |
| }, | |
| } | |
| } | |
| const PROBE_REQUEST_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); | |
| impl HiddenStateProbeScorer { | |
| /// Builds a scorer from startup-validated artifact and policy resources. | |
| pub(crate) fn new( | |
| base_url: impl Into<String>, | |
| model: impl Into<String>, | |
| hidden_states_dir: impl Into<PathBuf>, | |
| artifact: Arc<InferenceArtifact>, | |
| weak_head_index: usize, | |
| strong_head_index: usize, | |
| policy: CostAwareRoutingPolicy, | |
| ) -> Result<Self> { | |
| let base_url = base_url.into(); | |
| let client = reqwest::Client::builder() | |
| .timeout(PROBE_REQUEST_TIMEOUT) | |
| .build() | |
| .map_err(|error| probe_error(format!("probe client build failed: {error}")))?; | |
| Ok(Self { | |
| completions_url: format!("{}/chat/completions", base_url.trim_end_matches('/')), | |
| model: model.into(), | |
| hidden_states_dir: hidden_states_dir.into(), | |
| client, | |
| routing: LearnedRouting { | |
| artifact, | |
| weak_head_index, | |
| strong_head_index, | |
| policy, | |
| }, | |
| }) | |
| } | |
| } |
🤖 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/switchyard-components/src/prefill_probe/scorer.rs` around lines 95 -
119, Update HiddenStateProbeScorer::new and its callers to configure the reqwest
client with a finite probe timeout, preferably by threading a probe_timeout_s
value from PrefillProbeConfig. Ensure score() uses this client so stalled probe
requests terminate and the existing strong-tier fallback on probe failure is
preserved.
| wait_for_artifact(hidden_states_path).await?; | ||
|
|
||
| let raw_features = read_and_cleanup_hidden_states( | ||
| &self.hidden_states_dir, | ||
| hidden_states_path, | ||
| HiddenStateLayout::from_artifact(&self.routing.artifact), | ||
| )?; | ||
| self.routing.score(&raw_features) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Blocking file I/O and an advisory flock run on the async worker thread.
read_and_cleanup_hidden_states is synchronous: it opens the artifact, takes a blocking exclusive lock, reads the whole tensor file into memory, and unlinks it — all inside score().await on a tokio worker. Hidden-state artifacts are prompt_tokens × layers × hidden_size × 4 bytes (tens of MB for realistic prompts), and File::lock parks the thread until any other holder releases. Wrap the sync section in spawn_blocking so one slow artifact cannot stall unrelated requests on the same worker.
♻️ Proposed fix: move the sync section off the runtime worker
- let raw_features = read_and_cleanup_hidden_states(
- &self.hidden_states_dir,
- hidden_states_path,
- HiddenStateLayout::from_artifact(&self.routing.artifact),
- )?;
+ let root = self.hidden_states_dir.clone();
+ let artifact_path = hidden_states_path.to_path_buf();
+ let expected = HiddenStateLayout::from_artifact(&self.routing.artifact);
+ let raw_features = tokio::task::spawn_blocking(move || {
+ read_and_cleanup_hidden_states(&root, &artifact_path, expected)
+ })
+ .await
+ .map_err(|error| probe_error(format!("hidden-state read task failed: {error}")))??;
self.routing.score(&raw_features)Also applies to: 474-496
🤖 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/switchyard-components/src/prefill_probe/scorer.rs` around lines 170 -
177, Update the async scoring flow around read_and_cleanup_hidden_states so the
entire synchronous artifact read, lock, and cleanup operation runs inside
tokio::task::spawn_blocking rather than on the async worker. Preserve the
existing hidden_states_dir, hidden_states_path, and
HiddenStateLayout::from_artifact arguments, propagate the blocking task’s join
error and I/O result, then pass the returned features to self.routing.score.
| async fn select_for_input(&self, input: String) -> LlmTargetId { | ||
| if let Some(selected) = self.decision_cache.lock().get(&input).cloned() { | ||
| return selected; | ||
| } | ||
|
|
||
| match self.scorer.score(&input).await { | ||
| Ok(score) => { | ||
| let selected = if score >= LEARNED_SCORE_THRESHOLD { | ||
| self.weak_target_id.clone() | ||
| } else { | ||
| self.strong_target_id.clone() | ||
| }; | ||
| self.decision_cache.lock().insert(input, selected.clone()); | ||
| selected | ||
| } | ||
| Err(error) => { | ||
| tracing::warn!( | ||
| error = %error, | ||
| fallback_target = %self.strong_target_id, | ||
| "prefill probe failed; using uncached strong fallback" | ||
| ); | ||
| self.strong_target_id.clone() | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
decision_cache is unbounded and never evicted.
Every distinct resolved probe text (potentially a full first-user prompt) is retained for the lifetime of the process, keyed by the raw text. On a long-running proxy with high task cardinality this grows without bound, and it also keeps raw user prompt text resident indefinitely. Consider a capacity bound (LRU or capped insert) and/or hashing the key.
🛡️ Sketch: capacity-bounded insert
+const DECISION_CACHE_CAPACITY: usize = 4_096;
+
async fn select_for_input(&self, input: String) -> LlmTargetId {
if let Some(selected) = self.decision_cache.lock().get(&input).cloned() {
return selected;
}
@@
- self.decision_cache.lock().insert(input, selected.clone());
+ let mut cache = self.decision_cache.lock();
+ if cache.len() >= DECISION_CACHE_CAPACITY {
+ cache.clear();
+ }
+ cache.insert(input, selected.clone());
+ drop(cache);
selected📝 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.
| async fn select_for_input(&self, input: String) -> LlmTargetId { | |
| if let Some(selected) = self.decision_cache.lock().get(&input).cloned() { | |
| return selected; | |
| } | |
| match self.scorer.score(&input).await { | |
| Ok(score) => { | |
| let selected = if score >= LEARNED_SCORE_THRESHOLD { | |
| self.weak_target_id.clone() | |
| } else { | |
| self.strong_target_id.clone() | |
| }; | |
| self.decision_cache.lock().insert(input, selected.clone()); | |
| selected | |
| } | |
| Err(error) => { | |
| tracing::warn!( | |
| error = %error, | |
| fallback_target = %self.strong_target_id, | |
| "prefill probe failed; using uncached strong fallback" | |
| ); | |
| self.strong_target_id.clone() | |
| } | |
| } | |
| } | |
| const DECISION_CACHE_CAPACITY: usize = 4_096; | |
| async fn select_for_input(&self, input: String) -> LlmTargetId { | |
| if let Some(selected) = self.decision_cache.lock().get(&input).cloned() { | |
| return selected; | |
| } | |
| match self.scorer.score(&input).await { | |
| Ok(score) => { | |
| let selected = if score >= LEARNED_SCORE_THRESHOLD { | |
| self.weak_target_id.clone() | |
| } else { | |
| self.strong_target_id.clone() | |
| }; | |
| let mut cache = self.decision_cache.lock(); | |
| if cache.len() >= DECISION_CACHE_CAPACITY { | |
| cache.clear(); | |
| } | |
| cache.insert(input, selected.clone()); | |
| drop(cache); | |
| selected | |
| } | |
| Err(error) => { | |
| tracing::warn!( | |
| error = %error, | |
| fallback_target = %self.strong_target_id, | |
| "prefill probe failed; using uncached strong fallback" | |
| ); | |
| self.strong_target_id.clone() | |
| } | |
| } | |
| } |
🤖 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/switchyard-components/src/request_processors/prefill_probe_request_processor.rs`
around lines 125 - 149, Update the decision_cache used by select_for_input to
prevent unbounded retention of raw probe text: introduce a capacity-bounded
eviction strategy such as an LRU or capped cache, and use hashed keys if needed
to avoid storing full prompts. Preserve cache lookup behavior and scoring-based
strong/weak target selection while ensuring insertions cannot grow the cache
indefinitely.
| def _prefill_probe_switchyard( | ||
| route: Mapping[str, object], | ||
| *, | ||
| target_defaults: Mapping[str, object], | ||
| stats: StatsAccumulator, | ||
| pre_routing_request_processors: Sequence[Any] = (), | ||
| extra_response_processors: Sequence[Any] = (), | ||
| ) -> ChainRuntime: | ||
| """Build a learned prefill-probe routing profile from route-bundle data.""" | ||
| config = PrefillProbeConfig.model_validate( | ||
| _route_config(route, target_defaults, ("probe", "strong", "weak")) | ||
| ) | ||
| return ProfileSwitchyard( | ||
| PrefillProbeProfileConfig.from_config(config) | ||
| .build() | ||
| .with_runtime_components( | ||
| stats_accumulator=stats, | ||
| enable_stats=config.enable_stats, | ||
| pre_request_processors=pre_routing_request_processors, | ||
| response_processors=extra_response_processors, | ||
| ) | ||
| ) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check whether sibling route helpers wrap pydantic validation errors and include model_id.
rg -nP -C8 'def _(deterministic|escalation_router|stage_router|plan_execute)_switchyard\b' switchyard/cli/route_bundle.py
rg -nP -C4 'ValidationError' switchyard/cli/route_bundle.pyRepository: NVIDIA-NeMo/Switchyard
Length of output: 3664
🏁 Script executed:
#!/bin/bash
sed -n '900,980p' switchyard/cli/route_bundle.py
printf '\n----\n'
sed -n '1220,1325p' switchyard/cli/route_bundle.py
printf '\n----\n'
sed -n '1,120p' switchyard/cli/route_bundle.pyRepository: NVIDIA-NeMo/Switchyard
Length of output: 12565
🏁 Script executed:
#!/bin/bash
rg -n -C3 'PrefillProbeConfig|_prefill_probe_switchyard|RouteBundleConfigError|ValidationError' switchyard/cli/route_bundle.py
printf '\n----\n'
sed -n '520,760p' switchyard/cli/route_bundle.py
printf '\n----\n'
sed -n '760,930p' switchyard/cli/route_bundle.pyRepository: NVIDIA-NeMo/Switchyard
Length of output: 28951
🏁 Script executed:
#!/bin/bash
rg -n -C4 'def _route_config\b|model_validate\(' switchyard/cli/route_bundle.py
printf '\n----\n'
sed -n '1420,1475p' switchyard/cli/route_bundle.py
printf '\n----\n'
sed -n '980,1060p' switchyard/cli/route_bundle.pyRepository: NVIDIA-NeMo/Switchyard
Length of output: 8365
Wrap prefill-probe validation errors in RouteBundleConfigError
PrefillProbeConfig.model_validate(...) still raises a raw ValidationError; thread model_id into this helper and re-raise with route context so malformed prefill-probe YAML points at the offending route.
🤖 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 `@switchyard/cli/route_bundle.py` around lines 929 - 950, Update
_prefill_probe_switchyard to accept model_id, wrap
PrefillProbeConfig.model_validate(...) in a ValidationError handler, and
re-raise RouteBundleConfigError that includes the model_id and original
validation details. Update its callers to pass the route’s model_id so malformed
prefill-probe configuration identifies the offending route.
1 participant
What
Adds a learned
prefill_proberouting profile using Switchyard’s supported v1 component chain.prefill_probein the current route-bundle configuration and model discovery paths while keeping the probe model internal.The checkpoint remains external deployment data. Switchyard does not download checkpoints, fit PCA, train models, or perform online learning.
Why
The earlier prefill-router implementation targeted the removed components-v2 architecture and is not present on canonical
main. This ports the learned router onto the supported v1 request-component and deterministic-backend chain.The router enables task-level strong/weak selection using learned prompt complexity while preserving the original completion request and Switchyard’s existing backend, statistics, translation, and provider-cache behavior.
Related to #97 , #13 , and #119 .
How tested
uv run ruff check .cleanuv run mypy switchyardcleanuv run pytest tests/greenAdditional validation:
cargo fmt --all --checkcargo clippy --workspace --all-targets -- -D warningscargo test --workspaceuv run --only-group docs mkdocs build --strict[prompt_tokens, layers, hidden_size]artifact handling, PCA-200 checkpoint inference, and hidden-state file cleanup.15,420cache-creation tokens on the first request and15,420cached tokens on the second.mteb-leaderboardtask atlambda: 0.75with stock Terminus 2. The task made 14 successful weak-tier Nemotron calls from one probe decision, recorded zero endpoint errors, and cleaned up the hidden-state artifact. The verifier ran successfully; the task reward was0.0because the agent produced an incorrect answer.Checklist
snake_caseof the primary class.switchyard/__init__.py.__all__if intended for downstream use.--helpupdated if customer-facing surface changed.Signed-off-by: Your Name <email>) per the DCO.Notes for reviewers
lambdais the only continuous routing knob, an exact utility tie selects weak, and checkpoint probabilities remain private.build_tier_backend, which preserves standard Anthropic cache wrapping and statistics.Intentional minimal-release trade-offs:
Summary by CodeRabbit