Make request payload size configurable (#2444)

Signed-off-by: Ata Fatahi <immrata@gmail.com>

rust/src/lib.rs

@@ -22,6 +22,7 @@ struct Router {
     balance_rel_threshold: f32,
     eviction_interval_secs: u64,
     max_tree_size: usize,
+    max_payload_size: usize,
     verbose: bool,
 }
 
@@ -38,6 +39,7 @@ impl Router {
         balance_rel_threshold = 1.0001,
         eviction_interval_secs = 60,
         max_tree_size = 2usize.pow(24),
+        max_payload_size = 4 * 1024 * 1024,
         verbose = false
     ))]
     fn new(
@@ -50,6 +52,7 @@ impl Router {
         balance_rel_threshold: f32,
         eviction_interval_secs: u64,
         max_tree_size: usize,
+        max_payload_size: usize,
         verbose: bool,
     ) -> PyResult<Self> {
         Ok(Router {
@@ -62,6 +65,7 @@ impl Router {
             balance_rel_threshold,
             eviction_interval_secs,
             max_tree_size,
+            max_payload_size,
             verbose,
         })
     }
@@ -86,6 +90,7 @@ impl Router {
                 worker_urls: self.worker_urls.clone(),
                 policy_config,
                 verbose: self.verbose,
+                max_payload_size: self.max_payload_size,
             })
             .await
             .unwrap();

rust/src/server.rs

@@ -127,6 +127,7 @@ pub struct ServerConfig {
     pub worker_urls: Vec<String>,
     pub policy_config: PolicyConfig,
     pub verbose: bool,
+    pub max_payload_size: usize,
 }
 
 pub async fn startup(config: ServerConfig) -> std::io::Result<()> {
@@ -164,10 +165,16 @@ pub async fn startup(config: ServerConfig) -> std::io::Result<()> {
     info!("✅ Starting router on {}:{}", config.host, config.port);
     info!("✅ Serving Worker URLs: {:?}", config.worker_urls);
     info!("✅ Policy Config: {:?}", config.policy_config);
+    info!(
+        "✅ Max payload size: {} MB",
+        config.max_payload_size / (1024 * 1024)
+    );
 
     HttpServer::new(move || {
         App::new()
             .app_data(app_state.clone())
+            .app_data(web::JsonConfig::default().limit(config.max_payload_size))
+            .app_data(web::PayloadConfig::default().limit(config.max_payload_size))
             .service(generate)
             .service(v1_chat_completions)
             .service(v1_completions)