From 30148b776b5e54f07984af8d34d29e8a00077638 Mon Sep 17 00:00:00 2001
From: OlivierDehaene <23298448+OlivierDehaene@users.noreply.github.com>
Date: Tue, 4 Apr 2023 14:51:31 +0200
Subject: [PATCH] fix instrumentation

---
 router/src/infer.rs      | 6 +++++-
 router/src/validation.rs | 7 ++-----
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/router/src/infer.rs b/router/src/infer.rs
index 6d13400f..a13c464d 100644
--- a/router/src/infer.rs
+++ b/router/src/infer.rs
@@ -87,7 +87,11 @@ impl Infer {
             })?;
 
         // Validate request
-        let valid_request = self.validation.validate(request).await?;
+        let valid_request = self.validation.validate(request).await.map_err(|err| {
+            metrics::increment_counter!("tgi_request_failure", "err" => "validation");
+            tracing::error!("{err}");
+            err
+        })?;
 
         // MPSC channel to communicate with the background batching task
         let (response_tx, response_rx) = flume::unbounded();
diff --git a/router/src/validation.rs b/router/src/validation.rs
index b781210c..c3fd152f 100644
--- a/router/src/validation.rs
+++ b/router/src/validation.rs
@@ -59,6 +59,7 @@ impl Validation {
         }
     }
 
+    #[instrument(skip_all)]
     async fn validate_input(
         &self,
         inputs: String,
@@ -270,11 +271,7 @@ fn tokenizer_worker(tokenizer: Tokenizer, receiver: flume::Receiver<TokenizerReq
     while let Ok(((inputs, truncate), response_tx, parent_span)) = receiver.recv() {
         parent_span.in_scope(|| {
             response_tx
-                .send(prepare_input(inputs, truncate, &tokenizer).map_err(|err| {
-                    metrics::increment_counter!("tgi_request_failure", "err" => "validation");
-                    tracing::error!("{err}");
-                    err
-                }))
+                .send(prepare_input(inputs, truncate, &tokenizer))
                 .unwrap_or(())
         })
     }