Skip to content

Commit a6002de

Browse files
authored
[tool] fix: load MCP tools in async rollout mode (verl-project#2821)
### What does this PR do? Currently, the tool registry isn't aware of an event loop already existing, so it may fail when using the new async rollout architecture. This PR allows `initialize_tools_from_config` to load MCP tools when using the async architecture by spawning a new, temporary event loop in a separate thread to load from config. There is also a minor bugfix to `mcp_base_tool` which fixes a possibility of concatenating a string to None. NOTE: in the future, we should use async methods entirely, since this fix is not the most elegant. This fix works for now as verl is transitioning to a full async architecture. ### Checklist Before Starting - [x] Search for similar PRs. Paste at least one query link here: ... - [x] Format the PR title as `[{modules}] {type}: {description}` (This will be checked by the CI) - `{modules}` include `fsdp`, `megatron`, `sglang`, `vllm`, `rollout`, `trainer`, `ci`, `training_utils`, `recipe`, `hardware`, `deployment`, `ray`, `worker`, `single_controller`, `misc`, `perf`, `model`, `algo`, `env`, `tool`, `ckpt`, `doc`, `data` - If this PR involves multiple modules, separate them with `,` like `[megatron, fsdp, doc]` - `{type}` is in `feat`, `fix`, `refactor`, `chore`, `test` - If this PR breaks any API (CLI arguments, config, function signature, etc.), add `[BREAKING]` to the beginning of the title. - Example: `[BREAKING][fsdp, megatron] feat: dynamic batching` ### Design & Code Changes > Demonstrate the high-level design if this PR is complex, and list the specific changes. ### Checklist Before Submitting > [!IMPORTANT] > Please check all the following items before requesting a review, otherwise the reviewer might deprioritize this PR for review. - [x] Read the [Contribute Guide](https://github.com/volcengine/verl/blob/main/CONTRIBUTING.md). - [x] Apply [pre-commit checks](https://github.com/volcengine/verl/blob/main/CONTRIBUTING.md#code-linting-and-formatting): `pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always` - [x] Add / Update [the documentation](https://github.com/volcengine/verl/tree/main/docs). - [x] Add unit or end-to-end test(s) to [the CI workflow](https://github.com/volcengine/verl/tree/main/.github/workflows) to cover all the code. If not feasible, explain why: ... - [x] Once your PR is ready for CI, send a message in [the `ci-request` channel](https://verl-project.slack.com/archives/C091TCESWB1) in [the `verl` Slack workspace](https://join.slack.com/t/verl-project/shared_invite/zt-3855yhg8g-CTkqXu~hKojPCmo7k_yXTQ). (If not accessible, please try [the Feishu group (飞书群)](https://applink.larkoffice.com/client/chat/chatter/add_by_link?link_token=772jd4f1-cd91-441e-a820-498c6614126a).)
1 parent 0e14d81 commit a6002de

2 files changed

Lines changed: 47 additions & 24 deletions

File tree

verl/tools/mcp_base_tool.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,7 @@ async def _call_tool(self, instance_id, parameters) -> tuple[str, dict]:
7373

7474
logger.debug(f"Tool result for instance {instance_id} with tool {self.name}: {call_tool_result.content}")
7575
result, metadata = self._parse_tool_result(call_tool_result.content)
76-
metadata["api_request_error"] += err_msg
76+
metadata["api_request_error"] = None if not err_msg else err_msg
7777
return result, metadata
7878

7979
@rollout_trace_op

verl/tools/utils/tool_registry.py

Lines changed: 46 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import logging
1818
import os
1919
import sys
20+
import threading
2021
from enum import Enum
2122

2223
from omegaconf import OmegaConf
@@ -81,27 +82,49 @@ def get_tool_class(cls_name):
8182
def initialize_tools_from_config(tools_config_file):
8283
tools_config = OmegaConf.load(tools_config_file)
8384
tool_list = []
84-
for tool_config in tools_config.tools:
85-
cls_name = tool_config.class_name
86-
tool_type = ToolType(tool_config.config.type)
87-
tool_cls = get_tool_class(cls_name)
88-
89-
match tool_type:
90-
case ToolType.NATIVE:
91-
if tool_config.get("tool_schema", None) is None:
92-
tool_schema = None
93-
else:
94-
tool_schema_dict = OmegaConf.to_container(tool_config.tool_schema, resolve=True)
95-
tool_schema = OpenAIFunctionToolSchema.model_validate(tool_schema_dict)
96-
tool = tool_cls(
97-
config=OmegaConf.to_container(tool_config.config, resolve=True),
98-
tool_schema=tool_schema,
99-
)
100-
tool_list.append(tool)
101-
case ToolType.MCP:
102-
loop = asyncio.get_event_loop()
103-
mcp_tools = loop.run_until_complete(initialize_mcp_tool(tool_cls, tool_config))
104-
tool_list.extend(mcp_tools)
105-
case _:
106-
raise NotImplementedError
85+
86+
# Use a temporary event loop in a new thread because event
87+
# loop may already exist in new async architecture while retaining
88+
# backwards compatibility
89+
tmp_event_loop = asyncio.new_event_loop()
90+
thread = threading.Thread(target=tmp_event_loop.run_forever, name="mcp tool list fetcher", daemon=True)
91+
92+
def run_coroutine(coroutine):
93+
if not thread.is_alive():
94+
thread.start()
95+
96+
future = asyncio.run_coroutine_threadsafe(coroutine, tmp_event_loop)
97+
return future.result()
98+
99+
async def stop_loop():
100+
tmp_event_loop.stop()
101+
102+
try:
103+
for tool_config in tools_config.tools:
104+
cls_name = tool_config.class_name
105+
tool_type = ToolType(tool_config.config.type)
106+
tool_cls = get_tool_class(cls_name)
107+
108+
match tool_type:
109+
case ToolType.NATIVE:
110+
if tool_config.get("tool_schema", None) is None:
111+
tool_schema = None
112+
else:
113+
tool_schema_dict = OmegaConf.to_container(tool_config.tool_schema, resolve=True)
114+
tool_schema = OpenAIFunctionToolSchema.model_validate(tool_schema_dict)
115+
tool = tool_cls(
116+
config=OmegaConf.to_container(tool_config.config, resolve=True),
117+
tool_schema=tool_schema,
118+
)
119+
tool_list.append(tool)
120+
case ToolType.MCP:
121+
mcp_tools = run_coroutine(initialize_mcp_tool(tool_cls, tool_config))
122+
tool_list.extend(mcp_tools)
123+
case _:
124+
raise NotImplementedError
125+
finally:
126+
if thread.is_alive():
127+
asyncio.run_coroutine_threadsafe(stop_loop(), tmp_event_loop)
128+
thread.join()
129+
107130
return tool_list

0 commit comments

Comments
 (0)