Bound VM Ranker DPP work by the server pool limit

vm-ranker/args.rs

@@ -21,7 +21,11 @@ pub struct Args {
     #[arg(long, default_value_t = 0.5)]
     pub dpp_theta: f64,
 
-    #[arg(long, default_value_t = 100)]
+    #[arg(
+        long,
+        default_value_t = 100,
+        help = "Maximum DPP candidate pool size; requests may select a smaller pool"
+    )]
     pub dpp_max_selected_rank: usize,
 
     #[arg(long, default_value_t = 1024)]

vm-ranker/scoring/mod.rs

@@ -4,7 +4,7 @@ use std::sync::Arc;
 
 use log::error;
 
-use xai_vm_ranker_proto::{RankRequest, RankedCandidate};
+use xai_vm_ranker_proto::{DppParams, RankRequest, RankedCandidate};
 
 use crate::dpp::DppConfig;
 use crate::embedding_store::EmbeddingStore;
@@ -15,6 +15,16 @@ pub struct DppContext {
     pub config: DppConfig,
 }
 
+fn apply_request_dpp_params(config: &mut DppConfig, params: &DppParams) {
+    if params.theta != 0.0 {
+        config.theta = params.theta;
+    }
+    if params.max_selected_rank != 0 {
+        config.max_selected_rank =
+            (params.max_selected_rank as usize).min(config.max_selected_rank);
+    }
+}
+
 pub async fn rank(
     req: RankRequest,
     dpp: Option<&DppContext>,
@@ -29,12 +39,7 @@ pub async fn rank(
         let mut ctx = ctx.clone();
 
         if let Some(params) = &req.dpp_params {
-            if params.theta != 0.0 {
-                ctx.config.theta = params.theta;
-            }
-            if params.max_selected_rank != 0 {
-                ctx.config.max_selected_rank = params.max_selected_rank as usize;
-            }
+            apply_request_dpp_params(&mut ctx.config, params);
         }
 
         let dpp_result = tokio::task::spawn_blocking(move || dpp_model::rank(&req, &ctx))
@@ -52,3 +57,56 @@ pub async fn rank(
         })
         .collect())
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn config_with_pool_limit(max_selected_rank: usize) -> DppConfig {
+        DppConfig {
+            top_k: 50,
+            theta: 0.5,
+            max_selected_rank,
+            debug_viewer_id: 0,
+        }
+    }
+
+    #[test]
+    fn request_cannot_raise_server_dpp_pool_limit() {
+        let mut config = config_with_pool_limit(100);
+        let params = DppParams {
+            theta: 0.0,
+            max_selected_rank: u32::MAX,
+        };
+
+        apply_request_dpp_params(&mut config, &params);
+
+        assert_eq!(config.max_selected_rank, 100);
+    }
+
+    #[test]
+    fn request_can_select_a_smaller_dpp_pool() {
+        let mut config = config_with_pool_limit(100);
+        let params = DppParams {
+            theta: 0.0,
+            max_selected_rank: 40,
+        };
+
+        apply_request_dpp_params(&mut config, &params);
+
+        assert_eq!(config.max_selected_rank, 40);
+    }
+
+    #[test]
+    fn zero_request_value_keeps_server_dpp_pool_limit() {
+        let mut config = config_with_pool_limit(100);
+        let params = DppParams {
+            theta: 0.0,
+            max_selected_rank: 0,
+        };
+
+        apply_request_dpp_params(&mut config, &params);
+
+        assert_eq!(config.max_selected_rank, 100);
+    }
+}