diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
index ac2ca806a..565b83a56 100644
--- a/.buildkite/pipeline.yml
+++ b/.buildkite/pipeline.yml
@@ -58,7 +58,7 @@ steps:
python:3.11 bash -c '
set -euo pipefail
pip install -q torch --index-url https://download.pytorch.org/whl/cpu
- pip install -q pytest numpy packaging pyyaml omegaconf tqdm httpx requests ray pybase64 pylatexenc sympy aiohttp pillow safetensors transformers cloudpickle blake3 xxhash zstandard psutil
+ pip install -q pytest numpy packaging pyyaml omegaconf tqdm httpx requests ray pybase64 pylatexenc sympy aiohttp pillow safetensors transformers cloudpickle blake3 xxhash zstandard psutil jinja2
pip install -q -e . --no-deps
python tests/test_megatron_argument_validation.py
python tests/test_value_temperature.py
@@ -85,6 +85,7 @@ steps:
python tests/test_empty_colocated_weight_bucket.py
python tests/test_reloadable_process_group_memory_check.py
python tests/test_ppo_logprob_entropy.py
+ python tests/test_retool_generate.py
python tests/utils/test_hf_checkpoint_saver.py
'
diff --git a/examples/retool/README.md b/examples/retool/README.md
new file mode 100644
index 000000000..8d555c797
--- /dev/null
+++ b/examples/retool/README.md
@@ -0,0 +1,135 @@
+# Retool: from SFT to RL
+
+This example demonstrates how to use the retool functionality for tool-enabled language model generation.
+
+## Overview
+
+The retool example provides:
+- Safe Python code execution in a sandbox environment
+- Tool registry for managing available tools
+- Integration with language model generation
+- Reward calculation for tool usage
+
+## Files
+
+- `generate_with_retool.py`: Main generation function with tool support
+- `tool_sandbox.py`: Tool execution and safety management
+- `sft_data_processing.py`: Process SFT dataset
+- `rl_data_preprocess.py`: Process the RL (DAPO-Math-17k) dataset
+
+## Usage
+
+1. Setup and download datasets:
+```bash
+cd vime
+pip install -e . --no-deps
+pip install -r examples/retool/requirements.txt
+# For SFT part, you can use later model to RL directly and skip SFT.
+hf download --repo-type dataset JoeYing/ReTool-SFT --local-dir /root/JoeYing/ReTool-SFT
+hf download Qwen/Qwen3-4B-Instruct-2507 --local-dir /root/Qwen/Qwen3-4B-Instruct-2507
+
+# For RL part
+hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/dapo-math-17k
+hf download --repo-type dataset zhuzilin/aime-2024 --local-dir /root/aime-2024
+# download our SFT model if you want to skip SFT
+hf download font-info/qwen3-4b-sft-SGLang-RL --local-dir /root/font-info/qwen3-4b-sft
+```
+
+2. Create torch dist
+
+Both checkpoints use rope theta `5e6`, which differs from the `1e6` default in
+`scripts/models/qwen3-4B.sh`. Override it with `MODEL_ARGS_ROTARY_BASE` so the
+conversion and the training scripts agree.
+
+For SFT
+```bash
+MODEL_ARGS_ROTARY_BASE=5000000 source scripts/models/qwen3-4B.sh
+PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \
+ ${MODEL_ARGS[@]} \
+ --hf-checkpoint /root/Qwen/Qwen3-4B-Instruct-2507 \
+ --save /root/Qwen/Qwen3-4B-Instruct-2507_torch_dist
+```
+
+Or RL only
+```bash
+MODEL_ARGS_ROTARY_BASE=5000000 source scripts/models/qwen3-4B.sh
+PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \
+ ${MODEL_ARGS[@]} \
+ --hf-checkpoint /root/font-info/qwen3-4b-sft \
+ --save /root/font-info/qwen3-4b-sft_torch_dist
+```
+
+3. SFT:
+```bash
+python examples/retool/sft_data_processing.py
+bash examples/retool/retool_qwen3_4b_sft.sh
+```
+
+4. RL:
+```bash
+bash examples/retool/retool_qwen3_4b_rl.sh
+```
+
+5. Use in your training scripts by importing the generate function:
+```python
+from generate_with_retool import generate, reward_func
+```
+
+The RL script wires these up with:
+```bash
+--custom-generate-function-path generate_with_retool.generate
+--custom-rm-path generate_with_retool.reward_func
+```
+`generate_with_retool` is resolved as a top-level module, so the example
+directory is added to `PYTHONPATH` in the script's Ray runtime env (this is also
+what lets it import its sibling `tool_sandbox`).
+
+## Tool Format
+
+The system uses the following tool format:
+
+```
+You may call one or more functions to assist with the user query.
+
+You are provided with function signatures within XML tags:
+
+{"type": "function", "function": {"name": "code_interpreter", "description": "A tool for executing code.", "parameters": {"type": "object", "properties": {"code": {"type": "string", "description": "The code to execute."}}, "required": ["code"]}}}
+
+
+For each function call, return a json object with function name and arguments within XML tags:
+
+{"name": , "arguments": }
+
+```
+
+## Safety Features
+
+- Code execution in isolated sandbox
+- Memory and time limits
+- Dangerous operation detection
+- Allowed module restrictions
+
+Note that `PythonSandbox._check_code_safety` is deliberately strict: it allows
+only the stdlib modules in `PythonSandbox.allowed_modules` (`math`, `random`,
+`statistics`, `decimal`, `fractions`, …) and rejects `eval`/`exec`/`open`,
+dunder access, and imports outside that set. Widen `allowed_modules` if your task
+needs more (e.g. `sympy` or `numpy`).
+
+## Notes on the vLLM port
+
+This example was ported from slime's SGLang implementation. The rollout loop
+talks to vime's vLLM router at `/inference/v1/generate` with a
+`{"model", "token_ids", "sampling_params"}` body, and reads back
+`choices[0].token_ids` plus `choices[0].logprobs.content[i].logprob`.
+
+Two behaviours differ from `vime.rollout.vllm_rollout.generate` on purpose:
+
+- When the engine returns tokens but no usable per-token logprobs, this example
+ marks the sample `ABORTED` instead of substituting zeros. Zero-filled logprobs
+ would desync `rollout_log_probs` from the response tokens and silently corrupt
+ the importance ratio, so the sample is returned to the buffer for retry.
+- The tool-concurrency limit is taken exactly once, inside
+ `ToolRegistry.execute_tool`. `tool_sandbox.SEMAPHORE` is a plain
+ `asyncio.Semaphore` and is not reentrant, so acquiring it in both the caller
+ and the registry needs two permits per tool call and hangs once enough calls
+ are in flight.
diff --git a/examples/retool/generate_with_retool.py b/examples/retool/generate_with_retool.py
new file mode 100644
index 000000000..d8304adda
--- /dev/null
+++ b/examples/retool/generate_with_retool.py
@@ -0,0 +1,485 @@
+# Adapted from https://github.com/volcengine/verl/blob/cb809d66e46dfd3342d008628891a14a054fa424/recipe/retool/retool.py
+# Ported from slime's SGLang-based example to vime's vLLM ``/inference/v1/generate`` path.
+import re
+from typing import Any
+
+try:
+ from jinja2 import Template
+except ImportError as e:
+ raise ImportError("Jinja2 is required. Please install it with: pip install jinja2") from e
+
+from vime.rollout.vllm_rollout import GenerateState, _build_inference_sampling_params
+from vime.utils.http_utils import post
+from vime.utils.types import Sample
+
+# Import reward models
+try:
+ from vime.rollout.rm_hub.math_dapo_utils import compute_score as math_dapo_compute_score
+except ImportError as e:
+ raise ImportError("MathDapo is not installed") from e
+
+# Import tool sandbox functionality
+from tool_sandbox import TOOL_CONFIGS, tool_registry
+
+# Jinja2 template for tool-enabled conversations
+TOOL_TEMPLATE = """<|im_start|>system
+{%- if messages[0]['role'] == 'system' %}
+{{- messages[0]['content'] }}
+{%- else %}
+You are a helpful assistant.
+{%- endif %}
+{%- if tools %}
+# Tools
+
+You may call one or more functions to assist with the user query.
+
+You are provided with function signatures within XML tags:
+
+{%- for tool in tools %}
+{{- tool | tojson }}
+{%- endfor %}
+
+
+For each function call, return a json object with function name and arguments within XML tags:
+
+{"name": , "arguments": }
+
+{%- endif %}
+<|im_end|>
+{%- for message in messages %}
+{%- if message['role'] == 'user' %}
+<|im_start|>user
+{{- message['content'] }}<|im_end|>
+{%- elif message['role'] == 'assistant' %}
+<|im_start|>assistant
+{{- message['content'] }}<|im_end|>
+{%- endif %}
+{%- endfor %}
+<|im_start|>assistant
+"""
+
+
+def split_prompt(prompt: Any) -> tuple[str | None, str]:
+ """Return ``(system_prompt, user_text)`` from a dataset prompt.
+
+ ``prompt`` is a plain string only when the data file stores one. Both datasets
+ this example uses -- DAPO-Math-17k and AIME-2024 -- store a list of chat
+ messages, and ``vime.utils.data._build_messages`` passes that list through
+ untouched unless ``--apply-chat-template`` is set. Rendering the list here
+ keeps the chat template applied exactly once (the flag would apply it a second
+ time in the data loader) and keeps ``reward_func`` working on a string.
+ """
+ if isinstance(prompt, str):
+ return None, prompt
+ if isinstance(prompt, list):
+ system = next(
+ (str(m.get("content", "")) for m in prompt if isinstance(m, dict) and m.get("role") == "system"),
+ None,
+ )
+ user = "\n".join(
+ str(m.get("content", "")) for m in prompt if isinstance(m, dict) and m.get("role") != "system"
+ )
+ return system, user
+ return None, str(prompt)
+
+
+def format_conversation_with_tools(
+ prompt: str, tools: list[dict[str, Any]] = None, system_prompt: str = None, messages: list[dict[str, Any]] = None
+) -> str:
+ """Format conversation using Jinja2 template with tool support"""
+ template = Template(TOOL_TEMPLATE)
+
+ # Prepare messages
+ messages_to_render = []
+
+ # Always add system message - use provided one or default
+ if system_prompt:
+ system_content = system_prompt
+ else:
+ system_content = (
+ "You are a helpful assistant that can use Python "
+ "tools to solve mathematical problems. When you need "
+ "to perform calculations, use the code_interpreter "
+ "tool to execute code and get results."
+ )
+
+ messages_to_render.append({"role": "system", "content": system_content})
+
+ # Add user message if provided
+ if prompt:
+ messages_to_render.append({"role": "user", "content": prompt})
+
+ # Add assistant responses from previous turns if provided
+ if messages:
+ messages_to_render.extend(messages)
+
+ # Render template
+ formatted_text = template.render(messages=messages_to_render, tools=tools or [])
+
+ return formatted_text
+
+
+def postprocess_predictions(prediction: str):
+ """Extract action and content from prediction string"""
+ # Check for Answer: \boxed{...} format (only format we need for math_dapo)
+ # Use a more robust regex that handles nested braces
+ answer_pattern = r"Answer:\s*\\boxed\{((?:[^{}]|\{[^{}]*\})*)\}"
+ answer_match = re.search(answer_pattern, prediction, re.DOTALL)
+ if answer_match:
+ content = answer_match.group(1).strip()
+ return "answer", content
+
+ # Then check for tags (new format from Jinja2 template)
+ tool_call_pattern = r"\s*(\{.*?\})\s*"
+ tool_call_match = re.search(tool_call_pattern, prediction, re.DOTALL)
+ if tool_call_match:
+ try:
+ import json
+
+ json_str = tool_call_match.group(1)
+ try:
+ tool_call_data = json.loads(json_str)
+ except json.JSONDecodeError:
+ # Raw newlines inside the "code" string are invalid JSON; escaping
+ # recovers them. Only as a fallback -- escaping unconditionally
+ # breaks pretty-printed JSON, whose newlines are between tokens.
+ tool_call_data = json.loads(json_str.replace("\n", "\\n"))
+ tool_name = tool_call_data.get("name")
+ arguments = tool_call_data.get("arguments", {})
+
+ if tool_name == "code_interpreter":
+ code = arguments.get("code", "")
+ if code.strip():
+ return "code", code
+ except (json.JSONDecodeError, KeyError, AttributeError):
+ pass
+
+ # Then check for tags
+ code_pattern = r"(.*?)"
+ code_match = re.search(code_pattern, prediction, re.DOTALL)
+ if code_match:
+ content = code_match.group(1).strip()
+ return "code", content
+
+ # Finally check for ```python code blocks (lowest priority)
+ python_code_pattern = r"```python\s*(.*?)\s*```"
+ python_code_match = re.search(python_code_pattern, prediction, re.DOTALL)
+ if python_code_match:
+ content = python_code_match.group(1).strip()
+ return "code", content
+
+ return None, ""
+
+
+def postprocess_responses(resp: str) -> str:
+ """Post-process response to ensure tag completeness"""
+ # Handle tags (new format from Jinja2 template)
+ if "" in resp:
+ # Find the last occurrence of ...
+ tool_call_pattern = r"\s*\{.*?\}\s*"
+ matches = list(re.finditer(tool_call_pattern, resp, re.DOTALL))
+ if matches:
+ last_match = matches[-1]
+ return resp[: last_match.end()]
+
+ # Handle tags
+ if "" in resp:
+ return resp.split("")[0] + ""
+
+ # Handle ```python code blocks
+ if "```python" in resp:
+ # Find the last occurrence of ```python...```
+ python_pattern = r"```python\s*.*?```"
+ matches = list(re.finditer(python_pattern, resp, re.DOTALL))
+ if matches:
+ last_match = matches[-1]
+ return resp[: last_match.end()]
+
+ # Handle Answer: \boxed{...} format (only format we need for math_dapo)
+ if "Answer:" in resp and "\\boxed{" in resp:
+ # Find the last occurrence of Answer: \boxed{...} with nested braces support
+ answer_pattern = r"Answer:\s*\\boxed\{((?:[^{}]|\{[^{}]*\})*)\}"
+ matches = list(re.finditer(answer_pattern, resp, re.DOTALL))
+ if matches:
+ last_match = matches[-1]
+ return resp[: last_match.end()]
+
+ return resp
+
+
+async def execute_predictions(prediction: str) -> str:
+ """Execute predictions and return results"""
+ action, content = postprocess_predictions(prediction)
+
+ if action == "code":
+ # Content is already the Python code (extracted by
+ # postprocess_predictions)
+ code = content.strip()
+ if code:
+ # No SEMAPHORE acquire here: ``execute_tool`` already takes the same
+ # non-reentrant semaphore, and taking it twice deadlocks.
+ result = await tool_registry.execute_tool("code_interpreter", {"code": code})
+ next_obs = f"\n\n\n{result}\n\n\n"
+ done = False
+ else:
+ next_obs = "\n\n\nError: No Python code found" "\n\n\n"
+ done = False
+ elif action == "answer":
+ next_obs = ""
+ done = True
+ else:
+ next_obs = (
+ "\nMy previous action is invalid. "
+ "If I want to execute code, I should put the code between "
+ " and . "
+ "If I want to give the final answer, I should use the format "
+ "'Answer: \\boxed{answer}'. Let me try again.\n"
+ )
+ done = False
+
+ return next_obs, done
+
+
+def _parse_vllm_choice(choice: dict[str, Any]) -> tuple[list[int], list[float], dict[str, Any]]:
+ """Parse one vLLM ``/inference/v1/generate`` choice into tokens, logprobs and meta.
+
+ Returns ``log_probs=[]`` when the engine reports no per-token logprobs, so the
+ caller can abort instead of training on fabricated values.
+ """
+ tokens = [int(t) for t in (choice.get("token_ids") or [])]
+
+ log_probs: list[float] = []
+ lp = choice.get("logprobs")
+ if isinstance(lp, dict):
+ content_items = lp.get("content") or []
+ log_probs = [float(item.get("logprob", 0.0)) if isinstance(item, dict) else 0.0 for item in content_items]
+
+ # Normalize the bare vLLM ``finish_reason`` string into the nested shape.
+ fr = choice.get("finish_reason") or "stop"
+ if isinstance(fr, dict):
+ finish = fr
+ elif fr == "length":
+ finish = {"type": "length"}
+ elif fr in ("abort", "cancelled"):
+ finish = {"type": "abort"}
+ else:
+ finish = {"type": "stop"}
+
+ return tokens, log_probs, {"finish_reason": finish}
+
+
+async def generate(args, sample: Sample, sampling_params) -> Sample:
+ """Custom generation function supporting tool calls"""
+ assert not args.partial_rollout, "Partial rollout is not supported for " "this function at the moment."
+
+ # Retried samples (previously aborted / partial) arrive here with stale
+ # rollout state from the first attempt. Clear it so this generation starts
+ # clean; otherwise the concatenation below appends new tokens to old ones
+ # and downstream `slice_log_prob_with_cp` sees a length mismatch.
+ sample.rollout_log_probs = None
+ sample.rollout_top_p_token_ids = None
+ sample.rollout_top_p_token_offsets = None
+ sample.response = ""
+ sample.response_length = 0
+ sample.loss_mask = []
+
+ state = GenerateState(args)
+ url = f"http://{args.vllm_router_ip}:{args.vllm_router_port}/inference/v1/generate"
+
+ # Set up the initial prompt with system prompt and tools (outside the loop)
+ tool_specs = tool_registry.get_tool_specs()
+ system_prompt, user_text = split_prompt(sample.prompt)
+ prompt = format_conversation_with_tools(prompt=user_text, tools=tool_specs, system_prompt=system_prompt)
+
+ prompt_tokens_ids = state.tokenizer(prompt, add_special_tokens=False)["input_ids"]
+ sample.tokens = list(prompt_tokens_ids)
+ response = ""
+ response_token_ids = []
+ loss_masks = sample.loss_mask
+ tool_call_count = 0 # Track actual tool call rounds
+
+ if args.rollout_max_context_len is not None:
+ max_context_length = args.rollout_max_context_len
+ else:
+ max_context_length = args.context_parallel_size * args.max_tokens_per_gpu
+
+ meta_info = {"finish_reason": {"type": "stop"}}
+
+ for turn in range(TOOL_CONFIGS["max_turns"]):
+ # Check if total length exceeds max context length
+ total_length = len(prompt_tokens_ids) + len(response_token_ids)
+ if total_length >= max_context_length:
+ sample.status = Sample.Status.TRUNCATED
+ break
+
+ # Clamp per-turn max_new_tokens to the remaining context budget so a
+ # single turn cannot push total_length past max_context_length. Without
+ # this, a turn can append up to rollout_max_response_len tokens on top
+ # of a total that was just barely under the cap, producing samples
+ # that exceed the training-side max_tokens_per_gpu * cp_size budget
+ # and crash the partition/batch code (asserts or OOMs on an oversized
+ # partition).
+ remaining_budget = max_context_length - total_length
+ per_turn_sampling_params = dict(sampling_params)
+ per_turn_sampling_params["max_new_tokens"] = min(
+ sampling_params.get("max_new_tokens", remaining_budget),
+ remaining_budget,
+ )
+
+ current_token_ids = prompt_tokens_ids + response_token_ids
+ payload = {
+ "model": args.hf_checkpoint,
+ "token_ids": current_token_ids,
+ "sampling_params": _build_inference_sampling_params(per_turn_sampling_params),
+ }
+
+ # Log payload to wandb for debugging
+ try:
+ import wandb
+
+ if wandb.run is not None:
+ # Count available tools (from tool_specs)
+ available_tools = len(tool_specs)
+ # Count tools used in the current response
+ tools_used = response.count("")
+
+ wandb.log(
+ {
+ "debug/payload_length": len(prompt + response),
+ "debug/available_tools": available_tools,
+ "debug/tools_used": tools_used,
+ "debug/turn": turn,
+ }
+ )
+ except ImportError:
+ pass # wandb not available
+
+ output = await post(url, payload)
+ cur_response_token_ids, cur_log_probs, meta_info = _parse_vllm_choice(output["choices"][0])
+
+ # Handle abort
+ if meta_info["finish_reason"]["type"] == "abort":
+ sample.status = Sample.Status.ABORTED
+ return sample
+
+ if not cur_log_probs or len(cur_log_probs) != len(cur_response_token_ids):
+ # Unlike `vllm_rollout.generate`, do NOT substitute zeros: fabricated
+ # logprobs silently corrupt the importance ratio. Abort so the group
+ # goes back to the buffer for retry.
+ sample.status = Sample.Status.ABORTED
+ return sample
+
+ cur_response = state.tokenizer.decode(cur_response_token_ids)
+
+ response += cur_response
+ response_token_ids += cur_response_token_ids
+ sample.append_response_tokens(
+ args,
+ tokens=cur_response_token_ids,
+ log_probs=cur_log_probs,
+ trainable=True,
+ meta_info=meta_info,
+ )
+
+ # Check length limit
+ if meta_info["finish_reason"]["type"] == "length":
+ break
+
+ next_obs, done = await execute_predictions(cur_response)
+ if done:
+ break
+
+ # Count tool calls (when we get interpreter output, it means a tool
+ # was called)
+ if "" in next_obs:
+ tool_call_count += 1
+
+ assert next_obs != "", "Next observation should not be empty."
+ obs_tokens_ids = state.tokenizer(next_obs, add_special_tokens=False)["input_ids"]
+ overflow = len(prompt_tokens_ids) + len(response_token_ids) + len(obs_tokens_ids) - max_context_length
+ truncated_by_observation = overflow > 0
+ if truncated_by_observation:
+ obs_tokens_ids = obs_tokens_ids[: max(0, len(obs_tokens_ids) - overflow)]
+
+ # Add dummy log probs for observation tokens (they won't be used due to loss_mask=0)
+ # Check if maximum tool call count reached
+ response_token_ids += obs_tokens_ids
+ sample.append_response_tokens(args, tokens=obs_tokens_ids, trainable=False)
+
+ if sample.rollout_log_probs is not None:
+ assert len(response_token_ids) == len(
+ sample.rollout_log_probs
+ ), f"Token/logp length mismatch at turn {turn}: {len(response_token_ids)} tokens vs {len(sample.rollout_log_probs)} logps"
+
+ # Tool output is appended verbatim and can push total_length past
+ # max_context_length (the per-turn generation was clamped to the
+ # remaining budget, but tool output is unconstrained). Trim tail
+ # tokens so the final sample fits the training budget exactly.
+ if truncated_by_observation:
+ # Resync the text field from the trimmed token list so
+ # reward_func's `sample.prompt + sample.response` matches what
+ # the model was actually trained on. decode(tokenize(text)) can
+ # be lossy on some tokenizers (whitespace / special-token
+ # collapse), but reward_func's regex is whitespace-robust and
+ # the trainer sees tokens, not text — so the drift is safe.
+ response = state.tokenizer.decode(response_token_ids)
+ sample.status = Sample.Status.TRUNCATED
+ break
+ response += next_obs
+
+ if tool_call_count >= TOOL_CONFIGS["max_tool_calls"]:
+ break
+
+ # Set sample attributes
+ sample.tokens = prompt_tokens_ids + response_token_ids
+ sample.response_length = len(response_token_ids)
+ sample.response = response
+ sample.loss_mask = loss_masks
+
+ # Store payload information for wandb logging
+ sample.payload_text = prompt + response
+ sample.payload_has_system = "<|im_start|>system" in prompt + response
+ sample.payload_has_tools = "# Tools" in prompt + response
+
+ # Store tool call count for reward calculation
+ sample.tool_call_count = tool_call_count
+
+ # Set status
+ if sample.status is Sample.Status.PENDING:
+ match meta_info["finish_reason"]["type"]:
+ case "length":
+ sample.status = Sample.Status.TRUNCATED
+ case "abort":
+ sample.status = Sample.Status.ABORTED
+ case "stop":
+ sample.status = Sample.Status.COMPLETED
+
+ return sample
+
+
+async def reward_func(args, sample, **kwargs):
+ """Tool call reward function using math_dapo as primary reward model"""
+ if not isinstance(sample, Sample):
+ raise TypeError("Sample must be an instance of Sample class.")
+
+ # Build complete solution string. sample.prompt may be a chat-message list.
+ solution_str = split_prompt(sample.prompt)[1] + sample.response
+
+ # Get ground truth answer - label is a string, not a dict
+ ground_truth = sample.label if sample.label is not None else ""
+
+ # Get tool call count as num_turns
+ num_turns = getattr(sample, "tool_call_count", 0)
+
+ # use \\boxed{...} answer
+ result = math_dapo_compute_score(solution_str, ground_truth, strict_box_verify=True)
+
+ # encourage model to call tools
+ if result["score"] < 0:
+ tool_call_reward = (num_turns - 2) / 2 * 0.1
+ result["score"] = min(-0.6, result["score"] + tool_call_reward)
+
+ if result["pred"] is None:
+ result["pred"] = ""
+
+ return result
diff --git a/examples/retool/requirements.txt b/examples/retool/requirements.txt
new file mode 100644
index 000000000..154747336
--- /dev/null
+++ b/examples/retool/requirements.txt
@@ -0,0 +1,3 @@
+jinja2>=3.0.0
+psutil>=5.8.0
+pytest>=7.0.0
diff --git a/examples/retool/retool_qwen3_4b_rl.sh b/examples/retool/retool_qwen3_4b_rl.sh
new file mode 100644
index 000000000..add9f99a4
--- /dev/null
+++ b/examples/retool/retool_qwen3_4b_rl.sh
@@ -0,0 +1,159 @@
+#!/bin/bash
+
+# for rerun the task
+pkill -9 -f '[v]llm serve|VLL[M]::'
+sleep 3
+ray stop --force
+pkill -9 ray
+pkill -9 python
+sleep 3
+pkill -9 ray
+pkill -9 python
+
+set -ex
+
+# will prevent ray from buffering stdout/stderr
+export PYTHONUNBUFFERED=1
+
+NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l)
+if [ "$NVLINK_COUNT" -gt 0 ]; then
+ HAS_NVLINK=1
+else
+ HAS_NVLINK=0
+fi
+echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)"
+
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
+REPO_ROOT="$(cd -- "${SCRIPT_DIR}/../.." &>/dev/null && pwd)"
+# ReTool's SFT checkpoint uses rope theta 5e6, not qwen3-4B.sh's 1e6 default.
+MODEL_ARGS_ROTARY_BASE=5000000 source "${REPO_ROOT}/scripts/models/qwen3-4B.sh"
+
+CKPT_ARGS=(
+ --hf-checkpoint /root/font-info/qwen3-4b-sft
+ --ref-load /root/font-info/qwen3-4b-sft_torch_dist
+ # --load /root/Qwen3-4B_vime/
+ --save /root/font-info/qwen3-4b-sft/qwen3-4b-sft-multi-turn/
+ --save-interval 20
+)
+
+ROLLOUT_ARGS=(
+ --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl
+ --input-key prompt
+ --label-key label
+ # No --apply-chat-template: generate_with_retool renders the conversation
+ # itself, so enabling it here double-wraps the prompt.
+ --rollout-shuffle
+ --reward-key score
+ --num-rollout 3000
+ --rollout-batch-size 32
+ --n-samples-per-prompt 8
+ --rollout-max-response-len 8192
+ --rollout-temperature 1
+
+ --global-batch-size 256
+ --balance-data
+)
+
+EVAL_ARGS=(
+ --eval-interval 20
+ --eval-prompt-data aime /root/aime-2024/aime-2024.jsonl
+ --n-samples-per-eval-prompt 16
+ --eval-max-response-len 16384
+ --eval-top-p 1
+)
+
+PERF_ARGS=(
+ --tensor-model-parallel-size 2
+ --sequence-parallel
+ --pipeline-model-parallel-size 1
+ --context-parallel-size 1
+ --expert-model-parallel-size 1
+ --expert-tensor-parallel-size 1
+
+ --recompute-granularity full
+ --recompute-method uniform
+ --recompute-num-layers 1
+
+ # --micro-batch-size 1
+ --use-dynamic-batch-size
+ --max-tokens-per-gpu 9216
+)
+
+GRPO_ARGS=(
+ --advantage-estimator grpo
+ --use-kl-loss
+ --kl-loss-coef 0.00
+ --kl-loss-type low_var_kl
+ --entropy-coef 0.00
+ --eps-clip 0.2
+ --eps-clip-high 0.28
+)
+
+OPTIMIZER_ARGS=(
+ --optimizer adam
+ --lr 1e-6
+ --lr-decay-style constant
+ --weight-decay 0.1
+ --adam-beta1 0.9
+ --adam-beta2 0.98
+)
+
+WANDB_ARGS=(
+ --use-wandb
+ --wandb-project vime-dapo
+ --wandb-group qwen3-4B-test-multi-turn
+ --wandb-key "${WANDB_KEY}"
+)
+
+VLLM_ARGS=(
+ --rollout-num-gpus-per-engine 2
+ --vllm-gpu-memory-utilization 0.7
+)
+
+MISC_ARGS=(
+ # default dropout in megatron is 0.1
+ --attention-dropout 0.0
+ --hidden-dropout 0.0
+ # should be good for model performance
+ --accumulate-allreduce-grads-in-fp32
+ --attention-softmax-in-fp32
+ # need to comment this when using model with MLA
+ --attention-backend flash
+)
+
+CUSTOM_ARGS=(
+ --custom-generate-function-path generate_with_retool.generate
+ --custom-rm-path generate_with_retool.reward_func
+)
+
+# launch the master node of ray in container
+export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
+ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 4 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265
+
+# SCRIPT_DIR is on PYTHONPATH so `generate_with_retool` resolves as a top-level
+# module and can import its sibling `tool_sandbox` (as examples/tau-bench does).
+RUNTIME_ENV_JSON="{
+ \"env_vars\": {
+ \"PYTHONPATH\": \"/root/Megatron-LM/:${SCRIPT_DIR}:${REPO_ROOT}\",
+ \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\",
+ \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\"
+ }
+}"
+
+ray job submit --address="http://127.0.0.1:8265" \
+ --runtime-env-json="${RUNTIME_ENV_JSON}" \
+ -- python3 "${REPO_ROOT}/train.py" \
+ --actor-num-nodes 1 \
+ --actor-num-gpus-per-node 4 \
+ --colocate \
+ ${MODEL_ARGS[@]} \
+ ${CKPT_ARGS[@]} \
+ ${ROLLOUT_ARGS[@]} \
+ ${OPTIMIZER_ARGS[@]} \
+ ${GRPO_ARGS[@]} \
+ ${WANDB_ARGS[@]} \
+ ${PERF_ARGS[@]} \
+ ${EVAL_ARGS[@]} \
+ ${VLLM_ARGS[@]} \
+ ${MISC_ARGS[@]} \
+ ${CUSTOM_ARGS[@]}
diff --git a/examples/retool/retool_qwen3_4b_sft.sh b/examples/retool/retool_qwen3_4b_sft.sh
new file mode 100644
index 000000000..07317a448
--- /dev/null
+++ b/examples/retool/retool_qwen3_4b_sft.sh
@@ -0,0 +1,127 @@
+#!/bin/bash
+
+# for rerun the task
+pkill -9 -f '[v]llm serve|VLL[M]::'
+sleep 3
+ray stop --force
+pkill -9 ray
+pkill -9 python
+sleep 3
+pkill -9 ray
+pkill -9 python
+
+set -ex
+
+# will prevent ray from buffering stdout/stderr
+export PYTHONUNBUFFERED=1
+
+NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l)
+if [ "$NVLINK_COUNT" -gt 0 ]; then
+ HAS_NVLINK=1
+else
+ HAS_NVLINK=0
+fi
+echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)"
+
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
+REPO_ROOT="$(cd -- "${SCRIPT_DIR}/../.." &>/dev/null && pwd)"
+# Qwen3-4B-Instruct-2507 uses rope theta 5e6, not qwen3-4B.sh's 1e6 default.
+MODEL_ARGS_ROTARY_BASE=5000000 source "${REPO_ROOT}/scripts/models/qwen3-4B.sh"
+
+CKPT_ARGS=(
+ --hf-checkpoint /root/Qwen/Qwen3-4B-Instruct-2507/
+ --ref-load /root/Qwen/Qwen3-4B-Instruct-2507_torch_dist
+# --load ./models/Qwen/Qwen3-4B-Instruct_vime/
+ --save /root/Qwen/Qwen3-4B-Instruct-2507_sft_vime/
+ --save-interval 1000
+)
+
+SFT_ARGS=(
+ --rollout-function-path vime.rollout.sft_rollout.generate_rollout
+ --prompt-data ./data/retool/ReTool-SFT.parquet
+ --input-key messages
+ --rollout-shuffle
+ --num-epoch 3
+ --rollout-batch-size 128
+ --global-batch-size 128
+
+ --loss-type sft_loss
+ --calculate-per-token-loss
+ --disable-compute-advantages-and-returns
+ --debug-train-only
+)
+
+PERF_ARGS=(
+ --tensor-model-parallel-size 1
+ --sequence-parallel
+ --pipeline-model-parallel-size 1
+ --context-parallel-size 1
+ --expert-model-parallel-size 1
+ --expert-tensor-parallel-size 1
+
+ --recompute-granularity full
+ --recompute-method uniform
+ --recompute-num-layers 1
+
+ # --micro-batch-size 1
+ --use-dynamic-batch-size
+ --max-tokens-per-gpu 9216
+)
+
+OPTIMIZER_ARGS=(
+ --optimizer adam
+ --lr 1e-5
+ --lr-decay-style cosine
+ --min-lr 1e-6
+ --lr-warmup-fraction 0.1
+ --weight-decay 0.1
+ --adam-beta1 0.9
+ --adam-beta2 0.95
+)
+
+WANDB_ARGS=(
+ --use-wandb
+ --wandb-project vime-dev
+ --wandb-group qwen3-4B-base-sft
+ --wandb-key "${WANDB_KEY}"
+)
+
+MISC_ARGS=(
+ # default dropout in megatron is 0.1
+ --attention-dropout 0.0
+ --hidden-dropout 0.0
+ # should be good for model performance
+ --accumulate-allreduce-grads-in-fp32
+ --attention-softmax-in-fp32
+ # need to comment this when using model with MLA
+ --attention-backend flash
+)
+
+# launch the master node of ray in container
+export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
+export no_proxy="127.0.0.1,${MASTER_ADDR}"
+ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265
+
+
+# Build the runtime environment JSON with proper variable substitution
+RUNTIME_ENV_JSON="{
+ \"env_vars\": {
+ \"PYTHONPATH\": \"/root/Megatron-LM/:${REPO_ROOT}\",
+ \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\",
+ \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\",
+ \"PYTORCH_CUDA_ALLOC_CONF\": \"expandable_segments:True\"
+ }
+}"
+
+ray job submit --address="http://127.0.0.1:8265" \
+ --runtime-env-json="${RUNTIME_ENV_JSON}" \
+ -- python3 "${REPO_ROOT}/train_async.py" \
+ --actor-num-nodes 1 \
+ --actor-num-gpus-per-node 8 \
+ ${MODEL_ARGS[@]} \
+ ${CKPT_ARGS[@]} \
+ ${SFT_ARGS[@]} \
+ ${OPTIMIZER_ARGS[@]} \
+ ${WANDB_ARGS[@]} \
+ ${PERF_ARGS[@]} \
+ ${MISC_ARGS[@]}
diff --git a/examples/retool/rl_data_preprocess.py b/examples/retool/rl_data_preprocess.py
new file mode 100644
index 000000000..548b74468
--- /dev/null
+++ b/examples/retool/rl_data_preprocess.py
@@ -0,0 +1,21 @@
+from datasets import load_dataset
+
+# Load the original dataset
+ds = load_dataset("BytedTsinghua-SIA/DAPO-Math-17k", split="train")
+
+
+# Map to extract the ground_truth from the reward_model dict and create a new 'label' field
+def transform(example):
+ return {
+ "prompt": example["prompt"][0]["content"] if example["prompt"] else None,
+ "label": example["reward_model"]["ground_truth"],
+ }
+
+
+ds2 = ds.map(transform, remove_columns=ds.column_names)
+
+# Optionally, verify the first few entries
+print(ds2[0])
+
+# save to jsonl
+ds2.to_json("/root/dapo-math-17k-processed/dapo_math_17k_cleaned.jsonl", orient="records", lines=True)
diff --git a/examples/retool/sft_data_processing.py b/examples/retool/sft_data_processing.py
new file mode 100644
index 000000000..96790af1e
--- /dev/null
+++ b/examples/retool/sft_data_processing.py
@@ -0,0 +1,31 @@
+from datasets import load_dataset
+
+ds = load_dataset("JoeYing/ReTool-SFT")["train"]
+
+
+def convert(sample):
+ conversations = sample["messages"]
+
+ def convert_role(role):
+ if role == "user":
+ return "user"
+ elif role == "assistant":
+ return "assistant"
+ elif role == "system":
+ return "system"
+ else:
+ raise ValueError(f"Unknown role: {role}")
+
+ messages = [
+ {
+ "role": convert_role(turn["role"]),
+ "content": turn["content"],
+ }
+ for turn in conversations
+ ]
+
+ return {"messages": messages}
+
+
+ds = ds.map(convert)
+ds.to_parquet("./data/retool/ReTool-SFT.parquet")
diff --git a/examples/retool/tool_sandbox.py b/examples/retool/tool_sandbox.py
new file mode 100644
index 000000000..cdf68aa02
--- /dev/null
+++ b/examples/retool/tool_sandbox.py
@@ -0,0 +1,360 @@
+"""
+Tool sandbox module for safe code execution and tool management.
+
+This module provides:
+- PythonSandbox: Safe Python code execution environment
+- ToolRegistry: Tool registration and execution management
+- Memory management utilities
+"""
+
+import asyncio
+import gc
+import os
+import re
+import subprocess
+import tempfile
+from contextlib import contextmanager
+from typing import Any
+
+import psutil
+
+# Configuration for tool execution
+TOOL_CONFIGS = {
+ "max_turns": 16,
+ "max_tool_calls": 16,
+ "tool_concurrency": 32, # Aggressive: 32 concurrent processes
+ # Python interpreter settings
+ "python_timeout": 120, # 2 minutes for complex calculations
+ "python_memory_limit": "4GB", # 4GB per Python process
+ "python_cpu_limit": 1,
+ # Memory management settings
+ "max_memory_usage": 12288, # 12GB total (75% of 16GB)
+ "cleanup_threshold": 6144, # 6GB
+ "aggressive_cleanup_threshold": 3072, # 3GB
+ "force_cleanup_threshold": 9216, # 9GB
+}
+
+# Global semaphore for controlling concurrent tool executions
+SEMAPHORE = asyncio.Semaphore(TOOL_CONFIGS["tool_concurrency"])
+
+
+def get_memory_usage() -> float:
+ """Get current memory usage in MB"""
+ process = psutil.Process()
+ return process.memory_info().rss / 1024 / 1024
+
+
+def cleanup_memory():
+ """Force garbage collection to free memory"""
+ gc.collect()
+
+
+def aggressive_cleanup_memory():
+ """More aggressive memory cleanup"""
+ # Force multiple garbage collection cycles
+ for _ in range(3):
+ gc.collect()
+
+ # Clear Python's internal caches
+ import sys
+
+ # Note: sys.intern doesn't have a clear method, so we skip this
+ # Clear module cache if possible
+ if hasattr(sys, "modules"):
+ # Don't clear all modules, but clear some common ones that might cache data
+ modules_to_clear = ["numpy", "pandas", "matplotlib", "scipy"]
+ for module_name in modules_to_clear:
+ if module_name in sys.modules:
+ module = sys.modules[module_name]
+ if hasattr(module, "clear_cache"):
+ module.clear_cache()
+
+
+def check_and_cleanup_memory():
+ """Check memory usage and perform appropriate cleanup"""
+ current_memory = get_memory_usage()
+
+ if current_memory > TOOL_CONFIGS["force_cleanup_threshold"]:
+ # Force aggressive cleanup
+ aggressive_cleanup_memory()
+ return f"Warning: High memory usage ({current_memory:.1f}MB), performed aggressive cleanup"
+ elif current_memory > TOOL_CONFIGS["cleanup_threshold"]:
+ # Normal cleanup
+ cleanup_memory()
+ return f"Info: Memory usage ({current_memory:.1f}MB), performed cleanup"
+ elif current_memory > TOOL_CONFIGS["aggressive_cleanup_threshold"]:
+ # Light cleanup
+ gc.collect()
+ return f"Info: Memory usage ({current_memory:.1f}MB), performed light cleanup"
+
+ return None
+
+
+class PythonSandbox:
+ """Python code sandbox, provides safe code execution environment"""
+
+ def __init__(self, timeout: int = 10, memory_limit: str = "100MB"):
+ self.timeout = timeout
+ self.memory_limit = memory_limit
+ self.allowed_modules = {
+ "math",
+ "random",
+ "datetime",
+ "collections",
+ "itertools",
+ "functools",
+ "operator",
+ "statistics",
+ "decimal",
+ "fractions",
+ }
+
+ def _check_code_safety(self, code: str) -> tuple[bool, str]:
+ """Check code safety by scanning for dangerous patterns"""
+ # Check for dangerous operations
+ dangerous_patterns = [
+ r"import\s+os",
+ r"import\s+sys",
+ r"import\s+subprocess",
+ r"import\s+shutil",
+ r"import\s+glob",
+ r"import\s+pathlib",
+ r"__import__",
+ r"eval\s*\(",
+ r"exec\s*\(",
+ r"open\s*\(",
+ r"file\s*\(",
+ r"input\s*\(",
+ r"raw_input\s*\(",
+ r"compile\s*\(",
+ r"execfile\s*\(",
+ r"getattr\s*\(",
+ r"setattr\s*\(",
+ r"delattr\s*\(",
+ r"hasattr\s*\(",
+ r"globals\s*\(",
+ r"locals\s*\(",
+ r"vars\s*\(",
+ r"dir\s*\(",
+ r"type\s*\(",
+ r"isinstance\s*\(",
+ r"issubclass\s*\(",
+ r"super\s*\(",
+ r"property\s*\(",
+ r"staticmethod\s*\(",
+ r"classmethod\s*\(",
+ r"__\w+__", # double underscore methods
+ ]
+
+ for pattern in dangerous_patterns:
+ if re.search(pattern, code, re.IGNORECASE):
+ return False, f"Code contains dangerous pattern: {pattern}"
+
+ # Check imported modules
+ import_pattern = r"import\s+(\w+)"
+ from_pattern = r"from\s+(\w+)"
+
+ imports = re.findall(import_pattern, code)
+ froms = re.findall(from_pattern, code)
+
+ all_imports = set(imports + froms)
+ for imp in all_imports:
+ if imp not in self.allowed_modules:
+ return False, f"Import of '{imp}' is not allowed"
+
+ return True, "Code is safe"
+
+ @contextmanager
+ def _create_safe_environment(self):
+ """Create safe execution environment with temporary directory"""
+ # Create temporary directory
+ temp_dir = tempfile.mkdtemp(prefix="python_sandbox_")
+
+ try:
+ # Create safe Python script
+ script_path = os.path.join(temp_dir, "code.py")
+
+ # Set environment variables
+ env = os.environ.copy()
+ env["PYTHONPATH"] = temp_dir
+ env["PYTHONUNBUFFERED"] = "1"
+
+ yield script_path, env, temp_dir
+
+ finally:
+ # Clean up temporary directory
+ try:
+ import shutil
+
+ shutil.rmtree(temp_dir)
+ except Exception:
+ pass
+
+ async def execute_code(self, code: str) -> str:
+ """Execute Python code in sandbox with safety checks"""
+ # Check memory usage before execution
+ current_memory = get_memory_usage()
+ if current_memory > TOOL_CONFIGS["max_memory_usage"]:
+ aggressive_cleanup_memory()
+ return "Error: Memory usage too high, please try again"
+
+ # Check code safety
+ is_safe, message = self._check_code_safety(code)
+ if not is_safe:
+ return f"Error: {message}"
+
+ # Add necessary wrapper code with memory limits
+ # Properly indent the user code within the try block
+ # Handle indentation properly by adding 4 spaces to each line
+ indented_code = "\n".join(" " + line for line in code.split("\n"))
+
+ wrapped_code = f"""import sys
+import traceback
+from io import StringIO
+import resource
+
+# Set memory limit (4GB)
+try:
+ resource.setrlimit(resource.RLIMIT_AS, (4 * 1024 * 1024 * 1024, -1))
+except Exception:
+ pass
+
+# Redirect stdout and stderr
+old_stdout = sys.stdout
+old_stderr = sys.stderr
+stdout_capture = StringIO()
+stderr_capture = StringIO()
+sys.stdout = stdout_capture
+sys.stderr = stderr_capture
+
+try:
+ # User code
+{indented_code}
+
+ # Get output
+ stdout_output = stdout_capture.getvalue()
+ stderr_output = stderr_capture.getvalue()
+
+ # Restore standard output
+ sys.stdout = old_stdout
+ sys.stderr = old_stderr
+
+ # Return result
+ result = ""
+ if stdout_output:
+ result += f"Output:\\n{{stdout_output}}"
+ if stderr_output:
+ result += f"\\nErrors:\\n{{stderr_output}}"
+
+ print(result)
+
+except Exception as e:
+ # Restore standard output
+ sys.stdout = old_stdout
+ sys.stderr = old_stderr
+
+ # Return error information
+ error_msg = f"Error: {{str(e)}}\\nTraceback:\\n{{traceback.format_exc()}}"
+ print(error_msg)"""
+
+ with self._create_safe_environment() as (script_path, env, temp_dir):
+ # Write code to file
+ with open(script_path, "w") as f:
+ f.write(wrapped_code)
+
+ try:
+ # Use subprocess to run code
+ process = subprocess.Popen(
+ ["python3", script_path],
+ stdout=subprocess.PIPE,
+ stderr=subprocess.PIPE,
+ env=env,
+ cwd=temp_dir,
+ text=True,
+ )
+
+ # Set timeout
+ try:
+ stdout, stderr = process.communicate(timeout=self.timeout)
+
+ if process.returncode == 0:
+ result = stdout.strip()
+ else:
+ result = f"Error: Process exited with code {process.returncode}\n{stderr}"
+
+ except subprocess.TimeoutExpired:
+ process.kill()
+ result = f"Error: Code execution timed out after {self.timeout} seconds"
+
+ except Exception as e:
+ result = f"Error: Failed to execute code: {str(e)}"
+
+ # Check memory usage after execution and cleanup if needed
+ cleanup_message = check_and_cleanup_memory()
+ if cleanup_message:
+ print(f"Memory cleanup: {cleanup_message}")
+
+ return result
+
+
+class ToolRegistry:
+ """Tool registry, manages available tools and their execution"""
+
+ def __init__(self):
+ self.tools = {}
+ self.python_sandbox = PythonSandbox(
+ timeout=TOOL_CONFIGS["python_timeout"], memory_limit=TOOL_CONFIGS["python_memory_limit"]
+ )
+ self._register_default_tools()
+
+ def _register_default_tools(self):
+ """Register default tools in the registry"""
+ # Python code interpreter
+ self.register_tool(
+ "code_interpreter",
+ {
+ "type": "function",
+ "function": {
+ "name": "code_interpreter",
+ "description": "A tool for executing Python code in a safe sandbox environment.",
+ "parameters": {
+ "type": "object",
+ "properties": {"code": {"type": "string", "description": "The Python code to execute"}},
+ "required": ["code"],
+ },
+ },
+ },
+ )
+
+ def register_tool(self, name: str, tool_spec: dict[str, Any]):
+ """Register a new tool in the registry"""
+ self.tools[name] = tool_spec
+
+ def get_tool_specs(self) -> list[dict[str, Any]]:
+ """Get all tool specifications as a list"""
+ return list(self.tools.values())
+
+ async def execute_tool(self, tool_name: str, arguments: dict[str, Any]) -> str:
+ """Execute a tool call with the given arguments"""
+ if tool_name not in self.tools:
+ return f"Error: Tool '{tool_name}' not found"
+
+ async with SEMAPHORE:
+ if tool_name == "code_interpreter":
+ return await self._execute_python(arguments)
+ else:
+ return f"Error: Tool '{tool_name}' not implemented"
+
+ async def _execute_python(self, arguments: dict[str, Any]) -> str:
+ """Execute Python code using the sandbox"""
+ code = arguments.get("code", "")
+ if not code.strip():
+ return "Error: No code provided"
+
+ # Execute code in sandbox
+ result = await self.python_sandbox.execute_code(code)
+ return result
+
+
+# Global tool registry instance
+tool_registry = ToolRegistry()
diff --git a/tests/test_retool_generate.py b/tests/test_retool_generate.py
new file mode 100644
index 000000000..50d25f5f2
--- /dev/null
+++ b/tests/test_retool_generate.py
@@ -0,0 +1,664 @@
+"""CPU unit tests for the ``examples/retool`` vLLM rollout port.
+
+Covers the parts of the example that the slime -> vime port actually changed:
+the ``/inference/v1/generate`` request body, the ``choices[0]`` response parse,
+the multi-turn tool loop, and the tool-concurrency limit. The engine is mocked,
+so no GPU or running router is required.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import sys
+import types
+from argparse import Namespace
+from pathlib import Path
+
+_tests_root = Path(__file__).resolve().parent
+if str(_tests_root) not in sys.path:
+ sys.path.insert(0, str(_tests_root))
+
+import _unit_stubs
+import pytest
+
+_unit_stubs.install_rollout_optional_stubs()
+
+if not _unit_stubs.real_module_available("psutil"):
+ # tool_sandbox uses psutil only for RSS-based cleanup heuristics.
+ _psutil = types.ModuleType("psutil")
+
+ class _FakeProcess:
+ def memory_info(self):
+ return types.SimpleNamespace(rss=64 * 1024 * 1024)
+
+ _psutil.Process = _FakeProcess
+ sys.modules["psutil"] = _psutil
+
+# The RL script puts the example dir on PYTHONPATH so `generate_with_retool`
+# resolves as a top-level module and can import its sibling `tool_sandbox`.
+_RETOOL_DIR = _tests_root.parent / "examples" / "retool"
+if str(_RETOOL_DIR) not in sys.path:
+ sys.path.insert(0, str(_RETOOL_DIR))
+
+import generate_with_retool as mod # noqa: E402
+import tool_sandbox # noqa: E402
+
+from vime.utils.types import Sample # noqa: E402
+
+NUM_GPUS = 0
+
+
+class _FakeTokenizer:
+ """Character-code tokenizer: reversible, so decode(encode(t)) == t."""
+
+ def __call__(self, text: str, add_special_tokens: bool = False):
+ assert add_special_tokens is False
+ return {"input_ids": [ord(c) for c in text]}
+
+ def decode(self, token_ids, skip_special_tokens: bool = True) -> str:
+ return "".join(chr(int(t)) for t in token_ids)
+
+
+class _FakeState:
+ def __init__(self, args):
+ self.tokenizer = _FakeTokenizer()
+ self.processor = None
+
+
+def _args(**overrides) -> Namespace:
+ args = Namespace(
+ partial_rollout=False,
+ hf_checkpoint="/fake/qwen3-4b",
+ vllm_router_ip="127.0.0.1",
+ vllm_router_port=3250,
+ rollout_max_context_len=4096,
+ context_parallel_size=1,
+ max_tokens_per_gpu=4096,
+ vllm_speculative_config=None,
+ num_layers=2,
+ moe_router_topk=1,
+ )
+ for key, value in overrides.items():
+ setattr(args, key, value)
+ return args
+
+
+def _sampling_params(**overrides) -> dict:
+ sp = {"max_new_tokens": 256, "temperature": 1.0, "top_p": 1.0}
+ sp.update(overrides)
+ return sp
+
+
+def _choice(text: str, finish_reason: str = "stop", *, logprobs: bool = True) -> dict:
+ """Build a vLLM ``/inference/v1/generate`` choice for `text`."""
+ token_ids = [ord(c) for c in text]
+ choice: dict = {"token_ids": token_ids, "finish_reason": finish_reason}
+ if logprobs:
+ choice["logprobs"] = {"content": [{"logprob": -0.5} for _ in token_ids]}
+ return choice
+
+
+def _pending_sample(prompt: str = "2+2?") -> Sample:
+ return Sample(prompt=prompt, label="4", status=Sample.Status.PENDING)
+
+
+def _prompt_len(prompt: str = "2+2?") -> int:
+ """Token length of the rendered tool-enabled prompt, per _FakeTokenizer."""
+ rendered = mod.format_conversation_with_tools(prompt=prompt, tools=mod.tool_registry.get_tool_specs())
+ return len(_FakeTokenizer()(rendered)["input_ids"])
+
+
+@pytest.fixture(autouse=True)
+def _patch_state(monkeypatch):
+ monkeypatch.setattr(mod, "GenerateState", _FakeState)
+
+
+@pytest.fixture(autouse=True)
+def _stub_tool_subprocess(monkeypatch):
+ """Keep the rollout tests off real `python3` subprocesses.
+
+ They exercise the turn loop, not the sandbox, and spawning a subprocess per
+ tool call makes them slow and dependent on the runner's environment.
+ `test_real_sandbox_executes_code` covers real execution explicitly.
+ """
+
+ async def fake_execute_code(code):
+ return "Output:\n4"
+
+ monkeypatch.setattr(tool_sandbox.tool_registry.python_sandbox, "execute_code", fake_execute_code)
+
+
+def _run_generate(monkeypatch, responses, *, args=None, sample=None, sampling_params=None):
+ """Drive mod.generate with a scripted list of engine `choices`, capturing payloads."""
+ payloads: list[dict] = []
+ queue = list(responses)
+
+ async def fake_post(url, payload, **kwargs):
+ payloads.append({"url": url, "payload": payload})
+ assert queue, "engine called more times than the test scripted"
+ return {"choices": [queue.pop(0)]}
+
+ monkeypatch.setattr(mod, "post", fake_post)
+ result = asyncio.run(
+ mod.generate(
+ args or _args(),
+ sample if sample is not None else _pending_sample(),
+ sampling_params or _sampling_params(),
+ )
+ )
+ return result, payloads
+
+
+# --------------------------------------------------------------------------
+# response parsing (the ported SGLang -> vLLM surface)
+# --------------------------------------------------------------------------
+
+
+def test_parse_vllm_choice_reads_tokens_and_logprobs():
+ tokens, log_probs, meta = mod._parse_vllm_choice(_choice("hi"))
+ assert tokens == [ord("h"), ord("i")]
+ assert log_probs == [-0.5, -0.5]
+ assert meta == {"finish_reason": {"type": "stop"}}
+
+
+@pytest.mark.parametrize(
+ ("engine_finish_reason", "expected_type"),
+ [("stop", "stop"), ("length", "length"), ("abort", "abort"), ("cancelled", "abort"), (None, "stop")],
+)
+def test_parse_vllm_choice_normalizes_finish_reason(engine_finish_reason, expected_type):
+ choice = _choice("x", finish_reason=engine_finish_reason)
+ _, _, meta = mod._parse_vllm_choice(choice)
+ assert meta["finish_reason"] == {"type": expected_type}
+
+
+def test_parse_vllm_choice_passes_through_nested_finish_reason():
+ """Defensive: a dict finish_reason (SGLang shape) is used as-is."""
+ _, _, meta = mod._parse_vllm_choice({"token_ids": [1], "finish_reason": {"type": "length"}})
+ assert meta["finish_reason"] == {"type": "length"}
+
+
+def test_parse_vllm_choice_reports_missing_logprobs_instead_of_zero_filling():
+ tokens, log_probs, _ = mod._parse_vllm_choice(_choice("hi", logprobs=False))
+ assert tokens == [ord("h"), ord("i")]
+ # Empty, NOT [0.0, 0.0] -- generate() must abort rather than train on fakes.
+ assert log_probs == []
+
+
+# --------------------------------------------------------------------------
+# request body
+# --------------------------------------------------------------------------
+
+
+def test_generate_posts_token_ids_body_to_inference_endpoint(monkeypatch):
+ _, payloads = _run_generate(monkeypatch, [_choice("Answer: \\boxed{4}")])
+
+ assert len(payloads) == 1
+ assert payloads[0]["url"] == "http://127.0.0.1:3250/inference/v1/generate"
+ body = payloads[0]["payload"]
+ assert body["model"] == "/fake/qwen3-4b"
+ assert isinstance(body["token_ids"], list) and body["token_ids"]
+ # SGLang's `input_ids` / `return_logprob` must not survive the port.
+ assert "input_ids" not in body
+ assert "return_logprob" not in body
+ # _build_inference_sampling_params renames max_new_tokens and asks for logprobs.
+ assert body["sampling_params"]["max_tokens"] == 256
+ assert body["sampling_params"]["logprobs"] == 1
+ assert "max_new_tokens" not in body["sampling_params"]
+
+
+def test_generate_prompt_includes_tool_specs(monkeypatch):
+ sample = _pending_sample()
+ _, payloads = _run_generate(monkeypatch, [_choice("Answer: \\boxed{4}")], sample=sample)
+
+ prompt_text = _FakeTokenizer().decode(payloads[0]["payload"]["token_ids"])
+ assert "# Tools" in prompt_text
+ assert "code_interpreter" in prompt_text
+ assert sample.payload_has_tools is True
+ assert sample.payload_has_system is True
+
+
+# DAPO-Math-17k and AIME-2024 both store `prompt` as a chat-message list, and
+# `_build_messages` passes a list through untouched without --apply-chat-template.
+LIST_PROMPT = [{"role": "user", "content": "What is 2+2?"}]
+LIST_PROMPT_WITH_SYSTEM = [
+ {"role": "system", "content": "You are terse."},
+ {"role": "user", "content": "What is 2+2?"},
+]
+
+
+@pytest.mark.parametrize(
+ ("prompt", "expected"),
+ [
+ ("plain string", (None, "plain string")),
+ (LIST_PROMPT, (None, "What is 2+2?")),
+ (LIST_PROMPT_WITH_SYSTEM, ("You are terse.", "What is 2+2?")),
+ ([], (None, "")),
+ ],
+)
+def test_split_prompt_handles_both_dataset_shapes(prompt, expected):
+ assert mod.split_prompt(prompt) == expected
+
+
+def test_generate_accepts_a_chat_message_list_prompt(monkeypatch):
+ """The real datasets ship list prompts; rendering must still be single-templated."""
+ sample = Sample(prompt=LIST_PROMPT, label="4", status=Sample.Status.PENDING)
+ _, payloads = _run_generate(monkeypatch, [_choice("Answer: \\boxed{4}")], sample=sample)
+
+ rendered = _FakeTokenizer().decode(payloads[0]["payload"]["token_ids"])
+ assert "What is 2+2?" in rendered
+ assert rendered.count("<|im_start|>user") == 1, "list prompt must not double-wrap"
+ assert rendered.count("<|im_start|>assistant") == 1
+ # Passing the list straight to the template renders its repr, which still
+ # *contains* the question -- so assert the repr artifacts are absent instead.
+ assert "'role'" not in rendered, f"prompt list was rendered as a repr: {rendered[-200:]}"
+ assert "{'" not in rendered
+
+
+def test_generate_uses_a_system_message_from_the_prompt_list(monkeypatch):
+ sample = Sample(prompt=LIST_PROMPT_WITH_SYSTEM, label="4", status=Sample.Status.PENDING)
+ _, payloads = _run_generate(monkeypatch, [_choice("Answer: \\boxed{4}")], sample=sample)
+
+ rendered = _FakeTokenizer().decode(payloads[0]["payload"]["token_ids"])
+ assert "You are terse." in rendered
+ assert rendered.count("<|im_start|>system") == 1
+
+
+def test_reward_func_accepts_a_chat_message_list_prompt():
+ """Regression: `sample.prompt + sample.response` raised TypeError on the AIME
+ eval set, crashing RolloutManager.eval() at the first --eval-interval."""
+ sample = Sample(prompt=LIST_PROMPT, label="4", status=Sample.Status.COMPLETED)
+ sample.response = " Answer: \\boxed{4}"
+ result = asyncio.run(mod.reward_func(_args(), sample))
+ assert result["score"] > 0
+
+
+def test_prompt_has_exactly_one_conversation_structure():
+ """Guards against the nested-`user` prompt that --apply-chat-template produces."""
+ rendered = mod.format_conversation_with_tools(prompt="2+2?", tools=mod.tool_registry.get_tool_specs())
+
+ assert rendered.count("<|im_start|>system") == 1
+ assert rendered.count("<|im_start|>user") == 1
+ assert rendered.count("<|im_start|>assistant") == 1
+ # Jinja strips the template's trailing newline, so the open generation turn is
+ # `<|im_start|>assistant` with no "\n" (upstream behaviour, preserved).
+ assert rendered.endswith("<|im_start|>assistant")
+ assert rendered.count("<|im_end|>") == 2
+
+
+def test_retool_rl_script_does_not_apply_chat_template():
+ script = (_tests_root.parent / "examples" / "retool" / "retool_qwen3_4b_rl.sh").read_text()
+ active = [ln for ln in script.splitlines() if ln.strip().startswith("--apply-chat-template")]
+ assert not active, "retool renders its own chat template; --apply-chat-template double-wraps the prompt"
+
+
+def test_generate_clamps_per_turn_budget_to_remaining_context(monkeypatch):
+ """A single turn must not be allowed to exceed the remaining context budget."""
+ headroom = 48
+ args = _args(rollout_max_context_len=_prompt_len() + headroom)
+ _, payloads = _run_generate(
+ monkeypatch,
+ [_choice("Answer: \\boxed{4}")],
+ args=args,
+ sampling_params=_sampling_params(max_new_tokens=10_000),
+ )
+ assert payloads[0]["payload"]["sampling_params"]["max_tokens"] == headroom
+
+
+# --------------------------------------------------------------------------
+# turn loop
+# --------------------------------------------------------------------------
+
+
+def test_generate_completes_on_boxed_answer(monkeypatch):
+ sample, payloads = _run_generate(monkeypatch, [_choice("Answer: \\boxed{4}")])
+
+ assert len(payloads) == 1, "an answer must end the loop"
+ assert sample.status is Sample.Status.COMPLETED
+ assert sample.response == "Answer: \\boxed{4}"
+ assert sample.tool_call_count == 0
+ assert sample.response_length == len(sample.response)
+ assert sample.loss_mask == [1] * sample.response_length
+ assert len(sample.rollout_log_probs) == sample.response_length
+
+
+def test_generate_runs_tool_then_answers(monkeypatch):
+ """A code turn feeds output back and continues to a second turn."""
+ sample, payloads = _run_generate(
+ monkeypatch,
+ [
+ _choice("print(2+2)"),
+ _choice("Answer: \\boxed{4}"),
+ ],
+ )
+
+ assert len(payloads) == 2, "tool turn must trigger a follow-up generation"
+ assert sample.status is Sample.Status.COMPLETED
+ assert sample.tool_call_count == 1
+ assert "" in sample.response
+ assert "4" in sample.response.split("")[1]
+
+ # Turn 2 must resend prompt + everything generated/observed so far.
+ assert len(payloads[1]["payload"]["token_ids"]) > len(payloads[0]["payload"]["token_ids"])
+
+ # Tool tokens are masked out; model tokens are trainable.
+ assert len(sample.loss_mask) == sample.response_length
+ assert set(sample.loss_mask) == {0, 1}
+ assert len(sample.rollout_log_probs) == sample.response_length
+
+
+def test_generate_masks_only_the_observation_tokens(monkeypatch):
+ sample, _ = _run_generate(
+ monkeypatch,
+ [_choice("print(2+2)"), _choice("Answer: \\boxed{4}")],
+ )
+ observation = sample.response[sample.response.index("\n\n") :]
+ observation = observation[: observation.index("") + len("") + 2]
+ assert sample.loss_mask.count(0) == len(observation), "exactly the tool output is masked"
+
+
+def test_generate_truncates_on_length_finish_reason(monkeypatch):
+ sample, payloads = _run_generate(monkeypatch, [_choice("thinking hard", finish_reason="length")])
+
+ assert len(payloads) == 1, "length stop must end the loop"
+ assert sample.status is Sample.Status.TRUNCATED
+
+
+def test_generate_aborts_on_abort_finish_reason(monkeypatch):
+ sample, _ = _run_generate(monkeypatch, [_choice("partial", finish_reason="abort")])
+
+ assert sample.status is Sample.Status.ABORTED
+ assert sample.response == "", "aborted sample carries no trainable response"
+
+
+def test_generate_aborts_when_engine_omits_logprobs(monkeypatch):
+ """Must not zero-fill: that would desync rollout_log_probs from the tokens."""
+ sample, _ = _run_generate(monkeypatch, [_choice("hello", logprobs=False)])
+
+ assert sample.status is Sample.Status.ABORTED
+ assert sample.rollout_log_probs is None
+
+
+def test_generate_aborts_on_logprob_length_mismatch(monkeypatch):
+ bad = _choice("hello")
+ bad["logprobs"]["content"] = bad["logprobs"]["content"][:2] # 5 tokens, 2 logprobs
+ sample, _ = _run_generate(monkeypatch, [bad])
+
+ assert sample.status is Sample.Status.ABORTED
+
+
+def test_generate_stops_at_max_tool_calls(monkeypatch):
+ max_calls = tool_sandbox.TOOL_CONFIGS["max_tool_calls"]
+ # Always emit code, never an answer: the loop must stop itself.
+ sample, payloads = _run_generate(
+ monkeypatch,
+ [_choice("print(1)") for _ in range(max_calls + 5)],
+ args=_args(rollout_max_context_len=200_000),
+ )
+ assert sample.tool_call_count == max_calls
+ assert len(payloads) <= max_calls + 1
+
+
+def test_generate_resets_stale_state_from_a_retried_sample(monkeypatch):
+ """Aborted/partial samples come back with state from the first attempt."""
+ sample = _pending_sample()
+ sample.response = "stale text"
+ sample.response_length = 3
+ sample.rollout_log_probs = [-1.0, -1.0, -1.0]
+ sample.loss_mask = [1, 1, 1]
+ sample.tokens = [1, 2, 3]
+ sample.status = Sample.Status.PENDING
+
+ sample, _ = _run_generate(monkeypatch, [_choice("Answer: \\boxed{4}")], sample=sample)
+
+ assert "stale text" not in sample.response
+ assert sample.response_length == len(sample.response)
+ assert len(sample.rollout_log_probs) == sample.response_length
+ assert len(sample.loss_mask) == sample.response_length
+
+
+def test_generate_truncates_when_observation_overflows_context(monkeypatch):
+ """Tool output is unbounded, so it must be trimmed to the context budget."""
+ code = "print(2+2)"
+ # Leave room for the code turn but not for the whole block.
+ args = _args(rollout_max_context_len=_prompt_len() + len(code) + 10)
+
+ sample, _ = _run_generate(monkeypatch, [_choice(code), _choice("Answer: \\boxed{4}")], args=args)
+
+ assert sample.status is Sample.Status.TRUNCATED
+ assert len(sample.tokens) <= args.rollout_max_context_len
+ # response text is resynced from the trimmed tokens
+ assert sample.response_length == len(sample.response)
+ assert len(sample.loss_mask) == sample.response_length
+
+
+def test_generate_rejects_partial_rollout(monkeypatch):
+ with pytest.raises(AssertionError):
+ _run_generate(monkeypatch, [_choice("x")], args=_args(partial_rollout=True))
+
+
+# --------------------------------------------------------------------------
+# prediction parsing
+# --------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ ("text", "action", "content"),
+ [
+ ("Answer: \\boxed{4}", "answer", "4"),
+ ("Answer: \\boxed{\\frac{1}{2}}", "answer", "\\frac{1}{2}"),
+ ("print(1)", "code", "print(1)"),
+ ('{"name": "code_interpreter", "arguments": {"code": "print(1)"}}', "code", "print(1)"),
+ ("```python\nprint(1)\n```", "code", "print(1)"),
+ ("just some prose", None, ""),
+ ],
+)
+def test_postprocess_predictions(text, action, content):
+ assert mod.postprocess_predictions(text) == (action, content)
+
+
+def test_postprocess_predictions_parses_pretty_printed_tool_call():
+ """Newlines between JSON tokens must not be escaped -- that is a parse error,
+ and the dropped tool call silently degrades into the "invalid action" reprompt."""
+ text = '\n{\n "name": "code_interpreter",\n "arguments": {"code": "print(1)"}\n}\n'
+ assert mod.postprocess_predictions(text) == ("code", "print(1)")
+
+
+def test_postprocess_predictions_recovers_raw_newlines_inside_code():
+ """Raw newlines *inside* the code string are invalid JSON; escaping recovers them."""
+ text = '{"name": "code_interpreter", "arguments": {"code": "import math\nprint(math.sqrt(16))"}}'
+ action, code = mod.postprocess_predictions(text)
+ assert action == "code"
+ assert code == "import math\nprint(math.sqrt(16))"
+
+
+def test_postprocess_predictions_prefers_answer_over_code():
+ text = "print(1)\nAnswer: \\boxed{7}"
+ assert mod.postprocess_predictions(text) == ("answer", "7")
+
+
+@pytest.mark.parametrize(
+ ("text", "expected"),
+ [
+ ("print(1)trailing junk", "print(1)"),
+ ("Answer: \\boxed{4} and then rambling", "Answer: \\boxed{4}"),
+ ("```python\nprint(1)\n```junk", "```python\nprint(1)\n```"),
+ ("nothing to trim", "nothing to trim"),
+ ],
+)
+def test_postprocess_responses_trims_after_last_complete_tag(text, expected):
+ assert mod.postprocess_responses(text) == expected
+
+
+def test_execute_predictions_invalid_action_reprompts():
+ next_obs, done = asyncio.run(mod.execute_predictions("prose with no action"))
+ assert done is False
+ assert "previous action is invalid" in next_obs
+
+
+def test_execute_predictions_answer_is_terminal():
+ next_obs, done = asyncio.run(mod.execute_predictions("Answer: \\boxed{4}"))
+ assert (next_obs, done) == ("", True)
+
+
+# --------------------------------------------------------------------------
+# tool concurrency
+# --------------------------------------------------------------------------
+
+
+class _CountingSemaphore:
+ """asyncio.Semaphore that records how many times it was acquired."""
+
+ def __init__(self, value: int):
+ self._sem = asyncio.Semaphore(value)
+ self.acquires = 0
+
+ async def __aenter__(self):
+ self.acquires += 1
+ await self._sem.acquire()
+ return self
+
+ async def __aexit__(self, *exc_info):
+ self._sem.release()
+
+
+def _install_semaphore(monkeypatch, value: int) -> _CountingSemaphore:
+ """Point every reference to the tool semaphore at one counting instance.
+
+ ``generate_with_retool`` may hold its own ``from tool_sandbox import SEMAPHORE``
+ alias, so patching only ``tool_sandbox.SEMAPHORE`` would leave a second,
+ independent semaphore behind and hide a double-acquire.
+ """
+ sem = _CountingSemaphore(value)
+ monkeypatch.setattr(tool_sandbox, "SEMAPHORE", sem)
+ monkeypatch.setattr(mod, "SEMAPHORE", sem, raising=False)
+ return sem
+
+
+def _stub_sandbox(monkeypatch, on_execute=None):
+ async def fake_execute_code(code):
+ if on_execute is not None:
+ await on_execute()
+ return "Output:\n4"
+
+ monkeypatch.setattr(tool_sandbox.tool_registry.python_sandbox, "execute_code", fake_execute_code)
+
+
+def test_execute_predictions_takes_the_tool_semaphore_exactly_once(monkeypatch):
+ """Pinned to 1 permit, a double-acquire self-deadlocks."""
+ sem = _install_semaphore(monkeypatch, 1)
+ _stub_sandbox(monkeypatch)
+
+ async def run():
+ return await asyncio.wait_for(mod.execute_predictions("print(2+2)"), timeout=5)
+
+ next_obs, done = asyncio.run(run())
+ assert done is False
+ assert "" in next_obs and "4" in next_obs
+ assert sem.acquires == 1, f"tool semaphore acquired {sem.acquires}x per call, expected 1"
+
+
+def test_concurrent_tool_calls_reach_the_configured_concurrency(monkeypatch):
+ """All `tool_concurrency` calls must be able to run at once.
+
+ Gated on a barrier rather than a sleep: each call blocks inside the critical
+ section until `limit` of them are in there together. That makes the peak an
+ invariant instead of a scheduling race -- and a double-acquire, which only
+ fits limit//2 callers, can never fill the barrier and trips the timeout.
+ """
+ limit = 4
+ sem = _install_semaphore(monkeypatch, limit)
+
+ live = 0
+ peak = 0
+ barrier = asyncio.Event()
+
+ async def track():
+ nonlocal live, peak
+ live += 1
+ peak = max(peak, live)
+ if live >= limit:
+ barrier.set()
+ await asyncio.wait_for(barrier.wait(), timeout=10)
+ live -= 1
+
+ _stub_sandbox(monkeypatch, on_execute=track)
+
+ async def run():
+ tasks = [mod.execute_predictions("print(2+2)") for _ in range(limit * 3)]
+ return await asyncio.wait_for(asyncio.gather(*tasks), timeout=30)
+
+ results = asyncio.run(run())
+ assert len(results) == limit * 3
+ assert peak == limit, f"expected {limit} concurrent tool executions, saw {peak}"
+ assert sem.acquires == limit * 3, f"expected 1 acquire per call, got {sem.acquires} for {limit * 3} calls"
+
+
+# --------------------------------------------------------------------------
+# sandbox
+# --------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ "code",
+ ["import os", "eval('1')", "open('/etc/passwd')", "__import__('os')", "import numpy"],
+)
+def test_sandbox_rejects_unsafe_code(code):
+ ok, _ = tool_sandbox.tool_registry.python_sandbox._check_code_safety(code)
+ assert ok is False
+
+
+@pytest.mark.parametrize("code", ["print(2+2)", "import math\nprint(math.sqrt(16))", "x = sum(range(10))"])
+def test_sandbox_allows_plain_math(code):
+ ok, message = tool_sandbox.tool_registry.python_sandbox._check_code_safety(code)
+ assert ok is True, message
+
+
+def test_real_sandbox_executes_code():
+ """The one test that actually spawns the sandbox subprocess.
+
+ A fresh PythonSandbox sidesteps the autouse stub on the registry's instance.
+ """
+ sandbox = tool_sandbox.PythonSandbox(timeout=60, memory_limit="1GB")
+ out = asyncio.run(sandbox.execute_code("print(2 + 2)"))
+ assert "4" in out, out
+
+
+def test_real_sandbox_reports_rejected_code():
+ sandbox = tool_sandbox.PythonSandbox(timeout=60, memory_limit="1GB")
+ out = asyncio.run(sandbox.execute_code("import os\nprint(os.getcwd())"))
+ assert "Error" in out and "not allowed" in out.lower() or "dangerous" in out.lower(), out
+
+
+def test_unknown_tool_is_reported_not_raised():
+ result = asyncio.run(tool_sandbox.tool_registry.execute_tool("nope", {}))
+ assert "not found" in result
+
+
+# --------------------------------------------------------------------------
+# reward
+# --------------------------------------------------------------------------
+
+
+def test_reward_func_scores_correct_answer():
+ sample = _pending_sample()
+ sample.response = " Answer: \\boxed{4}"
+ result = asyncio.run(mod.reward_func(_args(), sample))
+ assert result["score"] > 0
+
+
+def test_reward_func_wrong_answer_gets_tool_use_bonus_but_stays_negative():
+ sample = _pending_sample()
+ sample.response = " Answer: \\boxed{5}"
+ sample.tool_call_count = 8
+ result = asyncio.run(mod.reward_func(_args(), sample))
+ assert result["score"] <= -0.6
+ assert result["pred"] is not None
+
+
+def test_reward_func_rejects_non_sample():
+ with pytest.raises(TypeError):
+ asyncio.run(mod.reward_func(_args(), {"response": "x"}))
+
+
+if __name__ == "__main__":
+ raise SystemExit(pytest.main([__file__]))