Update arch stats (#250)

This commit is contained in:
Aayush 2024-11-12 15:03:26 -08:00 committed by GitHub
parent 30647fd508
commit 5993e36f22
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 1443 additions and 17 deletions

View file

@ -80,7 +80,7 @@ impl RecordingMetric for Gauge {}
/// For offset deltas
impl IncrementingMetric for Gauge {}
#[derive(Copy, Clone)]
#[derive(Copy, Clone, Debug)]
pub struct Histogram {
id: u32,
}

View file

@ -5,6 +5,7 @@ use common::llm_providers::LlmProviders;
use common::ratelimit;
use common::stats::Counter;
use common::stats::Gauge;
use common::stats::Histogram;
use log::debug;
use proxy_wasm::traits::*;
use proxy_wasm::types::*;
@ -16,6 +17,10 @@ use std::rc::Rc;
pub struct WasmMetrics {
pub active_http_calls: Gauge,
pub ratelimited_rq: Counter,
pub time_to_first_token: Histogram,
pub request_latency: Histogram,
pub output_sequence_length: Histogram,
pub input_sequence_length: Histogram,
}
impl WasmMetrics {
@ -23,6 +28,10 @@ impl WasmMetrics {
WasmMetrics {
active_http_calls: Gauge::new(String::from("active_http_calls")),
ratelimited_rq: Counter::new(String::from("ratelimited_rq")),
time_to_first_token: Histogram::new(String::from("time_to_first_token")),
request_latency: Histogram::new(String::from("request_latency")),
output_sequence_length: Histogram::new(String::from("output_sequence_length")),
input_sequence_length: Histogram::new(String::from("input_sequence_length")),
}
}
}

View file

@ -20,7 +20,10 @@ use proxy_wasm::types::*;
use std::num::NonZero;
use std::rc::Rc;
use common::stats::IncrementingMetric;
use common::stats::{IncrementingMetric, RecordingMetric};
use proxy_wasm::hostcalls::get_current_time;
use std::time::{Duration, SystemTime};
pub struct StreamContext {
context_id: u32,
@ -32,6 +35,8 @@ pub struct StreamContext {
llm_providers: Rc<LlmProviders>,
llm_provider: Option<Rc<LlmProvider>>,
request_id: Option<String>,
start_time: Option<SystemTime>,
ttft_duration: Option<Duration>, // Store the duration directly
}
impl StreamContext {
@ -46,6 +51,8 @@ impl StreamContext {
llm_providers,
llm_provider: None,
request_id: None,
start_time: None,
ttft_duration: None,
}
}
fn llm_provider(&self) -> &LlmProvider {
@ -120,16 +127,27 @@ impl StreamContext {
model: &str,
json_string: &str,
) -> Result<(), ratelimit::Error> {
// Tokenize and record token count.
let token_count = tokenizer::token_count(model, json_string).unwrap_or(0);
// Record the token count to metrics.
self.metrics
.input_sequence_length
.record(token_count as u64);
log::debug!("Recorded input token count: {}", token_count);
// Check if rate limiting needs to be applied.
if let Some(selector) = self.ratelimit_selector.take() {
// Tokenize and Ratelimit.
if let Ok(token_count) = tokenizer::token_count(model, json_string) {
ratelimit::ratelimits(None).read().unwrap().check_limit(
model.to_owned(),
selector,
NonZero::new(token_count as u32).unwrap(),
)?;
}
log::debug!("Applying ratelimit for model: {}", model);
ratelimit::ratelimits(None).read().unwrap().check_limit(
model.to_owned(),
selector,
NonZero::new(token_count as u32).unwrap(),
)?;
} else {
log::debug!("No rate limit applied for model: {}", model);
}
Ok(())
}
}
@ -158,6 +176,12 @@ impl HttpContext for StreamContext {
);
self.request_id = self.get_http_request_header(REQUEST_ID_HEADER);
//start the timing for the request using get_current_time()
let current_time = get_current_time().unwrap();
self.start_time = Some(current_time);
self.ttft_duration = None;
Action::Continue
}
@ -226,9 +250,15 @@ impl HttpContext for StreamContext {
});
}
// only use the tokens from the messages, excluding the metadata and json tags
let input_tokens_str = deserialized_body
.messages
.iter()
.fold(String::new(), |acc, m| {
acc + " " + m.content.as_ref().unwrap_or(&String::new())
});
// enforce ratelimits on ingress
if let Err(e) =
self.enforce_ratelimits(&deserialized_body.model, &chat_completion_request_str)
if let Err(e) = self.enforce_ratelimits(&deserialized_body.model, input_tokens_str.as_str())
{
self.send_server_error(
ServerError::ExceededRatelimit(e),
@ -254,10 +284,33 @@ impl HttpContext for StreamContext {
return Action::Continue;
}
let body = if self.streaming_response {
if end_of_stream && body_size == 0 {
return Action::Continue;
let current_time = get_current_time().unwrap();
if end_of_stream && body_size == 0 {
// All streaming responses end with bytes=0 and end_stream=true
// Record the latency for the request
if let Some(start_time) = self.start_time {
match current_time.duration_since(start_time) {
Ok(duration) => {
// Convert the duration to milliseconds
let duration_ms = duration.as_millis();
debug!("Total latency: {} milliseconds", duration_ms);
// Record the latency to the latency histogram
self.metrics.request_latency.record(duration_ms as u64);
}
Err(e) => {
warn!("SystemTime error: {:?}", e);
}
}
}
// Record the output sequence length
self.metrics
.output_sequence_length
.record(self.response_tokens as u64);
return Action::Continue;
}
let body = if self.streaming_response {
let chunk_start = 0;
let chunk_size = body_size;
debug!(
@ -344,6 +397,26 @@ impl HttpContext for StreamContext {
}
};
self.response_tokens += token_count;
// Compute TTFT if not already recorded
if self.ttft_duration.is_none() {
if let Some(start_time) = self.start_time {
let current_time = get_current_time().unwrap();
match current_time.duration_since(start_time) {
Ok(duration) => {
let duration_ms = duration.as_millis();
debug!("Time to First Token (TTFT): {} milliseconds", duration_ms);
self.ttft_duration = Some(duration);
self.metrics.time_to_first_token.record(duration_ms as u64);
}
Err(e) => {
warn!("SystemTime error: {:?}", e);
}
}
} else {
warn!("Start time was not recorded");
}
}
} else {
debug!("non streaming response");
let chat_completions_response: ChatCompletionsResponse =

View file

@ -51,6 +51,8 @@ fn request_headers_expectations(module: &mut Tester, http_context: i32) {
.expect_log(Some(LogLevel::Debug), None)
.expect_get_header_map_value(Some(MapType::HttpRequestHeaders), Some("x-request-id"))
.returning(None)
.expect_get_current_time_nanos()
.returning(Some(0))
.execute_and_expect(ReturnType::Action(Action::Continue))
.unwrap();
}
@ -72,6 +74,10 @@ fn setup_filter(module: &mut Tester, config: &str) -> i32 {
.call_proxy_on_context_create(filter_context, 0)
.expect_metric_creation(MetricType::Gauge, "active_http_calls")
.expect_metric_creation(MetricType::Counter, "ratelimited_rq")
.expect_metric_creation(MetricType::Histogram, "time_to_first_token")
.expect_metric_creation(MetricType::Histogram, "request_latency")
.expect_metric_creation(MetricType::Histogram, "output_sequence_length")
.expect_metric_creation(MetricType::Histogram, "input_sequence_length")
.execute_and_expect(ReturnType::None)
.unwrap();
@ -209,6 +215,9 @@ fn llm_gateway_successful_request_to_open_ai_chat_completions() {
.returning(Some(chat_completions_request_body))
.expect_log(Some(LogLevel::Trace), None)
.expect_log(Some(LogLevel::Debug), None)
.expect_metric_record("input_sequence_length", 21)
.expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Debug), None)
.expect_set_buffer_bytes(Some(BufferType::HttpRequestBody), None)
.execute_and_expect(ReturnType::Action(Action::Continue))
@ -310,7 +319,7 @@ fn llm_gateway_request_ratelimited() {
},\
{\
\"role\": \"user\",\
\"content\": \"Compose a poem that explains the concept of recursion in programming. Compose a poem that explains the concept of recursion in programming. Compose a poem that explains the concept of recursion in programming. And also summarize it how a 4th graded would understand it.\"\
\"content\": \"Compose a poem that explains the concept of recursion in programming. Compose a poem that explains the concept of recursion in programming. Compose a poem that explains the concept of recursion in programming. And also summarize it how a 4th graded would understand it. Compose a poem that explains the concept of recursion in programming. And also summarize it how a 4th graded would understand it.\"\
}\
],\
\"model\": \"gpt-4\"\
@ -327,9 +336,11 @@ fn llm_gateway_request_ratelimited() {
// The actual call is not important in this test, we just need to grab the token_id
.expect_log(Some(LogLevel::Trace), None)
.expect_log(Some(LogLevel::Debug), None)
.expect_metric_record("input_sequence_length", 107)
.expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Debug), None)
// .expect_metric_increment("active_http_calls", 1)
.expect_send_local_response(
Some(StatusCode::TOO_MANY_REQUESTS.as_u16().into()),
None,
@ -391,6 +402,9 @@ fn llm_gateway_request_not_ratelimited() {
// The actual call is not important in this test, we just need to grab the token_id
.expect_log(Some(LogLevel::Trace), None)
.expect_log(Some(LogLevel::Debug), None)
.expect_metric_record("input_sequence_length", 29)
.expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Debug), None)
.expect_set_buffer_bytes(Some(BufferType::HttpRequestBody), None)
.execute_and_expect(ReturnType::Action(Action::Continue))