f57123aa2d · tduncan/Nvidia-Optimization-Stack-VLLM

rust/proto/inference.proto

+2

@@ -45,6 +45,8 @@ message GenerateRequest {

45 uint32 truncate_prompt_tokens = 11;

46

47 int32 priority = 12;

48 +

49 + optional string session_id = 13;

50 }

51

52 message RandomSampling {

rust/src/chat/src/lib.rs

+1

@@ -168,6 +168,7 @@ impl ChatRequestProcessor {

168 cache_salt: request.cache_salt,

169 add_special_tokens: request.add_special_tokens,

170 data_parallel_rank: request.data_parallel_rank,

171 + session_id: request.session_id,

172 reasoning_parser_kwargs,

173 lora_request: request.lora_request,

174 arrival_time: Some(arrival_time),

rust/src/chat/src/request.rs

+4

@@ -491,6 +491,9 @@ pub struct ChatRequest {

491 /// Override data parallel rank.

492 #[serde(default)]

493 pub data_parallel_rank: Option<u32>,

494 + /// Stable session identity shared by related requests.

495 + #[serde(default)]

496 + pub session_id: Option<String>,

497 /// LoRA adapter selected for this request.

498 #[serde(default)]

499 pub lora_request: Option<LoraRequest>,

@@ -514,6 +517,7 @@ impl ChatRequest {

517 cache_salt: None,

518 add_special_tokens: false,

519 data_parallel_rank: None,

520 + session_id: None,

521 lora_request: None,

522 }

523 }

rust/src/engine-core-client/src/protocol/request.rs

+6 -1

@@ -124,6 +124,9 @@ pub struct EngineCoreRequest {

124 /// standard `request_finished` hook.

125 #[serde(default)]

126 pub abort_immediately: bool,

127 + /// Stable session identity shared by related requests.

128 + #[serde(default)]

129 + pub session_id: Option<String>,

130 }

131

132 impl EngineCoreRequest {

@@ -175,6 +178,7 @@ mod tests {

178 }),

179 arrival_time: 1234.5,

180 client_index: 7,

181 + session_id: Some("session-1".to_string()),

182 ..EngineCoreRequest::default()

183 };

184

@@ -185,12 +189,13 @@ mod tests {

189 other => panic!("expected array, got {other:?}"),

190 };

191

188 - assert_eq!(array.len(), 20);

192 + assert_eq!(array.len(), 21);

193 assert_eq!(array[0], Value::from("req-1"));

194 assert_eq!(array[2], Value::Nil);

195 assert_eq!(array[4], Value::Nil);

196 assert_eq!(array[10], Value::Nil);

197 assert_eq!(array[11], Value::from(7));

198 + assert_eq!(array[20], Value::from("session-1"));

199 }

200

201 #[test]

rust/src/engine-core-client/src/tests/client.rs

+1

@@ -162,6 +162,7 @@ fn sample_request_with_id(request_id: &str) -> EngineCoreRequest {

162 ..EngineCoreSamplingParams::for_test()

163 }),

164 arrival_time: 42.5,

165 + session_id: Some("session-1".to_string()),

166 ..EngineCoreRequest::default()

167 }

168 }

rust/src/engine-core-client/src/tests/python_compat.py

+2

@@ -75,6 +75,7 @@ class EngineCoreRequest(

75 reasoning_ended: bool | None = None

76 reasoning_parser_kwargs: dict[str, object] | None = None

77 abort_immediately: bool = False

78 + session_id: str | None = None

79

80

81 class EngineCoreOutput(

@@ -137,6 +138,7 @@ request = EngineCoreRequest(

138 pooling_params=None,

139 arrival_time=42.5,

140 client_index=0,

141 + session_id="session-1",

142 )

143

144 # All defaults -> empty map. Regression guard for the sparse-map decode.

rust/src/llm/examples/external_engine_smoke.rs

+1

@@ -59,6 +59,7 @@ fn build_request(request_id: String, max_tokens: u32) -> GenerateRequest {

59 trace_headers: None,

60 priority: 0,

61 data_parallel_rank: None,

62 + session_id: None,

63 reasoning_parser_kwargs: None,

64 lora_request: None,

65 }

rust/src/llm/src/request.rs

+6

@@ -46,6 +46,8 @@ pub struct GenerateRequest {

46 pub priority: i32,

47 /// Optional data-parallel rank override for routing this request.

48 pub data_parallel_rank: Option<u32>,

49 + /// Stable session identity shared by related requests.

50 + pub session_id: Option<String>,

51 /// Optional reasoning-parser kwargs forwarded to engine-side structured

52 /// output logic.

53 pub reasoning_parser_kwargs: Option<ReasoningParserKwargs>,

@@ -76,6 +78,7 @@ impl GenerateRequest {

78 trace_headers,

79 priority,

80 data_parallel_rank,

81 + session_id,

82 reasoning_parser_kwargs,

83 lora_request,

84 } = self;

@@ -105,6 +108,7 @@ impl GenerateRequest {

108 priority,

109 trace_headers,

110 resumable: false,

111 + session_id,

112 external_req_id: Some(external_request_id),

113 // Rust parser doesn't expose this information, leave it unset and let the

114 // reasoning logic in engine-sided structured output manager handle it.

@@ -150,6 +154,7 @@ mod tests {

154 )])),

155 priority: 3,

156 data_parallel_rank: Some(2),

157 + session_id: Some("session-1".to_string()),

158 reasoning_parser_kwargs: Some(ReasoningParserKwargs {

159 chat_template_kwargs: [(

160 "chat_template_kwargs".to_string(),

@@ -177,6 +182,7 @@ mod tests {

182 assert_eq!(request.arrival_time, 42.5);

183 assert_eq!(request.cache_salt.as_deref(), Some("salt"));

184 assert_eq!(request.data_parallel_rank, Some(2));

185 + assert_eq!(request.session_id.as_deref(), Some("session-1"));

186 assert_eq!(

187 request.trace_headers,

188 Some(BTreeMap::from([(

rust/src/llm/tests/generate.rs

+1

@@ -188,6 +188,7 @@ fn sample_generate_request(request_id: &str, max_tokens: u32) -> GenerateRequest

188 trace_headers: None,

189 priority: 0,

190 data_parallel_rank: None,

191 + session_id: None,

192 reasoning_parser_kwargs: None,

193 lora_request: None,

194 }

rust/src/server/src/grpc/convert.rs

+2

@@ -53,6 +53,7 @@ pub fn to_text_request(

53 } else {

54 req.request_id

55 };

56 + let session_id = req.session_id.filter(|s| !s.is_empty());

57

58 let sampling = req.sampling.as_ref();

59 let decoding = req.decoding.as_ref();

@@ -100,6 +101,7 @@ pub fn to_text_request(

101 cache_salt: kv.map(|k| &k.cache_salt).filter(|s| !s.is_empty()).cloned(),

102 add_special_tokens: true,

103 data_parallel_rank: None,

104 + session_id,

105 reasoning_parser_kwargs: None,

106 lora_request: None,

107 arrival_time: None,

rust/src/server/src/routes/inference/generate/convert.rs

+1

@@ -75,6 +75,7 @@ pub(super) fn prepare_generate_request(

75 cache_salt: request.cache_salt,

76 add_special_tokens: false,

77 data_parallel_rank: ctx.data_parallel_rank,

78 + session_id: ctx.session_id,

79 reasoning_parser_kwargs: None,

80 lora_request: lora_resolution.lora_request.clone(),

81 arrival_time: None,

rust/src/server/src/routes/openai/chat_completions/convert.rs

+50

@@ -18,6 +18,7 @@ use crate::routes::openai::utils::types::{

18 };

19 use crate::utils::{

20 ResolvedRequestContext, convert_logit_bias, merge_ec_transfer_params, merge_kv_transfer_params,

21 + resolve_session_id,

22 };

23

24 /// Lowered chat request plus the public response metadata carried by every SSE

@@ -123,6 +124,11 @@ pub(super) fn prepare_chat_request(

124 request.response_format.as_ref(),

125 &request.structured_outputs,

126 )?;

127 + let session_id = resolve_session_id(

128 + &ctx,

129 + request.session_id.as_deref(),

130 + request.vllm_xargs.as_ref(),

131 + );

132

133 let chat_request = ChatRequest {

134 request_id: request_id.clone(),

@@ -177,6 +183,7 @@ pub(super) fn prepare_chat_request(

183 cache_salt: request.cache_salt,

184 add_special_tokens: request.add_special_tokens,

185 data_parallel_rank: ctx.data_parallel_rank,

186 + session_id,

187 lora_request: lora_resolution.lora_request.clone(),

188 };

189

@@ -1241,6 +1248,49 @@ mod tests {

1248 assert_eq!(prepared.chat_request.data_parallel_rank, Some(7));

1249 }

1250

1251 + #[test]

1252 + fn prepare_chat_request_threads_header_session_id() {

1253 + let mut headers = HeaderMap::new();

1254 + headers.insert("X-Session-ID", "header-session".parse().unwrap());

1255 + let prepared = prepare_chat_request(

1256 + base_request(),

1257 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),

1258 + request_context(&headers, None),

1259 + )

1260 + .expect("request is valid");

1261 + assert_eq!(

1262 + prepared.chat_request.session_id.as_deref(),

1263 + Some("header-session")

1264 + );

1265 + }

1266 +

1267 + #[test]

1268 + fn prepare_chat_request_ignores_correlation_header() {

1269 + let mut headers = HeaderMap::new();

1270 + headers.insert("X-Correlation-ID", "correlation-header".parse().unwrap());

1271 + let prepared = prepare_chat_request(

1272 + base_request(),

1273 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),

1274 + request_context(&headers, None),

1275 + )

1276 + .expect("request is valid");

1277 + assert_eq!(prepared.chat_request.session_id, None);

1278 + }

1279 +

1280 + #[test]

1281 + fn prepare_chat_request_ignores_empty_session_header_without_fallback() {

1282 + let mut headers = HeaderMap::new();

1283 + headers.insert("X-Session-ID", "".parse().unwrap());

1284 + headers.insert("X-Correlation-ID", "correlation-header".parse().unwrap());

1285 + let prepared = prepare_chat_request(

1286 + base_request(),

1287 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),

1288 + request_context(&headers, None),

1289 + )

1290 + .expect("request is valid");

1291 + assert_eq!(prepared.chat_request.session_id, None);

1292 + }

1293 +

1294 #[test]

1295 fn prepare_chat_request_leaves_data_parallel_rank_none_when_absent() {

1296 let prepared = prepare_chat_request(

rust/src/server/src/routes/openai/chat_completions/types.rs

+4

@@ -222,6 +222,9 @@ pub struct ChatCompletionRequest {

222 /// External request ID used for response correlation.

223 pub request_id: Option<String>,

224

225 + /// Stable session identity shared by related requests.

226 + pub session_id: Option<String>,

227 +

228 /// Tokens represented as strings of the form 'token_id:{token_id}' in

229 /// logprobs

230 pub return_tokens_as_token_ids: Option<bool>,

@@ -301,6 +304,7 @@ impl Default for ChatCompletionRequest {

304 structured_outputs: None,

305 priority: None,

306 request_id: None,

307 + session_id: None,

308 return_tokens_as_token_ids: None,

309 return_token_ids: None,

310 cache_salt: None,

rust/src/server/src/routes/openai/completions/convert.rs

+77

@@ -12,6 +12,7 @@ use crate::routes::openai::completions::validate;

12 use crate::routes::openai::utils::structured_outputs::convert_from_response_format_value;

13 use crate::utils::{

14 ResolvedRequestContext, convert_logit_bias, merge_ec_transfer_params, merge_kv_transfer_params,

15 + resolve_session_id,

16 };

17

18 /// Lowered completion request plus the public response metadata carried by

@@ -103,6 +104,11 @@ pub(super) fn prepare_completion_request(

104

105 let structured_outputs =

106 convert_from_response_format_value(&request.response_format, &request.structured_outputs)?;

107 + let session_id = resolve_session_id(

108 + &ctx,

109 + request.session_id.as_deref(),

110 + request.vllm_xargs.as_ref(),

111 + );

112

113 let text_request = TextRequest {

114 request_id: request_id.clone(),

@@ -147,6 +153,7 @@ pub(super) fn prepare_completion_request(

153 cache_salt: request.cache_salt,

154 add_special_tokens: request.add_special_tokens,

155 data_parallel_rank: ctx.data_parallel_rank,

156 + session_id,

157 reasoning_parser_kwargs: None,

158 lora_request: lora_resolution.lora_request.clone(),

159 arrival_time: None,

@@ -586,6 +593,76 @@ mod tests {

593 assert_eq!(prepared.text_request.data_parallel_rank, Some(3));

594 }

595

596 + #[test]

597 + fn prepare_completion_request_threads_body_session_id() {

598 + let request: CompletionRequest = serde_json::from_value(json!({

599 + "model": "Qwen/Qwen1.5-0.5B-Chat",

600 + "prompt": "hello",

601 + "stream": false,

602 + "session_id": "body-session",

603 + "vllm_xargs": {"session_id": "xargs-session"},

604 + }))

605 + .expect("parse request");

606 +

607 + let mut headers = HeaderMap::new();

608 + headers.insert("X-Session-ID", "header-session".parse().unwrap());

609 + let prepared = prepare_completion_request(

610 + request,

611 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),

612 + request_context(&headers, None),

613 + &test_tokenizer(),

614 + )

615 + .expect("prepare");

616 + assert_eq!(

617 + prepared.text_request.session_id.as_deref(),

618 + Some("body-session")

619 + );

620 + }

621 +

622 + #[test]

623 + fn prepare_completion_request_uses_vllm_xargs_session_id_fallback() {

624 + let request: CompletionRequest = serde_json::from_value(json!({

625 + "model": "Qwen/Qwen1.5-0.5B-Chat",

626 + "prompt": "hello",

627 + "stream": false,

628 + "vllm_xargs": {"session_id": "xargs-session"},

629 + }))

630 + .expect("parse request");

631 +

632 + let prepared = prepare_completion_request(

633 + request,

634 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),

635 + ResolvedRequestContext::default(),

636 + &test_tokenizer(),

637 + )

638 + .expect("prepare");

639 + assert_eq!(

640 + prepared.text_request.session_id.as_deref(),

641 + Some("xargs-session")

642 + );

643 + }

644 +

645 + #[test]

646 + fn prepare_completion_request_ignores_empty_and_non_string_session_id_values() {

647 + let request: CompletionRequest = serde_json::from_value(json!({

648 + "model": "Qwen/Qwen1.5-0.5B-Chat",

649 + "prompt": "hello",

650 + "stream": false,

651 + "session_id": "",

652 + "vllm_xargs": {"session_id": 7},

653 + }))

654 + .expect("parse request");

655 +

656 + let prepared = prepare_completion_request(

657 + request,

658 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),

659 + ResolvedRequestContext::default(),

660 + &test_tokenizer(),

661 + )

662 + .expect("prepare");

663 + assert_eq!(prepared.text_request.session_id, None);

664 + }

665 +

666 #[test]

667 fn prepare_completion_request_leaves_data_parallel_rank_none_when_absent() {

668 let request: CompletionRequest = serde_json::from_value(json!({

rust/src/server/src/routes/openai/completions/types.rs

+3

@@ -164,6 +164,9 @@ pub struct CompletionRequest {

164 /// External request ID used for response correlation.

165 pub request_id: Option<String>,

166

167 + /// Stable session identity shared by related requests.

168 + pub session_id: Option<String>,

169 +

170 /// Tokens represented as strings of the form 'token_id:{token_id}' in

171 /// logprobs

172 pub return_tokens_as_token_ids: Option<bool>,

rust/src/server/src/routes/tokenize/types.rs

+1

@@ -95,6 +95,7 @@ impl TokenizeChatRequest {

95 cache_salt: None,

96 add_special_tokens: self.add_special_tokens,

97 data_parallel_rank: None,

98 + session_id: None,

99 lora_request: None,

100 })

101 }

rust/src/server/src/utils.rs

+27 -2

@@ -15,6 +15,7 @@ use crate::error::ApiError;

15 pub struct ResolvedRequestContext {

16 pub request_id: String,

17 pub data_parallel_rank: Option<u32>,

18 + pub session_id: Option<String>,

19 }

20

21 /// Return the current Unix timestamp in seconds for OpenAI response objects.

@@ -66,6 +67,24 @@ pub fn merge_ec_transfer_params(

67 xargs

68 }

69

70 +pub fn resolve_session_id(

71 + ctx: &ResolvedRequestContext,

72 + request_session_id: Option<&str>,

73 + xargs: Option<&HashMap<String, Value>>,

74 +) -> Option<String> {

75 + request_session_id

76 + .filter(|value| !value.is_empty())

77 + .map(str::to_owned)

78 + .or_else(|| ctx.session_id.clone())

79 + .or_else(|| {

80 + xargs

81 + .and_then(|map| map.get("session_id"))

82 + .and_then(Value::as_str)

83 + .filter(|value| !value.is_empty())

84 + .map(str::to_owned)

85 + })

86 +}

87 +

88 /// Convert OpenAI-style `logit_bias` with string token-ID keys into the

89 /// internal `HashMap<u32, f32>` representation, validating that every key

90 /// parses as a `u32`.

@@ -91,8 +110,8 @@ pub fn convert_logit_bias(

110 .transpose()

111 }

112

94 -/// Extract common request metadata from HTTP headers: the external request ID

95 -/// and the optional data-parallel rank used for engine routing.

113 +/// Extract common request metadata from HTTP headers: the external request ID,

114 +/// session ID, and the optional data-parallel rank used for engine routing.

115 pub fn resolve_request_context(

116 headers: &HeaderMap,

117 request_id: Option<&str>,

@@ -106,10 +125,16 @@ pub fn resolve_request_context(

125 // Extract request id from header.

126 let request_id_header = headers.get("X-Request-Id").and_then(|value| value.to_str().ok());

127 let request_id = resolve_base_request_id(request_id_header, request_id);

128 + let session_id = headers

129 + .get("X-Session-ID")

130 + .and_then(|value| value.to_str().ok())

131 + .filter(|value| !value.is_empty())

132 + .map(str::to_owned);

133

134 ResolvedRequestContext {

135 request_id,

136 data_parallel_rank,

137 + session_id,

138 }

139 }

140

rust/src/text/src/lower.rs

+1

@@ -59,6 +59,7 @@ pub fn lower_text_request(

59 cache_salt: request.cache_salt.clone(),

60 priority: request.priority,

61 data_parallel_rank: request.data_parallel_rank,

62 + session_id: request.session_id.clone(),

63 reasoning_parser_kwargs: request.reasoning_parser_kwargs.clone(),

64 lora_request: request.lora_request.clone(),

65 arrival_time: request.arrival_time,

rust/src/text/src/request.rs

+4

@@ -183,6 +183,9 @@ pub struct TextRequest {

183 /// Override data parallel rank.

184 #[serde(default)]

185 pub data_parallel_rank: Option<u32>,

186 + /// Stable session identity shared by related requests.

187 + #[serde(default)]

188 + pub session_id: Option<String>,

189 /// Optional reasoning-parser kwargs forwarded to engine-side structured

190 /// output logic.

191 #[serde(default)]

@@ -212,6 +215,7 @@ impl TextRequest {

215 cache_salt: None,

216 add_special_tokens: false,

217 data_parallel_rank: None,

218 + session_id: None,

219 reasoning_parser_kwargs: None,

220 lora_request: None,

221 arrival_time: None,

tests/entrypoints/openai/test_session_id.py new

+97

@@ -0,0 +1,97 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +import pytest

5 +from starlette.requests import Request

6 +

7 +from vllm.entrypoints.generate.base.serving import GenerateBaseServing

8 +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest

9 +from vllm.entrypoints.openai.completion.protocol import CompletionRequest

10 +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest

11 +

12 +

13 +def _raw_request(headers: dict[str, str]) -> Request:

14 + return Request(

15 + {

16 + "type": "http",

17 + "method": "POST",

18 + "path": "/v1/chat/completions",

19 + "headers": [

20 + (key.lower().encode("latin-1"), value.encode("latin-1"))

21 + for key, value in headers.items()

22 + ],

23 + }

24 + )

25 +

26 +

27 +@pytest.mark.parametrize(

28 + "openai_request",

29 + [

30 + ChatCompletionRequest(

31 + model="test-model",

32 + messages=[{"role": "user", "content": "hi"}],

33 + ),

34 + CompletionRequest(model="test-model", prompt="hi"),

35 + ResponsesRequest(model="test-model", input="hi"),

36 + ],

37 +)

38 +def test_get_session_id_accepts_body_field(openai_request):

39 + openai_request.session_id = "body-session"

40 +

41 + session_id = GenerateBaseServing._get_session_id(

42 + openai_request,

43 + _raw_request({"X-Session-ID": "header-session"}),

44 + )

45 +

46 + assert session_id == "body-session"

47 +

48 +

49 +def test_get_session_id_accepts_session_header():

50 + request = CompletionRequest(model="test-model", prompt="hi")

51 +

52 + session_id = GenerateBaseServing._get_session_id(

53 + request,

54 + _raw_request({"X-Session-ID": "header-session"}),

55 + )

56 +

57 + assert session_id == "header-session"

58 +

59 +

60 +def test_get_session_id_ignores_correlation_header():

61 + request = CompletionRequest(

62 + model="test-model",

63 + prompt="hi",

64 + vllm_xargs={"session_id": "xargs-session"},

65 + )

66 +

67 + session_id = GenerateBaseServing._get_session_id(

68 + request,

69 + _raw_request({"X-Correlation-ID": "correlation-session"}),

70 + )

71 +

72 + assert session_id == "xargs-session"

73 +

74 +

75 +def test_get_session_id_keeps_vllm_xargs_as_compatibility_fallback():

76 + request = CompletionRequest(

77 + model="test-model",

78 + prompt="hi",

79 + vllm_xargs={"session_id": "xargs-session"},

80 + )

81 +

82 + session_id = GenerateBaseServing._get_session_id(request, None)

83 +

84 + assert session_id == "xargs-session"

85 +

86 +

87 +def test_get_session_id_ignores_empty_and_non_string_values():

88 + request = CompletionRequest(

89 + model="test-model",

90 + prompt="hi",

91 + session_id="",

92 + vllm_xargs={"session_id": 7},

93 + )

94 +

95 + session_id = GenerateBaseServing._get_session_id(request, None)

96 +

97 + assert session_id is None

tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py

+26

@@ -203,6 +203,32 @@ async def test_serve_tokens_skips_mm_cache_for_remote_engine_execution():

203 )

204

205

206 +@pytest.mark.asyncio

207 +async def test_serve_tokens_threads_session_id_header_to_engine():

208 + engine = _mock_engine()

209 +

210 + async def mock_generate(*args, **kwargs):

211 + yield _make_request_output(

212 + "req-1", token_ids=[10], finish_reason="stop", finished=True

213 + )

214 +

215 + engine.generate = MagicMock(side_effect=mock_generate)

216 + serving = _build_serving_tokens(engine)

217 +

218 + request = GenerateRequest(

219 + token_ids=[1, 2, 3],

220 + sampling_params=SamplingParams(max_tokens=1),

221 + model=MODEL_NAME,

222 + stream=False,

223 + )

224 + raw_request = MagicMock()

225 + raw_request.headers = {"X-Session-ID": "header-session"}

226 +

227 + await serving.serve_tokens(request, raw_request)

228 +

229 + assert engine.generate.call_args.kwargs["session_id"] == "header-session"

230 +

231 +

232 @pytest.mark.asyncio

233 async def test_stream_basic():

234 """Streaming returns SSE chunks with correct token_ids and ends with [DONE]."""

tests/v1/engine/test_parallel_sampling.py

+16

@@ -1,6 +1,8 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 +from copy import copy

5 +

6 from vllm import SamplingParams

7 from vllm.outputs import CompletionOutput

8 from vllm.sampling_params import RequestOutputKind

@@ -68,6 +70,20 @@ def test_parent_request_to_output_final_only() -> None:

70 )

71

72

73 +def test_parallel_sampling_child_requests_preserve_session_id() -> None:

74 + request = make_request(SamplingParams(n=2))

75 + request.session_id = "session-1"

76 + parent_request = ParentRequest(request)

77 +

78 + for idx in range(parent_request.n):

79 + request_id, child_params = parent_request.get_child_info(idx)

80 + child_request = request if idx == parent_request.n - 1 else copy(request)

81 + child_request.request_id = request_id

82 + child_request.sampling_params = child_params

83 +

84 + assert child_request.session_id == "session-1"

85 +

86 +

87 def make_request(sampling_params: SamplingParams) -> EngineCoreRequest:

88 return EngineCoreRequest(

89 request_id="parent_id",

tests/v1/test_request.py

+23 -1

@@ -1,6 +1,9 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -from vllm.v1.request import RequestStatus

3 +

4 +from vllm import SamplingParams

5 +from vllm.v1.engine import EngineCoreRequest

6 +from vllm.v1.request import Request, RequestStatus

7

8

9 def test_request_status_fmt_str():

@@ -18,3 +21,22 @@ def test_request_status_fmt_str():

21 assert f"{RequestStatus.FINISHED_LENGTH_CAPPED}" == "FINISHED_LENGTH_CAPPED"

22 assert f"{RequestStatus.FINISHED_ABORTED}" == "FINISHED_ABORTED"

23 assert f"{RequestStatus.FINISHED_IGNORED}" == "FINISHED_IGNORED"

24 +

25 +

26 +def test_request_copies_session_id_from_engine_core_request():

27 + engine_request = EngineCoreRequest(

28 + request_id="request-1",

29 + prompt_token_ids=[1, 2, 3],

30 + mm_features=None,

31 + sampling_params=SamplingParams(max_tokens=1),

32 + pooling_params=None,

33 + arrival_time=0.0,

34 + lora_request=None,

35 + cache_salt=None,

36 + data_parallel_rank=None,

37 + session_id="session-1",

38 + )

39 +

40 + request = Request.from_engine_core_request(engine_request, block_hasher=None)

41 +

42 + assert request.session_id == "session-1"

vllm/engine/protocol.py

+1

@@ -78,6 +78,7 @@ class EngineClient(ABC):

78 trace_headers: Mapping[str, str] | None = None,

79 priority: int = 0,

80 data_parallel_rank: int | None = None,

81 + session_id: str | None = None,

82 reasoning_ended: bool | None = None,

83 reasoning_parser_kwargs: dict[str, Any] | None = None,

84 ) -> AsyncGenerator[RequestOutput, None]:

vllm/entrypoints/generate/base/serving.py

+24

@@ -41,6 +41,7 @@ logger = init_logger(__name__)

41

42 RequestT = TypeVar("RequestT", bound=AnyRequest)

43 _T = TypeVar("_T")

44 +SESSION_ID_HEADER = "X-Session-ID"

45

46

47 def build_per_request_timing_metrics(

@@ -215,6 +216,29 @@ class GenerateBaseServing(BaseServing, BeamSearchOnlineMixin):

216 except ValueError:

217 return None

218

219 + @staticmethod

220 + def _get_session_id_from_headers(raw_request: Request | None) -> str | None:

221 + if raw_request is None:

222 + return None

223 + if value := raw_request.headers.get(SESSION_ID_HEADER):

224 + return value

225 + return None

226 +

227 + @staticmethod

228 + def _get_session_id(

229 + request: ChatCompletionRequest | CompletionRequest | ResponsesRequest,

230 + raw_request: Request | None,

231 + ) -> str | None:

232 + if request.session_id:

233 + return request.session_id

234 + if value := GenerateBaseServing._get_session_id_from_headers(raw_request):

235 + return value

236 + if request.vllm_xargs:

237 + session_id = request.vllm_xargs.get("session_id")

238 + if isinstance(session_id, str) and session_id:

239 + return session_id

240 + return None

241 +

242 async def _with_kv_transfer_rejection_cleanup(

243 self,

244 awaitable: Awaitable[_T],

vllm/entrypoints/generate/beam_search/online.py

+2

@@ -32,6 +32,7 @@ class BeamSearchOnlineMixin(ABC):

32 params: BeamSearchParams,

33 lora_request: LoRARequest | None = None,

34 trace_headers: Mapping[str, str] | None = None,

35 + session_id: str | None = None,

36 ) -> AsyncGenerator[RequestOutput, None]:

37 beam_width = params.beam_width

38 max_tokens = params.max_tokens

@@ -90,6 +91,7 @@ class BeamSearchOnlineMixin(ABC):

91 request_id_item,

92 lora_request=lora_request_item,

93 trace_headers=trace_headers,

94 + session_id=session_id,

95 )

96 )

97 )

vllm/entrypoints/openai/chat_completion/batch_serving.py

+2

@@ -174,6 +174,7 @@ class OpenAIServingChatBatch(OpenAIServingChat):

174 if raw_request is None

175 else await self._get_trace_headers(raw_request.headers)

176 )

177 + session_id = self._get_session_id(single_request, raw_request)

178 generators.append(

179 self.engine_client.generate(

180 engine_prompt,

@@ -183,6 +184,7 @@ class OpenAIServingChatBatch(OpenAIServingChat):

184 trace_headers=trace_headers,

185 priority=request.priority,

186 data_parallel_rank=data_parallel_rank,

187 + session_id=session_id,

188 reasoning_ended=None,

189 )

190 )

vllm/entrypoints/openai/chat_completion/protocol.py

+8

@@ -394,6 +394,14 @@ class ChatCompletionRequest(OpenAIBaseModel):

394 "through out the inference process and return in response."

395 ),

396 )

397 + session_id: str | None = Field(

398 + default=None,

399 + description=(

400 + "Stable session identity shared by related requests. Unlike "

401 + "request_id, this value is expected to remain stable across "

402 + "multiple requests in the same conversation or agent session."

403 + ),

404 + )

405

406 return_tokens_as_token_ids: bool | None = Field(

407 default=None,

vllm/entrypoints/openai/chat_completion/serving.py

+3

@@ -318,6 +318,7 @@ class OpenAIServingChat(GenerateBaseServing):

318 if raw_request is None

319 else await self._get_trace_headers(raw_request.headers)

320 )

321 + session_id = self._get_session_id(request, raw_request)

322

323 if isinstance(sampling_params, BeamSearchParams):

324 generator = self.beam_search(

@@ -326,6 +327,7 @@ class OpenAIServingChat(GenerateBaseServing):

327 params=sampling_params,

328 lora_request=lora_request,

329 trace_headers=trace_headers,

330 + session_id=session_id,

331 )

332 else:

333 if not request.include_reasoning:

@@ -348,6 +350,7 @@ class OpenAIServingChat(GenerateBaseServing):

350 trace_headers=trace_headers,

351 priority=request.priority,

352 data_parallel_rank=data_parallel_rank,

353 + session_id=session_id,

354 reasoning_ended=reasoning_ended,

355 reasoning_parser_kwargs={

356 "chat_template_kwargs": chat_template_kwargs,

vllm/entrypoints/openai/completion/protocol.py

+8

@@ -147,6 +147,14 @@ class CompletionRequest(OpenAIBaseModel):

147 "through out the inference process and return in response."

148 ),

149 )

150 + session_id: str | None = Field(

151 + default=None,

152 + description=(

153 + "Stable session identity shared by related requests. Unlike "

154 + "request_id, this value is expected to remain stable across "

155 + "multiple requests in the same conversation or agent session."

156 + ),

157 + )

158

159 return_tokens_as_token_ids: bool | None = Field(

160 default=None,

vllm/entrypoints/openai/completion/serving.py

+3

@@ -194,6 +194,7 @@ class OpenAIServingCompletion(GenerateBaseServing):

194 if raw_request is None

195 else await self._get_trace_headers(raw_request.headers)

196 )

197 + session_id = self._get_session_id(request, raw_request)

198

199 if isinstance(sampling_params, BeamSearchParams):

200 generator = self.beam_search(

@@ -202,6 +203,7 @@ class OpenAIServingCompletion(GenerateBaseServing):

203 params=sampling_params,

204 lora_request=lora_request,

205 trace_headers=trace_headers,

206 + session_id=session_id,

207 )

208 else:

209 generator = self.engine_client.generate(

@@ -212,6 +214,7 @@ class OpenAIServingCompletion(GenerateBaseServing):

214 trace_headers=trace_headers,

215 priority=request.priority,

216 data_parallel_rank=data_parallel_rank,

217 + session_id=session_id,

218 )

219

220 generators.append(generator)

vllm/entrypoints/openai/responses/protocol.py

+8

@@ -220,6 +220,14 @@ class ResponsesRequest(OpenAIBaseModel):

220 "through out the inference process and return in response."

221 ),

222 )

223 + session_id: str | None = Field(

224 + default=None,

225 + description=(

226 + "Stable session identity shared by related requests. Unlike "

227 + "request_id, this value is expected to remain stable across "

228 + "multiple requests in the same conversation or agent session."

229 + ),

230 + )

231 media_io_kwargs: dict[str, dict[str, Any]] | None = Field(

232 default=None,

233 description=(

vllm/entrypoints/openai/responses/serving.py

+4

@@ -451,6 +451,7 @@ class OpenAIServingResponses(GenerateBaseServing):

451 if raw_request is None

452 else await self._get_trace_headers(raw_request.headers)

453 )

454 + session_id = self._get_session_id(request, raw_request)

455

456 chat_template_kwargs = self._effective_chat_template_kwargs(request)

457 response_parser = self._make_response_parser(

@@ -516,6 +517,7 @@ class OpenAIServingResponses(GenerateBaseServing):

517 lora_request=lora_request,

518 priority=request.priority,

519 trace_headers=trace_headers,

520 + session_id=session_id,

521 reasoning_parser_kwargs=reasoning_parser_kwargs

522 if self.parser and self.parser.reasoning_parser_cls is not None

523 else None,

@@ -669,6 +671,7 @@ class OpenAIServingResponses(GenerateBaseServing):

671 lora_request: LoRARequest | None = None,

672 priority: int = 0,

673 trace_headers: Mapping[str, str] | None = None,

674 + session_id: str | None = None,

675 reasoning_parser_kwargs: dict[str, Any] | None = None,

676 ):

677 max_model_len = self.model_config.max_model_len

@@ -693,6 +696,7 @@ class OpenAIServingResponses(GenerateBaseServing):

696 lora_request=lora_request,

697 trace_headers=trace_headers,

698 priority=priority,

699 + session_id=session_id,

700 reasoning_parser_kwargs=reasoning_parser_kwargs,

701 )

702

vllm/entrypoints/scale_out/token_in_token_out/serving.py

+2

@@ -215,6 +215,7 @@ class ServingTokens(GenerateBaseServing):

215

216 # Extract data_parallel_rank from header (router can inject it)

217 data_parallel_rank = self._get_data_parallel_rank(raw_request)

218 + session_id = self._get_session_id_from_headers(raw_request)

219

220 result_generator = self.engine_client.generate(

221 engine_input,

@@ -224,6 +225,7 @@ class ServingTokens(GenerateBaseServing):

225 trace_headers=trace_headers,

226 priority=request.priority,

227 data_parallel_rank=data_parallel_rank,

228 + session_id=session_id,

229 )

230

231 assert result_generator is not None

vllm/v1/engine/__init__.py

+2

@@ -145,6 +145,8 @@ class EngineCoreRequest(

145 # KV-transfer request is rejected on the D node before engine admission.

146 abort_immediately: bool = False

147

148 + session_id: str | None = None

149 +

150 @property

151 def params(self) -> SamplingParams | PoolingParams:

152 """Return the processed params (sampling or pooling)."""

vllm/v1/engine/async_llm.py

+8

@@ -294,6 +294,7 @@ class AsyncLLM(EngineClient):

294 trace_headers: Mapping[str, str] | None = None,

295 priority: int = 0,

296 data_parallel_rank: int | None = None,

297 + session_id: str | None = None,

298 prompt_text: str | None = None,

299 reasoning_ended: bool | None = None,

300 reasoning_parser_kwargs: dict[str, Any] | None = None,

@@ -331,6 +332,7 @@ class AsyncLLM(EngineClient):

332 trace_headers,

333 priority,

334 data_parallel_rank,

335 + session_id,

336 )

337

338 # Convert Input --> Request.

@@ -362,6 +364,7 @@ class AsyncLLM(EngineClient):

364 trace_headers=trace_headers,

365 priority=priority,

366 data_parallel_rank=data_parallel_rank,

367 + session_id=session_id,

368 )

369 else:

370 # Raw prompts require tokenization and possibly multimodal

@@ -377,6 +380,7 @@ class AsyncLLM(EngineClient):

380 trace_headers=trace_headers,

381 priority=priority,

382 data_parallel_rank=data_parallel_rank,

383 + session_id=session_id,

384 )

385 prompt_text, _, _ = extract_prompt_components(self.model_config, prompt)

386

@@ -445,6 +449,7 @@ class AsyncLLM(EngineClient):

449 trace_headers: Mapping[str, str] | None = None,

450 priority: int = 0,

451 data_parallel_rank: int | None = None,

452 + session_id: str | None = None,

453 ) -> RequestOutputCollector:

454 self._validate_streaming_input_sampling_params(sampling_params)

455

@@ -456,6 +461,7 @@ class AsyncLLM(EngineClient):

461 trace_headers=trace_headers,

462 priority=priority,

463 data_parallel_rank=data_parallel_rank,

464 + session_id=session_id,

465 )

466

467 if not sampling_params.skip_clone:

@@ -556,6 +562,7 @@ class AsyncLLM(EngineClient):

562 trace_headers: Mapping[str, str] | None = None,

563 priority: int = 0,

564 data_parallel_rank: int | None = None,

565 + session_id: str | None = None,

566 reasoning_ended: bool | None = None,

567 reasoning_parser_kwargs: dict[str, Any] | None = None,

568 ) -> AsyncGenerator[RequestOutput, None]:

@@ -585,6 +592,7 @@ class AsyncLLM(EngineClient):

592 trace_headers=trace_headers,

593 priority=priority,

594 data_parallel_rank=data_parallel_rank,

595 + session_id=session_id,

596 prompt_text=prompt_text,

597 reasoning_ended=reasoning_ended,

598 reasoning_parser_kwargs=reasoning_parser_kwargs,

vllm/v1/engine/input_processor.py

+2

@@ -261,6 +261,7 @@ class InputProcessor:

261 priority: int = 0,

262 data_parallel_rank: int | None = None,

263 resumable: bool = False,

264 + session_id: str | None = None,

265 ) -> EngineCoreRequest:

266 self._validate_params(params, supported_tasks)

267 self._validate_lora(lora_request)

@@ -391,6 +392,7 @@ class InputProcessor:

392 data_parallel_rank=data_parallel_rank,

393 trace_headers=trace_headers,

394 resumable=resumable,

395 + session_id=session_id,

396 )

397

398 def _validate_prompt_len(

vllm/v1/engine/llm_engine.py

+2

@@ -225,6 +225,7 @@ class LLMEngine:

225 tokenization_kwargs: dict[str, Any] | None = None,

226 trace_headers: Mapping[str, str] | None = None,

227 priority: int = 0,

228 + session_id: str | None = None,

229 prompt_text: str | None = None,

230 ) -> str:

231 # Validate the request_id type.

@@ -257,6 +258,7 @@ class LLMEngine:

258 tokenization_kwargs=tokenization_kwargs,

259 trace_headers=trace_headers,

260 priority=priority,

261 + session_id=session_id,

262 )

263 prompt_text, _, _ = extract_prompt_components(self.model_config, prompt)

264

vllm/v1/request.py

+3

@@ -74,6 +74,7 @@ class Request:

74 trace_headers: Mapping[str, str] | None = None,

75 block_hasher: Callable[["Request"], list["BlockHash"]] | None = None,

76 resumable: bool = False,

77 + session_id: str | None = None,

78 reasoning_ended: bool | None = None,

79 reasoning_parser_kwargs: dict[str, Any] | None = None,

80 abort_immediately: bool = False,

@@ -183,6 +184,7 @@ class Request:

184 self.all_token_ids = ConstantList(self._all_token_ids)

185 # trace_headers

186 self.trace_headers = trace_headers

187 + self.session_id = session_id

188

189 # True if this request is scheduled as a non-final prefill chunk.

190 self.is_prefill_chunk = False

@@ -241,6 +243,7 @@ class Request:

243 trace_headers=request.trace_headers,

244 block_hasher=block_hasher,

245 resumable=request.resumable,

246 + session_id=request.session_id,

247 reasoning_ended=request.reasoning_ended,

248 reasoning_parser_kwargs=request.reasoning_parser_kwargs,

249 abort_immediately=request.abort_immediately,

feat(frontend): session id plumbing into requests (#48048)