From c60d114c140c4bab9cfae9bc0c3fdd7297d3bcc8 Mon Sep 17 00:00:00 2001 From: Hangrui Cao Date: Thu, 13 Aug 2026 20:08:22 +0000 Subject: [PATCH 1/5] [Example] Add retool tool-calling RL/SFT example (port slime -> vime vLLM) Re-adds `examples/retool`, the ReTool (tool-calling math RL, SFT -> RL) end-to-end example, ported from slime's SGLang implementation to vime's vLLM rollout path. #74 removed it as SGLang-only; this brings it back on `/inference/v1/generate`. The port itself is confined to the engine call in `generate_with_retool.py`: SGLang vime / vLLM ------------------------------------ ------------------------------------- /generate /inference/v1/generate {"input_ids", "sampling_params", {"model", "token_ids", "return_logprob": true} "sampling_params"} (the builder emits logprobs:1 and renames max_new_tokens -> max_tokens) meta_info.output_token_logprobs choices[0].token_ids + choices[0].logprobs.content[i] meta_info.finish_reason.type choices[0].finish_reason (bare string, normalized to the nested shape) The choice parse is inlined as `_parse_vllm_choice` rather than imported: #178 removed the shared `_inference_generate_tokens_and_logprobs` / `_vllm_meta_from_generate_choice` helpers from `vime.rollout.vllm_rollout`, and #184 established that callers parse the choice locally (as `vllm_streaming_rollout` and `vime/agent/adapters/common.py` both do). `tool_sandbox.py`, `requirements.txt`, `rl_data_preprocess.py` and `sft_data_processing.py` are engine-agnostic and vendored unchanged. Shell scripts get the usual sglang->vllm treatment (pkill pattern, VLLM_ARGS, --vllm-gpu-memory-utilization, SCRIPT_DIR/REPO_ROOT, vime.rollout.sft_rollout, wandb/ckpt renames) plus `MODEL_ARGS_ROTARY_BASE=5000000` instead of a duplicate `--rotary-base`, and drops a dead `${EVAL_ARGS[@]}` the upstream SFT script expands but never defines. Adding the directory also un-dangles the `examples/README.md` link that has pointed at nothing since #74. Three fixes on top of a straight port: 1. Drop `--apply-chat-template` from the RL script. `format_conversation_with_tools` renders the whole `<|im_start|>...<|im_end|>` conversation itself, so templating in the data loader too made the example wrap an already-templated prompt in its own `user` turn. Rendered against Qwen3-4B's real template, the model received nested `<|im_start|>user<|im_start|>user` plus a spurious empty assistant turn. 2. Take the tool-concurrency semaphore exactly once. `tool_sandbox.SEMAPHORE` is a plain, non-reentrant `asyncio.Semaphore`, and upstream acquires it in both `execute_predictions` and `ToolRegistry.execute_tool` -- two permits per tool call, which hangs the tool path (a single call self-deadlocks at `tool_concurrency == 1`). The limit now lives only in the registry, its natural owner. 3. Abort instead of zero-filling missing logprobs. `vllm_rollout.generate` substitutes `[0.0] * len(tokens)` when the engine reports no per-token logprobs. Here a missing or length-mismatched logprob array marks the sample ABORTED so the group returns to the buffer, rather than desyncing `rollout_log_probs` from the response tokens and silently corrupting the importance ratio. Adds `tests/test_retool_generate.py` (51 tests, mocked engine, no GPU) covering the request body, all five finish_reason mappings, the multi-turn tool loop, loss-mask/logprob alignment, context-budget clamping, observation-overflow truncation, stale-state reset on retry, prompt structure, prediction parsing, sandbox safety and reward shaping. Registered in the existing CPU CI step, with jinja2 added to that step's installs rather than relying on it arriving transitively via torch. Each of the three fixes above was mutation-tested: reintroducing the bug makes the corresponding tests fail, and the first version of the concurrency test did not (`from tool_sandbox import SEMAPHORE` binds an alias that patching `tool_sandbox.SEMAPHORE` never reaches), so it was rewritten around a counting semaphore installed on both names. Co-Authored-By: Claude Opus 5 Signed-off-by: Hangrui Cao --- .buildkite/pipeline.yml | 3 +- examples/retool/README.md | 135 ++++++ examples/retool/generate_with_retool.py | 474 ++++++++++++++++++++ examples/retool/requirements.txt | 3 + examples/retool/retool_qwen3_4b_rl.sh | 163 +++++++ examples/retool/retool_qwen3_4b_sft.sh | 127 ++++++ examples/retool/rl_data_preprocess.py | 21 + examples/retool/sft_data_processing.py | 31 ++ examples/retool/tool_sandbox.py | 360 +++++++++++++++ tests/test_retool_generate.py | 563 ++++++++++++++++++++++++ 10 files changed, 1879 insertions(+), 1 deletion(-) create mode 100644 examples/retool/README.md create mode 100644 examples/retool/generate_with_retool.py create mode 100644 examples/retool/requirements.txt create mode 100644 examples/retool/retool_qwen3_4b_rl.sh create mode 100644 examples/retool/retool_qwen3_4b_sft.sh create mode 100644 examples/retool/rl_data_preprocess.py create mode 100644 examples/retool/sft_data_processing.py create mode 100644 examples/retool/tool_sandbox.py create mode 100644 tests/test_retool_generate.py 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..721a6498d --- /dev/null +++ b/examples/retool/generate_with_retool.py @@ -0,0 +1,474 @@ +# 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 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 + + # Clean up the JSON string by removing newlines and extra + # whitespace + json_str = tool_call_match.group(1) + # Replace newlines in string values with \n + json_str = json_str.replace("\n", "\\n") + tool_call_data = json.loads(json_str) + 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: ``tool_registry.execute_tool`` already + # takes the same non-reentrant ``tool_sandbox.SEMAPHORE``. Acquiring + # it in both places needs 2 permits per call and hangs once enough + # tool calls are in flight (a single call self-deadlocks when + # ``tool_concurrency == 1``). The limit is owned by the registry. + 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. + + Inlined rather than imported: #178 removed the shared + ``_inference_generate_tokens_and_logprobs`` / ``_vllm_meta_from_generate_choice`` + helpers from ``vime.rollout.vllm_rollout``, and #184 established that callers + parse the choice locally (as ``vllm_streaming_rollout`` and + ``vime/agent/adapters/common.py`` both do). + + 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 slime's nested shape + # so the turn loop below reads the same as the upstream example. + 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() + prompt = format_conversation_with_tools(prompt=sample.prompt, tools=tool_specs) + + 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, + ) + + # Use token IDs instead of text. ``_build_inference_sampling_params`` + # maps ``max_new_tokens`` -> ``max_tokens`` and requests ``logprobs: 1``, + # which replaces SGLang's ``return_logprob`` flag. + 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): + # The engine returned tokens but no usable per-token logprobs. Unlike + # `vllm_rollout.generate`, do NOT substitute zeros: that would desync + # rollout_log_probs from response_token_ids and silently corrupt the + # importance ratio (and blow up `slice_log_prob_with_cp` downstream). + # Abort so the rollout manager returns the whole group to the buffer + # for retry instead of poisoning the trainer. + 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 + solution_str = sample.prompt + 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..4cac515fd --- /dev/null +++ b/examples/retool/retool_qwen3_4b_rl.sh @@ -0,0 +1,163 @@ +#!/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.format_conversation_with_tools + # renders the full <|im_start|>...<|im_end|> conversation itself. Enabling the + # flag templates the prompt in the data loader too, and the example then wraps + # that inside its own `user` turn -- producing nested `<|im_start|>user + # <|im_start|>user` plus an empty assistant turn. + --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 + +# Build the runtime environment JSON with proper variable substitution. +# SCRIPT_DIR is on PYTHONPATH so `generate_with_retool` resolves as a top-level +# module (and can import its sibling `tool_sandbox`), matching examples/tau-bench. +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..3c7b01425 --- /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..fe177aedb --- /dev/null +++ b/tests/test_retool_generate.py @@ -0,0 +1,563 @@ +"""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) + + +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 + + +def test_prompt_has_exactly_one_conversation_structure(): + """The example renders the whole conversation itself, so the RL script must not + also pass --apply-chat-template. Guards against the nested + `<|im_start|>user<|im_start|>user` + empty-assistant-turn prompt that 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 + # exactly one trailing generation turn, and nothing after it + 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") + # system + user are closed; the assistant turn is left open for generation + 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_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): + """`tool_sandbox.SEMAPHORE` is a plain asyncio.Semaphore, so nesting the + acquire in both the caller and ToolRegistry.execute_tool needs 2 permits + per call. 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. + + A double-acquire burns two permits per call and halves the peak. + """ + limit = 4 + sem = _install_semaphore(monkeypatch, limit) + + live = 0 + peak = 0 + + async def track(): + nonlocal live, peak + live += 1 + peak = max(peak, live) + await asyncio.sleep(0.01) + 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=10) + + 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_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__])) From 1345afc5def3950abdbb2b9282d84c862163f7f6 Mon Sep 17 00:00:00 2001 From: Hangrui Cao Date: Thu, 13 Aug 2026 21:54:47 +0000 Subject: [PATCH 2/5] Address high-priority review findings Three high-severity findings from the automated review on #382, each verified before fixing: 1. `postprocess_predictions` escaped newlines unconditionally before `json.loads`. That rescues raw newlines *inside* the "code" string (invalid JSON), but corrupts pretty-printed tool calls, where the newlines sit *between* tokens -- a backslash there is a parse error. Confirmed: a pretty-printed `` failed to parse and silently degraded into the "My previous action is invalid" reprompt, wasting a turn. Now parses first and escapes only as a fallback, so both shapes work. 2/3. `--wandb-key ${WANDB_KEY}` was unquoted in both scripts. With WANDB_KEY unset the word disappears, so `--wandb-key` swallows the next flag as its value. Verified in bash: the array expands to `--use-wandb --wandb-key --tensor-model-parallel-size 2`. Quoting yields an empty-string argument instead. (The same latent issue exists in other examples, e.g. eval_multi_task and geo3k_vlm; left alone here.) Adds two tests pinning both JSON shapes. Mutation-tested: restoring the unconditional escape fails the pretty-printed test. Co-Authored-By: Claude Opus 5 Signed-off-by: Hangrui Cao --- examples/retool/generate_with_retool.py | 14 +++++++++----- examples/retool/retool_qwen3_4b_rl.sh | 4 +++- examples/retool/retool_qwen3_4b_sft.sh | 4 +++- tests/test_retool_generate.py | 16 ++++++++++++++++ 4 files changed, 31 insertions(+), 7 deletions(-) diff --git a/examples/retool/generate_with_retool.py b/examples/retool/generate_with_retool.py index 721a6498d..ac9d5a461 100644 --- a/examples/retool/generate_with_retool.py +++ b/examples/retool/generate_with_retool.py @@ -112,12 +112,16 @@ def postprocess_predictions(prediction: str): try: import json - # Clean up the JSON string by removing newlines and extra - # whitespace json_str = tool_call_match.group(1) - # Replace newlines in string values with \n - json_str = json_str.replace("\n", "\\n") - tool_call_data = json.loads(json_str) + try: + tool_call_data = json.loads(json_str) + except json.JSONDecodeError: + # Some models emit raw newlines inside the "code" string, which is + # not valid JSON; escaping them recovers the call. Do this only as + # a fallback -- escaping unconditionally corrupts pretty-printed + # JSON, whose newlines sit *between* tokens rather than inside a + # string, and a backslash there is a parse error. + tool_call_data = json.loads(json_str.replace("\n", "\\n")) tool_name = tool_call_data.get("name") arguments = tool_call_data.get("arguments", {}) diff --git a/examples/retool/retool_qwen3_4b_rl.sh b/examples/retool/retool_qwen3_4b_rl.sh index 4cac515fd..e0367d2e4 100644 --- a/examples/retool/retool_qwen3_4b_rl.sh +++ b/examples/retool/retool_qwen3_4b_rl.sh @@ -105,7 +105,9 @@ WANDB_ARGS=( --use-wandb --wandb-project vime-dapo --wandb-group qwen3-4B-test-multi-turn - --wandb-key ${WANDB_KEY} + # Quoted: unquoted+unset expands to nothing, so --wandb-key would swallow + # the next flag (e.g. --tensor-model-parallel-size) as its value. + --wandb-key "${WANDB_KEY}" ) VLLM_ARGS=( diff --git a/examples/retool/retool_qwen3_4b_sft.sh b/examples/retool/retool_qwen3_4b_sft.sh index 3c7b01425..efd6fbc22 100644 --- a/examples/retool/retool_qwen3_4b_sft.sh +++ b/examples/retool/retool_qwen3_4b_sft.sh @@ -83,7 +83,9 @@ WANDB_ARGS=( --use-wandb --wandb-project vime-dev --wandb-group qwen3-4B-base-sft - --wandb-key ${WANDB_KEY} + # Quoted: unquoted+unset expands to nothing, so --wandb-key would swallow + # the next flag (e.g. --tensor-model-parallel-size) as its value. + --wandb-key "${WANDB_KEY}" ) MISC_ARGS=( diff --git a/tests/test_retool_generate.py b/tests/test_retool_generate.py index fe177aedb..f66e9be22 100644 --- a/tests/test_retool_generate.py +++ b/tests/test_retool_generate.py @@ -390,6 +390,22 @@ 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: a backslash outside a + string is a JSON parse error, and the dropped tool call would silently + degrade 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") From d28d4ee90a4af8274c80c17b6676eaf2208254f2 Mon Sep 17 00:00:00 2001 From: Hangrui Cao Date: Thu, 13 Aug 2026 22:04:55 +0000 Subject: [PATCH 3/5] Trim redundant comments Drops comments that restated the code they sat above or carried PR archaeology into the source: - `_parse_vllm_choice`: the #178/#184 rationale for inlining belongs in the PR, not the docstring. - The payload comment re-described `_build_inference_sampling_params`. - Shortened the three notes documenting deliberate divergence from upstream (semaphore, logprob abort, no --apply-chat-template) to the reason only; a future sync still needs the why, just not the walkthrough. - Dropped the justification for quoting ${WANDB_KEY} -- quoting is idiomatic. - Test docstrings that repeated the assertion beneath them. Upstream's own comments are left untouched to keep the vendored files diffable against slime. Net -33 comment lines; still 53 tests passing, and the three fixes remain mutation-guarded (reintroducing any of them fails its tests). Co-Authored-By: Claude Opus 5 Signed-off-by: Hangrui Cao --- examples/retool/generate_with_retool.py | 36 +++++++------------------ examples/retool/retool_qwen3_4b_rl.sh | 12 +++------ examples/retool/retool_qwen3_4b_sft.sh | 2 -- tests/test_retool_generate.py | 15 +++-------- 4 files changed, 16 insertions(+), 49 deletions(-) diff --git a/examples/retool/generate_with_retool.py b/examples/retool/generate_with_retool.py index ac9d5a461..9c244e829 100644 --- a/examples/retool/generate_with_retool.py +++ b/examples/retool/generate_with_retool.py @@ -116,11 +116,9 @@ def postprocess_predictions(prediction: str): try: tool_call_data = json.loads(json_str) except json.JSONDecodeError: - # Some models emit raw newlines inside the "code" string, which is - # not valid JSON; escaping them recovers the call. Do this only as - # a fallback -- escaping unconditionally corrupts pretty-printed - # JSON, whose newlines sit *between* tokens rather than inside a - # string, and a backslash there is a parse error. + # 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", {}) @@ -194,11 +192,8 @@ async def execute_predictions(prediction: str) -> str: # postprocess_predictions) code = content.strip() if code: - # No SEMAPHORE acquire here: ``tool_registry.execute_tool`` already - # takes the same non-reentrant ``tool_sandbox.SEMAPHORE``. Acquiring - # it in both places needs 2 permits per call and hangs once enough - # tool calls are in flight (a single call self-deadlocks when - # ``tool_concurrency == 1``). The limit is owned by the registry. + # 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 @@ -224,12 +219,6 @@ async def execute_predictions(prediction: str) -> str: 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. - Inlined rather than imported: #178 removed the shared - ``_inference_generate_tokens_and_logprobs`` / ``_vllm_meta_from_generate_choice`` - helpers from ``vime.rollout.vllm_rollout``, and #184 established that callers - parse the choice locally (as ``vllm_streaming_rollout`` and - ``vime/agent/adapters/common.py`` both do). - Returns ``log_probs=[]`` when the engine reports no per-token logprobs, so the caller can abort instead of training on fabricated values. """ @@ -241,8 +230,7 @@ def _parse_vllm_choice(choice: dict[str, Any]) -> tuple[list[int], list[float], 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 slime's nested shape - # so the turn loop below reads the same as the upstream example. + # Normalize the bare vLLM ``finish_reason`` string into the nested shape. fr = choice.get("finish_reason") or "stop" if isinstance(fr, dict): finish = fr @@ -313,9 +301,6 @@ async def generate(args, sample: Sample, sampling_params) -> Sample: remaining_budget, ) - # Use token IDs instead of text. ``_build_inference_sampling_params`` - # maps ``max_new_tokens`` -> ``max_tokens`` and requests ``logprobs: 1``, - # which replaces SGLang's ``return_logprob`` flag. current_token_ids = prompt_tokens_ids + response_token_ids payload = { "model": args.hf_checkpoint, @@ -353,12 +338,9 @@ async def generate(args, sample: Sample, sampling_params) -> Sample: return sample if not cur_log_probs or len(cur_log_probs) != len(cur_response_token_ids): - # The engine returned tokens but no usable per-token logprobs. Unlike - # `vllm_rollout.generate`, do NOT substitute zeros: that would desync - # rollout_log_probs from response_token_ids and silently corrupt the - # importance ratio (and blow up `slice_log_prob_with_cp` downstream). - # Abort so the rollout manager returns the whole group to the buffer - # for retry instead of poisoning the trainer. + # 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 diff --git a/examples/retool/retool_qwen3_4b_rl.sh b/examples/retool/retool_qwen3_4b_rl.sh index e0367d2e4..add9f99a4 100644 --- a/examples/retool/retool_qwen3_4b_rl.sh +++ b/examples/retool/retool_qwen3_4b_rl.sh @@ -40,11 +40,8 @@ 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.format_conversation_with_tools - # renders the full <|im_start|>...<|im_end|> conversation itself. Enabling the - # flag templates the prompt in the data loader too, and the example then wraps - # that inside its own `user` turn -- producing nested `<|im_start|>user - # <|im_start|>user` plus an empty assistant turn. + # 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 @@ -105,8 +102,6 @@ WANDB_ARGS=( --use-wandb --wandb-project vime-dapo --wandb-group qwen3-4B-test-multi-turn - # Quoted: unquoted+unset expands to nothing, so --wandb-key would swallow - # the next flag (e.g. --tensor-model-parallel-size) as its value. --wandb-key "${WANDB_KEY}" ) @@ -135,9 +130,8 @@ CUSTOM_ARGS=( 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 -# Build the runtime environment JSON with proper variable substitution. # SCRIPT_DIR is on PYTHONPATH so `generate_with_retool` resolves as a top-level -# module (and can import its sibling `tool_sandbox`), matching examples/tau-bench. +# 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}\", diff --git a/examples/retool/retool_qwen3_4b_sft.sh b/examples/retool/retool_qwen3_4b_sft.sh index efd6fbc22..07317a448 100644 --- a/examples/retool/retool_qwen3_4b_sft.sh +++ b/examples/retool/retool_qwen3_4b_sft.sh @@ -83,8 +83,6 @@ WANDB_ARGS=( --use-wandb --wandb-project vime-dev --wandb-group qwen3-4B-base-sft - # Quoted: unquoted+unset expands to nothing, so --wandb-key would swallow - # the next flag (e.g. --tensor-model-parallel-size) as its value. --wandb-key "${WANDB_KEY}" ) diff --git a/tests/test_retool_generate.py b/tests/test_retool_generate.py index f66e9be22..96bb7fd7b 100644 --- a/tests/test_retool_generate.py +++ b/tests/test_retool_generate.py @@ -203,19 +203,15 @@ def test_generate_prompt_includes_tool_specs(monkeypatch): def test_prompt_has_exactly_one_conversation_structure(): - """The example renders the whole conversation itself, so the RL script must not - also pass --apply-chat-template. Guards against the nested - `<|im_start|>user<|im_start|>user` + empty-assistant-turn prompt that produces.""" + """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 - # exactly one trailing generation turn, and nothing after it 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") - # system + user are closed; the assistant turn is left open for generation assert rendered.count("<|im_end|>") == 2 @@ -391,9 +387,8 @@ def test_postprocess_predictions(text, action, content): def test_postprocess_predictions_parses_pretty_printed_tool_call(): - """Newlines *between* JSON tokens must not be escaped: a backslash outside a - string is a JSON parse error, and the dropped tool call would silently - degrade into the "invalid action" reprompt.""" + """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)") @@ -479,9 +474,7 @@ async def fake_execute_code(code): def test_execute_predictions_takes_the_tool_semaphore_exactly_once(monkeypatch): - """`tool_sandbox.SEMAPHORE` is a plain asyncio.Semaphore, so nesting the - acquire in both the caller and ToolRegistry.execute_tool needs 2 permits - per call. Pinned to 1 permit, a double-acquire self-deadlocks.""" + """Pinned to 1 permit, a double-acquire self-deadlocks.""" sem = _install_semaphore(monkeypatch, 1) _stub_sandbox(monkeypatch) From 1758d8e91226934cb6590bc8533161b12efdd77c Mon Sep 17 00:00:00 2001 From: Hangrui Cao Date: Fri, 14 Aug 2026 18:13:39 +0000 Subject: [PATCH 4/5] Make the tool-concurrency tests deterministic `test_concurrent_tool_calls_reach_the_configured_concurrency` asserted `peak == limit` while the tasks were held open only by `asyncio.sleep(0.01)`. On a loaded runner the first sleep can expire before the last task starts, so the peak lands below the limit and the test fails spuriously. It now gates on an `asyncio.Event` barrier: each call blocks inside the critical section until `limit` of them are in there together, which makes the peak an invariant rather than a scheduling race. The double-acquire bug still trips it -- only limit//2 callers fit, so the barrier can never fill and the wait times out. Also stops the rollout tests spawning real `python3` subprocesses. They exercise the turn loop, not the sandbox, and a subprocess per tool call made them slow and dependent on the runner (`test_generate_stops_at_max_tool_calls` alone spawned 16). An autouse fixture stubs the sandbox; two new tests cover real execution and real rejection explicitly, using a fresh PythonSandbox to sidestep the stub. All test phases now sum to 0.06s, down from ~0.6s. 55 tests, and both concurrency tests still fail if the double-acquire is reintroduced. Co-Authored-By: Claude Opus 5 Signed-off-by: Hangrui Cao --- tests/test_retool_generate.py | 43 ++++++++++++++++++++++++++++++++--- 1 file changed, 40 insertions(+), 3 deletions(-) diff --git a/tests/test_retool_generate.py b/tests/test_retool_generate.py index 96bb7fd7b..52d898b27 100644 --- a/tests/test_retool_generate.py +++ b/tests/test_retool_generate.py @@ -113,6 +113,21 @@ 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] = [] @@ -490,26 +505,32 @@ async def run(): def test_concurrent_tool_calls_reach_the_configured_concurrency(monkeypatch): """All `tool_concurrency` calls must be able to run at once. - A double-acquire burns two permits per call and halves the peak. + 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) - await asyncio.sleep(0.01) + 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=10) + return await asyncio.wait_for(asyncio.gather(*tasks), timeout=30) results = asyncio.run(run()) assert len(results) == limit * 3 @@ -537,6 +558,22 @@ def test_sandbox_allows_plain_math(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 From cec3b8ac1b239baf662cf36e22e0dcde94ff188c Mon Sep 17 00:00:00 2001 From: Hangrui Cao Date: Fri, 14 Aug 2026 18:47:27 +0000 Subject: [PATCH 5/5] Handle chat-message-list prompts (fixes eval crash) The bounded RL run crashed at the first AIME eval: File "examples/retool/generate_with_retool.py", in reward_func solution_str = sample.prompt + sample.response TypeError: can only concatenate list (not "str") to list Both datasets this example wires up store `prompt` as a list of chat messages, not a string -- DAPO-Math-17k for training and AIME-2024 for eval -- and `vime.utils.data._build_messages` passes a list through untouched when `--apply-chat-template` is off. This also corrects the earlier commit that dropped `--apply-chat-template`. That removal was right about the symptom (the flag templates in the data loader and the example then wraps the result in its own `user` turn, so the model sees nested `<|im_start|>user<|im_start|>user`) but wrong about the remedy: with the flag off, `sample.prompt` stays a list, which renders as its own repr in the prompt and raises TypeError in the reward. The earlier verification used a synthetic dataset whose `prompt` was a string, so neither failure showed up. `split_prompt` now normalizes both shapes, feeding the user text to `format_conversation_with_tools` and any system message through its existing `system_prompt` parameter. The chat template is applied exactly once and `reward_func` always gets a string, so the flag stays off. Adds 5 tests over the real data shapes. Mutation-tested: reverting either call site fails its test. The first version of the generate-side test did not catch it -- passing the list to the template renders a repr that still contains the question text -- so it now asserts the repr artifacts are absent. Co-Authored-By: Claude Opus 5 Signed-off-by: Hangrui Cao --- examples/retool/generate_with_retool.py | 31 ++++++++++++-- tests/test_retool_generate.py | 55 +++++++++++++++++++++++++ 2 files changed, 83 insertions(+), 3 deletions(-) diff --git a/examples/retool/generate_with_retool.py b/examples/retool/generate_with_retool.py index 9c244e829..d8304adda 100644 --- a/examples/retool/generate_with_retool.py +++ b/examples/retool/generate_with_retool.py @@ -59,6 +59,30 @@ """ +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: @@ -264,7 +288,8 @@ async def generate(args, sample: Sample, sampling_params) -> Sample: # Set up the initial prompt with system prompt and tools (outside the loop) tool_specs = tool_registry.get_tool_specs() - prompt = format_conversation_with_tools(prompt=sample.prompt, tools=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) @@ -437,8 +462,8 @@ async def reward_func(args, sample, **kwargs): if not isinstance(sample, Sample): raise TypeError("Sample must be an instance of Sample class.") - # Build complete solution string - solution_str = sample.prompt + sample.response + # 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 "" diff --git a/tests/test_retool_generate.py b/tests/test_retool_generate.py index 52d898b27..50d25f5f2 100644 --- a/tests/test_retool_generate.py +++ b/tests/test_retool_generate.py @@ -217,6 +217,61 @@ def test_generate_prompt_includes_tool_specs(monkeypatch): 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())