Add prometheus metrics for tokenize API

Author: njhill

README.md

@@ -154,3 +154,7 @@ They are all prefixed with `tgi_`. Descriptions will be added to the table below
 | `tgi_prefill_weight_limit_exceeded`        | `counter`   |                                                     |              |
 | `tgi_prompt_load_failure`                  | `counter`   |                                                     |              |
 | `tgi_prompt_load_duration`                 | `histogram` |                                                     |              |
+| `tgi_tokenize_request_count`               | `counter`   |                                                     |              |
+| `tgi_tokenize_request_input_count`         | `counter`   |                                                     |              |
+| `tgi_tokenize_request_tokens`              | `histogram` |                                                     |              |
+| `tgi_tokenize_request_duration`            | `histogram` |                                                     |              |

router/src/grpc_server.rs

@@ -60,6 +60,7 @@ pub(crate) async fn start_grpc_server<F: Future<Output = ()> + Send +'static> (
         state: shared_state,
         tokenizer,
         input_counter: metrics::register_counter!("tgi_request_input_count"),
+        tokenize_input_counter: metrics::register_counter!("tgi_tokenize_request_input_count"),
     };
     let grpc_server = builder
         .add_service(GenerationServiceServer::new(grpc_service))
@@ -82,6 +83,7 @@ pub struct GenerationServicer {
     state: ServerState,
     tokenizer: Tokenizer,
     input_counter: metrics::Counter,
+    tokenize_input_counter: metrics::Counter,
 }
 
 #[tonic::async_trait]
@@ -245,15 +247,25 @@ impl GenerationService for GenerationServicer {
         &self, request: Request<BatchedTokenizeRequest>
     ) -> Result<Response<BatchedTokenizeResponse>, Status> {
         let br = request.into_inner();
+        metrics::increment_counter!("tgi_tokenize_request_count");
+        let start_time = Instant::now();
+        self.tokenize_input_counter.increment(br.requests.len() as u64);
 
+        let mut token_total = 0;
         let responses = self.tokenizer.encode_batch(
             br.requests.into_iter().map(|tr| tr.text).collect(), true
         )
-            .map_err(Status::from_error)?
-            .into_iter().map(|e| TokenizeResponse {
-                token_count: e.len() as u32,
-                tokens: if br.return_tokens { e.get_tokens().to_vec() } else { vec![] },
-        }).collect();
+            .map_err(Status::from_error)?.into_iter()
+            .map(|e| {
+                let token_count = e.len() as u32;
+                token_total += token_count;
+                TokenizeResponse {
+                    token_count,
+                    tokens: if br.return_tokens { e.get_tokens().to_vec() } else { vec![] },
+                }
+            }).collect();
+        metrics::histogram!("tgi_tokenize_request_tokens", token_total as f64);
+        metrics::histogram!("tgi_tokenize_request_duration", start_time.elapsed().as_secs_f64());
 
         Ok(Response::new(BatchedTokenizeResponse { responses }))
     }

router/src/server.rs

@@ -343,6 +343,9 @@ async fn do_run<B: BatchType>(
         value *= 1.5;
         duration_buckets.push(value);
     }
+    // Tokenization token count buckets
+    let tokenized_tokens_matcher = Matcher::Full(String::from("tgi_tokenize_request_tokens"));
+    let tokenized_tokens_buckets: Vec<f64> = (6..20).map(|x| (1 << x) as f64).collect();
     // Input Length buckets
     let input_length_matcher = Matcher::Full(String::from("tgi_request_input_length"));
     let max_sequence_length_buckets: Vec<f64> = (0..64)
@@ -364,20 +367,14 @@ async fn do_run<B: BatchType>(
 
     // Prometheus handler
     let builder = PrometheusBuilder::new()
-        .set_buckets_for_metric(duration_matcher, &duration_buckets)
-        .unwrap()
-        .set_buckets_for_metric(input_length_matcher, &max_sequence_length_buckets)
-        .unwrap()
-        .set_buckets_for_metric(generated_tokens_matcher, &max_new_tokens_buckets)
-        .unwrap()
-        .set_buckets_for_metric(max_new_tokens_matcher, &max_new_tokens_buckets)
-        .unwrap()
-        .set_buckets_for_metric(total_tokens_matcher, &max_sequence_length_buckets)
-        .unwrap()
-        .set_buckets_for_metric(batch_size_matcher, &batch_size_buckets)
-        .unwrap()
-        .set_buckets_for_metric(batch_inference_size_matcher, &batch_size_buckets)
-        .unwrap();
+        .set_buckets_for_metric(duration_matcher, &duration_buckets).unwrap()
+        .set_buckets_for_metric(tokenized_tokens_matcher, &tokenized_tokens_buckets).unwrap()
+        .set_buckets_for_metric(input_length_matcher, &max_sequence_length_buckets).unwrap()
+        .set_buckets_for_metric(generated_tokens_matcher, &max_new_tokens_buckets).unwrap()
+        .set_buckets_for_metric(max_new_tokens_matcher, &max_new_tokens_buckets).unwrap()
+        .set_buckets_for_metric(total_tokens_matcher, &max_sequence_length_buckets).unwrap()
+        .set_buckets_for_metric(batch_size_matcher, &batch_size_buckets).unwrap()
+        .set_buckets_for_metric(batch_inference_size_matcher, &batch_size_buckets).unwrap();
     let prom_handle = builder
         .install_recorder()
         .expect("failed to install metrics recorder");