diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-admission-control.md b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-admission-control.md index 4434abd02bf5..6915cd7ccfc0 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-admission-control.md +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-admission-control.md @@ -78,6 +78,7 @@ use std::time::Duration; use dynamo_kv_router::plugins::request_classifier::{ ClassifierError, ClassifyFuture, ClassifyRequest, RequestClassifier, }; +use dynamo_kv_router::plugins::worker_selection::WorkerCandidates; use dynamo_runtime::error::{DynamoError, ErrorClass}; use tokio::time::Instant; @@ -86,7 +87,7 @@ struct QueueBudget { } impl RequestClassifier for QueueBudget { - fn classify(&mut self, mut request: ClassifyRequest) -> ClassifyFuture { + fn classify(&mut self, mut request: ClassifyRequest, _workers: WorkerCandidates<'_>) -> ClassifyFuture { let deadline = request.ingress_at() + self.budget; Box::pin(async move { if Instant::now() >= deadline { @@ -230,7 +231,47 @@ Each entry returned by `workers()` is a `RequestClassifierWorker`: | `worker()` | Worker ID and data-parallel rank, returned as `WorkerWithDpRank` | | `total_kv_blocks()` | Advertised total KV-block capacity for this rank, or `None` when unavailable; this is not the number of free blocks | -Classifiers do not currently receive the per-worker cache and load views exposed to worker-selection plugins through `WorkerInputs`. +### Worker Inputs + +A classifier can read the per-worker cache and load data that [worker-selection plugins](custom-worker-selection.mdx) receive. Declare the groups in `required_worker_inputs`; `classify` receives them as `workers`: + +```rust +use dynamo_kv_router::plugins::request_classifier::{ + ClassifyFuture, ClassifyRequest, RequestClassifier, +}; +use dynamo_kv_router::plugins::worker_selection::{WorkerCandidates, WorkerInputs}; + +struct LoadAwareClass { + max_active_requests: usize, +} + +impl RequestClassifier for LoadAwareClass { + fn required_worker_inputs(&self) -> WorkerInputs { + WorkerInputs::LOAD + } + + fn classify(&mut self, mut request: ClassifyRequest, workers: WorkerCandidates<'_>) -> ClassifyFuture { + // Count only observed loads, and require at least one. + let mut loads = workers + .iter() + .filter_map(|worker| worker.load()) + .filter(|load| load.is_available()) + .peekable(); + let saturated = loads.peek().is_some() + && loads.all(|load| load.active_requests() >= self.max_active_requests); + if saturated { + request.set_policy_class("batch"); + } + Box::pin(async move { Ok(request) }) + } +} +``` + +Dynamo builds these rows with the same code that worker selection uses, so `CACHE`, `LOAD`, and `PREFERRED_TAINT` have the same meaning in both plugins. Dynamo reads `required_worker_inputs` once, when it installs the classifier. With the default `WorkerInputs::NONE`, Dynamo builds no rows and `workers` is empty. + +- Rows cover each worker rank that is eligible for the request, including overloaded workers. Row order is unspecified. +- `workers` is a snapshot taken for this `classify` call and is lent only for its duration. Dynamo frees the rows when `classify` returns, so a deferred classification holds none of them. Copy into the returned future whatever it needs, and do not treat that copy as current load when the request is admitted later. +- Worker selection takes its own snapshot when it picks a worker, so it can pick a different worker than the classifier expected. Use `set_worker_selection_target` to state a preference. ### Admission Decisions diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx index 782cbc620140..30b1d4809512 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx @@ -604,11 +604,12 @@ For example, a classifier can inspect the sequence while retaining the request i use dynamo_kv_router::plugins::request_classifier::{ ClassifyFuture, ClassifyRequest, RequestClassifier, }; +use dynamo_kv_router::plugins::worker_selection::WorkerCandidates; struct InspectPrefixes; impl RequestClassifier for InspectPrefixes { - fn classify(&mut self, request: ClassifyRequest) -> ClassifyFuture { + fn classify(&mut self, request: ClassifyRequest, _workers: WorkerCandidates<'_>) -> ClassifyFuture { Box::pin(async move { let _tracked_blocks = request.sequence_hashes().map(|hashes| hashes.len()); Ok(request) diff --git a/lib/kv-router/src/plugins/request_classifier.rs b/lib/kv-router/src/plugins/request_classifier.rs index 68b31d8fe992..903108248f4c 100644 --- a/lib/kv-router/src/plugins/request_classifier.rs +++ b/lib/kv-router/src/plugins/request_classifier.rs @@ -27,6 +27,7 @@ use async_trait::async_trait; use dynamo_tokens::SequenceHash; use tokio::time::Instant; +use crate::plugins::worker_selection::{WorkerCandidates, WorkerInputs}; use crate::protocols::{WorkerAffinityTarget, WorkerWithDpRank}; use crate::scheduling::{SessionContext, policy_queue::QueueSnapshot}; @@ -270,7 +271,7 @@ pub type ClassifyFuture = /// } /// /// impl RequestClassifier for Pauser { -/// fn classify(&mut self, request: ClassifyRequest) -> ClassifyFuture { +/// fn classify(&mut self, request: ClassifyRequest, _workers: WorkerCandidates<'_>) -> ClassifyFuture { /// // Clone the handle into the future; `self` cannot be borrowed. /// let state = Arc::clone(&self.state); /// Box::pin(async move { @@ -293,7 +294,29 @@ pub type ClassifyFuture = /// the request ends; events still queued at router shutdown are dropped. #[async_trait] pub trait RequestClassifier: Send + 'static { - fn classify(&mut self, request: ClassifyRequest) -> ClassifyFuture { + /// Declare the worker-signal groups the `workers` view passed to [`Self::classify`] + /// carries. The host reads this once at install and builds rows only when it is not + /// [`WorkerInputs::NONE`]. + fn required_worker_inputs(&self) -> WorkerInputs { + WorkerInputs::NONE + } + + /// `workers` covers every worker rank eligible for this request under its pin, allowed + /// workers, availability, and routing constraints, in unspecified order, with the groups + /// [`Self::required_worker_inputs`] declares; it is empty when that is + /// [`WorkerInputs::NONE`]. Overloaded workers are included. Values come from the row + /// builder worker selection uses, so each group means the same thing here as in a filter + /// or scorer. + /// + /// The view is a snapshot taken for this call and is lent only for its duration: the host + /// frees the rows when `classify` returns, so a deferred classification does not hold them. + /// Copy into the returned future anything it needs. Worker selection takes its own + /// snapshot when it later picks a worker, so it can pick a different worker. + fn classify( + &mut self, + request: ClassifyRequest, + _workers: WorkerCandidates<'_>, + ) -> ClassifyFuture { Box::pin(async move { Ok(request) }) } diff --git a/lib/kv-router/src/plugins/worker_selection.rs b/lib/kv-router/src/plugins/worker_selection.rs index 6288c7eaf4c9..7b1aff913397 100644 --- a/lib/kv-router/src/plugins/worker_selection.rs +++ b/lib/kv-router/src/plugins/worker_selection.rs @@ -14,7 +14,7 @@ pub use runtime::{ }; pub use context::WorkerSelectionContext; -pub(crate) use inputs::{CacheSnapshot, CandidateData, WorkerCacheData}; +pub(crate) use inputs::{CacheSnapshot, CandidateData, WorkerCacheData, WorkerInputRows}; pub use inputs::{ ScoredWorkerCandidate, WorkerCacheInput, WorkerCacheInputs, WorkerCandidate, WorkerCandidates, WorkerInputView, WorkerInputs, WorkerLoadInput, WorkerTransferInput, diff --git a/lib/kv-router/src/plugins/worker_selection/inputs.rs b/lib/kv-router/src/plugins/worker_selection/inputs.rs index aace9783c6f1..b0f64a2ac197 100644 --- a/lib/kv-router/src/plugins/worker_selection/inputs.rs +++ b/lib/kv-router/src/plugins/worker_selection/inputs.rs @@ -183,6 +183,66 @@ pub(crate) struct CacheSnapshot<'a> { pub(crate) has_tier_matches: bool, } +/// Rows and request-wide cache facts owned by a reader outside worker selection. +pub(crate) struct WorkerInputRows { + inputs: WorkerInputs, + rows: Vec, + shared_hits: Option, + has_tier_matches: bool, +} + +impl WorkerInputRows { + pub(crate) fn new( + inputs: WorkerInputs, + rows: Vec, + cache_snapshot: CacheSnapshot<'_>, + ) -> Self { + Self { + inputs, + rows, + shared_hits: if inputs.contains(WorkerInputs::CACHE) { + cache_snapshot.shared_hits.cloned() + } else { + None + }, + has_tier_matches: cache_snapshot.has_tier_matches, + } + } + + /// Rows for a reader that declared no inputs. + pub(crate) fn empty() -> Self { + Self { + inputs: WorkerInputs::NONE, + rows: Vec::new(), + shared_hits: None, + has_tier_matches: false, + } + } + + /// Lend the rows as a view restricted to the declared inputs. + pub(crate) fn read(&self, read: impl FnOnce(WorkerCandidates<'_>) -> R) -> R { + let cache_snapshot = CacheSnapshot { + shared_hits: self.shared_hits.as_ref(), + has_tier_matches: self.has_tier_matches, + }; + read(WorkerCandidates::new( + &self.rows, + self.inputs, + &cache_snapshot, + )) + } +} + +impl std::fmt::Debug for WorkerInputRows { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("WorkerInputRows") + .field("inputs", &self.inputs) + .field("rows", &self.rows.len()) + .finish_non_exhaustive() + } +} + /// Numeric cache row retained in host buffers between selections. #[derive(Clone, Copy, Default)] pub(crate) struct WorkerCacheData { diff --git a/lib/kv-router/src/scheduling/local.rs b/lib/kv-router/src/scheduling/local.rs index 42285b0160fe..3ea3096e6249 100644 --- a/lib/kv-router/src/scheduling/local.rs +++ b/lib/kv-router/src/scheduling/local.rs @@ -22,7 +22,7 @@ use super::queue::{ SchedulerQueue, }; use super::request_classifier::{RequestClassifierRuntime, RequestLifecycle}; -use super::selector::WorkerSelector; +use super::selector::{WorkerInputs, WorkerSelector}; use super::types::{ AdmissionAttempt, AdmittedSchedulingResponse, AdvisorySchedulingResponse, AttemptId, KvSchedulerError, NonMaxOverlapSelectionObserver, OverloadedWorkerProvider, PotentialLoad, @@ -373,9 +373,9 @@ where ) -> Result<(AdmittedSchedulingResponse, Option), KvSchedulerError> { let (resp_tx, resp_rx) = tokio::sync::oneshot::channel(); let (attempt_tx, attempt_rx) = tokio::sync::oneshot::channel(); - let (request, block_hashes) = self.make_scheduling_request(request, Some(resp_tx)); + let (mut request, block_hashes) = self.make_scheduling_request(request, Some(resp_tx)); let tracked = request.mode.is_tracked(); - let classified_request = self.classify_request(&request, ingress_at).await?; + let classified_request = self.classify_request(&mut request, ingress_at).await?; let lifecycle_lease = self .queue .new_request_lifecycle_lease(request.mode.lifecycle_request_id()); @@ -414,7 +414,7 @@ where async fn classify_request( &self, - request: &SchedulingRequest, + request: &mut SchedulingRequest, ingress_at: Instant, ) -> Result, KvSchedulerError> { let Some(classifier) = self.request_classifier.get() else { @@ -427,11 +427,31 @@ where // `best_worker`, `RouterRequest::New`) never emits lifecycle events, so // classifying it would corrupt plugin bookkeeping. It uses the default // queue inputs, exactly as if no classifier were installed. - if !classifier.has_request(request_id) { + let Some(overrides_cached) = classifier.cached_overrides(request_id) else { return Ok(None); + }; + // A retry whose lifecycle already carries overrides returns from `classify_with` + // before the plugin runs, so it gets no load projection and no rows. + let worker_inputs = if overrides_cached { + WorkerInputs::NONE + } else { + classifier.worker_inputs() + }; + if worker_inputs.contains(WorkerInputs::LOAD) { + self.slots.project_worker_loads_into( + request.token_seq.as_deref(), + Instant::now(), + &mut request.worker_loads, + ); } + let (classification, workers) = + self.queue + .build_classify_request(request, ingress_at, worker_inputs); + // The rows own their copy. Free the projection so a queued request does not retain + // one entry per worker; the actor projects load again before selection. + request.worker_loads = FxHashMap::default(); classifier - .classify_with(self.queue.build_classify_request(request, ingress_at)) + .classify_with(classification, workers) .await .map(Some) } @@ -938,6 +958,7 @@ mod tests { use tokio::sync::{mpsc, watch}; use super::*; + use crate::plugins::worker_selection::WorkerCandidates; use crate::protocols::{ActiveSequenceEvent, ActiveSequenceEventData}; use crate::scheduling::PrefillLoadEstimator; use crate::scheduling::request_classifier::ClassifyFuture; @@ -1343,8 +1364,13 @@ mod tests { } impl RequestClassifier for CountingClassifier { - fn classify(&mut self, request: ClassifyRequest) -> ClassifyFuture { + fn classify( + &mut self, + request: ClassifyRequest, + workers: WorkerCandidates<'_>, + ) -> ClassifyFuture { self.calls.fetch_add(1, Ordering::Relaxed); + assert!(workers.is_empty()); Box::pin(async move { Ok(request) }) } } @@ -1432,7 +1458,11 @@ mod tests { async fn booking_admission_applies_classifier_deadline_from_original_ingress() { struct DeadlineClassifier(Instant); impl RequestClassifier for DeadlineClassifier { - fn classify(&mut self, mut request: ClassifyRequest) -> ClassifyFuture { + fn classify( + &mut self, + mut request: ClassifyRequest, + _workers: WorkerCandidates<'_>, + ) -> ClassifyFuture { assert_eq!(request.ingress_at(), self.0); request.set_due_at(self.0 + Duration::from_secs(1)); Box::pin(async move { Ok(request) }) @@ -1468,6 +1498,91 @@ mod tests { cancel_token.cancel(); } + #[tokio::test] + async fn classifier_reads_declared_worker_inputs() { + type Row = (WorkerWithDpRank, f64, bool, usize); + struct InputsClassifier(Arc>>>); + impl RequestClassifier for InputsClassifier { + fn required_worker_inputs(&self) -> WorkerInputs { + WorkerInputs::CACHE | WorkerInputs::LOAD + } + + fn classify( + &mut self, + request: ClassifyRequest, + workers: WorkerCandidates<'_>, + ) -> ClassifyFuture { + *self.0.lock() = Some( + workers + .iter() + .map(|candidate| { + let load = candidate.load().expect("LOAD was declared"); + ( + candidate.worker(), + candidate + .cache() + .expect("CACHE was declared") + .device_overlap_blocks(), + load.is_available(), + load.active_requests(), + ) + }) + .collect(), + ); + Box::pin(async move { Ok(request) }) + } + } + + let cached = WorkerWithDpRank::new(0, 0); + let busy = WorkerWithDpRank::new(1, 0); + let workers = HashMap::from([ + (0, SimpleWorkerConfig::default()), + (1, SimpleWorkerConfig::default()), + ]); + let (scheduler, _slots, _cfg_tx, cancel_token) = make_scheduler(workers, None, true, None); + let rows = Arc::new(parking_lot::Mutex::new(None)); + assert!(scheduler.install_request_classifier( + Box::new(InputsClassifier(Arc::clone(&rows))), + cancel_token.clone(), + )); + let _busy = scheduler + .add_request_if_registered_guarded(SequenceRequest { + request_id: "busy".into(), + token_sequence: None, + track_prefill_tokens: true, + expected_output_tokens: None, + prefill_load_hint: None, + worker: busy, + lora_name: None, + }) + .unwrap(); + let _lifecycle = scheduler + .begin_request_lifecycle("classified") + .unwrap() + .unwrap(); + let mut classified = request(ScheduleMode::TrackedWithLifecycle { + request_id: "classified".to_string(), + }); + classified + .overlap + .tier_overlap_blocks + .device + .insert(cached, 2); + + scheduler + .schedule_request_with_booking(classified) + .await + .unwrap(); + + let mut rows = rows + .lock() + .take() + .expect("declared inputs must reach the classifier"); + rows.sort_by_key(|row| row.0); + assert_eq!(rows, vec![(cached, 2.0, true, 0), (busy, 0.0, true, 1)]); + cancel_token.cancel(); + } + #[tokio::test] async fn legacy_tracked_request_without_id_fails_before_enqueue() { let workers = HashMap::from([(0, SimpleWorkerConfig::default())]); diff --git a/lib/kv-router/src/scheduling/queue.rs b/lib/kv-router/src/scheduling/queue.rs index 0a5af3c3fc69..db7738db692a 100644 --- a/lib/kv-router/src/scheduling/queue.rs +++ b/lib/kv-router/src/scheduling/queue.rs @@ -26,12 +26,13 @@ use super::policy_queue::{PolicyQueue, QueueMetadata, QueueSnapshot}; use super::prefill_load::{PrefillLoadEstimator, effective_prefill_tokens}; use super::queue_admission::WorkerPlacement; use super::request_classifier::{ClassificationOverrides, ClassifyRequest}; -use super::selector::{WorkerSelectionInput, WorkerSelector}; +use super::selector::{WorkerInputs, WorkerSelectionInput, WorkerSelector, worker_input_rows}; use super::types::{ AdvisorySchedulingResponse, AdvisoryWorkerLoad, AttemptId, KvSchedulerError, NonMaxOverlapSelection, NonMaxOverlapSelectionObserver, OverloadedWorkerProvider, SchedulingContext, SchedulingRequest, SchedulingResponse, WorkerAvailabilityProvider, }; +use crate::plugins::worker_selection::WorkerInputRows; use crate::plugins::worker_selection::{ RuntimePolicyConfig, RuntimePolicyError, RuntimePolicyUpdate, }; @@ -848,15 +849,16 @@ impl< } } - /// Build the classifier's input. Uses `request.isl_tokens` — the - /// routing-token basis the queue's bucketing, limits, and DRR cost already - /// use — so a pass-through classifier is behavior-identical to no + /// Build the classifier's input and the rows for its declared worker inputs. Uses + /// `request.isl_tokens` — the routing-token basis the queue's bucketing, limits, and DRR + /// cost already use — so a pass-through classifier is behavior-identical to no /// classifier. pub(crate) fn build_classify_request( &self, request: &SchedulingRequest, ingress_at: Instant, - ) -> ClassifyRequest { + worker_inputs: WorkerInputs, + ) -> (ClassifyRequest, WorkerInputRows) { let available = self .available_worker_provider .as_ref() @@ -878,7 +880,15 @@ impl< if let Some(session_context) = request.session_context.clone() { classification = classification.with_session_context(session_context); } - classification + let rows = if worker_inputs == WorkerInputs::NONE { + WorkerInputRows::empty() + } else { + let eligibility = request + .eligibility() + .with_available_workers(available.as_deref()); + worker_input_rows(&workers, request, eligibility, worker_inputs) + }; + (classification, rows) } /// Enqueue a request with `due_at` for actor-expiry tests. @@ -2052,6 +2062,7 @@ mod tests { use super::*; use crate::kv_hints::KvTransferCandidates; + use crate::plugins::worker_selection::WorkerCandidates; use crate::protocols::{ ActiveSequenceEvent, ExternalSequenceBlockHash, WorkerSelectionResult, WorkerWithDpRank, }; @@ -2742,7 +2753,11 @@ mod tests { struct InspectHashes(Option>, WorkerWithDpRank); impl RequestClassifier for InspectHashes { - fn classify(&mut self, request: ClassifyRequest) -> ClassifyFuture { + fn classify( + &mut self, + request: ClassifyRequest, + _workers: WorkerCandidates<'_>, + ) -> ClassifyFuture { let expected = self.0.clone(); let pin = self.1; Box::pin(async move { @@ -2761,10 +2776,11 @@ mod tests { let (mut request, _rx) = make_request("prefix-input", 193); request.token_seq = hashes.clone(); request.pinned_worker = Some(pin); - let classified = queue.build_classify_request(&request, Instant::now()); + let (classified, workers) = + queue.build_classify_request(&request, Instant::now(), WorkerInputs::NONE); assert_eq!(request.token_seq, hashes); let mut classifier = InspectHashes(hashes, pin); - let pending = classifier.classify(classified); + let pending = workers.read(|workers| classifier.classify(classified, workers)); // The independently pollable future must not borrow actor-owned data. request.token_seq = Some(vec![99]); request.pinned_worker = None; @@ -2786,7 +2802,11 @@ mod tests { struct Park; impl RequestClassifier for Park { - fn classify(&mut self, request: ClassifyRequest) -> ClassifyFuture { + fn classify( + &mut self, + request: ClassifyRequest, + _workers: WorkerCandidates<'_>, + ) -> ClassifyFuture { Box::pin(async move { std::future::pending::<()>().await; Ok(black_box(request)) @@ -2829,8 +2849,12 @@ mod tests { for _ in 0..repeats { let start = WallClock::now(); for (request, _) in &sources { - let input = queue.build_classify_request(black_box(request), ingress_at); - let mut future = classifier.classify(input); + let (input, workers) = queue.build_classify_request( + black_box(request), + ingress_at, + WorkerInputs::NONE, + ); + let mut future = workers.read(|workers| classifier.classify(input, workers)); assert!(future.as_mut().poll(&mut cx).is_pending()); futures.push(future); } @@ -2891,7 +2915,9 @@ policy_classes: request.policy_class = Some("latency".to_string()); let ingress_at = Instant::now(); let due_at = ingress_at + Duration::from_secs(20); - let mut classified = queue.build_classify_request(&request, ingress_at); + let mut classified = queue + .build_classify_request(&request, ingress_at, WorkerInputs::NONE) + .0; assert_eq!(classified.ingress_at(), ingress_at); classified.set_policy_class("bulk"); classified.set_due_at(due_at); @@ -2907,21 +2933,27 @@ policy_classes: assert_eq!(metadata.snapshot.scheduling_cost_tokens, 7); assert_eq!(metadata.due_at, Some(due_at)); - let mut invalid = queue.build_classify_request(&request, ingress_at); + let mut invalid = queue + .build_classify_request(&request, ingress_at, WorkerInputs::NONE) + .0; invalid.set_policy_class("missing"); assert!(matches!( queue.validate_classification(&mut request, invalid, ingress_at), Err(KvSchedulerError::InvalidClassificationMetadata(_)) )); - let mut physical = queue.build_classify_request(&request, ingress_at); + let mut physical = queue + .build_classify_request(&request, ingress_at, WorkerInputs::NONE) + .0; physical.set_policy_class("latency_cached"); assert!(matches!( queue.validate_classification(&mut request, physical, ingress_at), Err(KvSchedulerError::InvalidClassificationMetadata(_)) )); - let mut zero_cost = queue.build_classify_request(&request, ingress_at); + let mut zero_cost = queue + .build_classify_request(&request, ingress_at, WorkerInputs::NONE) + .0; zero_cost.set_scheduling_cost_tokens(0); assert!(matches!( queue.validate_classification(&mut request, zero_cost, ingress_at), @@ -2937,14 +2969,18 @@ policy_classes: let pin = WorkerWithDpRank::new(0, 0); let now = Instant::now(); - let mut classified = queue.build_classify_request(&request, now); + let mut classified = queue + .build_classify_request(&request, now, WorkerInputs::NONE) + .0; classified.set_worker_selection_target(worker); queue .validate_classification(&mut request, classified, now) .unwrap(); assert_eq!(request.affinity_target, Some(worker.into())); - let mut classified = queue.build_classify_request(&request, now); + let mut classified = queue + .build_classify_request(&request, now, WorkerInputs::NONE) + .0; classified.clear_worker_selection_target(); queue .validate_classification(&mut request, classified, now) @@ -2952,7 +2988,9 @@ policy_classes: assert!(request.affinity_target.is_none()); request.pinned_worker = Some(pin); - let mut classified = queue.build_classify_request(&request, now); + let mut classified = queue + .build_classify_request(&request, now, WorkerInputs::NONE) + .0; assert_eq!(classified.pinned_worker(), Some(pin)); classified.set_worker_selection_target(worker); queue @@ -2962,7 +3000,9 @@ policy_classes: assert!(request.affinity_target.is_none()); request.pinned_worker = None; request.allowed_worker_ids = Some(HashSet::from([0])); - let mut classified = queue.build_classify_request(&request, now); + let mut classified = queue + .build_classify_request(&request, now, WorkerInputs::NONE) + .0; classified.set_worker_selection_target(worker); queue .validate_classification(&mut request, classified, now) @@ -3002,7 +3042,9 @@ policy_classes: .insert(WorkerWithDpRank::new(0, 0), 32); let ingress_at = Instant::now(); - let classified = queue.build_classify_request(&request, ingress_at); + let classified = queue + .build_classify_request(&request, ingress_at, WorkerInputs::NONE) + .0; // The only cache-bearing worker disappears while the classification is // pending; enqueue-time state governs, so the request buckets as @@ -3142,7 +3184,9 @@ policy_classes: let mut receivers = Vec::new(); for (name, hash) in [("stalled", 41), ("swept", 42)] { let (queued, queued_rx) = make_request(name, isl); - let mut classified = queue.build_classify_request(&queued, ingress_at); + let mut classified = queue + .build_classify_request(&queued, ingress_at, WorkerInputs::NONE) + .0; classified.set_due_at(due_at); queue .enqueue_admitted_with_block_hashes_and_lease( @@ -4662,7 +4706,8 @@ policy_classes: ); assert_eq!( queue - .build_classify_request(&request, Instant::now()) + .build_classify_request(&request, Instant::now(), WorkerInputs::NONE) + .0 .scheduling_cost_tokens(), 224 ); @@ -4677,7 +4722,8 @@ policy_classes: ); assert_eq!( queue - .build_classify_request(&request, Instant::now()) + .build_classify_request(&request, Instant::now(), WorkerInputs::NONE) + .0 .scheduling_cost_tokens(), 256 ); diff --git a/lib/kv-router/src/scheduling/request_classifier.rs b/lib/kv-router/src/scheduling/request_classifier.rs index 891d99713177..05b817e8c816 100644 --- a/lib/kv-router/src/scheduling/request_classifier.rs +++ b/lib/kv-router/src/scheduling/request_classifier.rs @@ -16,6 +16,7 @@ use tokio::task::JoinHandle; use tokio_util::sync::CancellationToken; use super::types::KvSchedulerError; +use crate::plugins::worker_selection::{WorkerInputRows, WorkerInputs}; use crate::protocols::WorkerWithDpRank; static NEXT_CLASSIFICATION_ID: AtomicU64 = AtomicU64::new(1); @@ -81,6 +82,8 @@ pub(crate) struct RequestClassifierRuntime { // keeps serving: one request's failure must not disable classification // router-wide, and the plugin owns its own state across an unwind. classifier: Arc>>, + // Read once at install so classification never takes the classifier lock to ask. + worker_inputs: WorkerInputs, live_requests: Mutex>, // Shared with the delivery task, which releases each terminal event once // `on_event` has returned for it. @@ -105,6 +108,7 @@ impl RequestClassifierRuntime { shutdown: CancellationToken, ) -> Arc { let (events, mut receiver) = mpsc::unbounded_channel(); + let worker_inputs = classifier.required_worker_inputs(); let classifier = Arc::new(AsyncMutex::new(classifier)); let pending_terminals = Arc::new(PendingTerminals::default()); let delivery_classifier = Arc::clone(&classifier); @@ -151,6 +155,7 @@ impl RequestClassifierRuntime { }); Arc::new(Self { classifier, + worker_inputs, live_requests: Mutex::new(HashMap::new()), pending_terminals, events, @@ -159,9 +164,16 @@ impl RequestClassifierRuntime { }) } + pub(crate) fn worker_inputs(&self) -> WorkerInputs { + self.worker_inputs + } + + /// `workers` holds the rows for the classifier's declared inputs. They are lent to + /// `classify` and freed when it returns, before the classification future is awaited. pub(crate) async fn classify_with( &self, mut request: ClassifyRequest, + workers: WorkerInputRows, ) -> Result { if self.shutdown.is_cancelled() { return Err(KvSchedulerError::SubscriberShutdown); @@ -205,9 +217,12 @@ impl RequestClassifierRuntime { request_id.to_owned(), )); } - catch_unwind(AssertUnwindSafe(|| classifier.classify(request))).map_err(|panic| { - KvSchedulerError::RequestClassifierPanicked(panic_message(panic)) - })? + let classification = catch_unwind(AssertUnwindSafe(|| { + workers.read(|workers| classifier.classify(request, workers)) + })) + .map_err(|panic| KvSchedulerError::RequestClassifierPanicked(panic_message(panic)))?; + drop(workers); + classification }; let classification = AssertUnwindSafe(classification).catch_unwind(); @@ -279,8 +294,14 @@ impl RequestClassifierRuntime { }) } - pub(crate) fn has_request(&self, request_id: &str) -> bool { - self.live_requests.lock().contains_key(request_id) + /// Whether `request_id` has a live lifecycle and, if so, whether that lifecycle already + /// carries classification overrides. `Some(true)` means `classify_with` will return + /// them without running the plugin, so the caller can skip building its inputs. + pub(crate) fn cached_overrides(&self, request_id: &str) -> Option { + self.live_requests + .lock() + .get(request_id) + .map(|live| live.overrides.is_some()) } /// Take the classifier lock for a `classify` call. A reused id must not @@ -524,6 +545,7 @@ mod tests { use tokio_util::sync::CancellationToken; use super::*; + use crate::plugins::worker_selection::WorkerCandidates; use crate::protocols::WorkerWithDpRank; use crate::scheduling::KvSchedulerError; @@ -536,7 +558,11 @@ mod tests { } impl RequestClassifier for SynchronousPanicOnce { - fn classify(&mut self, request: ClassifyRequest) -> ClassifyFuture { + fn classify( + &mut self, + request: ClassifyRequest, + _workers: WorkerCandidates<'_>, + ) -> ClassifyFuture { if self.calls.fetch_add(1, Ordering::Relaxed) == 0 { panic!("synchronous classifier panic"); } @@ -555,7 +581,7 @@ mod tests { ); let error = runtime - .classify_with(ClassifyRequest::new(1, 0)) + .classify_with(ClassifyRequest::new(1, 0), WorkerInputRows::empty()) .await .unwrap_err(); assert!(matches!( @@ -564,7 +590,7 @@ mod tests { if message == "synchronous classifier panic" )); runtime - .classify_with(ClassifyRequest::new(2, 0)) + .classify_with(ClassifyRequest::new(2, 0), WorkerInputRows::empty()) .await .unwrap(); assert_eq!(calls.load(Ordering::Relaxed), 2); @@ -575,7 +601,11 @@ mod tests { } impl RequestClassifier for FuturePanicOnce { - fn classify(&mut self, request: ClassifyRequest) -> ClassifyFuture { + fn classify( + &mut self, + request: ClassifyRequest, + _workers: WorkerCandidates<'_>, + ) -> ClassifyFuture { let should_panic = self.calls.fetch_add(1, Ordering::Relaxed) == 0; Box::pin(async move { assert!(!should_panic, "classifier future panic"); @@ -595,7 +625,7 @@ mod tests { ); let error = runtime - .classify_with(ClassifyRequest::new(1, 0)) + .classify_with(ClassifyRequest::new(1, 0), WorkerInputRows::empty()) .await .unwrap_err(); assert!(matches!( @@ -604,7 +634,7 @@ mod tests { if message == "classifier future panic" )); runtime - .classify_with(ClassifyRequest::new(2, 0)) + .classify_with(ClassifyRequest::new(2, 0), WorkerInputRows::empty()) .await .unwrap(); assert_eq!(calls.load(Ordering::Relaxed), 2); @@ -617,7 +647,11 @@ mod tests { struct FailingClassifier; impl RequestClassifier for FailingClassifier { - fn classify(&mut self, _request: ClassifyRequest) -> ClassifyFuture { + fn classify( + &mut self, + _request: ClassifyRequest, + _workers: WorkerCandidates<'_>, + ) -> ClassifyFuture { Box::pin(async { Err(Box::new(TestClassifierError) as Box) }) } } @@ -628,7 +662,7 @@ mod tests { RequestClassifierRuntime::new(Box::new(FailingClassifier), CancellationToken::new()); let error = runtime - .classify_with(ClassifyRequest::new(1, 0)) + .classify_with(ClassifyRequest::new(1, 0), WorkerInputRows::empty()) .await .unwrap_err(); assert!(matches!( @@ -641,7 +675,11 @@ mod tests { struct ReplacingClassifier; impl RequestClassifier for ReplacingClassifier { - fn classify(&mut self, _request: ClassifyRequest) -> ClassifyFuture { + fn classify( + &mut self, + _request: ClassifyRequest, + _workers: WorkerCandidates<'_>, + ) -> ClassifyFuture { Box::pin(async { Ok(ClassifyRequest::new(2, 0)) }) } } @@ -652,7 +690,7 @@ mod tests { RequestClassifierRuntime::new(Box::new(ReplacingClassifier), CancellationToken::new()); let error = runtime - .classify_with(ClassifyRequest::new(1, 0)) + .classify_with(ClassifyRequest::new(1, 0), WorkerInputRows::empty()) .await .unwrap_err(); assert!(matches!( @@ -685,7 +723,10 @@ mod tests { .with_request_id("request-1") .with_initial_policy_class("latency"); - let result = PassThrough.classify(request).await.unwrap(); + let result = WorkerInputRows::empty() + .read(|workers| PassThrough.classify(request, workers)) + .await + .unwrap(); assert_eq!(result.request_id(), Some("request-1")); assert_eq!(result.policy_class(), Some("latency")); @@ -704,7 +745,11 @@ mod tests { #[async_trait] impl RequestClassifier for EventReleasedClassifier { - fn classify(&mut self, request: ClassifyRequest) -> ClassifyFuture { + fn classify( + &mut self, + request: ClassifyRequest, + _workers: WorkerCandidates<'_>, + ) -> ClassifyFuture { let entered = Arc::clone(&self.entered); let released = Arc::clone(&self.released); Box::pin(async move { @@ -736,7 +781,7 @@ mod tests { let pending_runtime = Arc::clone(&runtime); let pending = tokio::spawn(async move { pending_runtime - .classify_with(ClassifyRequest::new(1, 1)) + .classify_with(ClassifyRequest::new(1, 1), WorkerInputRows::empty()) .await }); entered.notified().await; @@ -753,7 +798,11 @@ mod tests { #[async_trait] impl RequestClassifier for PendingClassifier { - fn classify(&mut self, _request: ClassifyRequest) -> ClassifyFuture { + fn classify( + &mut self, + _request: ClassifyRequest, + _workers: WorkerCandidates<'_>, + ) -> ClassifyFuture { self.entered.fetch_add(1, Ordering::Relaxed); Box::pin(std::future::pending()) } @@ -781,7 +830,10 @@ mod tests { let cancelled = tokio::spawn(async move { let _lifecycle = cancelled_runtime.begin_request("cancelled").unwrap(); cancelled_runtime - .classify_with(ClassifyRequest::new(1, 0).with_request_id("cancelled")) + .classify_with( + ClassifyRequest::new(1, 0).with_request_id("cancelled"), + WorkerInputRows::empty(), + ) .await }); while entered.load(Ordering::Relaxed) < 1 { @@ -818,7 +870,10 @@ mod tests { ); let lifecycle = runtime.begin_request("off-runtime-drop").unwrap(); runtime - .classify_with(ClassifyRequest::new(1, 0).with_request_id("off-runtime-drop")) + .classify_with( + ClassifyRequest::new(1, 0).with_request_id("off-runtime-drop"), + WorkerInputRows::empty(), + ) .await .unwrap(); @@ -831,7 +886,11 @@ mod tests { } impl RequestClassifier for CountingClassifier { - fn classify(&mut self, mut request: ClassifyRequest) -> ClassifyFuture { + fn classify( + &mut self, + mut request: ClassifyRequest, + _workers: WorkerCandidates<'_>, + ) -> ClassifyFuture { self.calls.fetch_add(1, Ordering::Relaxed); request.set_scheduling_cost_tokens(7); request.set_worker_selection_target(WorkerWithDpRank::new(9, 1)); @@ -839,6 +898,41 @@ mod tests { } } + #[tokio::test] + async fn cached_overrides_follow_the_lifecycle() { + let runtime = RequestClassifierRuntime::new( + Box::new(CountingClassifier { + calls: Arc::new(AtomicUsize::new(0)), + }), + CancellationToken::new(), + ); + assert_eq!(runtime.cached_overrides("request-1"), None); + + let mut lifecycle = runtime.begin_request("request-1").unwrap(); + assert_eq!(runtime.cached_overrides("request-1"), Some(false)); + + runtime + .classify_with( + ClassifyRequest::new(10, 10).with_request_id("request-1"), + WorkerInputRows::empty(), + ) + .await + .unwrap(); + assert_eq!(runtime.cached_overrides("request-1"), Some(true)); + lifecycle.prepare_retry(); + assert_eq!( + runtime.cached_overrides("request-1"), + Some(true), + "a retry reuses the overrides, so its inputs need not be built" + ); + + lifecycle.abort(None); + assert_eq!(runtime.cached_overrides("request-1"), None); + let next_lifecycle = runtime.begin_request("request-1").unwrap(); + assert_eq!(runtime.cached_overrides("request-1"), Some(false)); + drop(next_lifecycle); + } + #[tokio::test] async fn retry_reuses_classification_overrides_without_reinvoking_plugin() { let calls = Arc::new(AtomicUsize::new(0)); @@ -851,7 +945,10 @@ mod tests { let mut lifecycle = runtime.begin_request("request-1").unwrap(); let first = runtime - .classify_with(ClassifyRequest::new(10, 10).with_request_id("request-1")) + .classify_with( + ClassifyRequest::new(10, 10).with_request_id("request-1"), + WorkerInputRows::empty(), + ) .await .unwrap(); assert_eq!(first.progress().context_tokens(), 10); @@ -860,7 +957,10 @@ mod tests { assert_eq!(first.progress().context_tokens(), 25); lifecycle.prepare_retry(); let retry = runtime - .classify_with(ClassifyRequest::new(20, 20).with_request_id("request-1")) + .classify_with( + ClassifyRequest::new(20, 20).with_request_id("request-1"), + WorkerInputRows::empty(), + ) .await .unwrap(); @@ -881,7 +981,10 @@ mod tests { let mut next_lifecycle = runtime.begin_request("request-1").unwrap(); let next = runtime - .classify_with(ClassifyRequest::new(3, 0).with_request_id("request-1")) + .classify_with( + ClassifyRequest::new(3, 0).with_request_id("request-1"), + WorkerInputRows::empty(), + ) .await .unwrap(); next_lifecycle.observe_context_tokens(3); @@ -908,7 +1011,10 @@ mod tests { ); let result = runtime - .classify_with(ClassifyRequest::new(4, 0).with_request_id("ghost")) + .classify_with( + ClassifyRequest::new(4, 0).with_request_id("ghost"), + WorkerInputRows::empty(), + ) .await; assert!(matches!( @@ -925,7 +1031,11 @@ mod tests { } impl RequestClassifier for GatedClassifier { - fn classify(&mut self, mut request: ClassifyRequest) -> ClassifyFuture { + fn classify( + &mut self, + mut request: ClassifyRequest, + _workers: WorkerCandidates<'_>, + ) -> ClassifyFuture { let call = self.calls.fetch_add(1, Ordering::Relaxed) + 1; let entered = Arc::clone(&self.entered); let release = Arc::clone(&self.release); @@ -956,7 +1066,10 @@ mod tests { let stale_runtime = Arc::clone(&runtime); let stale = tokio::spawn(async move { stale_runtime - .classify_with(ClassifyRequest::new(1, 0).with_request_id("reused")) + .classify_with( + ClassifyRequest::new(1, 0).with_request_id("reused"), + WorkerInputRows::empty(), + ) .await }); entered.notified().await; @@ -976,7 +1089,10 @@ mod tests { // classification reaches the plugin instead of reusing overrides. release.notify_one(); let second = runtime - .classify_with(ClassifyRequest::new(1, 0).with_request_id("reused")) + .classify_with( + ClassifyRequest::new(1, 0).with_request_id("reused"), + WorkerInputRows::empty(), + ) .await .unwrap(); assert_eq!(calls.load(Ordering::Relaxed), 2); @@ -994,8 +1110,10 @@ mod tests { ); let lifecycle = runtime.begin_request("reused").unwrap(); let guard = runtime.classifier.lock().await; - let mut stale = - Box::pin(runtime.classify_with(ClassifyRequest::new(1, 0).with_request_id("reused"))); + let mut stale = Box::pin(runtime.classify_with( + ClassifyRequest::new(1, 0).with_request_id("reused"), + WorkerInputRows::empty(), + )); // Poll once to guarantee the old call is waiting for the lock, then // replace its lifecycle before the plugin can be invoked. @@ -1012,7 +1130,10 @@ mod tests { // The replacement must reach the plugin, without stale overrides. let replacement = runtime - .classify_with(ClassifyRequest::new(100, 0).with_request_id("reused")) + .classify_with( + ClassifyRequest::new(100, 0).with_request_id("reused"), + WorkerInputRows::empty(), + ) .await .unwrap(); assert_eq!(replacement.input_tokens(), 100); @@ -1037,7 +1158,10 @@ mod tests { let ended_runtime = Arc::clone(&runtime); let ended = tokio::spawn(async move { ended_runtime - .classify_with(ClassifyRequest::new(1, 0).with_request_id("ended")) + .classify_with( + ClassifyRequest::new(1, 0).with_request_id("ended"), + WorkerInputRows::empty(), + ) .await }); tokio::task::yield_now().await; @@ -1067,7 +1191,10 @@ mod tests { let ended_runtime = Arc::clone(&runtime); let ended = tokio::spawn(async move { ended_runtime - .classify_with(ClassifyRequest::new(1, 0).with_request_id("ended")) + .classify_with( + ClassifyRequest::new(1, 0).with_request_id("ended"), + WorkerInputRows::empty(), + ) .await }); entered.notified().await; @@ -1081,7 +1208,7 @@ mod tests { ended.await.unwrap(), Err(KvSchedulerError::ClassificationLifecycleEnded(_)) )); - assert!(!runtime.has_request("ended")); + assert_eq!(runtime.cached_overrides("ended"), None); } #[derive(Debug, PartialEq, Eq)] @@ -1128,7 +1255,10 @@ mod tests { let worker = WorkerWithDpRank::new(7, 2); let mut lifecycle = runtime.begin_request("request-1").unwrap(); runtime - .classify_with(ClassifyRequest::new(40, 0).with_request_id("request-1")) + .classify_with( + ClassifyRequest::new(40, 0).with_request_id("request-1"), + WorkerInputRows::empty(), + ) .await .unwrap(); @@ -1219,7 +1349,11 @@ mod tests { #[async_trait] impl RequestClassifier for OrderingClassifier { - fn classify(&mut self, request: ClassifyRequest) -> ClassifyFuture { + fn classify( + &mut self, + request: ClassifyRequest, + _workers: WorkerCandidates<'_>, + ) -> ClassifyFuture { let _ = self .observed .send(format!("classify:{}", request.request_id().unwrap_or(""))); @@ -1274,7 +1408,10 @@ mod tests { let classify_runtime = Arc::clone(&runtime); let classify = tokio::spawn(async move { classify_runtime - .classify_with(ClassifyRequest::new(1, 0).with_request_id("reused")) + .classify_with( + ClassifyRequest::new(1, 0).with_request_id("reused"), + WorkerInputRows::empty(), + ) .await }); tokio::task::yield_now().await; diff --git a/lib/kv-router/src/scheduling/selector/mod.rs b/lib/kv-router/src/scheduling/selector/mod.rs index 56b18f908932..a607a6283ae0 100644 --- a/lib/kv-router/src/scheduling/selector/mod.rs +++ b/lib/kv-router/src/scheduling/selector/mod.rs @@ -26,7 +26,7 @@ use crate::kv_hints::{ KvTransferScope, KvTransferSourceChoice, KvTransferSourceSummary, KvTransferTarget, }; use crate::plugins::worker_selection::{ - CacheSnapshot, CandidateData, WorkerCacheData, WorkerTransferInput, + CacheSnapshot, CandidateData, WorkerCacheData, WorkerInputRows, WorkerTransferInput, }; use crate::plugins::worker_selection::{ RuntimePolicyConfig, RuntimePolicyError, RuntimePolicyUpdate, @@ -182,12 +182,7 @@ impl<'a> MaterializedSelectionInput<'a> { fn new(request: &'a SchedulingRequest, block_size: u32) -> Self { Self { request, - cache_snapshot: CacheSnapshot { - shared_hits: request.shared_cache_hits.as_ref(), - has_tier_matches: !request.overlap.tier_overlap_blocks.device.is_empty() - || !request.overlap.tier_overlap_blocks.host_pinned.is_empty() - || !request.overlap.tier_overlap_blocks.disk.is_empty(), - }, + cache_snapshot: cache_snapshot(request), context: WorkerSelectionContext { request, request_blocks: request.request_blocks(block_size), @@ -296,75 +291,131 @@ impl<'a> MaterializedSelectionInput<'a> { inputs: WorkerInputs, select_device_overlap: impl FnOnce(f64, f64) -> f64, ) -> CandidateData { - let cached_tokens = if inputs.contains(WorkerInputs::CACHE) { - self.request.effective_cached_tokens_for(worker) - } else { - 0 - }; - let worker_load = if inputs.contains(WorkerInputs::LOAD) { - self.request.worker_loads.get(&worker).copied() - } else { - None - }; - let cache = if inputs.contains(WorkerInputs::CACHE) { - let effective_overlap_blocks = self.request.effective_overlap_blocks_for(worker); - let reported_device_overlap_blocks = self - .request + worker_input_row( + self.request, + worker, + preferred_taint_multiplier, + inputs, + select_device_overlap, + ) + } +} + +/// Request-wide cache facts shared by every row built from `request`. +fn cache_snapshot(request: &SchedulingRequest) -> CacheSnapshot<'_> { + CacheSnapshot { + shared_hits: request.shared_cache_hits.as_ref(), + has_tier_matches: !request.overlap.tier_overlap_blocks.device.is_empty() + || !request.overlap.tier_overlap_blocks.host_pinned.is_empty() + || !request.overlap.tier_overlap_blocks.disk.is_empty(), + } +} + +/// Build one worker's input row. Worker selection and request classification share this +/// builder, so each input group carries the same values in both. +#[inline(always)] +fn worker_input_row( + request: &SchedulingRequest, + worker: WorkerWithDpRank, + preferred_taint_multiplier: Option, + inputs: WorkerInputs, + select_device_overlap: impl FnOnce(f64, f64) -> f64, +) -> CandidateData { + let cached_tokens = if inputs.contains(WorkerInputs::CACHE) { + request.effective_cached_tokens_for(worker) + } else { + 0 + }; + let worker_load = if inputs.contains(WorkerInputs::LOAD) { + request.worker_loads.get(&worker).copied() + } else { + None + }; + let cache = if inputs.contains(WorkerInputs::CACHE) { + let effective_overlap_blocks = request.effective_overlap_blocks_for(worker); + let reported_device_overlap_blocks = request + .overlap + .tier_overlap_blocks + .device + .get(&worker) + .copied() + .map(|blocks| blocks as f64) + .unwrap_or(0.0); + let device_overlap_blocks = + select_device_overlap(effective_overlap_blocks, reported_device_overlap_blocks); + WorkerCacheData { + effective_overlap_blocks, + estimated_cached_tokens: cached_tokens, + device_overlap_blocks, + host_overlap_blocks: request .overlap .tier_overlap_blocks - .device + .host_pinned .get(&worker) .copied() - .map(|blocks| blocks as f64) - .unwrap_or(0.0); - let device_overlap_blocks = - select_device_overlap(effective_overlap_blocks, reported_device_overlap_blocks); - WorkerCacheData { - effective_overlap_blocks, - estimated_cached_tokens: cached_tokens, - device_overlap_blocks, - host_overlap_blocks: self - .request - .overlap - .tier_overlap_blocks - .host_pinned - .get(&worker) - .copied() - .unwrap_or(0) as f64, - disk_overlap_blocks: self - .request - .overlap - .tier_overlap_blocks - .disk - .get(&worker) - .copied() - .unwrap_or(0) as f64, - } - } else { - WorkerCacheData::default() - }; - let load = if inputs.contains(WorkerInputs::LOAD) { - let available = worker_load.is_some(); - let worker_load = worker_load.unwrap_or_else(|| self.request.worker_load_for(worker)); - WorkerLoadInput { - available, - active_prefill_tokens: worker_load.active_prefill_tokens, - decode_cost_blocks: worker_load.potential_decode_blocks() as f64, - active_requests: worker_load.active_requests, - } + .unwrap_or(0) as f64, + disk_overlap_blocks: request + .overlap + .tier_overlap_blocks + .disk + .get(&worker) + .copied() + .unwrap_or(0) as f64, + } + } else { + WorkerCacheData::default() + }; + let load = if inputs.contains(WorkerInputs::LOAD) { + let available = worker_load.is_some(); + let worker_load = worker_load.unwrap_or_else(|| request.worker_load_for(worker)); + WorkerLoadInput { + available, + active_prefill_tokens: worker_load.active_prefill_tokens, + decode_cost_blocks: worker_load.potential_decode_blocks() as f64, + active_requests: worker_load.active_requests, + } + } else { + WorkerLoadInput::default() + }; + + CandidateData { + worker, + inputs, + cache, + load, + transfer: None, + preferred_taint_multiplier, + } +} + +/// Build owned input rows for every worker rank that `eligibility` admits. +pub(crate) fn worker_input_rows( + workers: &HashMap, + request: &SchedulingRequest, + eligibility: RoutingEligibility<'_>, + inputs: WorkerInputs, +) -> WorkerInputRows { + let materialize_preferred_taint = + eligibility.pinned_worker().is_none() && inputs.contains(WorkerInputs::PREFERRED_TAINT); + let mut rows = Vec::new(); + eligibility.any_eligible_worker_rank(workers, |worker, config| { + let preferred_taint_multiplier = if materialize_preferred_taint { + request + .routing_constraints + .preferred_taint_multiplier(config.taints()) } else { - WorkerLoadInput::default() + None }; - - CandidateData { + rows.push(worker_input_row( + request, worker, - inputs, - cache, - load, - transfer: None, preferred_taint_multiplier, - } - } + inputs, + |_, device_overlap_blocks| device_overlap_blocks, + )); + false + }); + WorkerInputRows::new(inputs, rows, cache_snapshot(request)) } fn source_prefix_blocks(source: KvTransferSourceChoice) -> u32 { diff --git a/lib/kv-router/tests/plugin_api_compat.rs b/lib/kv-router/tests/plugin_api_compat.rs index 98d9871855eb..ad698f6ce3f9 100644 --- a/lib/kv-router/tests/plugin_api_compat.rs +++ b/lib/kv-router/tests/plugin_api_compat.rs @@ -74,7 +74,11 @@ impl WorkerPicker for LegacyPolicy { struct LegacyClassifier; impl RequestClassifier for LegacyClassifier { - fn classify(&mut self, request: ClassifyRequest) -> ClassifyFuture { + fn classify( + &mut self, + request: ClassifyRequest, + _workers: WorkerCandidates<'_>, + ) -> ClassifyFuture { let _: Option<&str> = request.request_id(); let _: usize = request.input_tokens(); let _: Option<&[u64]> = request.sequence_hashes(); diff --git a/lib/llm/src/discovery/watcher.rs b/lib/llm/src/discovery/watcher.rs index a76ef90cb01e..eeb3fb71958c 100644 --- a/lib/llm/src/discovery/watcher.rs +++ b/lib/llm/src/discovery/watcher.rs @@ -2200,6 +2200,7 @@ mod tests { crate::test_utils::assert_isolated_success(&output); return; } + use dynamo_kv_router::plugins::worker_selection::WorkerCandidates; use dynamo_kv_router::scheduling::{ ClassifyEvent, ClassifyFuture, ClassifyRequest, RequestClassifier, }; @@ -2212,7 +2213,11 @@ mod tests { } #[async_trait] impl RequestClassifier for RejectingClassifier { - fn classify(&mut self, _request: ClassifyRequest) -> ClassifyFuture { + fn classify( + &mut self, + _request: ClassifyRequest, + _workers: WorkerCandidates<'_>, + ) -> ClassifyFuture { self.calls += 1; self.events .send((self.instance, "classify", self.calls)) diff --git a/lib/llm/src/kv_router/routing_host/tests.rs b/lib/llm/src/kv_router/routing_host/tests.rs index 356f9f13faac..6d2be227fb65 100644 --- a/lib/llm/src/kv_router/routing_host/tests.rs +++ b/lib/llm/src/kv_router/routing_host/tests.rs @@ -11,6 +11,7 @@ use std::{ }; use crate::kv_router::SelectionPolicySource; +use dynamo_kv_router::plugins::worker_selection::WorkerCandidates; use dynamo_kv_router::{ config::KvRouterConfig, protocols::RoutingConstraints, @@ -1798,7 +1799,11 @@ struct RecordingClassifier { #[async_trait::async_trait] impl RequestClassifier for RecordingClassifier { - fn classify(&mut self, request: ClassifyRequest) -> ClassifyFuture { + fn classify( + &mut self, + request: ClassifyRequest, + _workers: WorkerCandidates<'_>, + ) -> ClassifyFuture { self.calls.fetch_add(1, Ordering::Relaxed); Box::pin(async move { Ok(request) }) } @@ -1831,7 +1836,11 @@ struct RejectingClassifier { #[async_trait::async_trait] impl RequestClassifier for RejectingClassifier { - fn classify(&mut self, _request: ClassifyRequest) -> ClassifyFuture { + fn classify( + &mut self, + _request: ClassifyRequest, + _workers: WorkerCandidates<'_>, + ) -> ClassifyFuture { self.calls.fetch_add(1, Ordering::Relaxed); Box::pin(async { Err(Box::new(ClassifierRejected) as Box) }) } @@ -1901,7 +1910,11 @@ struct TypedRejectingClassifier { #[async_trait::async_trait] impl RequestClassifier for TypedRejectingClassifier { - fn classify(&mut self, _request: ClassifyRequest) -> ClassifyFuture { + fn classify( + &mut self, + _request: ClassifyRequest, + _workers: WorkerCandidates<'_>, + ) -> ClassifyFuture { self.calls.fetch_add(1, Ordering::Relaxed); let error_type = self.error_type; Box::pin(async move { @@ -2100,7 +2113,11 @@ struct PausingClassifier { } impl RequestClassifier for PausingClassifier { - fn classify(&mut self, request: ClassifyRequest) -> ClassifyFuture { + fn classify( + &mut self, + request: ClassifyRequest, + _workers: WorkerCandidates<'_>, + ) -> ClassifyFuture { let entered = Arc::clone(&self.entered); let resumed = Arc::clone(&self.resumed); Box::pin(async move { @@ -2215,7 +2232,11 @@ struct TokenContractClassifier { #[async_trait::async_trait] impl RequestClassifier for TokenContractClassifier { - fn classify(&mut self, request: ClassifyRequest) -> ClassifyFuture { + fn classify( + &mut self, + request: ClassifyRequest, + _workers: WorkerCandidates<'_>, + ) -> ClassifyFuture { self.classified .send((request.input_tokens(), request.progress().clone())) .unwrap(); diff --git a/lib/router-plugins/builtin/src/thunderagent/request_classifier/mod.rs b/lib/router-plugins/builtin/src/thunderagent/request_classifier/mod.rs index 2179fd73d3a1..d56fcd2e88ed 100644 --- a/lib/router-plugins/builtin/src/thunderagent/request_classifier/mod.rs +++ b/lib/router-plugins/builtin/src/thunderagent/request_classifier/mod.rs @@ -17,6 +17,7 @@ use dynamo_kv_router::plugins::request_classifier::{ RequestClassifierContext, RequestClassifierFactory, RequestClassifierParameters, RequestClassifierProviderError, RequestClassifierRegistryError, RequestProgress, }; +use dynamo_kv_router::plugins::worker_selection::WorkerCandidates; use parking_lot::Mutex; use thiserror::Error; use tokio::sync::Notify; @@ -283,7 +284,11 @@ impl ThunderAgentClassifier { #[async_trait] impl RequestClassifier for ThunderAgentClassifier { - fn classify(&mut self, request: ClassifyRequest) -> ClassifyFuture { + fn classify( + &mut self, + request: ClassifyRequest, + _workers: WorkerCandidates<'_>, + ) -> ClassifyFuture { let Some(session) = request.session_context() else { return Box::pin(async move { Ok(request) }); }; diff --git a/lib/router-plugins/builtin/tests/plugin_registry.rs b/lib/router-plugins/builtin/tests/plugin_registry.rs index b63d13842bd8..03940820be15 100644 --- a/lib/router-plugins/builtin/tests/plugin_registry.rs +++ b/lib/router-plugins/builtin/tests/plugin_registry.rs @@ -7,12 +7,17 @@ use dynamo_kv_router::plugins::RouterPlugins; use dynamo_kv_router::plugins::request_classifier::{ ClassifyFuture, ClassifyRequest, RequestClassifier, }; +use dynamo_kv_router::plugins::worker_selection::WorkerCandidates; use dynamo_kv_router::{KvRouterConfig, RoutingPartitionRef, WorkerType}; struct PassThrough; impl RequestClassifier for PassThrough { - fn classify(&mut self, request: ClassifyRequest) -> ClassifyFuture { + fn classify( + &mut self, + request: ClassifyRequest, + _workers: WorkerCandidates<'_>, + ) -> ClassifyFuture { Box::pin(async move { Ok(request) }) } }