fix: endpoint changes should be prioritized over new requests in kv scheduler (#2071)

Signed-off-by: Yan Ru Pei <yanrpei@gmail.com>

Author: PeaBrane

lib/llm/src/kv_router/scheduler.rs

@@ -142,24 +142,20 @@ impl KvScheduler {
                 request = tokio::select! {
                     biased;
 
-                    new_request = request_rx.recv() => {
-                        match new_request {
-                            Some(new_request) => {
-                                tracing::trace!("received request to be scheduled");
-                                new_request
-                            },
-                            None => {
-                                tracing::trace!("scheduler shutdown");
-                                break 'outer;
-                            }
-                        }
-                    }
-
                     _ = endpoints_rx.changed() => {
                         endpoints = endpoints_rx.borrow_and_update().clone();
                         pending_endpoint_update = Some(endpoints.worker_ids());
                         continue 'outer;
                     }
+
+                    maybe_new_request = request_rx.recv() => {
+                        let Some(new_request) = maybe_new_request else {
+                            tracing::warn!("scheduler shutdown");
+                            break 'outer;
+                        };
+                        tracing::trace!("received request to be scheduled");
+                        new_request
+                    }
                 };
 
                 loop {
@@ -181,6 +177,7 @@ impl KvScheduler {
                             request.respond(response);
                             continue 'outer;
                         }
+                        // TODO: this is not actually hooked up
                         Err(KvSchedulerError::AllWorkersBusy) => {
                             tracing::trace!("all workers busy; waiting for more capacity");
                             tokio::time::sleep(Duration::from_millis(5)).await;