huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-12 04:44:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: avoid unwrap and pre allocate future vec

Browse Source
This commit is contained in:
drbh 2024-08-30 16:47:24 +00:00
parent b5dd58f73b
commit ed8c7726ba
1 changed files with 116 additions and 109 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

29
router/src/server.rs
View File

@ -1407,10 +1407,9 @@ async fn vertex_compatibility(
}
// Prepare futures for all instances
let futures: Vec<_> = req
.instances
.iter()
.map(|instance| {
let mut futures = Vec::with_capacity(req.instances.len());
for instance in req.instances.iter() {
let generate_request = match instance {
VertexInstance::Generate(instance) => GenerateRequest {
inputs: instance.inputs.clone(),
@ -1456,7 +1455,7 @@ async fn vertex_compatibility(
Some(temperature) if temperature == 0.0 => (false, None),
other => (true, other),
};
let (inputs, grammar, _using_tools) = prepare_chat_input(
let (inputs, grammar, _using_tools) = match prepare_chat_input(
&infer,
response_format,
tools,
@ -1464,10 +1463,19 @@ async fn vertex_compatibility(
&tool_prompt,
guideline,
messages,
)
.unwrap();
) {
Ok(result) => result,
Err(e) => {
return Err((
StatusCode::BAD_REQUEST,
Json(ErrorResponse {
error: format!("Failed to prepare chat input: {}", e),
error_type: "Input preparation error".to_string(),
}),
));
}
};
// build the request passing some parameters
GenerateRequest {
inputs: inputs.to_string(),
add_special_tokens: false,
@ -1500,7 +1508,7 @@ async fn vertex_compatibility(
let compute_type_clone = compute_type.clone();
let span_clone = span.clone();
async move {
futures.push(async move {
generate_internal(
Extension(infer_clone),
compute_type_clone,
@ -1518,9 +1526,8 @@ async fn vertex_compatibility(
}),
)
})
});
}
})
.collect();
// execute all futures in parallel, collect results, returning early if any error occurs
let results = futures::future::join_all(futures).await;