[Sync][G] streaming rollout + retool example (slime #1921/#1861)

[aoshen02]

PR-G — streaming + retool examples (RFC #107). Stacked on PR-A (#145) — base is sync/slime-mega-A so the diff is G-only.

Included

• #1921 streaming → vime/rollout/vllm_streaming_rollout.py + tests/test_qwen3_4B_streaming_partial_rollout.py. Translates the sglang streaming path to vLLM's /inference/v1/generate SSE. Load-bearing translation: sglang streams cumulative output_token_logprobs (overwrite per chunk); vLLM SSE sends delta token_ids+logprobs per chunk → accumulate (+=) and re-decode accumulated tokens each chunk (multi-byte-safe), finalize via _vllm_meta_from_generate_choice. routed_experts applied only from the terminal choice. MM raises NotImplementedError (the two-call render→generate flow doesn't fit a single stream). Streaming file is the reviewed sync/slime-pr-1921 (draft [Rollout] Add streaming vLLM rollout (port slime #1921) #117) design.
• #1861 retool → re-added examples/retool/generate_with_retool.py (vime lacked examples/retool/) from the slime base + #1861's multi-turn/retry-desync hardening (if not meta_info.get("output_token_logprobs") guard after routing through the vLLM field map).

Notes / follow-ups

• Retool example is partial: it depends on a sibling tool_sandbox.py not yet ported (out of #1861's scope; vime never had examples/retool/). The example won't run end-to-end until that lands — flagged for a follow-up.
• Two small correctness guards vs slime in retool (reviewed): post-loop status won't clobber a budget-driven TRUNCATED; finish_type tracked locally since vLLM responses are per-call.

CI

• test_qwen3_4B_streaming_partial_rollout registered in the e2e-test-megatron matrix (.j2 + generated .yml). py_compile clean.

🤖 Generated with Claude Code

---

Fidelity ledger (audited 2026-06-04)

op	detail
Ported	#1861 retool multi-turn/retry hardening ; #1921 streaming example
Translated (sglang→vllm)	streaming meta_info from sglang cumulative form → vLLM delta-accumulation
Dropped/Deferred	none
Verified	green on rebuilt FCBAGDE stack

[gemini-code-assist]

Code Review

This pull request introduces a custom tool-enabled generation script, a streaming vLLM rollout implementation to handle mid-generation aborts, and a corresponding CI smoke test. The reviewer identified two critical issues: first, a global newline replacement in the tool-call JSON parser that can break structural newlines and cause decoding errors, which can be resolved by using json.loads(..., strict=False); second, several vulnerabilities in the streaming finalization logic, including a lack of exception handling for stream disconnects, incorrect status assignment on stream failure, and potential aborted status overwrites.

[gemini-code-assist]

Replacing all newlines (\n) with escaped newlines (\\n) globally on the JSON string breaks structural newlines in formatted JSON, leading to JSONDecodeError during json.loads. Use json.loads(json_str, strict=False) instead, which is the standard way to allow raw control characters (such as newlines) inside JSON string values without breaking structural whitespace.

Suggested change

	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)
	json_str = tool_call_match.group(1)
	tool_call_data = json.loads(json_str, strict=False)

[gemini-code-assist]

There are three critical issues in the streaming finalization logic:

1. Lack of Exception Handling: Any transient network error or premature stream disconnect will raise an exception and crash the entire rollout task. Wrapping the stream block in a try...except block ensures robustness.
2. Incorrect Status on Stream Failure: If the stream fails or disconnects prematurely without being aborted, the sample status remains PENDING. It should be marked as FAILED to prevent training on incomplete data.
3. Aborted Status Overwrite: If state.aborted is True, the sample status must be set to ABORTED even if finish_reason is set, to ensure proper recycling of the sample.

    try:
        with trace_span(
            sample, "vllm_inference_generate_stream", attrs={"max_new_tokens": params["max_new_tokens"]}
        ) as span:
            async with client.stream("POST", url, json=payload, headers=headers) as response:
                response.raise_for_status()
                async for raw_line in response.aiter_lines():
                    if not raw_line or not raw_line.startswith("data:"):
                        continue
                    data_str = raw_line[len("data:") :].strip()
                    if not data_str or data_str == "[DONE]":
                        continue
                    try:
                        chunk = json.loads(data_str)
                    except json.JSONDecodeError:
                        logger.warning("vllm_streaming: skipping non-JSON chunk: %r", data_str[:120])
                        continue

                    choices = chunk.get("choices") or []
                    if not choices:
                        # usage-only / keepalive chunk
                        if chunk.get("usage"):
                            last_usage = chunk["usage"]
                        continue
                    choice = choices[0]
                    last_choice = choice
                    if chunk.get("usage"):
                        last_usage = chunk["usage"]
                    if choice.get("finish_reason"):
                        finish_reason = choice["finish_reason"]

                    # Each streamed choice carries only this chunk's *delta* tokens
                    # (GenerateResponseStreamChoice), so accumulate.
                    delta_tokens, delta_log_probs = _inference_generate_tokens_and_logprobs(choice)
                    if delta_tokens:
                        call_tokens += delta_tokens
                        call_log_probs += delta_log_probs

                    # Surface partial state on the sample immediately. If the outer
                    # abort path cuts us, whatever we've written so far is what
                    # survives. Decode the *accumulated* tokens (not the per-chunk
                    # delta) so multi-token characters straddling a chunk boundary
                    # decode correctly.
                    sample.tokens = base_tokens + call_tokens
                    sample.response = base_response + (
                        state.tokenizer.decode(call_tokens, skip_special_tokens=skip_decode) if call_tokens else ""
                    )
                    sample.response_length = base_response_length + len(call_tokens)
                    sample.rollout_log_probs = base_log_probs + call_log_probs
                    if base_loss_mask is not None:
                        assert args.partial_rollout and args.mask_offpolicy_in_partial_rollout
                        sample.loss_mask = base_loss_mask + [1] * len(call_tokens)

                    if state.aborted:
                        break

            if finish_reason and last_choice is not None:
                span.update(build_vllm_meta_trace_attrs({"choices": [last_choice], "usage": last_usage}))
    except Exception as e:
        logger.warning("vllm_streaming: stream failed with exception: %r", e)

    if state.aborted:
        sample.status = Sample.Status.ABORTED
    elif finish_reason and last_choice is not None:
        # Finalize exactly like the non-streaming path: align logprobs to tokens,
        # rebuild meta + output_token_logprobs, then let Sample own status.
        new_response_tokens, new_response_log_probs = _align_engine_tokens_and_logprobs(call_tokens, call_log_probs)

        meta = _vllm_meta_from_generate_choice(args, last_choice, last_usage)
        if new_response_tokens:
            meta["output_token_logprobs"] = [
                [float(lp), int(tid)] for lp, tid in zip(new_response_log_probs, new_response_tokens, strict=True)
            ]

        sample.tokens = base_tokens + new_response_tokens
        sample.response = base_response + (
            state.tokenizer.decode(new_response_tokens, skip_special_tokens=skip_decode) if new_response_tokens else ""
        )
        sample.response_length = base_response_length + len(new_response_tokens)
        sample.rollout_log_probs = base_log_probs + new_response_log_probs
        if base_loss_mask is not None:
            assert args.partial_rollout and args.mask_offpolicy_in_partial_rollout
            sample.loss_mask = base_loss_mask + [1] * len(new_response_tokens)

        sample.update_from_meta_info(args, meta)
        # MoE routing replay (when requested) ships on the terminal choice.
        _apply_vllm_routed_experts(args, sample, last_choice)
    else:
        sample.status = Sample.Status.FAILED

[CalvinXKY]

LGTM