@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 170129eb28

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[gpt-oss] Harmony changes with container tool support (#23386)

Signed-off-by: zhiweiz <zhiweiz@fb.com> Signed-off-by: Aaron Pham <contact@aarnphm.xyz> Signed-off-by: Lu Fang <30275821+houseroad@users.noreply.github.com> Co-authored-by: zhiweiz <zhiweiz@fb.com> Co-authored-by: Aaron Pham <contact@aarnphm.xyz> Co-authored-by: Simon Mo <simon.mo@hey.com> Co-authored-by: Lu Fang <30275821+houseroad@users.noreply.github.com>

zhiweiz committed Sep 8, 2025 at 19:03 UTC 170129eb283aa225f20666a434c02822f643242e

5 files changed +170 -27

vllm/entrypoints/context.py

+79 -8

@@ -1,5 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 +import asyncio

4 +import contextlib

5 import json

6 import logging

7 from abc import ABC, abstractmethod

@@ -57,9 +59,14 @@ class ConversationContext(ABC):

60 @abstractmethod

61 async def init_tool_sessions(self, tool_server: Optional[ToolServer],

60 - exit_stack: AsyncExitStack) -> None:

62 + exit_stack: AsyncExitStack,

63 + request_id: str) -> None:

64 pass

66 + @abstractmethod

67 + async def cleanup_session(self) -> None:

68 + raise NotImplementedError("Should not be called.")

69 +

71 class SimpleContext(ConversationContext):

@@ -89,9 +96,13 @@ class SimpleContext(ConversationContext):

96 raise NotImplementedError("Should not be called.")

98 async def init_tool_sessions(self, tool_server: Optional[ToolServer],

92 - exit_stack: AsyncExitStack) -> None:

99 + exit_stack: AsyncExitStack,

100 + request_id: str) -> None:

101 pass

102

103 + async def cleanup_session(self) -> None:

104 + raise NotImplementedError("Should not be called.")

105 +

106

107 class HarmonyContext(ConversationContext):

108

@@ -103,6 +114,7 @@ class HarmonyContext(ConversationContext):

114 self._messages = messages

115 self.available_tools = available_tools

116 self._tool_sessions: dict[str, Union[ClientSession, Tool]] = {}

117 + self.called_tools: set[str] = set()

118

119 self.parser = get_streamable_parser_for_assistant()

120 self.num_init_messages = len(messages)

@@ -234,7 +246,8 @@ class HarmonyContext(ConversationContext):

246 last_msg = self.messages[-1]

247 recipient = last_msg.recipient

248 return recipient is not None and (recipient.startswith("browser.")

237 - or recipient.startswith("python"))

249 + or recipient.startswith("python") or

250 + recipient.startswith("container."))

251

252 async def call_tool(self) -> list[Message]:

253 if not self.messages:

@@ -248,6 +261,9 @@ class HarmonyContext(ConversationContext):

261 elif recipient.startswith("python"):

262 return await self.call_python_tool(

263 self._tool_sessions["python"], last_msg)

264 + elif recipient.startswith("container."):

265 + return await self.call_container_tool(

266 + self._tool_sessions["container"], last_msg)

267 raise ValueError("No tool call found")

268

269 def render_for_completion(self) -> list[int]:

@@ -256,6 +272,7 @@ class HarmonyContext(ConversationContext):

272 async def call_search_tool(self, tool_session: Union["ClientSession",

273 Tool],

274 last_msg: Message) -> list[Message]:

275 + self.called_tools.add("browser")

276 if isinstance(tool_session, Tool):

277 return await tool_session.get_result(self)

278 tool_name = last_msg.recipient.split(".")[1]

@@ -265,12 +282,16 @@ class HarmonyContext(ConversationContext):

282 content = TextContent(text=result_str)

283 author = Author(role=Role.TOOL, name=last_msg.recipient)

284 return [

268 - Message(author=author, content=[content], recipient=Role.ASSISTANT)

285 + Message(author=author,

286 + content=[content],

287 + recipient=Role.ASSISTANT,

288 + channel=last_msg.channel)

289 ]

290

291 async def call_python_tool(self, tool_session: Union["ClientSession",

292 Tool],

293 last_msg: Message) -> list[Message]:

294 + self.called_tools.add("python")

295 if isinstance(tool_session, Tool):

296 return await tool_session.get_result(self)

297 param = {

@@ -290,13 +311,63 @@ class HarmonyContext(ConversationContext):

311 ]

312

313 async def init_tool_sessions(self, tool_server: Optional[ToolServer],

293 - exit_stack: AsyncExitStack) -> None:

314 + exit_stack: AsyncExitStack,

315 + request_id: str) -> None:

316 if tool_server:

317 for tool_name in self.available_tools:

318 if tool_name not in self._tool_sessions:

297 - self._tool_sessions[

298 - tool_name] = await exit_stack.enter_async_context(

299 - tool_server.new_session(tool_name))

319 + tool_session = await exit_stack.enter_async_context(

320 + tool_server.new_session(tool_name, request_id))

321 + self._tool_sessions[tool_name] = tool_session

322 + exit_stack.push_async_exit(self.cleanup_session)

323 +

324 + async def call_container_tool(self, tool_session: Union["ClientSession",

325 + Tool],

326 + last_msg: Message) -> list[Message]:

327 + """

328 + Call container tool. Expect this to be run in a stateful docker

329 + with command line terminal.

330 + The official container tool would at least

331 + expect the following format:

332 + - for tool name: exec

333 + - args:

334 + {

335 + "cmd":List[str] "command to execute",

336 + "workdir":optional[str] "current working directory",

337 + "env":optional[object/dict] "environment variables",

338 + "session_name":optional[str] "session name",

339 + "timeout":optional[int] "timeout in seconds",

340 + "user":optional[str] "user name",

341 + }

342 + """

343 + self.called_tools.add("container")

344 + if isinstance(tool_session, Tool):

345 + return await tool_session.get_result(self)

346 + tool_name = last_msg.recipient.split(".")[1].split(" ")[0]

347 + args = json.loads(last_msg.content[0].text)

348 + result = await tool_session.call_tool(tool_name, args)

349 + result_str = result.content[0].text

350 + content = TextContent(text=result_str)

351 + author = Author(role=Role.TOOL, name=last_msg.recipient)

352 + return [

353 + Message(author=author,

354 + content=[content],

355 + recipient=Role.ASSISTANT,

356 + channel=last_msg.channel)

357 + ]

358 +

359 + async def cleanup_session(self, *args, **kwargs) -> None:

360 + """Can be used as coro to used in __aexit__"""

361 +

362 + async def cleanup_tool_session(tool_session):

363 + if not isinstance(tool_session, Tool):

364 + logger.info("Cleaning up tool session for %s",

365 + tool_session._client_info)

366 + with contextlib.suppress(Exception):

367 + await tool_session.call_tool("cleanup_session", {})

368 +

369 + await asyncio.gather(*(cleanup_tool_session(self._tool_sessions[tool])

370 + for tool in self.called_tools))

371

372

373 class StreamingHarmonyContext(HarmonyContext):

vllm/entrypoints/harmony_utils.py

+44 -6

@@ -16,11 +16,13 @@ from openai.types.responses.response_function_web_search import (

16 from openai.types.responses.response_reasoning_item import (

17 Content as ResponseReasoningTextContent)

18 from openai.types.responses.tool import Tool

19 -from openai_harmony import (Author, Conversation, DeveloperContent,

20 - HarmonyEncodingName, Message, ReasoningEffort,

21 - Role, StreamableParser, SystemContent, TextContent,

22 - ToolDescription, load_harmony_encoding)

19 +from openai_harmony import (Author, ChannelConfig, Conversation,

20 + DeveloperContent, HarmonyEncodingName, Message,

21 + ReasoningEffort, Role, StreamableParser,

22 + SystemContent, TextContent, ToolDescription,

23 + load_harmony_encoding)

25 +from vllm import envs

26 from vllm.entrypoints.openai.protocol import (ChatCompletionToolsParam,

27 ResponseInputOutputItem)

28 from vllm.utils import random_uuid

@@ -33,6 +35,20 @@ REASONING_EFFORT = {

36 _harmony_encoding = None

38 +# Builtin tools that should be included in the system message when

39 +# they are available and requested by the user.

40 +# Tool args are provided by MCP tool descriptions. Output

41 +# of the tools are stringified.

42 +BUILTIN_TOOLS = {

43 + "web_search_preview",

44 + "code_interpreter",

45 + "container",

46 +}

47 +

48 +

49 +def has_custom_tools(tool_types: list[str]) -> bool:

50 + return not set(tool_types).issubset(BUILTIN_TOOLS)

51 +

53 def get_encoding():

54 global _harmony_encoding

@@ -48,10 +64,19 @@ def get_system_message(

64 start_date: Optional[str] = None,

65 browser_description: Optional[str] = None,

66 python_description: Optional[str] = None,

67 + container_description: Optional[str] = None,

68 + instructions: Optional[str] = None,

69 + with_custom_tools: bool = False,

70 ) -> Message:

71 sys_msg_content = SystemContent.new()

72 if model_identity is not None:

73 sys_msg_content = sys_msg_content.with_model_identity(model_identity)

74 + if (instructions is not None

75 + and envs.VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS):

76 + current_identity = sys_msg_content.model_identity

77 + new_identity = (f'{current_identity}\n{instructions}'

78 + if current_identity else instructions)

79 + sys_msg_content = sys_msg_content.with_model_identity(new_identity)

80 if reasoning_effort is not None:

81 sys_msg_content = sys_msg_content.with_reasoning_effort(

82 REASONING_EFFORT[reasoning_effort])

@@ -63,6 +88,14 @@ def get_system_message(

88 sys_msg_content = sys_msg_content.with_tools(browser_description)

89 if python_description is not None:

90 sys_msg_content = sys_msg_content.with_tools(python_description)

91 + if container_description is not None:

92 + sys_msg_content = sys_msg_content.with_tools(container_description)

93 + if not with_custom_tools:

94 + channel_config = sys_msg_content.channel_config

95 + invalid_channel = "commentary"

96 + new_config = ChannelConfig.require_channels(

97 + [c for c in channel_config.valid_channels if c != invalid_channel])

98 + sys_msg_content = sys_msg_content.with_channel_config(new_config)

99 sys_msg = Message.from_role_and_content(Role.SYSTEM, sys_msg_content)

100 return sys_msg

101

@@ -86,14 +119,17 @@ def get_developer_message(

119 tools: Optional[list[Union[Tool, ChatCompletionToolsParam]]] = None,

120 ) -> Message:

121 dev_msg_content = DeveloperContent.new()

89 - if instructions is not None:

122 + if (instructions is not None

123 + and not envs.VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS):

124 dev_msg_content = dev_msg_content.with_instructions(instructions)

125 if tools is not None:

126 function_tools: list[Union[Tool, ChatCompletionToolsParam]] = []

127 for tool in tools:

94 - if tool.type in ("web_search_preview", "code_interpreter"):

128 + if tool.type in ("web_search_preview", "code_interpreter",

129 + "container"):

130 # These are built-in tools that are added to the system message.

131 pass

132 +

133 elif tool.type == "function":

134 function_tools.append(tool)

135 else:

@@ -136,6 +172,8 @@ def parse_response_input(

172 TextContent(text=text_prefix + c["text"]) for c in content

173 ]

174 msg = Message.from_role_and_contents(role, contents)

175 + if role == "assistant":

176 + msg = msg.with_channel("final")

177 elif response_msg["type"] == "function_call_output":

178 call_id = response_msg["call_id"]

179 call_response: Optional[ResponseFunctionToolCall] = None

vllm/entrypoints/openai/serving_responses.py

+26 -7

@@ -44,8 +44,9 @@ from vllm.entrypoints.context import (ConversationContext, HarmonyContext,

44 SimpleContext, StreamingHarmonyContext)

45 from vllm.entrypoints.harmony_utils import (

46 get_developer_message, get_stop_tokens_for_assistant_actions,

47 - get_system_message, get_user_message, parse_output_message,

48 - parse_remaining_state, parse_response_input, render_for_completion)

47 + get_system_message, get_user_message, has_custom_tools,

48 + parse_output_message, parse_remaining_state, parse_response_input,

49 + render_for_completion)

50 from vllm.entrypoints.logger import RequestLogger

51 # yapf conflicts with isort for this block

52 # yapf: disable

@@ -266,6 +267,8 @@ class OpenAIServingResponses(OpenAIServing):

267 builtin_tool_list.append("browser")

268 if self.tool_server.has_tool("python"):

269 builtin_tool_list.append("python")

270 + if self.tool_server.has_tool("container"):

271 + builtin_tool_list.append("container")

272

273 if self.tool_server is not None:

274 available_tools = builtin_tool_list

@@ -448,7 +451,8 @@ class OpenAIServingResponses(OpenAIServing):

451

452 async with AsyncExitStack() as exit_stack:

453 try:

451 - await context.init_tool_sessions(self.tool_server, exit_stack)

454 + await context.init_tool_sessions(self.tool_server, exit_stack,

455 + request.request_id)

456 async for _ in result_generator:

457 pass

458 except asyncio.CancelledError:

@@ -710,13 +714,21 @@ class OpenAIServingResponses(OpenAIServing):

714 # New conversation.

715 reasoning_effort = (request.reasoning.effort

716 if request.reasoning else None)

717 + # Temporary: OpenAI types doesn't have container tool

718 + # so we used MCP to cover that, up for change

719 tool_types = [tool.type for tool in request.tools]

720 + if envs.VLLM_GPT_OSS_USE_CONTAINER_TOOL:

721 + tool_types.append("container")

722 enable_browser = ("web_search_preview" in tool_types

723 and self.tool_server is not None

724 and self.tool_server.has_tool("browser"))

725 enable_code_interpreter = ("code_interpreter" in tool_types

726 and self.tool_server is not None

727 and self.tool_server.has_tool("python"))

728 + enable_container = ("container" in tool_types

729 + and self.tool_server is not None

730 + and self.tool_server.has_tool("container"))

731 + with_custom_tools = has_custom_tools(tool_types)

732 sys_msg = get_system_message(

733 reasoning_effort=reasoning_effort,

734 browser_description=self.tool_server.get_tool_description(

@@ -725,11 +737,17 @@ class OpenAIServingResponses(OpenAIServing):

737 python_description=self.tool_server.get_tool_description(

738 "python") if enable_code_interpreter

739 and self.tool_server is not None else None,

740 + container_description=self.tool_server.get_tool_description(

741 + "container")

742 + if enable_container and self.tool_server is not None else None,

743 + instructions=request.instructions,

744 + with_custom_tools=with_custom_tools,

745 )

746 messages.append(sys_msg)

730 - dev_msg = get_developer_message(request.instructions,

731 - request.tools)

732 - messages.append(dev_msg)

747 + if with_custom_tools:

748 + dev_msg = get_developer_message(

749 + instructions=request.instructions, tools=request.tools)

750 + messages.append(dev_msg)

751 else:

752 # Continue the previous conversation.

753 # FIXME(woosuk): Currently, request params like reasoning and

@@ -1613,7 +1631,8 @@ class OpenAIServingResponses(OpenAIServing):

1631 async with AsyncExitStack() as exit_stack:

1632 processer = None

1633 if self.use_harmony:

1616 - await context.init_tool_sessions(self.tool_server, exit_stack)

1634 + await context.init_tool_sessions(self.tool_server, exit_stack,

1635 + request.request_id)

1636 processer = self._process_harmony_streaming_events

1637 else:

1638 processer = self._process_simple_streaming_events

vllm/entrypoints/tool_server.py

+10 -6

@@ -86,7 +86,8 @@ class ToolServer(ABC):

86 pass

88 @abstractmethod

89 - def new_session(self, tool_name: str) -> AbstractAsyncContextManager[Any]:

89 + def new_session(self, tool_name: str,

90 + session_id: str) -> AbstractAsyncContextManager[Any]:

91 """

92 Create a session for the tool.

93 """

@@ -124,7 +125,8 @@ class MCPToolServer(ToolServer):

125 description=tool.description,

126 parameters=tool.inputSchema)

127 for tool in list_tools_response.tools

127 - ])

128 + ],

129 + )

130 self.harmony_tool_descriptions[tool_from_mcp.name] = tool_from_mcp

131 if tool_from_mcp.name not in self.urls:

132 self.urls[tool_from_mcp.name] = url

@@ -142,14 +144,16 @@ class MCPToolServer(ToolServer):

144 return self.harmony_tool_descriptions.get(tool_name)

145

146 @asynccontextmanager

145 - async def new_session(self, tool_name: str):

147 + async def new_session(self, tool_name: str, session_id: str):

148 from mcp import ClientSession

149 from mcp.client.sse import sse_client

150 url = self.urls.get(tool_name)

151 + headers = {"x-session-id": session_id}

152 if not url:

153 raise KeyError(f"Tool '{tool_name}' is not supported")

151 - async with sse_client(url=url) as streams, ClientSession(

152 - *streams) as session:

154 + async with sse_client(url=url,

155 + headers=headers) as streams, ClientSession(

156 + *streams) as session:

157 await session.initialize()

158 yield session

159

@@ -182,7 +186,7 @@ class DemoToolServer(ToolServer):

186 raise ValueError(f"Unknown tool {tool_name}")

187

188 @asynccontextmanager

185 - async def new_session(self, tool_name: str):

189 + async def new_session(self, tool_name: str, session_id: str):

190 if tool_name not in self.tools:

191 raise KeyError(f"Tool '{tool_name}' is not supported")

192 yield self.tools[tool_name]

vllm/envs.py

+11

@@ -168,6 +168,8 @@ if TYPE_CHECKING:

168 VLLM_ALLREDUCE_USE_SYMM_MEM: bool = False

169 VLLM_TUNED_CONFIG_FOLDER: Optional[str] = None

170 VLLM_DISABLE_PAD_FOR_CUDAGRAPH: bool = False

171 + VLLM_GPT_OSS_USE_CONTAINER_TOOL: bool = False

172 + VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS: bool = False

173 VLLM_CUSTOM_SCOPES_FOR_PROFILING: bool = False

174

175

@@ -1201,6 +1203,15 @@ environment_variables: dict[str, Callable[[], Any]] = {

1203 "VLLM_TUNED_CONFIG_FOLDER":

1204 lambda: os.getenv("VLLM_TUNED_CONFIG_FOLDER", None),

1205

1206 + # Allows vllm use container tool

1207 + "VLLM_GPT_OSS_USE_CONTAINER_TOOL":

1208 + lambda: bool(int(os.getenv("VLLM_GPT_OSS_USE_CONTAINER_TOOL", "0"))),

1209 +

1210 + # Allows harmony instructions to be injected on system messages

1211 + "VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS":

1212 + lambda: bool(

1213 + int(os.getenv("VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS", "0"))),

1214 +

1215 # Add optional custom scopes for profiling, disable to avoid overheads

1216 "VLLM_CUSTOM_SCOPES_FOR_PROFILING":

1217 lambda: bool(int(os.getenv("VLLM_CUSTOM_SCOPES_FOR_PROFILING", "0"))),