propagate mark_prefill_completed and free results

Signed-off-by: PeaBrane <yanrpei@gmail.com>

Author: PeaBrane

lib/llm/src/kv_router.rs

@@ -399,11 +399,11 @@ impl KvRouter {
             .await;
     }
 
-    pub async fn mark_prefill_completed(&self, request_id: &str) {
+    pub async fn mark_prefill_completed(&self, request_id: &str) -> Result<()> {
         self.scheduler.mark_prefill_completed(request_id).await
     }
 
-    pub async fn free(&self, request_id: &str) {
+    pub async fn free(&self, request_id: &str) -> Result<()> {
         self.scheduler.free(request_id).await
     }
 
@@ -456,14 +456,12 @@ impl AsyncEngine<SingleIn<RouterRequest>, ManyOut<Annotated<RouterResponse>>, Er
                     overlap_blocks,
                 }
             }
-            RouterRequest::MarkPrefill => {
-                self.mark_prefill_completed(&context_id).await;
-                RouterResponse::PrefillMarked { success: true }
-            }
-            RouterRequest::MarkFree => {
-                self.free(&context_id).await;
-                RouterResponse::FreeMarked { success: true }
-            }
+            RouterRequest::MarkPrefill => RouterResponse::PrefillMarked {
+                success: self.mark_prefill_completed(&context_id).await.is_ok(),
+            },
+            RouterRequest::MarkFree => RouterResponse::FreeMarked {
+                success: self.free(&context_id).await.is_ok(),
+            },
         };
 
         let response = Annotated::from_data(response);
@@ -593,15 +591,19 @@ impl AsyncEngine<SingleIn<PreprocessedRequest>, ManyOut<Annotated<LLMEngineOutpu
 
                 let wrapped_stream = Box::pin(async_stream::stream! {
                     if let Some(first_item) = response_stream.next().await {
-                        chooser.mark_prefill_completed(&context_id).await;
+                        if let Err(e) = chooser.mark_prefill_completed(&context_id).await {
+                            tracing::warn!("Failed to mark prefill completed for request {context_id}: {e:?}");
+                        }
                         yield first_item;
                     }
 
                     while let Some(item) = response_stream.next().await {
                         yield item;
                     }
 
-                    chooser.free(&context_id).await;
+                    if let Err(e) = chooser.free(&context_id).await {
+                        tracing::warn!("Failed to free request {context_id}: {e:?}");
+                    }
                 });
                 Ok(ResponseStream::new(wrapped_stream, stream_context))
             }

lib/llm/src/kv_router/scheduler.rs

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use crate::local_model::runtime_config::ModelRuntimeConfig;
+use anyhow::Result;
 use dynamo_runtime::component::{Component, Instance};
 use dynamo_runtime::traits::DistributedRuntimeProvider;
 use dynamo_runtime::traits::events::EventPublisher;
@@ -329,15 +330,14 @@ impl KvScheduler {
             .await;
     }
 
-    pub async fn mark_prefill_completed(&self, request_id: &str) {
-        let _ = self
-            .slots
+    pub async fn mark_prefill_completed(&self, request_id: &str) -> Result<()> {
+        self.slots
             .mark_prefill_completed(&request_id.to_string())
-            .await;
+            .await
     }
 
-    pub async fn free(&self, request_id: &str) {
-        let _ = self.slots.free(&request_id.to_string()).await;
+    pub async fn free(&self, request_id: &str) -> Result<()> {
+        self.slots.free(&request_id.to_string()).await
     }
 
     pub async fn get_potential_loads(