feat(router): expose worker inputs to request classifiers

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
 

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)

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) })
     }
 

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,

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<CandidateData>,
+    shared_hits: Option<SharedCacheHits>,
+    has_tier_matches: bool,
+}
+
+impl WorkerInputRows {
+    pub(crate) fn new(
+        inputs: WorkerInputs,
+        rows: Vec<CandidateData>,
+        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<R>(&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 {

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<BookingHandle>), 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<Option<ClassifyRequest>, 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<parking_lot::Mutex<Option<Vec<Row>>>>);
+        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())]);