Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Stricter tracing inclusions #738

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Stricter tracing inclusions #738

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
70a7d7b
skipping a lot more input fields to the default tracing
noyoshi Jan 10, 2025
89ca296
added useful header logging
noyoshi Jan 10, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 7 additions & 7 deletions router/src/infer.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -276,7 +276,7 @@ impl Infer {
}

/// Add a new request to the queue and return a stream of InferStreamResponse
#[instrument(skip(self))]
#[instrument(skip_all,fields(parameters = ? request.parameters))]
pub(crate) async fn generate_stream(
&self,
request: GenerateRequest,
Expand Down Expand Up @@ -400,7 +400,7 @@ impl Infer {
}

/// Add a new request to the queue and return a InferResponse
#[instrument(skip(self))]
#[instrument(skip_all,fields(parameters = ? request.parameters))]
pub(crate) async fn generate(
&self,
request: GenerateRequest,
Expand Down Expand Up @@ -488,7 +488,7 @@ impl Infer {
}
}

#[instrument(skip(self))]
#[instrument(skip_all,fields(parameters = ? request.parameters))]
pub(crate) async fn embed(&self, request: EmbedRequest) -> Result<EmbedResponse, InferError> {
// Limit concurrent requests by acquiring a permit from the semaphore
let _permit = self
Expand Down Expand Up @@ -618,7 +618,7 @@ impl Infer {
}
}

#[instrument(skip(self))]
#[instrument(skip_all)]
pub(crate) async fn classify(
&self,
request: ClassifyRequest,
Expand Down Expand Up @@ -731,7 +731,7 @@ impl Infer {
}
}

#[instrument(skip(self))]
#[instrument(skip_all)]
pub(crate) async fn classify_batch(
&self,
request: BatchClassifyRequest,
Expand Down Expand Up @@ -861,7 +861,7 @@ impl Infer {
}
}

#[instrument(skip(self))]
#[instrument(skip_all,fields(parameters = ? request.parameters))]
pub(crate) async fn embed_batch(
&self,
request: BatchEmbedRequest,
Expand Down Expand Up @@ -996,7 +996,7 @@ impl Infer {

/// Add best_of new requests to the queue and return a InferResponse of the sequence with
/// the highest log probability per token
#[instrument(skip(self))]
#[instrument(skip_all,fields(parameters = ? request.parameters, best_of, prefix_caching))]
pub(crate) async fn generate_best_of(
&self,
request: GenerateRequest,
Expand Down
38 changes: 35 additions & 3 deletions router/src/server.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,18 @@ example = json ! ({"error": "Input validation error"})),
example = json ! ({"error": "Incomplete generation"})),
)
)]
#[instrument(skip(infer, req))]
#[instrument(
skip_all,
fields(
parameters = ? req.0.parameters,
total_time,
validation_time,
queue_time,
inference_time,
time_per_token,
seed,
)
)]
async fn compat_generate(
default_return_full_text: Extension<bool>,
infer: Extension<Infer>,
Expand All @@ -91,6 +102,10 @@ async fn compat_generate(
req_headers: HeaderMap,
req: Json<CompatGenerateRequest>,
) -> Result<Response, (StatusCode, Json<ErrorResponse>)> {
// Log some useful headers to the span.
let span = tracing::Span::current();
trace_headers(req_headers, &span);

let mut req = req.0;

// default return_full_text given the pipeline_tag
Expand Down Expand Up @@ -147,7 +162,7 @@ example = json ! ({"error": "Input validation error"})),
example = json ! ({"error": "Incomplete generation"})),
)
)]
#[instrument(skip(infer, req))]
#[instrument(skip(infer, req, req_headers))]
async fn completions_v1(
default_return_full_text: Extension<bool>,
infer: Extension<Infer>,
Expand All @@ -158,6 +173,8 @@ async fn completions_v1(
req_headers: HeaderMap,
req: Json<CompletionRequest>,
) -> Result<Response, (StatusCode, Json<ErrorResponse>)> {
let span = tracing::Span::current();
trace_headers(req_headers, &span);
let mut req = req.0;
if req.model == info.model_id.as_str() {
// Allow user to specify the base model, but treat it as an empty adapter_id
Expand Down Expand Up @@ -232,7 +249,7 @@ example = json ! ({"error": "Input validation error"})),
example = json ! ({"error": "Incomplete generation"})),
)
)]
#[instrument(skip(infer, req))]
#[instrument(skip(infer, req, req_headers))]
async fn chat_completions_v1(
default_return_full_text: Extension<bool>,
infer: Extension<Infer>,
Expand All @@ -243,6 +260,8 @@ async fn chat_completions_v1(
req_headers: HeaderMap,
req: Json<ChatCompletionRequest>,
) -> Result<Response, (StatusCode, Json<ErrorResponse>)> {
let span = tracing::Span::current();
trace_headers(req_headers, &span);
let mut req = req.0;
let model_id = info.model_id.clone();
if req.model == info.model_id.as_str() {
Expand Down Expand Up @@ -632,6 +651,7 @@ async fn generate(
mut req: Json<GenerateRequest>,
) -> Result<(HeaderMap, Json<GenerateResponse>), (StatusCode, Json<ErrorResponse>)> {
let span = tracing::Span::current();
trace_headers(req_headers, &span);
let start_time = Instant::now();
metrics::increment_counter!("lorax_request_count");

Expand Down Expand Up @@ -2058,3 +2078,15 @@ async fn tokenize(
))
}
}

fn trace_headers(headers: HeaderMap, span: &tracing::Span) {
headers
.get("x-predibase-tenant")
.map(|value| span.record("x-predibase-tenant", value));
headers
.get("user-agent")
.map(|value| span.record("user-agent", value));
headers
.get("x-b3-traceid")
.map(|value| span.record("x-b3-traceid", value));
}
Loading