diff --git a/vm-ranker/args.rs b/vm-ranker/args.rs index eb2ac60c..633884e5 100644 --- a/vm-ranker/args.rs +++ b/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)] diff --git a/vm-ranker/scoring/mod.rs b/vm-ranker/scoring/mod.rs index 916b5bd8..859e6632 100644 --- a/vm-ranker/scoring/mod.rs +++ b/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, ¶ms); + + 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, ¶ms); + + 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, ¶ms); + + assert_eq!(config.max_selected_rank, 100); + } +}