[Example] Add retool tool-calling RL/SFT example (port slime → vime vLLM)

[DiegoCao]

What

Re-adds examples/retool — the ReTool tool-calling math RL example (SFT → RL) — ported from slime's SGLang implementation to vime's vLLM rollout path.

This is a re-add, not a new example: #74 deleted retool and search-r1 as SGLang-only after #39's sglang→vllm rename. Adding the directory back also un-dangles the examples/README.md link that has pointed at nothing since then.

The port

All the real work is the engine call in generate_with_retool.py:

SGLang (upstream)	vime / vLLM
<router>/generate	<router>/inference/v1/generate
{"input_ids", "sampling_params", "return_logprob": true}	{"model", "token_ids", "sampling_params"} — _build_inference_sampling_params emits logprobs: 1 and renames max_new_tokens→max_tokens
meta_info.output_token_logprobs	choices[0].token_ids + choices[0].logprobs.content[i].logprob
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.

Three fixes on top of a straight port

1. Dropped --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 meant the example wrapped an already-templated prompt in its own user turn. Rendered against Qwen3-4B's real chat template, the model received:

<|im_start|>user<|im_start|>user
What is 12345 * 6789?<|im_end|>
<|im_start|>assistant
<|im_end|>
<|im_start|>assistant

Nested user turns 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.

Tests

tests/test_retool_generate.py — 51 tests, mocked engine, no GPU. Covers the request body (asserting SGLang's input_ids/return_logprob did not survive), 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 was mutation-tested — reintroducing the bug makes the corresponding tests fail. That mattered: the first version of the concurrency test passed even with the bug restored, because from tool_sandbox import SEMAPHORE binds an alias that patching tool_sandbox.SEMAPHORE never reaches. It was rewritten around a counting semaphore installed on both names.

Validation

Unit tests above, plus an end-to-end smoke run on 8xH100 80GB (Qwen3-4B, GRPO,
2 rollouts x 8 prompts x 4 samples) against a live vLLM router — enough to
exercise the request/response contract and the tool loop, not a convergence run.

rollout 0: response_len/mean 844.5  rollout_time 12.9s  truncated_ratio 0.0
rollout 1: response_len/mean 777.9  rollout_time 10.7s  truncated_ratio 0.0
rollout/rollout_log_probs -0.0430   rollout/raw_reward -0.3922
grad_norm 2.2928 / 2.8861          train exit=0

What this confirms about the port specifically:

• The real /inference/v1/generate accepts the {model, token_ids, sampling_params}
  body and returns token_ids + logprobs.content[i].logprob in the shape
  _parse_vllm_choice expects.
• rollout/rollout_log_probs = -0.043 is non-degenerate, so per-token logprobs
  genuinely flow through to training. Had the parse failed, this would be exactly
  0.0 (the zero-fill path) or the samples would have aborted.
• The multi-turn tool loop runs end to end: the model emits a <tool_call>, the
  sandbox executes it, the <interpreter> block is appended as non-trainable
  tokens, and generation continues to Answer: \boxed{...}.
• reward_func scores the result and two optimizer steps complete.

One observation worth recording, which is upstream behaviour and deliberately
not changed here: the model naturally emits bare expressions
({"code": "64345 + 3436"}), but PythonSandbox runs code as a script rather
than a REPL, so an expression yields no stdout and the tool returns an empty
<interpreter> block. The model then falls back to its own arithmetic. Making
the sandbox echo a trailing expression would change the training signal and
diverge from the published ReTool recipe, so it is left alone.

Known gaps

• The SFT stage renders prompts with Qwen's official chat template (via sft_rollout + --input-key messages), while the RL stage uses this example's own template, which Jinja renders without the trailing newline after <|im_start|>assistant. That SFT/RL prompt mismatch is upstream's design and is preserved here rather than silently changed.
• --use-kl-loss with --kl-loss-coef 0.00 is upstream's config and is preserved. Note placement_group.py creates the reference model on kl_coef != 0 or use_kl_loss, so the ref model is loaded and forward-passed while its KL term is multiplied by zero — dropping the flag would reclaim ~4 GB/GPU at no change to the loss.

🤖 Generated with Claude Code

[read-the-docs-community]

Documentation build overview

📚 vime | 🛠️ Build #34072741 | 📁 Comparing cec3b8a against latest (8144096)

  🔍 Preview build  

46 files changed · + 1 added · ± 45 modified

+ Added

• _examples_synced/retool/README.html

± Modified

• genindex.html
• index.html
• search.html
• advanced/arch-support-beyond-megatron.html
• advanced/delta-weight-sync.html
• advanced/external-rollout-engines.html
• advanced/fault-tolerance.html
• advanced/low-precision.html
• advanced/megatron-config.html
• advanced/observability.html
• and 35 more...

[gemini-code-assist]

Code Review

This pull request introduces the 'retool' example, porting tool-enabled language model generation from SFT to RL using vime's vLLM router. It includes custom generation and reward functions, a safe Python execution sandbox, data preprocessing scripts, training shell scripts, and comprehensive unit tests. The code review highlights several critical issues and improvement opportunities: potential JSON corruption from global newline replacement, argument parsing bugs in shell scripts when WANDB_KEY is unset, portability and cross-platform compatibility concerns (such as hardcoded interpreter paths, Unix-specific modules, and absolute file paths), security vulnerabilities in the regex-based sandbox checks, and potential file descriptor leaks. Additionally, the reviewer suggested code cleanups, such as avoiding inline imports, removing redundant variables, optimizing loop-level imports, and correcting CLI commands in the documentation.

[gemini-code-assist]

The global replacement of \n with \\n on json_str will corrupt any multi-line JSON structure (e.g., JSON with newlines between keys or pretty-printed JSON), making it invalid JSON and causing json.loads to fail. If the intent is to only escape literal newlines within string values, a global replace is unsafe. Since standard-conforming LLMs should output valid JSON with escaped newlines, we should avoid global replacement or only apply it defensively when direct parsing fails.

[gemini-code-assist]

If WANDB_KEY is empty or unset, the unquoted ${WANDB_KEY} will expand to nothing. This causes --wandb-key to consume the next argument in the command line as its value, corrupting the argument parsing and causing the script to fail. We should double-quote it as "${WANDB_KEY}" to ensure it expands to an empty string if unset.

Suggested change

	--wandb-key ${WANDB_KEY}
	--wandb-key "${WANDB_KEY}"

[gemini-code-assist]

If WANDB_KEY is empty or unset, the unquoted ${WANDB_KEY} will expand to nothing. This causes --wandb-key to consume the next argument in the command line (e.g., from PERF_ARGS) as its value, corrupting the argument parsing and causing the script to fail. We should double-quote it as "${WANDB_KEY}" to ensure it expands to an empty string if unset.

Suggested change

	--wandb-key ${WANDB_KEY}
	--wandb-key "${WANDB_KEY}"

[gemini-code-assist]

Using "python3" hardcodes the interpreter executable, which might not be available or might point to a different Python environment than the one running the main process (especially in virtual environments or different OS platforms). We should use sys.executable instead of "python3" to ensure the subprocess runs with the same Python interpreter.

Suggested change

	["python3", script_path],
	[sys.executable, script_path],

[gemini-code-assist]

Standard library modules like shutil (line 188) and sys (line 59) should be imported at the top of the file rather than inline inside functions or context managers, adhering to PEP 8 guidelines.

[gemini-code-assist]

The output path /root/dapo-math-17k-processed/dapo_math_17k_cleaned.jsonl is hardcoded to an absolute path in the /root directory. This makes the script non-portable and will fail on systems where the user is not root or /root is not writable/present. We should use a relative path (similar to sft_data_processing.py) or allow configuring the output path via an argument.

[gemini-code-assist]

The convert_role function can be simplified by checking if the role is in a set of allowed roles, rather than using multiple if-elif branches that return the same value.

    def convert_role(role):
        if role not in {"user", "assistant", "system"}:
            raise ValueError(f"Unknown role: {role}")
        return role

[gemini-code-assist]

Importing wandb inside the for turn loop on every iteration is inefficient and unidiomatic. We should import wandb once at the module level (or at the start of the function) using a try...except ImportError block, and then simply check if it is not None inside the loop.

[gemini-code-assist]

The local variable loss_masks is assigned to sample.loss_mask here and then re-assigned back to sample.loss_mask at line 426, but it is never modified or used anywhere else in the function. This is redundant and can lead to bugs if sample.loss_mask is re-assigned inside append_response_tokens. We should remove loss_masks and let append_response_tokens manage sample.loss_mask directly.

[gemini-code-assist]

The command hf download is used throughout the setup instructions, but the standard Hugging Face CLI command is huggingface-cli download. Unless hf is a pre-configured alias on the target system, users will encounter a command not found error. We should update these to use huggingface-cli download.

[CalvinXKY]

For an example merge, the 2-rollout smoke is not enough, it only checks the vLLM contract and that one train step runs.
Please post a longer RL run with WandB (or equivalent) showing reward/eval converging (AIME is already wired at --eval-interval 20). Without that curve, we cannot tell whether the ported recipe actually trains.

[DiegoCao]

@CalvinXKY Yes, we can add test trajectory beyond that.

[aoshen02]

We don't need a test for it.

[aoshen02]

No need to add these, we port most of the code from slime

[aoshen02]

And here as well. It is very clear that vime is a fork of slime