[Rollout] Add streaming vLLM rollout (port slime #1921)

[aoshen02]

Ports THUDM/slime PR #1921 (streaming rollout) onto vime's vLLM rollout seam.

What

• vime/rollout/vllm_streaming_rollout.py — a drop-in --custom-generate-function-path that consumes /inference/v1/generate as an SSE stream (stream: true) and writes each chunk's tokens/logprobs onto the sample immediately. The win is on abort: a partial-rollout / weight-update abort mid-generation leaves a coherent partial sample behind, without depending on an /abort_request response.
• tests/test_qwen3_4B_streaming_partial_rollout.py — 8-GPU GRPO smoke (Qwen3-4B + dapo-math) with --over-sampling-batch-size > --rollout-batch-size + --partial-rollout, so the loop must abort in-flight requests every step and exercise the streaming abort path. Registered in the e2e-test-megatron matrix (j2 + rendered yml).

vime vs slime port notes

• slime's sglang stream sends cumulative meta_info.output_token_logprobs; vLLM's GenerateResponseStreamChoice chunks carry DELTA token_ids + logprobs, so this accumulates deltas (+=) per chunk. Each delta choice is the same shape as the non-streaming choice, so _inference_generate_tokens_and_logprobs parses it unchanged.
• Finalization mirrors vllm_rollout.generate: align logprobs → build meta → update_from_meta_info → _apply_vllm_routed_experts (routing replay).
• Partial-continuation max_new_tokens math + consistent_hash x-session-id header match the non-streaming path.
• Multimodal samples raise NotImplementedError (the MM render→generate flow is a two-call dance that doesn't map onto a single streamed call) rather than silently dropping images.
• Passthrough confirmed end-to-end in source: vLLM-fork engine streams /inference/v1/generate via SSE, and the Rust vllm_router forwards SSE transparently (Body::from_stream, text/event-stream).

Validation

• cpu gate (fresh --rm container from the latest cumulative CI image vime-vllm-cu129-sync1916): module import OK, test-file import OK, test_megatron_argument_validation 7 passed, plugin_contracts 34 passed (rollout 6 / runtime-hook 10 / path-loading 14 / generate 4).
• lint: pinned pre-commit (ruff/isort/black) clean on changed files.
• GPU CI: not run (per request).

Upstream: THUDM/slime#1921. Part of the slime→vime sync tracked in #107 (not ticking the box pending sign-off).

[gemini-code-assist]

Code Review

This pull request introduces a streaming vLLM rollout implementation along with a corresponding CI smoke test. The streaming rollout consumes vLLM's SSE stream incrementally to handle mid-generation aborts gracefully. Feedback on the implementation points out a lack of retry logic for HTTP connection and stream consumption, which could cause training runs to crash on transient network glitches. It is recommended to wrap the stream connection in a retry loop and restore the sample state on failure.

[gemini-code-assist]

The streaming rollout implementation lacks retry logic for the HTTP connection and stream consumption. In distributed LLM training, transient network glitches, router restarts, or temporary vLLM engine overloads are common. Without retries, any such transient error will immediately raise an exception, propagate through asyncio.gather, and crash the entire training run.

We should wrap the stream connection and consumption in a retry loop (e.g., up to 60 retries with a 1-second delay, matching the non-streaming post behavior). If a retry is triggered, we must restore the sample object to its pre-call snapshot state to avoid accumulating duplicate or corrupted tokens/responses. Additionally, if the stream finishes normally but finish_reason is not explicitly set in the chunks, we should default it to 'stop' to ensure the sample is properly finalized rather than being left in a PENDING state.

    base_tokens = list(sample.tokens)
    base_response = sample.response or ""
    base_response_length = sample.response_length
    base_log_probs = list(sample.rollout_log_probs) if sample.rollout_log_probs is not None else None
    base_loss_mask = list(sample.loss_mask) if sample.loss_mask is not None else None

    skip_sp = params.get("skip_special_tokens")
    skip_decode = True if skip_sp is None else bool(skip_sp)

    call_tokens: list[int] = []
    call_log_probs: list[float] = []
    last_choice: dict[str, Any] | None = None
    last_usage: dict[str, Any] | None = None
    finish_reason: Any = None

    client = http_utils._http_client
    assert client is not None, "http client not initialized; call init_http_client first"

    max_retries = 60
    retry_count = 0
    while retry_count < max_retries:
        call_tokens = []
        call_log_probs = []
        last_choice = None
        last_usage = None
        finish_reason = None

        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:
                            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"]

                        delta_tokens, delta_log_probs = _inference_generate_tokens_and_logprobs(choice)
                        if delta_tokens:
                            call_tokens += delta_tokens
                            call_log_probs += delta_log_probs

                        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 or []) + 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}))
            break
        except Exception as e:
            retry_count += 1
            if retry_count >= max_retries:
                logger.error("Max retries reached for streaming generate, failing...")
                raise e
            logger.warning("Streaming generate error: %s, retrying (%d/%d)...", e, retry_count, max_retries)
            sample.tokens = list(base_tokens)
            sample.response = base_response
            sample.response_length = base_response_length
            sample.rollout_log_probs = list(base_log_probs) if base_log_probs is not None else None
            if base_loss_mask is not None:
                sample.loss_mask = list(base_loss_mask)
            await asyncio.sleep(1)

    if not finish_reason and last_choice is not None and not state.aborted:
        finish_reason = "stop"

[aoshen02]

Valid. The streaming path lacks the non-streaming post() retry/backoff, so a transient router/engine glitch would propagate through asyncio.gather and crash the run; the snapshot-restore-on-retry and finish_reason default are both right.

Note: #117 is intentionally deferred to draft — there's no current streaming-partial-rollout need, and streaming can't carry R3 routing-replay (the vLLM GenerateResponseStreamChoice SSE schema has no routed_experts field, unlike sglang's terminal meta_info). When it's resumed, this retry+restore+finish_reason fix will be applied alongside the routing-replay/image → non-streaming fallbacks and the stream usage chunk.

[aoshen02]

Closing for now (slime→vime sync being restructured). Branch is preserved — can reopen anytime.