fix: do not leak inputs on error (#1228)

Close #1225
huggingface · Nov 20, 2023 · 3dbc649 · 3dbc649
1 parent 8acdc1f
commit 3dbc649
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/router/src/infer.rs b/router/src/infer.rs
@@ -82,7 +82,7 @@ impl Infer {
     }
 
     /// Add a new request to the queue and return a stream of InferStreamResponse
-    #[instrument(skip(self))]
+    #[instrument(skip_all)]
     pub(crate) async fn generate_stream(
         &self,
         request: GenerateRequest,
@@ -133,7 +133,7 @@ impl Infer {
     }
 
     /// Add a new request to the queue and return a InferResponse
-    #[instrument(skip(self))]
+    #[instrument(skip_all)]
     pub(crate) async fn generate(
         &self,
         request: GenerateRequest,
@@ -214,7 +214,7 @@ impl Infer {
     }
     /// Add best_of new requests to the queue and return a InferResponse of the sequence with
     /// the highest log probability per token
-    #[instrument(skip(self))]
+    #[instrument(skip(self, request))]
     pub(crate) async fn generate_best_of(
         &self,
         request: GenerateRequest,

diff --git a/router/src/validation.rs b/router/src/validation.rs
@@ -69,7 +69,7 @@ impl Validation {
         }
     }
 
-    #[instrument(skip_all)]
+    #[instrument(skip(self, inputs))]
     async fn validate_input(
         &self,
         inputs: String,