From 521f6203d14605c6b074456a298d291d9b04d1e9 Mon Sep 17 00:00:00 2001
From: OlivierDehaene <23298448+OlivierDehaene@users.noreply.github.com>
Date: Wed, 19 Apr 2023 18:39:44 +0200
Subject: [PATCH] add metrics

---
 router/src/infer.rs | 5 ++++-
 router/src/queue.rs | 4 +++-
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/router/src/infer.rs b/router/src/infer.rs
index 1eb37e6a..484720a0 100644
--- a/router/src/infer.rs
+++ b/router/src/infer.rs
@@ -446,7 +446,10 @@ fn filter_send_generations(generations: Vec<Generation>, entries: &mut IntMap<u6
         // Send generation responses back to the infer task
         // If the receive an error from the Flume channel, it means that the client dropped the
         // request and we need to stop generating hence why we unwrap_or(true)
-        let stopped = send_responses(generation, entry).unwrap_or(true);
+        let stopped = send_responses(generation, entry).map_err(|err| {
+            metrics::increment_counter!("tgi_request_failure", "err" => "dropped");
+            err
+        }).unwrap_or(true);
         if stopped {
             entries.remove(&id).expect("ID not found in entries. This is a bug.");
         }
diff --git a/router/src/queue.rs b/router/src/queue.rs
index c2220cb0..43651ff3 100644
--- a/router/src/queue.rs
+++ b/router/src/queue.rs
@@ -161,6 +161,7 @@ impl State {
             // Filter entries where the response receiver was dropped (== entries where the request
             // was dropped by the client)
             if entry.response_tx.is_disconnected() {
+                metrics::increment_counter!("tgi_request_failure", "err" => "dropped");
                 continue;
             }
 
@@ -190,6 +191,8 @@ impl State {
             }
         }
 
+        metrics::gauge!("tgi_queue_size", self.entries.len() as f64);
+
         // Maybe all entries were dropped because their channel were closed
         if batch_requests.is_empty() {
             return None;
@@ -207,7 +210,6 @@ impl State {
         // Increment batch id
         self.next_batch_id += 1;
 
-        metrics::gauge!("tgi_queue_size", self.entries.len() as f64);
         metrics::histogram!("tgi_batch_next_size", batch.size as f64);
         Some((batch_entries, batch, next_batch_span))
     }