Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
196 changes: 181 additions & 15 deletions lib/llm/src/http/service/openai.rs
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,7 @@ use crate::request_template::{RequestTemplate, resolve_request_model};
use crate::types::Annotated;
use dynamo_protocols::types::ChatCompletionMessageContent;
use dynamo_protocols::types::ChatCompletionMessageToolCallChunk;
use dynamo_protocols::types::ChatCompletionStreamOptions;
use dynamo_protocols::types::ChatCompletionStreamResponseDelta;
use dynamo_protocols::types::Choice;
use dynamo_protocols::types::responses::{
Expand Down Expand Up @@ -1030,7 +1031,8 @@ async fn handler_completions(
body: Body,
) -> Result<Response, ErrorResponse> {
let body = read_json_request_body(&headers, body).await?;
let mut request: NvCreateCompletionRequest = parse_json_request("completions", &body)?;
let mut request: NvCreateCompletionRequest =
parse_completion_json_request("completions", &body)?;
if *FORCE_INCLUDE_USAGE && request.inner.stream.unwrap_or(false) {
delta_common::force_include_usage(&mut request.inner.stream_options);
}
Expand Down Expand Up @@ -2317,7 +2319,8 @@ async fn handler_chat_completions(
body: Body,
) -> Result<Response, ErrorResponse> {
let body = read_json_request_body(&headers, body).await?;
let mut request: NvCreateChatCompletionRequest = parse_json_request("chat completions", &body)?;
let mut request: NvCreateChatCompletionRequest =
parse_completion_json_request("chat completions", &body)?;
if *FORCE_INCLUDE_USAGE && request.inner.stream.unwrap_or(false) {
delta_common::force_include_usage(&mut request.inner.stream_options);
}
Expand Down Expand Up @@ -2396,6 +2399,56 @@ async fn handler_chat_completions(
response
}

#[derive(Deserialize)]
struct CompletionRequestWithNullableStreamOptions<T> {
#[serde(flatten)]

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Avoid copying escaped request text through another flatten layer

This still buffers the request twice: NvCreate*Request already contains flattened fields, so the additional wrapper traverses Serde's outer Content through ContentRefDeserializer, copying owned escaped strings into another buffer. The wrapper runs even when stream_options is absent. At this head, a 15 MiB decoded prompt containing newlines (18 MiB JSON body) raises peak live parsing allocation from 46 MiB to 61 MiB compared with direct deserialization of the same request type. That is an extra 15 MiB per concurrently parsing request; the plain-text control was effectively unchanged. These are allocator measurements in a debug build, excluding the input body, not RSS or throughput estimates.

Normalize the two booleans in ChatCompletionStreamOptions with #[serde(default, deserialize_with = "super::deserialize_null_as_default")], reuse the existing null-default helper in the protocol crate, and remove this wrapper so both handlers use parse_json_request directly. I validated that change locally with a patched protocol dependency: 750 tests passed, and the null-containing stream-options probe returned identical parsed output with peak parsing allocation back at 46 MiB. Shipping it needs the corresponding protocol release and dependency bump.

Follow-up to the earlier allocation finding: the explicit JSON-value retry is gone, but the extra payload copy remains in this replacement.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the detailed measurements. Agreed—the wrapper should be removed. The protocol-side null handling is already available in dynamo-protocols 6.0.1 through frontend-crate #239, but this PR is still on 5.4.1 and #14755 owns the protocols 6 migration. I’ll wait for #14755 to land, then update this branch, remove the wrapper, restore direct parse_json_request usage in both handlers, and rerun the targeted tests.

request: T,
stream_options: Option<NullableStreamOptions>,
}

#[derive(Deserialize)]
struct NullableStreamOptions {
include_usage: Option<bool>,
continuous_usage_stats: Option<bool>,
}

trait SetStreamOptions {
fn set_stream_options(&mut self, stream_options: Option<ChatCompletionStreamOptions>);
}

impl SetStreamOptions for NvCreateChatCompletionRequest {
fn set_stream_options(&mut self, stream_options: Option<ChatCompletionStreamOptions>) {
self.inner.stream_options = stream_options;
}
}

impl SetStreamOptions for NvCreateCompletionRequest {
fn set_stream_options(&mut self, stream_options: Option<ChatCompletionStreamOptions>) {
self.inner.stream_options = stream_options;
}
}

// The protocol crate requires include_usage and uses non-nullable booleans,
// while clients may omit stream_options flags or serialize them as null.
fn parse_completion_json_request<T>(endpoint: &'static str, body: &[u8]) -> Result<T, ErrorResponse>
where
T: DeserializeOwned + SetStreamOptions,
{
let mut parsed: CompletionRequestWithNullableStreamOptions<T> =
parse_json_request(endpoint, body)?;
parsed
.request
.set_stream_options(
parsed
.stream_options
.map(|options| ChatCompletionStreamOptions {
include_usage: options.include_usage.unwrap_or(false),
continuous_usage_stats: options.continuous_usage_stats.unwrap_or(false),
}),
);
Ok(parsed.request)
}

fn parse_json_request<T>(endpoint: &'static str, body: &[u8]) -> Result<T, ErrorResponse>
where
T: DeserializeOwned,
Expand Down Expand Up @@ -5859,12 +5912,123 @@ mod tests {
assert_eq!(err.0, StatusCode::UNSUPPORTED_MEDIA_TYPE);
}

#[test]
fn test_parse_completion_stream_options_null_flags() {
for (options, expected) in [
(serde_json::json!(null), None),
(serde_json::json!({}), Some((false, false))),
(
serde_json::json!({"continuous_usage_stats": null}),
Some((false, false)),
),
(
serde_json::json!({"continuous_usage_stats": true}),
Some((false, true)),
),
(
serde_json::json!({"include_usage": null}),
Some((false, false)),
),
(
serde_json::json!({"include_usage": true, "continuous_usage_stats": null}),
Some((true, false)),
),
(
serde_json::json!({"include_usage": null, "continuous_usage_stats": true}),
Some((false, true)),
),
] {
let mut payload = serde_json::json!({
"model": "test-model", "stream": true, "stream_options": options,
"messages": [{"role": "user", "content": "hello"}],
});
let chat: NvCreateChatCompletionRequest = parse_completion_json_request(
"chat completions",
&serde_json::to_vec(&payload).unwrap(),
)
.unwrap();
payload.as_object_mut().unwrap().remove("messages");
payload["prompt"] = serde_json::json!("hello");
let completion: NvCreateCompletionRequest = parse_completion_json_request(
"completions",
&serde_json::to_vec(&payload).unwrap(),
)
.unwrap();
crate::engines::ValidateRequest::validate(&chat).unwrap();
crate::engines::ValidateRequest::validate(&completion).unwrap();
for parsed in [chat.inner.stream_options, completion.inner.stream_options] {
assert_eq!(
parsed.map(|opts| (opts.include_usage, opts.continuous_usage_stats)),
expected
);
}
}
}

#[test]
fn test_parse_completion_stream_options_rejects_invalid_types() {
for options in [
serde_json::json!(false),
serde_json::json!({"include_usage": "true", "continuous_usage_stats": null}),
serde_json::json!({"include_usage": null, "continuous_usage_stats": 0}),
] {
let body = serde_json::to_vec(&serde_json::json!({
"model": "test-model", "messages": [{"role": "user", "content": "hello"}],
"prompt": "hello", "stream_options": options,
}))
.unwrap();
assert_eq!(
parse_completion_json_request::<NvCreateChatCompletionRequest>(
"chat completions",
&body
)
.unwrap_err()
.0,
StatusCode::BAD_REQUEST
);
assert_eq!(
parse_completion_json_request::<NvCreateCompletionRequest>("completions", &body)
.unwrap_err()
.0,
StatusCode::BAD_REQUEST
);
}
let body =
br#"{"model":"test-model","messages":42,"stream_options":{"include_usage":null}}"#;
assert_eq!(
parse_completion_json_request::<NvCreateChatCompletionRequest>(
"chat completions",
body
)
.unwrap_err()
.0,
StatusCode::BAD_REQUEST
);
}

#[test]
fn test_parse_completion_stream_options_preserves_duplicate_field_errors() {
let chat_body = br#"{
"model":"first-model",
"model":"second-model",
"messages":[{"role":"user","content":"hello"}],
"stream_options":{"include_usage":null}
}"#;
let chat_error = parse_completion_json_request::<NvCreateChatCompletionRequest>(
"chat completions",
chat_body,
)
.unwrap_err();
assert_eq!(chat_error.0, StatusCode::BAD_REQUEST);
assert!(chat_error.1.message.contains("duplicate field `model`"));
}

#[test]
fn test_parse_chat_completion_request_escapes_control_chars_in_strings() {
let body = b"{\"model\":\"test-model\",\"messages\":[{\"role\":\"user\",\"content\":\"log \x1b[33mPK\x03\x04\"}]}";

let request: NvCreateChatCompletionRequest =
parse_json_request("chat completions", body).expect("request should parse");
parse_completion_json_request("chat completions", body).expect("request should parse");

let message = request
.inner
Expand All @@ -5885,7 +6049,7 @@ mod tests {
let body = b"{\"model\":\"test-model\",\"messages\":[{\"role\":\"user\",\"content\":\"raw \xff data\"}]}";

let request: NvCreateChatCompletionRequest =
parse_json_request("chat completions", body).expect("request should parse");
parse_completion_json_request("chat completions", body).expect("request should parse");

let message = request
.inner
Expand All @@ -5906,7 +6070,7 @@ mod tests {
let body = b"{\"model\":\"test-model\",\"messages\":[{\"role\":\"user\",\"content\":\"slash \\\nnext\"}]}";

let request: NvCreateChatCompletionRequest =
parse_json_request("chat completions", body).expect("request should parse");
parse_completion_json_request("chat completions", body).expect("request should parse");

let message = request
.inner
Expand All @@ -5926,11 +6090,13 @@ mod tests {
fn test_parse_chat_completion_request_keeps_schema_errors() {
let body = br#"{"model":"test-model","messages":[{"role":"assistant","content":[{"type":"thinking","thinking":"working"}]}]}"#;

let err =
match parse_json_request::<NvCreateChatCompletionRequest>("chat completions", body) {
Ok(_) => panic!("schema should still fail"),
Err(err) => err,
};
let err = match parse_completion_json_request::<NvCreateChatCompletionRequest>(
"chat completions",
body,
) {
Ok(_) => panic!("schema should still fail"),
Err(err) => err,
};

assert_eq!(err.0, StatusCode::BAD_REQUEST);
assert!(
Expand All @@ -5947,7 +6113,7 @@ mod tests {
let body = br#"{"model":"test-model","messages":[{"role":"user","content":[{"type":"image_url","image_url":{"url":""},"uuid":"image-42"}]}]}"#;

let request: NvCreateChatCompletionRequest =
parse_json_request("chat completions", body).expect("request should parse");
parse_completion_json_request("chat completions", body).expect("request should parse");
let request = serde_json::to_value(request).expect("request should serialize");
assert_eq!(request["messages"][0]["content"][0]["uuid"], "image-42");
assert_eq!(
Expand All @@ -5967,7 +6133,7 @@ mod tests {
);

let request: NvCreateChatCompletionRequest =
parse_json_request("chat completions", body.as_bytes())
parse_completion_json_request("chat completions", body.as_bytes())
.expect("request should parse");
let request = serde_json::to_value(request).expect("request should serialize");
assert_eq!(request["messages"][0]["content"][0]["uuid"], uuid);
Expand All @@ -5983,7 +6149,7 @@ mod tests {
let body = b"{\"model\":\"test-model\",\"messages\":[{\"role\":\"user\",\"content\":[{\"type\":\"text\",\"text\":\"raw \xff \x1b data\"},{\"type\":\"image_url\",\"image_url\":{\"url\":\"\"},\"uuid\":\"image-42\"}]}]}";

let request: NvCreateChatCompletionRequest =
parse_json_request("chat completions", body).expect("request should parse");
parse_completion_json_request("chat completions", body).expect("request should parse");
let request = serde_json::to_value(request).expect("request should serialize");
assert_eq!(
request["messages"][0]["content"][0]["text"],
Expand All @@ -6002,7 +6168,7 @@ mod tests {
b"{\"model\":\"test-model\",\"prompt\":\"log \x1b[33mPK\x03\x04\",\"max_tokens\":1}";

let request: NvCreateCompletionRequest =
parse_json_request("completions", body).expect("request should parse");
parse_completion_json_request("completions", body).expect("request should parse");

let Prompt::String(prompt) = &request.inner.prompt else {
panic!("expected string prompt");
Expand All @@ -6015,7 +6181,7 @@ mod tests {
let body = b"{\"model\":\"test-model\",\"prompt\":\"raw \xff data\",\"max_tokens\":1}";

let request: NvCreateCompletionRequest =
parse_json_request("completions", body).expect("request should parse");
parse_completion_json_request("completions", body).expect("request should parse");

let Prompt::String(prompt) = &request.inner.prompt else {
panic!("expected string prompt");
Expand Down
7 changes: 6 additions & 1 deletion lib/llm/src/protocols/openai/chat_completions.rs
Original file line number Diff line number Diff line change
Expand Up @@ -124,7 +124,12 @@ pub struct NvCreateChatCompletionRequest {
pub return_tokens_as_token_ids: Option<bool>,

/// Catch-all for unsupported fields - checked during validation
#[serde(flatten, default, skip_serializing)]
#[serde(
flatten,
default,
skip_serializing,
deserialize_with = "validate::deserialize_extra_fields"
)]
pub unsupported_fields: std::collections::HashMap<String, serde_json::Value>,
}

Expand Down
7 changes: 6 additions & 1 deletion lib/llm/src/protocols/openai/completions.rs
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,12 @@ pub struct NvCreateCompletionRequest {
pub return_tokens_as_token_ids: Option<bool>,

/// Catch-all for unsupported fields - checked during validation
#[serde(flatten, default, skip_serializing)]
#[serde(
flatten,
default,
skip_serializing,
deserialize_with = "validate::deserialize_extra_fields"
)]
pub unsupported_fields: std::collections::HashMap<String, serde_json::Value>,
}

Expand Down
Loading
Loading