feat: configurable payload limit (#210)

Author: OlivierDehaene

README.md

@@ -223,11 +223,18 @@ Options:
           [default: /tmp/text-embeddings-inference-server]
 
       --huggingface-hub-cache <HUGGINGFACE_HUB_CACHE>
-          The location of the huggingface hub cache. Used to override the location if you want to provide a mounted disk
-          for instance
+          The location of the huggingface hub cache. Used to override the location if you want to provide a mounted disk for instance
 
           [env: HUGGINGFACE_HUB_CACHE=/data]
 
+      --payload-limit <PAYLOAD_LIMIT>
+          Payload size limit in bytes
+
+          Default is 2MB
+
+          [env: PAYLOAD_LIMIT=]
+          [default: 2000000]
+
       --json-output
           Outputs the logs in JSON format (useful for telemetry)
 

router/src/http/server.rs

@@ -12,7 +12,7 @@ use crate::{
 };
 use ::http::HeaderMap;
 use anyhow::Context;
-use axum::extract::Extension;
+use axum::extract::{DefaultBodyLimit, Extension};
 use axum::http::HeaderValue;
 use axum::http::{Method, StatusCode};
 use axum::routing::{get, post};
@@ -1262,6 +1262,7 @@ pub async fn run(
     info: Info,
     addr: SocketAddr,
     prom_builder: PrometheusBuilder,
+    payload_limit: usize,
     cors_allow_origin: Option<Vec<String>>,
 ) -> Result<(), anyhow::Error> {
     // OpenAPI documentation
@@ -1370,7 +1371,8 @@ pub async fn run(
     };
 
     // Create router
-    let base_routes = Router::new()
+    let mut app = Router::new()
+        .layer(DefaultBodyLimit::max(payload_limit))
         .merge(SwaggerUi::new("/docs").url("/api-doc/openapi.json", doc))
         // Base routes
         .route("/info", get(get_model_info))
@@ -1393,8 +1395,6 @@ pub async fn run(
         // Prometheus metrics route
         .route("/metrics", get(metrics));
 
-    let mut app = Router::new().merge(base_routes);
-
     // Set default routes
     app = match &info.model_type {
         ModelType::Classifier(_) => {

router/src/lib.rs

@@ -56,6 +56,7 @@ pub async fn run(
     port: u16,
     uds_path: Option<String>,
     huggingface_hub_cache: Option<String>,
+    payload_limit: usize,
     otlp_endpoint: Option<String>,
     cors_allow_origin: Option<Vec<String>>,
 ) -> Result<()> {
@@ -268,7 +269,15 @@ pub async fn run(
     #[cfg(feature = "http")]
     {
         let server = tokio::spawn(async move {
-            http::server::run(infer, info, addr, prom_builder, cors_allow_origin).await
+            http::server::run(
+                infer,
+                info,
+                addr,
+                prom_builder,
+                payload_limit,
+                cors_allow_origin,
+            )
+            .await
         });
         tracing::info!("Ready");
         server.await??;

router/src/main.rs

@@ -96,6 +96,12 @@ struct Args {
     #[clap(long, env)]
     huggingface_hub_cache: Option<String>,
 
+    /// Payload size limit in bytes
+    ///
+    /// Default is 2MB
+    #[clap(default_value = "2000000", long, env)]
+    payload_limit: usize,
+
     /// Outputs the logs in JSON format (useful for telemetry)
     #[clap(long, env)]
     json_output: bool,
@@ -136,6 +142,7 @@ async fn main() -> Result<()> {
         args.port,
         Some(args.uds_path),
         args.huggingface_hub_cache,
+        args.payload_limit,
         args.otlp_endpoint,
         args.cors_allow_origin,
     )