Update various rust crates

Metrics update includes new macro syntax.

Author: njhill

Cargo.lock

@@ -6,9 +6,9 @@ authors = ["Olivier Dehaene"]
 description = "Text Generation Launcher"
 
 [dependencies]
-clap = { version = "4.4.11", features = ["derive", "env"] }
-ctrlc = { version = "3.4.1", features = ["termination"] }
+clap = { version = "4.4.12", features = ["derive", "env"] }
+ctrlc = { version = "3.4.2", features = ["termination"] }
 subprocess = "0.2.9"
 tracing = "0.1.40"
-tracing-subscriber = { version = "0.3.17", features = ["json"] }
+tracing-subscriber = { version = "0.3.18", features = ["json"] }
 

launcher/Cargo.toml

@@ -17,30 +17,30 @@ path = "src/main.rs"
 axum = { version = "0.6.20", features = ["json"] }
 text-generation-client = { path = "client" }
 clap = { version = "^4.4.11", features = ["derive", "env"] }
-futures = "^0.3.29"
+futures = "^0.3.30"
 flume = "^0.11.0"
-metrics = "0.21.1"
-metrics-exporter-prometheus = { version = "0.12.1", features = [] }
-moka = { version = "0.12.1", features = ["future"] }
+metrics = "0.22.0"
+metrics-exporter-prometheus = { version = "0.13.0", features = [] }
+moka = { version = "0.12.2", features = ["future"] }
 nohash-hasher = "^0.2.0"
 num = "^0.4.1"
 num_cpus = "^1.16.0"
-hyper = "^0.14.27" # Override to address CVE-2023-26964
-openssl = "^0.10.59" # Override to address WS-2023-0082, WS-2023-0083, WS-2023-0195
-openssl-sys = "^0.9.95" # Override to address WS-2023-0082, WS-2023-0083, WS-2023-0195
-rustls-webpki = "^0.101.7" # Override to address WS-2023-0305, CVE-2018-16875
+hyper = "^0.14.28" # Override to address CVE-2023-26964
+openssl = "^0.10.62" # Override to address WS-2023-0082, WS-2023-0083, WS-2023-0195
+openssl-sys = "^0.9.98" # Override to address WS-2023-0082, WS-2023-0083, WS-2023-0195
+rustls-webpki = "0.102.0" # Override to address WS-2023-0305, CVE-2018-16875
 rand = "^0.8.5"
-serde = "^1.0.193"
-serde_json = "^1.0.108"
+serde = "^1.0.194"
+serde_json = "^1.0.110"
 # Pin sct to avoid WS-2023-0094 in older 0.5.2 version of spin sub-dependency
 sct = "=0.7.1"
-thiserror = "^1.0.50"
+thiserror = "^1.0.56"
 tokenizers = "0.15.0"
-tokio = { version = "1.35.0", features = ["rt", "rt-multi-thread", "parking_lot", "signal", "sync", "fs"] }
+tokio = { version = "1.35.1", features = ["rt", "rt-multi-thread", "parking_lot", "signal", "sync", "fs"] }
 tokio-rustls = "^0.25.0"
-rustls = "0.21.8"
+rustls = "0.22.1"
 tracing = "^0.1.40"
-tracing-subscriber = { version = "0.3.17", features = ["json"] }
+tracing-subscriber = { version = "0.3.18", features = ["json"] }
 prost = "^0.12.3"
 tonic = { version = "^0.10.2", features = ["tls"] }
 tokio-stream ="^0.1.14"
@@ -49,4 +49,4 @@ unicode-truncate = "^0.2.0"
 
 [build-dependencies]
 tonic-build = "^0.10.2"
-tempfile = "^3.8.1" # Override 0.3.3 version from tonic-build/prost-build, due to RUSTSEC-2023-0018 and WS-2023-0366
+tempfile = "^3.9.0" # Override 0.3.3 version from tonic-build/prost-build, due to RUSTSEC-2023-0018 and WS-2023-0366

router/Cargo.toml

@@ -5,10 +5,10 @@ edition = "2021"
 build="build.rs"
 
 [dependencies]
-futures = "^0.3.29"
-prost = "^0.12.1"
-thiserror = "^1.0.50"
-tokio = { version = "1.35.0", features = ["sync"] }
+futures = "^0.3.30"
+prost = "^0.12.3"
+thiserror = "^1.0.56"
+tokio = { version = "1.35.1", features = ["sync"] }
 tonic = "^0.10.2"
 tower = "^0.4.13"
 tracing = "^0.1.40"

router/client/Cargo.toml

@@ -380,9 +380,9 @@ async fn batching_task<B: BatchType>(
                 batch_size,
             );
 
-            metrics::gauge!("tgi_batch_current_size", batch_size as f64);
-            metrics::gauge!("tgi_batch_input_tokens", batch_tokens as f64);
-            metrics::gauge!("tgi_batch_max_remaining_tokens", batch_max_remaining_tokens.unwrap() as f64);
+            metrics::gauge!("tgi_batch_current_size").set(batch_size as f64);
+            metrics::gauge!("tgi_batch_input_tokens").set(batch_tokens as f64);
+            metrics::gauge!("tgi_batch_max_remaining_tokens").set(batch_max_remaining_tokens.unwrap() as f64);
 
             // Don't interfere with current batch if it's about to complete
             if batch_max_remaining_tokens.unwrap() >= 2 {
@@ -436,7 +436,7 @@ async fn batching_task<B: BatchType>(
                             if added_batch_size > 0 {
                                 info!("Extending batch #{} of {} with additional batch #{} of {}",
                                 batch_id, batch_size, new_batch_id, added_batch_size);
-                                metrics::increment_counter!("tgi_batch_concatenation_count");
+                                metrics::counter!("tgi_batch_concatenation_count").increment(1);
                             }
                         } else {
                             combined_batch_id = new_batch_id;
@@ -463,9 +463,9 @@ async fn batching_task<B: BatchType>(
             }
         }
 
-        metrics::gauge!("tgi_batch_current_size", 0.0);
-        metrics::gauge!("tgi_batch_input_tokens", 0.0);
-        metrics::gauge!("tgi_batch_max_remaining_tokens", 0.0);
+        metrics::gauge!("tgi_batch_current_size").set(0.0);
+        metrics::gauge!("tgi_batch_input_tokens").set(0.0);
+        metrics::gauge!("tgi_batch_max_remaining_tokens").set(0.0);
     }
 
     info!("Batching loop exiting");
@@ -524,10 +524,8 @@ impl<'a> TokenProcessor<'a> {
         let batch_size = batch.requests.len();
         let batch_tokens = batch.total_tokens;
         let start_time = Instant::now();
-        metrics::histogram!("tgi_batch_next_tokens", batch_tokens as f64);
-        metrics::histogram!(
-            "tgi_batch_inference_batch_size", batch_size as f64, "method" => "prefill"
-        );
+        metrics::histogram!("tgi_batch_next_tokens").record(batch_tokens as f64);
+        metrics::histogram!("tgi_batch_inference_batch_size", "method" => "prefill").record(batch_size as f64);
         self._wrap_future(
             client.prefill(batch, to_prune).map(|r| {
                 info!(
@@ -543,9 +541,8 @@ impl<'a> TokenProcessor<'a> {
     async fn next_token<B: BatchType>(
         &mut self, client: &mut ShardedClient, batches: Vec<CachedBatch>, queue: &mut Queue<B>,
     ) -> Option<CachedBatch> {
-        metrics::histogram!(
-            "tgi_batch_inference_batch_size", self.entries.len() as f64, "method" => "next_token"
-        );
+        metrics::histogram!("tgi_batch_inference_batch_size", "method" => "next_token")
+            .record(self.entries.len() as f64);
         let start_time = Instant::now();
         self._wrap_future(
             client.next_token(batches), "next_token", start_time, None, queue
@@ -562,7 +559,7 @@ impl<'a> TokenProcessor<'a> {
         start_id: Option<u64>,
         queue: &mut Queue<B>,
     ) -> Option<CachedBatch> {
-        metrics::increment_counter!("tgi_batch_inference_count", "method" => method);
+        metrics::counter!("tgi_batch_inference_count", "method" => method).increment(1);
 
         // We process the shared queue while waiting for the response from the python shard(s)
         let queue_servicer = queue.service_queue().fuse();
@@ -576,7 +573,8 @@ impl<'a> TokenProcessor<'a> {
 
         match result {
             Ok(
-                Some((generated_tokens, input_tokens, errors, next_batch_id, forward_duration))
+                Some((generated_tokens, input_tokens,
+                         errors, next_batch_id, forward_duration))
             ) => {
                 let pre_token_process_time = Instant::now();
                 self.process_input_tokens(input_tokens);
@@ -585,27 +583,27 @@ impl<'a> TokenProcessor<'a> {
                 );
                 // Update health
                 self.generation_health.store(true, Ordering::SeqCst);
-                metrics::histogram!(
+                let histogram = metrics::histogram!(
                     "tgi_batch_inference_duration",
-                    start_time.elapsed().as_secs_f64(),
                     "method" => method,
                     "makeup" => "single_only", // later will possibly be beam_only or mixed
                 );
-                metrics::histogram!(
+                histogram.record(start_time.elapsed().as_secs_f64());
+                let histogram = metrics::histogram!(
                     "tgi_batch_inference_forward_duration",
-                    forward_duration,
                     "method" => method,
                     "makeup" => "single_only", // later will possibly be beam_only or mixed
                 );
-                metrics::histogram!(
+                histogram.record(forward_duration);
+                let histogram = metrics::histogram!(
                     "tgi_batch_inference_tokproc_duration",
-                    pre_token_process_time.elapsed().as_secs_f64(),
                     "method" => method,
                     "makeup" => "single_only", // later will possibly be beam_only or mixed
                 );
+                histogram.record(pre_token_process_time.elapsed().as_secs_f64());
                 // Probably don't need this additional counter because the duration histogram
                 // records a total count
-                metrics::increment_counter!("tgi_batch_inference_success", "method" => method);
+                metrics::counter!("tgi_batch_inference_success", "method" => method).increment(1);
                 Some(CachedBatch{
                     batch_id: next_batch_id,
                     status: completed_request_ids.map(|c| RequestsStatus{completed_ids: c}),
@@ -622,7 +620,8 @@ impl<'a> TokenProcessor<'a> {
                     ClientError::Connection(_) => "connection",
                     _ => "error"
                 };
-                metrics::increment_counter!("tgi_batch_inference_failure", "method" => method, "reason" => reason);
+                metrics::counter!("tgi_batch_inference_failure", "method" => method, "reason" => reason)
+                    .increment(1);
                 self.send_errors(err, start_id);
                 None
             },
@@ -819,7 +818,7 @@ impl<'a> TokenProcessor<'a> {
                     // If receiver closed (request cancelled), cancel this entry
                     let e = self.entries.remove(&request_id).unwrap();
                     stop_reason = Cancelled;
-                    metrics::increment_counter!("tgi_request_failure", "err" => "cancelled");
+                    metrics::counter!("tgi_request_failure", "err" => "cancelled").increment(1);
                     //TODO include request context in log message
                     warn!("Aborted streaming request {request_id} cancelled by client \
                         after generating {} token(s)", e.generated_tokens);
@@ -831,7 +830,7 @@ impl<'a> TokenProcessor<'a> {
                 // If receiver closed (request cancelled), cancel this entry
                 let e = self.entries.remove(&request_id).unwrap();
                 stop_reason = Cancelled;
-                metrics::increment_counter!("tgi_request_failure", "err" => "cancelled");
+                metrics::counter!("tgi_request_failure", "err" => "cancelled").increment(1);
                 //TODO include request context in log message
                 warn!("Aborted request {request_id} cancelled by client \
                     after generating {} token(s)", e.generated_tokens);

router/src/batcher.rs

@@ -59,8 +59,8 @@ pub(crate) async fn start_grpc_server<F: Future<Output = ()> + Send +'static> (
     let grpc_service = GenerationServicer {
         state: shared_state,
         tokenizer,
-        input_counter: metrics::register_counter!("tgi_request_input_count"),
-        tokenize_input_counter: metrics::register_counter!("tgi_tokenize_request_input_count"),
+        input_counter: metrics::counter!("tgi_request_input_count"),
+        tokenize_input_counter: metrics::counter!("tgi_tokenize_request_input_count"),
     };
     let grpc_server = builder
         .add_service(GenerationServiceServer::new(grpc_service))
@@ -104,7 +104,7 @@ impl GenerationService for GenerationServicer {
         let br = request.into_inner();
         let batch_size = br.requests.len();
         let kind = if batch_size == 1 { "single" } else { "batch" };
-        metrics::increment_counter!("tgi_request_count", "kind" => kind);
+        metrics::counter!("tgi_request_count", "kind" => kind).increment(1);
         if batch_size == 0 {
             return Ok(Response::new(BatchedGenerationResponse{ responses: vec![] }));
         }
@@ -113,7 +113,7 @@ impl GenerationService for GenerationServicer {
         let _permit = self.state.limit_concurrent_requests
             .try_acquire_many(batch_size as u32)
             .map_err(|_| {
-                metrics::increment_counter!("tgi_request_failure", "err" => "conc_limit");
+                metrics::counter!("tgi_request_failure", "err" => "conc_limit").increment(1);
                 tracing::error!("Model is overloaded");
                 Status::resource_exhausted("Model is overloaded")
             })?;
@@ -155,11 +155,11 @@ impl GenerationService for GenerationServicer {
             }
         }.map_err(|err| match err {
             InferError::RequestQueueFull() => {
-                metrics::increment_counter!("tgi_request_failure", "err" => "queue_full");
+                metrics::counter!("tgi_request_failure", "err" => "queue_full").increment(1);
                 Status::resource_exhausted(err.to_string())
             },
             _ => {
-                metrics::increment_counter!("tgi_request_failure", "err" => "generate");
+                metrics::counter!("tgi_request_failure", "err" => "generate").increment(1);
                 tracing::error!("{err}");
                 Status::from_error(Box::new(err))
             },
@@ -184,11 +184,11 @@ impl GenerationService for GenerationServicer {
         &self, request: Request<SingleGenerationRequest>
     ) -> Result<Response<Self::GenerateStreamStream>, Status> {
         let start_time = Instant::now();
-        metrics::increment_counter!("tgi_request_count", "kind" => "stream");
+        metrics::counter!("tgi_request_count", "kind" => "stream").increment(1);
         self.input_counter.increment(1);
         let permit = self.state.limit_concurrent_requests.clone()
             .try_acquire_owned().map_err(|_| {
-                metrics::increment_counter!("tgi_request_failure", "err" => "conc_limit");
+                metrics::counter!("tgi_request_failure", "err" => "conc_limit").increment(1);
                 tracing::error!("Model is overloaded");
                 Status::resource_exhausted("Model is overloaded")
         })?;
@@ -210,7 +210,7 @@ impl GenerationService for GenerationServicer {
             }, |ctx, count, reason, request_id, times, out, err| {
                 let _enter = ctx.span.enter();
                 if let Some(e) = err {
-                    metrics::increment_counter!("tgi_request_failure", "err" => "generate");
+                    metrics::counter!("tgi_request_failure", "err" => "generate").increment(1);
                     tracing::error!("Streaming response failed after {count} tokens, \
                         output so far: '{:?}': {e}", truncate(&out, 32));
                 } else {
@@ -229,11 +229,11 @@ impl GenerationService for GenerationServicer {
             .await
             .map_err(|err| match err {
                 InferError::RequestQueueFull() => {
-                    metrics::increment_counter!("tgi_request_failure", "err" => "queue_full");
+                    metrics::counter!("tgi_request_failure", "err" => "queue_full").increment(1);
                     Status::resource_exhausted(err.to_string())
                 },
                 _ => {
-                    metrics::increment_counter!("tgi_request_failure", "err" => "unknown");
+                    metrics::counter!("tgi_request_failure", "err" => "unknown").increment(1);
                     tracing::error!("{err}");
                     Status::from_error(Box::new(err))
                 },
@@ -247,7 +247,7 @@ impl GenerationService for GenerationServicer {
         &self, request: Request<BatchedTokenizeRequest>
     ) -> Result<Response<BatchedTokenizeResponse>, Status> {
         let br = request.into_inner();
-        metrics::increment_counter!("tgi_tokenize_request_count");
+        metrics::counter!("tgi_tokenize_request_count").increment(1);
         let start_time = Instant::now();
         self.tokenize_input_counter.increment(br.requests.len() as u64);
 
@@ -262,8 +262,8 @@ impl GenerationService for GenerationServicer {
             ))).map_err(Status::from_error).await?;
 
         let token_total: u32 = responses.iter().map(|tr| tr.token_count).sum();
-        metrics::histogram!("tgi_tokenize_request_tokens", token_total as f64);
-        metrics::histogram!("tgi_tokenize_request_duration", start_time.elapsed().as_secs_f64());
+        metrics::histogram!("tgi_tokenize_request_tokens").record(token_total as f64);
+        metrics::histogram!("tgi_tokenize_request_duration").record(start_time.elapsed().as_secs_f64());
 
         Ok(Response::new(BatchedTokenizeResponse { responses }))
     }
@@ -304,11 +304,11 @@ impl GenerationServicer {
             ).await,
             Err(err) => Err(err),
         }.map_err(|err| {
-            metrics::increment_counter!("tgi_request_failure", "err" => "validation");
+            metrics::counter!("tgi_request_failure", "err" => "validation").increment(1);
             tracing::error!("{err}");
             Status::invalid_argument(err.to_string())
         }).map(|requests| {
-            metrics::histogram!("tgi_request_validation_duration", start_time.elapsed().as_secs_f64());
+            metrics::histogram!("tgi_request_validation_duration").record(start_time.elapsed().as_secs_f64());
             requests
         })
     }
@@ -349,23 +349,21 @@ fn log_response(
         );
         _enter = span.enter();
 
-        metrics::histogram!("tgi_request_inference_duration", inference_time.as_secs_f64());
-        metrics::histogram!("tgi_request_mean_time_per_token_duration", time_per_token.as_secs_f64());
+        metrics::histogram!("tgi_request_inference_duration").record(inference_time.as_secs_f64());
+        metrics::histogram!("tgi_request_mean_time_per_token_duration").record(time_per_token.as_secs_f64());
     }
 
     // Metrics
     match reason {
-        Error => metrics::increment_counter!("tgi_request_failure", "err" => "generate"),
+        Error => metrics::counter!("tgi_request_failure", "err" => "generate").increment(1),
         Cancelled => (), // recorded where cancellation is detected
         _ => {
-            metrics::increment_counter!(
-                "tgi_request_success", "stop_reason" => reason.as_str_name(), "kind" => kind
-            );
-            metrics::histogram!("tgi_request_duration", total_time.as_secs_f64());
-            metrics::histogram!("tgi_request_generated_tokens", generated_tokens as f64);
-            metrics::histogram!(
-                "tgi_request_total_tokens", (generated_tokens as usize + input_tokens) as f64
-            );
+            metrics::counter!("tgi_request_success", "stop_reason" => reason.as_str_name(), "kind" => kind)
+                .increment(1);
+            metrics::histogram!("tgi_request_duration").record(total_time.as_secs_f64());
+            metrics::histogram!("tgi_request_generated_tokens").record(generated_tokens as f64);
+            metrics::histogram!("tgi_request_total_tokens")
+                .record((generated_tokens as usize + input_tokens) as f64);
         }
     }