Change routing policy API to be async to support more policies (#17048)

sgl-model-gateway/tests/mesh_integration_test.rs

@@ -372,13 +372,13 @@ async fn test_rate_limit_window_reset() {
     manager.sync_rate_limit_inc(GLOBAL_RATE_LIMIT_COUNTER_KEY.to_string(), 50);
     let value_before = manager.get_rate_limit_value(GLOBAL_RATE_LIMIT_COUNTER_KEY);
     // Value may be None if not owner, or Some if owner
-    if value_before.is_some() {
-        assert!(value_before.unwrap() > 0);
+    if let Some(val) = value_before {
+        assert!(val > 0);
 
         // Reset counter
         manager.reset_global_rate_limit_counter();
         let value_after = manager.get_rate_limit_value(GLOBAL_RATE_LIMIT_COUNTER_KEY);
         // Should be reset
-        assert!(value_after.is_none() || value_after.unwrap() <= 0);
+        assert!(value_after.is_none() || value_after.unwrap_or(0) <= 0);
     }
 }

sgl-model-gateway/tests/routing/cache_aware_backward_compat_test.rs

@@ -5,8 +5,8 @@ use smg::{
     policies::{CacheAwareConfig, CacheAwarePolicy, LoadBalancingPolicy, SelectWorkerInfo},
 };
 
-#[test]
-fn test_backward_compatibility_with_empty_model_id() {
+#[tokio::test]
+async fn test_backward_compatibility_with_empty_model_id() {
     let config = CacheAwareConfig {
         cache_threshold: 0.5,
         balance_abs_threshold: 2,
@@ -40,13 +40,15 @@ fn test_backward_compatibility_with_empty_model_id() {
     let workers: Vec<Arc<dyn Worker>> = vec![Arc::new(worker1.clone()), Arc::new(worker2.clone())];
 
     // Select worker - should work without errors
-    let selected = policy.select_worker(
-        &workers,
-        &SelectWorkerInfo {
-            request_text: Some("test request"),
-            ..Default::default()
-        },
-    );
+    let selected = policy
+        .select_worker(
+            &workers,
+            &SelectWorkerInfo {
+                request_text: Some("test request"),
+                ..Default::default()
+            },
+        )
+        .await;
     assert!(selected.is_some(), "Should select a worker");
 
     // Remove workers - should work without errors
@@ -54,8 +56,8 @@ fn test_backward_compatibility_with_empty_model_id() {
     policy.remove_worker(&worker2);
 }
 
-#[test]
-fn test_mixed_model_ids() {
+#[tokio::test]
+async fn test_mixed_model_ids() {
     let config = CacheAwareConfig {
         cache_threshold: 0.5,
         balance_abs_threshold: 2,
@@ -107,12 +109,12 @@ fn test_mixed_model_ids() {
         request_text: Some("test request"),
         ..Default::default()
     };
-    let selected = policy.select_worker(&default_workers, &info);
+    let selected = policy.select_worker(&default_workers, &info).await;
     assert!(selected.is_some(), "Should select from default workers");
 
     let llama_workers: Vec<Arc<dyn Worker>> =
         vec![Arc::new(worker2.clone()), Arc::new(worker4.clone())];
-    let selected = policy.select_worker(&llama_workers, &info);
+    let selected = policy.select_worker(&llama_workers, &info).await;
     assert!(selected.is_some(), "Should select from llama-3 workers");
 
     let all_workers: Vec<Arc<dyn Worker>> = vec![
@@ -121,12 +123,12 @@ fn test_mixed_model_ids() {
         Arc::new(worker3.clone()),
         Arc::new(worker4.clone()),
     ];
-    let selected = policy.select_worker(&all_workers, &info);
+    let selected = policy.select_worker(&all_workers, &info).await;
     assert!(selected.is_some(), "Should select from all workers");
 }
 
-#[test]
-fn test_remove_worker_by_url_backward_compat() {
+#[tokio::test]
+async fn test_remove_worker_by_url_backward_compat() {
     let config = CacheAwareConfig::default();
     let policy = CacheAwarePolicy::with_config(config);
 
@@ -154,12 +156,14 @@ fn test_remove_worker_by_url_backward_compat() {
     policy.remove_worker_by_url("http://worker1:8080");
 
     let workers: Vec<Arc<dyn Worker>> = vec![Arc::new(worker2.clone())];
-    let selected = policy.select_worker(
-        &workers,
-        &SelectWorkerInfo {
-            request_text: Some("test"),
-            ..Default::default()
-        },
-    );
+    let selected = policy
+        .select_worker(
+            &workers,
+            &SelectWorkerInfo {
+                request_text: Some("test"),
+                ..Default::default()
+            },
+        )
+        .await;
     assert_eq!(selected, Some(0), "Should only have worker2 left");
 }