feat(llm-client): add vllm hidden-state probe - #185
Conversation
Signed-off-by: asurla1998 <asurla1998@users.noreply.huggingface.co>
Signed-off-by: asurla1998 <asurla1998@users.noreply.huggingface.co>
WalkthroughAdds a vLLM hidden-state extraction probe and a learned prefill classifier. The classifier validates and runs router artifacts, applies cost-aware weak/strong selection, caches successful decisions, and integrates with existing classifier routing. ChangesPrefill routing pipeline
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: 2
🧹 Nitpick comments (9)
crates/libsy-llm-client/src/vllm_hidden_state_probe.rs (3)
149-152: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winHTTP client build failure should be a
Configurationerror.
map_reqwest_errorclassifies a builder failure asTransport(orTimeout), but nothing has been sent yet — this is purely a bad-config outcome, andConfigurationexists for exactly that. Callers matching onTransportmay wrongly retry.♻️ Suggested change
let client = reqwest::Client::builder() .timeout(config.request_timeout) .build() - .map_err(map_reqwest_error)?; + .map_err(|error| { + configuration_error(format!("failed to build probe HTTP client: {error}")) + })?;🤖 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-llm-client/src/vllm_hidden_state_probe.rs` around lines 149 - 152, Update the reqwest client construction in the probe initialization flow to map `.build()` failures directly to the existing `Configuration` error variant instead of using `map_reqwest_error`. Keep `map_reqwest_error` for runtime request failures and preserve the existing successful client construction path.
168-169: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winEvery probe pays a full directory scan before its request.
reap_stale_artifactswalks the entire hidden-state directory (andstats + lock-probes each.safetensorsentry) on the blocking pool ahead of eachextract. On a busy shared connector directory this is O(files) blocking I/O per routing decision, directly on the prefill-latency path. Consider amortizing it — a periodic background sweep, or gating on an elapsed-since-last-sweep timestamp — rather than running it inline per request.🤖 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-llm-client/src/vllm_hidden_state_probe.rs` around lines 168 - 169, The extract method currently performs the full reap_stale_artifacts directory scan inline for every request. Amortize stale-artifact cleanup by gating reap_stale_artifacts with an elapsed-since-last-sweep check or moving it to a periodic background task, while preserving cleanup behavior and avoiding blocking I/O on the extract request path.
21-23: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winConsider exposing the artifact wait timeout through
VllmHiddenStateProbeConfig.
artifact_wait_timeoutandstale_artifact_retentionare struct fields but are only ever set from the consts, so they are effectively unconfigurable. A hard 1s ceiling on waiting for vLLM to release the synchronization lock is the most likely first failure mode under prefill load, andrequest_timeoutis already caller-tunable — the wait bound deserves the same treatment (optional fields defaulting to these consts).Also applies to: 118-119
🤖 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-llm-client/src/vllm_hidden_state_probe.rs` around lines 21 - 23, Expose artifact_wait_timeout and stale_artifact_retention as optional fields on VllmHiddenStateProbeConfig, defaulting to ARTIFACT_WAIT_TIMEOUT and STALE_ARTIFACT_RETENTION when unset. Update configuration construction and the probe’s wait/retention logic to use the resolved config values, while preserving the existing constants as defaults.crates/libsy-llm-client/Cargo.toml (1)
21-21: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider workspace-backing
safetensors(andtempfile).Every other dependency added here is
*.workspace = true, andcrates/libsy/src/algorithms/prefill_probe/artifact.rsin this same stack also depends onsafetensors. Two independent pins can drift and pull in duplicate versions of the tensor parser.♻️ Suggested change
-safetensors = "0.4" +safetensors.workspace = trueAdd to the workspace manifest
[workspace.dependencies]:safetensors = "0.4"Also applies to: 29-29
🤖 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-llm-client/Cargo.toml` at line 21, Move the safetensors dependency to the workspace manifest’s [workspace.dependencies] with version 0.4, then update the libsy-llm-client dependency declaration to use workspace backing; apply the same workspace-backed dependency change to tempfile where it is declared.crates/libsy/src/algorithms/prefill_probe/artifact.rs (2)
588-604: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low valueConsider accumulating dense products in
f64for parity withproject_pca.
project_pca(Line 547) accumulates inf64whiledense_layersums 200/256/128-wide dot products inf32, so trunk logits can drift from the exported reference math for large activations. Cheap to align.♻️ Proposed change
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 dense_layer’s dot-product accumulation to use f64, matching project_pca’s reference-math precision, then convert the summed result back to f32 before adding the bias and applying the existing finite-value and ReLU handling.
181-227: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueTest helper can produce non-finite logits for degenerate probabilities.
(probability / (1.0 - probability)).ln()yieldsinfat1.0and-inf/NaN at0.0, which would then fail the finite checks inensemble_probabilities. Current callers pass interior values, but a debug assertion or clamp keeps the helper safe for future tests.🤖 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 181 - 227, Update with_test_probabilities so boundary probabilities cannot produce non-finite logits before being stored in ensemble output biases. Clamp or otherwise validate each probability to a finite interior range before calculating the logarithm, while preserving existing behavior for valid interior probabilities.crates/libsy/src/algorithms/prefill_probe.rs (2)
300-322: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low valueEmptiness check trims but the distinctness check does not.
"weak"and" weak"passstrong_target != weak_targetyet resolve to the same semantic target downstream (androuting_tiercompares raw strings). Compare trimmed values, or store trimmed targets.♻️ Proposed change
- if config.strong_target == config.weak_target { + if config.strong_target.trim() == config.weak_target.trim() { return Err(config_error( "strong_target and weak_target must be distinct", )); }🤖 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 300 - 322, Update validate_config to compare trimmed strong_target and weak_target values for distinctness, and ensure downstream routing_tier uses the same normalized target values or equivalent trimmed comparisons. Preserve the existing non-empty validation and error behavior while preventing whitespace-only differences from being treated as distinct targets.
341-360: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winTerminus task text is used verbatim without trimming.
terminus_task_instructionreturns the raw slice between headers, so leading/trailing whitespace variations of the same task produce different cache keys and different probe inputs. Trimming (and rejecting empty after trim) would make keys stable.♻️ Proposed change
fn terminus_task_instruction(instruction: &str) -> Option<&str> { let (_, task_and_terminal) = instruction.split_once(TERMINUS_TASK_DESCRIPTION_HEADER)?; let (task, _) = task_and_terminal.split_once(TERMINUS_TERMINAL_STATE_HEADER)?; - (!task.is_empty()).then_some(task) + let task = task.trim(); + (!task.is_empty()).then_some(task) }🤖 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 341 - 360, Trim the task slice returned by terminus_task_instruction before validating and returning it, so leading or trailing whitespace does not affect probe inputs or cache keys. Ensure whitespace-only tasks are rejected after trimming, while preserving the existing header parsing and fallback behavior in probe_input.crates/libsy/Cargo.toml (1)
19-25: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUpdate
lruandsafetensorspins
lru = "0.12"andsafetensors = "0.4"are behind current releases; refresh them unless you need to stay on these lines.🤖 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/Cargo.toml` around lines 19 - 25, Update the dependency pins for lru and safetensors in the crate manifest to current compatible releases, replacing the existing 0.12 and 0.4 constraints. Keep the remaining dependency configuration unchanged unless compatibility requires retaining the current major lines.
🤖 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-llm-client/src/vllm_hidden_state_probe.rs`:
- Around line 466-473: Update cleanup_artifact_files so removing the main
artifact treats a filesystem NotFound error as successful cleanup, matching the
existing lock-file handling for lock_path. Continue converting other removal
errors through artifact_error, while preserving the current cleanup flow.
- Around line 484-498: Update cleanup_stale_artifacts in
crates/libsy-llm-client/src/vllm_hidden_state_probe.rs at lines 484-498 to skip
entries when entry, file_type, metadata, or modified lookups fail due to
concurrent disappearance (at minimum ErrorKind::NotFound), continuing the sweep
instead of returning an error; also update remove_file handling at lines 466-473
to treat ErrorKind::NotFound as success, consistent with the lock-file handling.
---
Nitpick comments:
In `@crates/libsy-llm-client/Cargo.toml`:
- Line 21: Move the safetensors dependency to the workspace manifest’s
[workspace.dependencies] with version 0.4, then update the libsy-llm-client
dependency declaration to use workspace backing; apply the same workspace-backed
dependency change to tempfile where it is declared.
In `@crates/libsy-llm-client/src/vllm_hidden_state_probe.rs`:
- Around line 149-152: Update the reqwest client construction in the probe
initialization flow to map `.build()` failures directly to the existing
`Configuration` error variant instead of using `map_reqwest_error`. Keep
`map_reqwest_error` for runtime request failures and preserve the existing
successful client construction path.
- Around line 168-169: The extract method currently performs the full
reap_stale_artifacts directory scan inline for every request. Amortize
stale-artifact cleanup by gating reap_stale_artifacts with an
elapsed-since-last-sweep check or moving it to a periodic background task, while
preserving cleanup behavior and avoiding blocking I/O on the extract request
path.
- Around line 21-23: Expose artifact_wait_timeout and stale_artifact_retention
as optional fields on VllmHiddenStateProbeConfig, defaulting to
ARTIFACT_WAIT_TIMEOUT and STALE_ARTIFACT_RETENTION when unset. Update
configuration construction and the probe’s wait/retention logic to use the
resolved config values, while preserving the existing constants as defaults.
In `@crates/libsy/Cargo.toml`:
- Around line 19-25: Update the dependency pins for lru and safetensors in the
crate manifest to current compatible releases, replacing the existing 0.12 and
0.4 constraints. Keep the remaining dependency configuration unchanged unless
compatibility requires retaining the current major lines.
In `@crates/libsy/src/algorithms/prefill_probe.rs`:
- Around line 300-322: Update validate_config to compare trimmed strong_target
and weak_target values for distinctness, and ensure downstream routing_tier uses
the same normalized target values or equivalent trimmed comparisons. Preserve
the existing non-empty validation and error behavior while preventing
whitespace-only differences from being treated as distinct targets.
- Around line 341-360: Trim the task slice returned by terminus_task_instruction
before validating and returning it, so leading or trailing whitespace does not
affect probe inputs or cache keys. Ensure whitespace-only tasks are rejected
after trimming, while preserving the existing header parsing and fallback
behavior in probe_input.
In `@crates/libsy/src/algorithms/prefill_probe/artifact.rs`:
- Around line 588-604: Update dense_layer’s dot-product accumulation to use f64,
matching project_pca’s reference-math precision, then convert the summed result
back to f32 before adding the bias and applying the existing finite-value and
ReLU handling.
- Around line 181-227: Update with_test_probabilities so boundary probabilities
cannot produce non-finite logits before being stored in ensemble output biases.
Clamp or otherwise validate each probability to a finite interior range before
calculating the logarithm, while preserving existing behavior for valid interior
probabilities.
🪄 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: 93deeaed-1b3e-4ce8-904c-c696d85cb308
⛔ Files ignored due to path filters (1)
Cargo.lockis excluded by!**/*.lock,!Cargo.lock
📒 Files selected for processing (9)
crates/libsy-llm-client/Cargo.tomlcrates/libsy-llm-client/src/lib.rscrates/libsy-llm-client/src/vllm_hidden_state_probe.rscrates/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
| fn cleanup_artifact_files(path: &Path) -> VllmHiddenStateProbeResult<()> { | ||
| let lock_path = companion_lock_path(path); | ||
| std::fs::remove_file(path).map_err(|error| { | ||
| artifact_error(format!( | ||
| "hidden-state artifact cleanup error for {}: {error}", | ||
| path.display() | ||
| )) | ||
| })?; |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Treat a already-removed artifact as successful cleanup.
The lock file's NotFound is tolerated (Line 476) but the artifact's is not, so a concurrent reaper or a second probe instance removing the file first converts a fully successful extraction into an Artifact error at Line 458. Mirror the lock-file handling. (Shares a root cause with the sweep issue above — see the consolidated note.)
🤖 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-llm-client/src/vllm_hidden_state_probe.rs` around lines 466 -
473, Update cleanup_artifact_files so removing the main artifact treats a
filesystem NotFound error as successful cleanup, matching the existing lock-file
handling for lock_path. Continue converting other removal errors through
artifact_error, while preserving the current cleanup flow.
| fn cleanup_stale_artifacts(root: &Path, retention: Duration) -> VllmHiddenStateProbeResult<()> { | ||
| let entries = std::fs::read_dir(root).map_err(|error| { | ||
| artifact_error(format!( | ||
| "failed to scan hidden-state directory {}: {error}", | ||
| root.display() | ||
| )) | ||
| })?; | ||
| let now = SystemTime::now(); | ||
| for entry in entries { | ||
| let entry = entry.map_err(|error| { | ||
| artifact_error(format!( | ||
| "failed to inspect hidden-state directory {}: {error}", | ||
| root.display() | ||
| )) | ||
| })?; |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Filesystem races in the shared hidden-state directory are treated as hard failures. Both sites assume this probe is the only mutator of hidden_states_dir, but vLLM writes there and concurrent probe instances reap there, so a file disappearing mid-operation is expected rather than exceptional. Classify concurrent disappearance as benign in both places.
crates/libsy-llm-client/src/vllm_hidden_state_probe.rs#L484-L498: skip entries whoseentry/file_type/metadata/modifiedlookups fail (at minimumNotFound) and continue the sweep, instead of returningErrand failingextractbefore the HTTP request is sent.crates/libsy-llm-client/src/vllm_hidden_state_probe.rs#L466-L473: treatErrorKind::NotFoundonremove_file(path)as success, matching the lock-file handling immediately below it.
📍 Affects 1 file
crates/libsy-llm-client/src/vllm_hidden_state_probe.rs#L484-L498(this comment)crates/libsy-llm-client/src/vllm_hidden_state_probe.rs#L466-L473
🤖 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-llm-client/src/vllm_hidden_state_probe.rs` around lines 484 -
498, Update cleanup_stale_artifacts in
crates/libsy-llm-client/src/vllm_hidden_state_probe.rs at lines 484-498 to skip
entries when entry, file_type, metadata, or modified lookups fail due to
concurrent disappearance (at minimum ErrorKind::NotFound), continuing the sweep
instead of returning an error; also update remove_file handling at lines 466-473
to treat ErrorKind::NotFound as success, consistent with the lock-file handling.
1 participant
What
VllmHiddenStateProbefor vLLMExampleHiddenStatesConnectorartifacts.switchyard-llm-client.Why
Provides the transport implementation consumed by the transport-independent prefill classifier in #174.
Depends on #174. Until #174 merges, GitHub will also show its commit in this stacked PR.
How tested
cargo fmt --all --checkcargo clippy --workspace --all-targets -- -D warningscargo test --workspaceuv run ruff check .uv run mypy switchyarduv run pytest tests/ -v -m "not integration"— 1718 passed, 12 skipped, 28 deselectedChecklist
Notes for reviewers
The intended delta is the single
feat(llm-client)commit on top of #174. Token-aware pricing remains a separate follow-up.Summary by CodeRabbit