feat(libsy): add prefill-probe classifier - #174
Conversation
Signed-off-by: asurla1998 <asurla1998@users.noreply.huggingface.co>
e1e415f to
caed4cd
Compare
WalkthroughAdds a transport-independent ChangesPrefill Probe Routing
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.
🧹 Nitpick comments (5)
crates/libsy/src/algorithms/prefill_probe/artifact.rs (1)
588-604: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low valueConsider accumulating dense dot products in
f64for consistency withproject_pca.
project_pca(Line 547) widens tof64before summing, whiledense_layersums 200-term products inf32. Not a correctness bug at these widths, but the mixed precision strategy is inconsistent and slightly more prone to rounding drift versus the training-time (float32/float64 numpy) reference.♻️ Optional: widen accumulation
.map(|(row, bias)| { let value = row .iter() .zip(input) - .map(|(weight, input)| *weight * *input) - .sum::<f32>() - + *bias; + .map(|(weight, input)| f64::from(*weight) * f64::from(*input)) + .sum::<f64>() as f32 + + *bias;🤖 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/prefill_probe/artifact.rs` around lines 588 - 604, Update the dense_layer dot-product accumulation to widen each weight/input product to f64 and sum in f64, matching project_pca’s accumulation strategy; apply the bias consistently before converting the result back to the expected f32 output and retain the existing non-finite and ReLU handling.crates/libsy/src/algorithms/prefill_probe.rs (3)
210-225: 🚀 Performance & Scalability | 🔵 Trivial | ⚖️ Poor tradeoffNo in-flight deduplication: concurrent identical tasks each trigger a full probe extraction.
The cache is only consulted before and populated after the await, so a burst of requests for the same task all miss and each performs the probe round-trip plus inference. Once
PrefillProbeis backed by an HTTP client (later in the stack), this amplifies load on the probe endpoint proportional to burst concurrency. A single-flight map keyed bycache_key(e.g. per-key broadcast/notify) would collapse duplicates.Worth deciding whether the added complexity is justified at your expected traffic shape — flagging so it is a conscious choice rather than an oversight.
🤖 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/prefill_probe.rs` around lines 210 - 225, Update PrefillProbe::select_for_task to add per-cache_key in-flight deduplication, so concurrent identical tasks await one shared probe extraction and inference operation instead of starting duplicates. Reuse the existing decision_cache key and ensure the in-flight entry is removed after completion while preserving the current cache lookup and result behavior.
236-245: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winSustained probe outage will emit one
warnper request.When the probe endpoint or checkpoint inference is persistently unavailable, this logs at
warnon every single classification, flooding logs at request rate with no added diagnostic value. Consider emitting a counter/metric for the fallback and rate-limiting the warning (or logging the first failure atwarnand subsequent ones atdebuguntil recovery).🤖 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/prefill_probe.rs` around lines 236 - 245, Adjust the error branch around the prefill probe fallback to avoid emitting a warn log for every request during sustained outages. Keep recording each fallback with a counter/metric, and rate-limit the warning or emit it at warn only for the first failure, lowering subsequent failures to debug until probe recovery; preserve the existing strong_target fallback behavior.
256-265: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low value
Debugacquires the cache mutex; keep it out of any lock-holding path.
parking_lot::Mutexis not reentrant, so formatting aPrefillProbeClassifierwhile thedecision_cachelock is held self-deadlocks. No current call site does this, but a futuretracingfield likeclassifier = ?selfinsideselect_for_task's locked scopes would deadlock rather than fail loudly.try_lockwith a placeholder would makeDebugunconditionally safe.🛡️ Optional defensive change
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { formatter .debug_struct("PrefillProbeClassifier") .field("strong_target", &self.strong_target) .field("weak_target", &self.weak_target) - .field("cache_capacity", &self.decision_cache.lock().cap()) + .field( + "cache_capacity", + &self.decision_cache.try_lock().map(|cache| cache.cap()), + ) .finish_non_exhaustive() }🤖 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/prefill_probe.rs` around lines 256 - 265, Update PrefillProbeClassifier’s fmt implementation to avoid blocking on decision_cache: use try_lock and emit a placeholder when the mutex is already held, while preserving the capacity output when the lock is available. Keep the existing strong_target and weak_target debug fields unchanged.crates/libsy/src/algorithms/prefill_probe/policy.rs (1)
106-150: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd coverage for an intermediate
lambdawhere both terms matter.Tests exercise only the degenerate endpoints (
lambda0.0 and 1.0), where one term of the utility is zeroed out. A case likelambda = 0.5with unequal costs would pin the actual tradeoff arithmetic and catch a future sign or weight inversion inselect.🤖 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/prefill_probe/policy.rs` around lines 106 - 150, Add a test alongside lambda_zero_uses_cost_and_weak_wins_equal_utility and lambda_one_uses_correctness_probabilities covering CostAwareRoutingPolicy::select with lambda = 0.5 and unequal weak/strong costs. Use probabilities and costs where the weighted correctness-cost tradeoff produces distinct expected PrefillTier choices, validating both routing directions and the intermediate utility arithmetic.
🤖 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.
Nitpick comments:
In `@crates/libsy/src/algorithms/prefill_probe.rs`:
- Around line 210-225: Update PrefillProbe::select_for_task to add per-cache_key
in-flight deduplication, so concurrent identical tasks await one shared probe
extraction and inference operation instead of starting duplicates. Reuse the
existing decision_cache key and ensure the in-flight entry is removed after
completion while preserving the current cache lookup and result behavior.
- Around line 236-245: Adjust the error branch around the prefill probe fallback
to avoid emitting a warn log for every request during sustained outages. Keep
recording each fallback with a counter/metric, and rate-limit the warning or
emit it at warn only for the first failure, lowering subsequent failures to
debug until probe recovery; preserve the existing strong_target fallback
behavior.
- Around line 256-265: Update PrefillProbeClassifier’s fmt implementation to
avoid blocking on decision_cache: use try_lock and emit a placeholder when the
mutex is already held, while preserving the capacity output when the lock is
available. Keep the existing strong_target and weak_target debug fields
unchanged.
In `@crates/libsy/src/algorithms/prefill_probe/artifact.rs`:
- Around line 588-604: Update the dense_layer dot-product accumulation to widen
each weight/input product to f64 and sum in f64, matching project_pca’s
accumulation strategy; apply the bias consistently before converting the result
back to the expected f32 output and retain the existing non-finite and ReLU
handling.
In `@crates/libsy/src/algorithms/prefill_probe/policy.rs`:
- Around line 106-150: Add a test alongside
lambda_zero_uses_cost_and_weak_wins_equal_utility and
lambda_one_uses_correctness_probabilities covering
CostAwareRoutingPolicy::select with lambda = 0.5 and unequal weak/strong costs.
Use probabilities and costs where the weighted correctness-cost tradeoff
produces distinct expected PrefillTier choices, validating both routing
directions and the intermediate utility arithmetic.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d56e1e4e-22a0-4f52-8246-7bf63869f9be
⛔ Files ignored due to path filters (1)
Cargo.lockis excluded by!**/*.lock,!Cargo.lock
📒 Files selected for processing (6)
crates/libsy/Cargo.tomlcrates/libsy/src/algorithms.rscrates/libsy/src/algorithms/prefill_probe.rscrates/libsy/src/algorithms/prefill_probe/artifact.rscrates/libsy/src/algorithms/prefill_probe/policy.rscrates/libsy/src/lib.rs
|
@anniesurla can you rebase? |
2 participants
Summary
libsyas aClassifierPrefillProbeboundary for hidden-state extractionContext
This is the first of three smaller PRs replacing #140 in response to maintainer review:
libsyclassifier and checkpoint inference (this PR)This PR intentionally contains no HTTP client, filesystem hidden-state transport, Python binding, or server configuration.
Implementation notes
tokio::task::spawn_blocking.Test plan
cargo fmt --all --checkcargo clippy --workspace --all-targets -- -D warningscargo test -p switchyard-libsycargo test --workspaceThe live VPN-dependent test remains ignored by the workspace suite.
Summary by CodeRabbit