diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index cb7432f97..1c6abeabf 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -669,6 +669,10 @@ jobs: "num_gpus": 0, "test_file": "test_external_sglang_engines.py" }, + { + "num_gpus": 0, + "test_file": "test_streaming_rollout.py" + }, { "num_gpus": 0, "test_file": "test_empty_colocated_weight_bucket.py" diff --git a/.github/workflows/pr-test.yml.j2 b/.github/workflows/pr-test.yml.j2 index fba22ea91..f47d8a9f6 100644 --- a/.github/workflows/pr-test.yml.j2 +++ b/.github/workflows/pr-test.yml.j2 @@ -93,6 +93,7 @@ {'test_file': 'test_reloadable_process_group_world.py', 'num_gpus': 0}, {'test_file': 'test_placement_group.py', 'num_gpus': 0}, {'test_file': 'test_external_sglang_engines.py', 'num_gpus': 0}, + {'test_file': 'test_streaming_rollout.py', 'num_gpus': 0}, {'test_file': 'test_empty_colocated_weight_bucket.py', 'num_gpus': 0}, {'test_file': 'test_expert_routing.py', 'num_gpus': 0}, {'test_file': 'test_glm5_indexer_short_context.py', 'num_gpus': 0}, diff --git a/slime/rollout/sglang_rollout.py b/slime/rollout/sglang_rollout.py index b76013bfb..d6951ef71 100644 --- a/slime/rollout/sglang_rollout.py +++ b/slime/rollout/sglang_rollout.py @@ -5,7 +5,7 @@ import logging import uuid from argparse import Namespace -from collections.abc import Callable +from collections.abc import Awaitable, Callable from contextlib import contextmanager from typing import Any @@ -132,6 +132,8 @@ def reset(self) -> None: self.remaining_batch_size = 0 self.pendings = set() self.aborted = False + self.cancellable_tasks = set() + self.active_server_generations = 0 def submit_generate_tasks(self, samples: list[list[Sample]]) -> None: for group in samples: @@ -219,6 +221,36 @@ async def generate(args: Namespace, sample: Sample, sampling_params: dict[str, A return sample +async def _run_request_abortable_generate( + state: GenerateState, + sample: Sample, + generate_call: Awaitable[Sample | list[Sample]], +) -> Sample | list[Sample]: + task = asyncio.current_task() + assert task is not None + state.cancellable_tasks.add(task) + try: + return await generate_call + except asyncio.CancelledError: + if task in state.cancellable_tasks: + raise + sample.status = Sample.Status.ABORTED + return sample + finally: + state.cancellable_tasks.discard(task) + + +async def _run_server_abort_generate( + state: GenerateState, + generate_call: Awaitable[Sample | list[Sample]], +) -> Sample | list[Sample]: + state.active_server_generations += 1 + try: + return await generate_call + finally: + state.active_server_generations -= 1 + + @trace_function("generate_and_rm", target="sample") async def generate_and_rm( args: Namespace, @@ -246,18 +278,18 @@ async def generate_and_rm( return sample with state.dp_rank_context() as _: - # Check sample.generate_function_path for per-sample custom_generate_function_path (e.g., from eval dataset config) - custom_func_path = getattr(sample, "generate_function_path", None) or args.custom_generate_function_path - - if custom_func_path is not None: - custom_generate_func = load_function(custom_func_path) - # if signature has evaluation, pass evaluation - if "evaluation" in inspect.signature(custom_generate_func).parameters: - sample = await custom_generate_func(args, sample, sampling_params, evaluation=evaluation) - else: - sample = await custom_generate_func(args, sample, sampling_params) + custom_func_path = sample.generate_function_path or args.custom_generate_function_path + generate_func = load_function(custom_func_path) if custom_func_path is not None else generate + + if custom_func_path is not None and "evaluation" in inspect.signature(generate_func).parameters: + generate_call = generate_func(args, sample, sampling_params, evaluation=evaluation) + else: + generate_call = generate_func(args, sample, sampling_params) + + if getattr(generate_func, "abort_mode", None) == "request": + sample = await _run_request_abortable_generate(state, sample, generate_call) else: - sample = await generate(args, sample, sampling_params) + sample = await _run_server_abort_generate(state, generate_call) sample = await apply_rollout_sample_hooks(args, sample, evaluation=evaluation) @@ -341,10 +373,17 @@ async def abort(args: Namespace, rollout_id: int) -> list[list[Sample]]: assert not state.aborted state.aborted = True - response = await get(f"http://{args.sglang_router_ip}:{args.sglang_router_port}/workers") - urls = [worker["url"] for worker in response["workers"]] + cancellable_tasks = list(state.cancellable_tasks) + state.cancellable_tasks.difference_update(cancellable_tasks) + for task in cancellable_tasks: + task.cancel() + + if state.active_server_generations: + response = await get(f"http://{args.sglang_router_ip}:{args.sglang_router_port}/workers") + urls = [worker["url"] for worker in response["workers"]] + await abort_servers_until_idle(urls) - await abort_servers_until_idle(urls) + await asyncio.gather(*cancellable_tasks, return_exceptions=True) # make sure all the pending tasks are finished count = 0 @@ -357,6 +396,8 @@ async def abort(args: Namespace, rollout_id: int) -> list[list[Sample]]: # for partial rollout, collect the partial samples into the data buffer for task in done: group = task.result() + if not any(sample.status == Sample.Status.ABORTED and sample.response_length > 0 for sample in group): + continue for sample in group: if sample.response and "start_rollout_id" not in sample.metadata: sample.metadata["start_rollout_id"] = rollout_id diff --git a/slime/rollout/sglang_streaming_rollout.py b/slime/rollout/sglang_streaming_rollout.py index 12471148c..e7309514f 100644 --- a/slime/rollout/sglang_streaming_rollout.py +++ b/slime/rollout/sglang_streaming_rollout.py @@ -13,15 +13,19 @@ --rollout-function-path slime.rollout.sglang_rollout.generate_rollout \\ --custom-generate-function-path slime.rollout.sglang_streaming_rollout.generate_streaming +This generator selects request abort, so Slime cancels each active streaming +HTTP request instead of using server-wide abort. + +Request cancellation only preserves metadata received before disconnect. +SGLang emits top-p and routed-expert replay data on its terminal chunk, so use +server abort when those features must survive a partial rollout. + The outer rollout loop (semaphore, dp_rank balancing, abort orchestration, partial-rollout buffer hand-off) is still owned by ``sglang_rollout``; this file only replaces the inner HTTP call. -sglang's default streaming output is cumulative — server-side -``state.output_token_logprobs`` accumulates and every chunk references the -full list-so-far (see ``tokenizer_manager.py``). If anyone ever flips -``--incremental-streaming-output`` on the sglang server, the text/output_ids -deltas will need different handling here. +Both cumulative and incremental SGLang streams are accepted. The latter is used +when the server enables incremental streaming output. """ import json @@ -30,6 +34,7 @@ from typing import Any from slime.rollout.sglang_rollout import GenerateState, _prepare_prompt_ids +from slime.rollout.streaming_utils import SGLangStreamAccumulator from slime.utils import http_utils from slime.utils.processing_utils import encode_image_for_rollout_engine from slime.utils.trace_utils import build_sglang_meta_trace_attrs, trace_span @@ -43,8 +48,8 @@ async def generate_streaming(args: Namespace, sample: Sample, sampling_params: dict[str, Any]) -> Sample: """Streaming counterpart to :func:`slime.rollout.sglang_rollout.generate`. - Writes the cumulative state from each SSE chunk onto ``sample`` so an - abort that cuts the stream still leaves a coherent partial sample behind. + Applies each SSE chunk onto ``sample`` so an abort that cuts the stream + still leaves a coherent partial sample behind. """ if args.ci_test: assert isinstance(sample.prompt, str) @@ -88,78 +93,86 @@ async def generate_streaming(args: Namespace, sample: Sample, sampling_params: d if sample.session_id and getattr(args, "router_policy", None) == "consistent_hashing": headers = {"X-SMG-Routing-Key": sample.session_id} - # Snapshot pre-call sample state. sglang's SSE chunks are cumulative - # *within this call*; on each chunk we rebuild the post-call view of the - # sample = prior state + chunk delta. That way a mid-stream break leaves - # the sample exactly at the boundary of the last chunk we observed. + # Preserve the pre-call state so both stream formats also work when + # resuming a partial rollout. A terminal full top-p snapshot may replay + # the current call from this base to keep metadata aligned. base_tokens = list(sample.tokens) base_response = sample.response or "" base_response_length = sample.response_length base_log_probs = None if sample.rollout_log_probs is None else list(sample.rollout_log_probs) base_top_p_token_ids = sample.rollout_top_p_token_ids base_top_p_token_offsets = sample.rollout_top_p_token_offsets + base_routed_experts = sample.rollout_routed_experts base_loss_mask = list(sample.loss_mask) if sample.loss_mask is not None else None last_meta_info: dict[str, Any] = {} - call_tokens: list[int] = [] - call_log_probs: list[float] = [] - call_text: str = "" + stream = SGLangStreamAccumulator( + output_mode="incremental" if args.sglang_incremental_streaming_output else "cumulative" + ) client = http_utils._http_client assert client is not None, "http client not initialized; call init_http_client first" - with trace_span( - sample, "sglang_generate_stream", attrs={"max_new_tokens": sampling_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("sglang_streaming: skipping non-JSON chunk: %r", data_str[:120]) - continue - - meta = chunk.get("meta_info") or {} - last_meta_info = meta - - call_text = chunk.get("text", call_text) - if "output_token_logprobs" in meta: - call_tokens = [item[1] for item in meta["output_token_logprobs"]] - call_log_probs = [item[0] for item in meta["output_token_logprobs"]] - - # Surface partial state on the sample immediately. If the - # outer abort path cuts us, whatever we've written so far is - # what survives — no /abort_request round-trip needed. - sample.tokens = list(base_tokens) - sample.response = base_response - sample.response_length = base_response_length - sample.rollout_log_probs = None if base_log_probs is None else list(base_log_probs) - sample.rollout_top_p_token_ids = base_top_p_token_ids - sample.rollout_top_p_token_offsets = base_top_p_token_offsets - sample.loss_mask = None if base_loss_mask is None else list(base_loss_mask) - sample.append_response_tokens( - args, - tokens=call_tokens, - log_probs=call_log_probs, - trainable=True, - meta_info=meta, - text=call_text, - update_terminal_info=bool(meta.get("finish_reason")), - ) - - if state.aborted: - break - - if last_meta_info.get("finish_reason"): - span.update(build_sglang_meta_trace_attrs(last_meta_info)) + try: + with trace_span( + sample, "sglang_generate_stream", attrs={"max_new_tokens": sampling_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("sglang_streaming: skipping non-JSON chunk: %r", data_str[:120]) + continue + + update = stream.add(chunk) + last_meta_info = update.meta_info + + if update.replace_call_state: + sample.tokens = list(base_tokens) + sample.response = base_response + sample.response_length = base_response_length + sample.rollout_log_probs = None if base_log_probs is None else list(base_log_probs) + sample.rollout_top_p_token_ids = base_top_p_token_ids + sample.rollout_top_p_token_offsets = base_top_p_token_offsets + sample.rollout_routed_experts = base_routed_experts + sample.loss_mask = None if base_loss_mask is None else list(base_loss_mask) + + sample.append_response_tokens( + args, + tokens=update.tokens, + log_probs=update.log_probs, + trainable=True, + meta_info=last_meta_info, + text=None, + update_terminal_info=bool(last_meta_info.get("finish_reason")), + ) + + if state.aborted: + break + + if last_meta_info.get("finish_reason"): + span.update(build_sglang_meta_trace_attrs(last_meta_info)) + finally: + sample.response = base_response + stream.response_text( + lambda token_ids: state.tokenizer.decode( + token_ids, + skip_special_tokens=sampling_params.get("skip_special_tokens", True), + ) + ) if state.aborted and not last_meta_info.get("finish_reason"): sample.status = Sample.Status.ABORTED + elif not last_meta_info.get("finish_reason"): + raise RuntimeError("SGLang streaming response ended without a terminal finish_reason.") return sample + + +generate_streaming.abort_mode = "request" diff --git a/slime/rollout/streaming_utils.py b/slime/rollout/streaming_utils.py new file mode 100644 index 000000000..42793f471 --- /dev/null +++ b/slime/rollout/streaming_utils.py @@ -0,0 +1,129 @@ +"""Helpers for consuming SGLang streaming responses.""" + +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Any, Literal + +from slime.utils.misc import decode_int32_meta_array + +_TOP_P_TOKEN_ID_META_KEYS = ("top_p_token_ids", "top_p_kept_token_ids") +_TOP_P_TOKEN_OFFSET_META_KEYS = ("top_p_token_offsets", "top_p_kept_token_offsets") + + +def _has_full_response_top_p_metadata( + meta_info: dict[str, Any], + *, + new_token_count: int, + reported_length: int, +) -> bool: + token_ids = decode_int32_meta_array(meta_info, _TOP_P_TOKEN_ID_META_KEYS) + offsets = decode_int32_meta_array(meta_info, _TOP_P_TOKEN_OFFSET_META_KEYS) + if token_ids is None or offsets is None or offsets.numel() == new_token_count + 1: + return False + if offsets.numel() != reported_length + 1: + raise ValueError( + "Incremental SGLang top-p metadata must describe either the current chunk or the full response: " + f"offsets={offsets.numel()}, chunk_tokens={new_token_count}, reported_tokens={reported_length}." + ) + + return True + + +@dataclass(frozen=True) +class SGLangStreamUpdate: + """A normalized update ready to apply to the current HTTP call state.""" + + replace_call_state: bool + tokens: list[int] + log_probs: list[float] + meta_info: dict[str, Any] + + +@dataclass +class SGLangStreamAccumulator: + """Validate and accumulate configured SGLang stream chunks. + + SGLang's ``output_token_logprobs_length`` is cumulative in both modes. + In cumulative mode each chunk must contain that many logprob pairs; in + incremental mode the current and prior chunk lengths must sum to it. + """ + + output_mode: Literal["incremental", "cumulative"] + output_length: int = 0 + tokens: list[int] = field(default_factory=list) + log_probs: list[float] = field(default_factory=list) + _text_chunks: list[str] = field(default_factory=list) + _cumulative_text: str | None = None + _decode_text: bool = False + + def add(self, chunk: dict[str, Any]) -> SGLangStreamUpdate: + meta_info = chunk.get("meta_info") or {} + if "output_token_logprobs_length" not in meta_info: + raise ValueError("SGLang streaming responses must include output_token_logprobs_length.") + + pairs = meta_info.get("output_token_logprobs") or [] + chunk_tokens = [item[1] for item in pairs] + chunk_log_probs = [item[0] for item in pairs] + reported_length = int(meta_info["output_token_logprobs_length"]) + + previous_length = self.output_length + cumulative = self.output_mode == "cumulative" + expected_length = len(chunk_tokens) if cumulative else previous_length + len(chunk_tokens) + if expected_length != reported_length: + raise ValueError( + f"SGLang {self.output_mode} streaming output has inconsistent output_token_logprobs_length: " + f"received={len(chunk_tokens)}, previous={previous_length}, " + f"expected={expected_length}, reported={reported_length}." + ) + if reported_length < previous_length: + raise ValueError( + "SGLang cumulative streaming output length decreased: " + f"previous={previous_length}, reported={reported_length}." + ) + + if cumulative: + new_tokens = chunk_tokens[previous_length:] + new_log_probs = chunk_log_probs[previous_length:] + self.tokens.extend(new_tokens) + self.log_probs.extend(new_log_probs) + else: + new_tokens = chunk_tokens + new_log_probs = chunk_log_probs + self.tokens.extend(new_tokens) + self.log_probs.extend(new_log_probs) + + full_top_p_metadata = _has_full_response_top_p_metadata( + meta_info, + new_token_count=len(new_tokens), + reported_length=reported_length, + ) + + chunk_text = chunk.get("text") + if cumulative: + if chunk_text is not None: + self._cumulative_text = chunk_text + elif new_tokens: + self._decode_text = True + else: + if chunk_text is None: + self._decode_text = True + elif chunk_text: + self._text_chunks.append(chunk_text) + + self.output_length = reported_length + return SGLangStreamUpdate( + replace_call_state=full_top_p_metadata, + tokens=list(self.tokens) if full_top_p_metadata else new_tokens, + log_probs=list(self.log_probs) if full_top_p_metadata else new_log_probs, + meta_info=meta_info, + ) + + def response_text(self, decode: Callable[[list[int]], str]) -> str: + """Materialize response text once after completion or cancellation.""" + if self.output_mode == "cumulative": + if self._cumulative_text is not None and not self._decode_text: + return self._cumulative_text + return decode(self.tokens) + if self._decode_text: + return decode(self.tokens) + return "".join(self._text_chunks) diff --git a/slime/utils/arguments.py b/slime/utils/arguments.py index 7eda64eb7..a6016ef42 100644 --- a/slime/utils/arguments.py +++ b/slime/utils/arguments.py @@ -479,7 +479,9 @@ def add_rollout_arguments(parser): default=None, help=( "Only substitue the `def generate(args, sample, sampling_params)` function within the example rollout function. " - "This should be useful if you need to implement some special rollout logic, e.g. multi-turn, function calling." + "This should be useful if you need to implement some special rollout logic, e.g. multi-turn, function calling. " + "Set `abort_mode = 'request'` on the function when it implements request-level abort; otherwise " + "Slime uses server-wide abort." ), ) parser.add_argument( diff --git a/tests/plugin_contracts/test_plugin_generate_contracts.py b/tests/plugin_contracts/test_plugin_generate_contracts.py index 9cf7af3a3..ceb91ba1c 100644 --- a/tests/plugin_contracts/test_plugin_generate_contracts.py +++ b/tests/plugin_contracts/test_plugin_generate_contracts.py @@ -61,6 +61,8 @@ def __init__(self, args) -> None: self.pendings = set() self.remaining_batch_size = 0 self.aborted = False + self.cancellable_tasks = set() + self.active_server_generations = 0 self.group_sampling_seeds = [args.rollout_seed + i for i in range(args.n_samples_per_prompt)] @contextmanager diff --git a/tests/test_streaming_rollout.py b/tests/test_streaming_rollout.py new file mode 100644 index 000000000..7bec8bb30 --- /dev/null +++ b/tests/test_streaming_rollout.py @@ -0,0 +1,527 @@ +import asyncio +import base64 +import json +import struct +import sys +from contextlib import asynccontextmanager, nullcontext +from pathlib import Path +from types import SimpleNamespace + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +try: + import sglang_router # noqa: F401 +except ImportError: + sys.modules["sglang_router"] = SimpleNamespace(__version__="0.3.0") + +try: + import transformers # noqa: F401 +except ImportError: + sys.modules["transformers"] = SimpleNamespace( + AutoProcessor=object, + AutoTokenizer=object, + PreTrainedTokenizerBase=object, + ProcessorMixin=object, + ) + +from slime.rollout import sglang_rollout +from slime.rollout import sglang_streaming_rollout as streaming +from slime.rollout.streaming_utils import SGLangStreamAccumulator +from slime.utils.types import Sample + +NUM_GPUS = 0 + + +def _chunk(text, pairs, output_length, **meta_info): + return { + "text": text, + "meta_info": { + "output_token_logprobs": pairs, + "output_token_logprobs_length": output_length, + **meta_info, + }, + } + + +def _b64_int32(values): + return base64.b64encode(struct.pack(f"<{len(values)}i", *values)).decode("ascii") + + +@pytest.mark.parametrize( + ("output_mode", "chunks", "expected_updates", "expected"), + [ + ( + "cumulative", + [ + _chunk("a", [[-0.1, 11, None]], 1), + _chunk("ab", [[-0.1, 11, None], [-0.2, 12, None]], 2), + ], + [[11], [12]], + ([11, 12], [-0.1, -0.2], "ab"), + ), + ( + "incremental", + [ + _chunk("a", [[-0.1, 11, None]], 1), + _chunk("b", [[-0.2, 12, None]], 2), + _chunk("", [], 2), + ], + [[11], [12], []], + ([11, 12], [-0.1, -0.2], "ab"), + ), + ( + "incremental", + [ + _chunk(None, [[-0.1, 11, None]], 1), + _chunk(None, [[-0.2, 12, None]], 2), + ], + [[11], [12]], + ([11, 12], [-0.1, -0.2], "<11><12>"), + ), + ], + ids=["cumulative", "incremental", "null-text"], +) +def test_merge_stream_chunks(output_mode, chunks, expected_updates, expected): + accumulator = SGLangStreamAccumulator(output_mode=output_mode) + updates = [] + for chunk in chunks: + updates.append(accumulator.add(chunk).tokens) + + decode_calls = 0 + + def decode(token_ids): + nonlocal decode_calls + decode_calls += 1 + return "".join(f"<{token_id}>" for token_id in token_ids) + + result = (accumulator.tokens, accumulator.log_probs, accumulator.response_text(decode)) + + assert updates == expected_updates + assert result == expected + assert decode_calls == int(chunks[-1]["text"] is None) + + +def test_merge_stream_rejects_inconsistent_length(): + accumulator = SGLangStreamAccumulator( + output_mode="incremental", + output_length=1, + tokens=[11], + ) + with pytest.raises(ValueError, match="output_token_logprobs_length"): + accumulator.add(_chunk("bc", [[-0.2, 12, None], [-0.3, 13, None]], 4)) + + +def _generation_state(): + return SimpleNamespace( + tokenizer=SimpleNamespace( + decode=lambda token_ids, skip_special_tokens=False: "".join(f"<{token_id}>" for token_id in token_ids) + ), + processor=None, + aborted=False, + cancellable_tasks=set(), + active_server_generations=0, + pendings=set(), + semaphore=asyncio.Semaphore(8), + dp_rank_context=lambda: nullcontext(None), + ) + + +def _streaming_args(): + return SimpleNamespace( + ci_test=False, + sglang_router_ip="frontend", + sglang_router_port=8000, + use_rollout_routing_replay=False, + sglang_incremental_streaming_output=True, + router_policy=None, + sglang_speculative_algorithm=False, + num_layers=1, + moe_router_topk=1, + partial_rollout=False, + mask_offpolicy_in_partial_rollout=False, + custom_generate_function_path="slime.rollout.sglang_streaming_rollout.generate_streaming", + group_rm=True, + rollout_sample_hook_path=None, + ) + + +def _patch_streaming(monkeypatch, state, stream): + monkeypatch.setattr(streaming, "GenerateState", lambda _args: state) + monkeypatch.setattr(sglang_rollout, "GenerateState", lambda _args: state) + monkeypatch.setattr(sglang_rollout, "load_function", lambda _path: streaming.generate_streaming) + monkeypatch.setattr(streaming, "_prepare_prompt_ids", lambda *_args: [1, 2]) + monkeypatch.setattr(streaming.http_utils, "_http_client", SimpleNamespace(stream=stream)) + monkeypatch.setattr( + streaming, + "trace_span", + lambda *_args, **_kwargs: nullcontext(SimpleNamespace(update=lambda *_args, **_kwargs: None)), + ) + + +def _http_stream(*chunks, lines=None, expected_payload=None, on_close=None): + async def chunk_lines(): + for chunk in chunks: + yield "data: " + json.dumps(chunk) + + response = SimpleNamespace(raise_for_status=lambda: None, aiter_lines=lines or chunk_lines) + + @asynccontextmanager + async def stream(method, url, json, headers): + assert (method, url, json["stream"]) == ("POST", "http://frontend:8000/generate", True) + for key, value in (expected_payload or {}).items(): + assert json[key] == value + try: + yield response + finally: + if on_close is not None: + on_close() + + return stream + + +def test_streaming_generate_selects_request_abort(): + assert streaming.generate_streaming.abort_mode == "request" + + +def test_stream_accumulator_requires_reported_length(): + with pytest.raises(ValueError, match="must include output_token_logprobs_length"): + SGLangStreamAccumulator(output_mode="incremental").add( + {"text": "x", "meta_info": {"output_token_logprobs": [[-0.1, 11, None]]}}, + ) + + +def test_stream_accumulator_rejects_output_incompatible_with_configured_mode(): + accumulator = SGLangStreamAccumulator(output_mode="incremental") + accumulator.add(_chunk("a", [[-0.1, 11, None]], 1)) + + with pytest.raises(ValueError, match="incremental streaming output has inconsistent"): + accumulator.add( + _chunk( + "ab", + [[-0.1, 11, None], [-0.2, 12, None]], + 2, + ), + ) + + +@pytest.mark.parametrize("stream_interval", [1, 20, 64]) +@pytest.mark.parametrize("incremental", [True, False], ids=["incremental", "cumulative"]) +@pytest.mark.parametrize("top_p_layout", ["final_full", "per_chunk"]) +def test_generate_streaming_preserves_metadata_across_stream_intervals( + monkeypatch, stream_interval, incremental, top_p_layout +): + # Keep the response longer than the largest interval so the second chunk + # makes cumulative versus incremental output observable on the wire. + response_tokens = list(range(11, 76)) + chunks = [] + previous = 0 + for end in range(stream_interval, len(response_tokens) + stream_interval, stream_interval): + end = min(end, len(response_tokens)) + start = previous if incremental else 0 + token_slice = response_tokens[start:end] + is_final = end == len(response_tokens) + top_p_tokens = token_slice if top_p_layout == "per_chunk" else response_tokens + top_p_metadata = ( + { + "top_p_token_ids": [token_id + 1000 for token_id in top_p_tokens], + "top_p_token_offsets": list(range(len(top_p_tokens) + 1)), + } + if top_p_layout == "per_chunk" or is_final + else {} + ) + chunks.append( + _chunk( + "".join(f"<{token_id}>" for token_id in token_slice), + [[-float(token_id), token_id, None] for token_id in token_slice], + end, + finish_reason={"type": "stop"} if is_final else None, + **top_p_metadata, + ) + ) + previous = end + if end == len(response_tokens): + break + + state = _generation_state() + args = _streaming_args() + args.sglang_incremental_streaming_output = incremental + _patch_streaming(monkeypatch, state, _http_stream(*chunks)) + + result = asyncio.run( + streaming.generate_streaming( + args, + Sample(prompt="hello"), + {"max_new_tokens": len(response_tokens), "skip_special_tokens": False}, + ) + ) + + assert result.status == Sample.Status.COMPLETED + assert result.tokens == [1, 2, *response_tokens] + assert result.response_length == len(response_tokens) + assert result.rollout_log_probs == [-float(token_id) for token_id in response_tokens] + assert result.rollout_top_p_token_ids.tolist() == [token_id + 1000 for token_id in response_tokens] + assert result.rollout_top_p_token_offsets.tolist() == list(range(len(response_tokens) + 1)) + + +@pytest.mark.parametrize("incremental", [True, False], ids=["incremental", "cumulative"]) +def test_generate_streaming_handles_encoded_terminal_metadata(monkeypatch, incremental): + first = _chunk("λ", [[-0.1, 11, None]], 1) + second = _chunk( + "<|eot_id|>" if incremental else "λ<|eot_id|>", + [[-0.2, 12, None]] if incremental else [[-0.1, 11, None], [-0.2, 12, None]], + 2, + ) + terminal = _chunk( + "" if incremental else "λ<|eot_id|>", + [] if incremental else [[-0.1, 11, None], [-0.2, 12, None]], + 2, + finish_reason={"type": "length"}, + top_p_token_ids=_b64_int32([1011, 1012]), + top_p_token_offsets=_b64_int32([0, 1, 2]), + routed_experts=_b64_int32([101, 102, 103]), + ) + + state = _generation_state() + args = _streaming_args() + args.use_rollout_routing_replay = True + args.sglang_incremental_streaming_output = incremental + stream = _http_stream(first, second, terminal, expected_payload={"return_routed_experts": True}) + _patch_streaming(monkeypatch, state, stream) + + result = asyncio.run( + streaming.generate_streaming( + args, + Sample(prompt="hello"), + {"max_new_tokens": 2, "skip_special_tokens": False}, + ) + ) + + assert result.status == Sample.Status.TRUNCATED + assert result.response == "λ<|eot_id|>" + assert result.tokens == [1, 2, 11, 12] + assert result.rollout_top_p_token_ids.tolist() == [1011, 1012] + assert result.rollout_top_p_token_offsets.tolist() == [0, 1, 2] + assert result.rollout_routed_experts.tolist() == [[[101]], [[102]], [[103]]] + + +def test_streaming_generator_fails_closed_on_unexpected_eof(monkeypatch): + stream = _http_stream(_chunk("a", [[-0.1, 11, None]], 1)) + _patch_streaming(monkeypatch, _generation_state(), stream) + + with pytest.raises(RuntimeError, match="without a terminal finish_reason"): + asyncio.run( + streaming.generate_streaming( + _streaming_args(), + Sample(prompt="hello"), + {"max_new_tokens": 2}, + ) + ) + + +def test_server_abort_keeps_last_observed_prefix(monkeypatch): + state = _generation_state() + + async def lines(): + state.aborted = True + yield "data: " + json.dumps(_chunk("a", [[-0.1, 11, None]], 1)) + raise AssertionError("server-aborted stream should stop after the next observed chunk") + + _patch_streaming(monkeypatch, state, _http_stream(lines=lines)) + result = asyncio.run( + streaming.generate_streaming( + _streaming_args(), + Sample(prompt="hello"), + {"max_new_tokens": 2}, + ) + ) + + assert result.status == Sample.Status.ABORTED + assert (result.tokens, result.rollout_log_probs, result.response) == ([1, 2, 11], [-0.1], "a") + assert not state.cancellable_tasks + + +def test_stream_cancellation_closes_request_and_keeps_prefix(monkeypatch): + first_chunk_seen = asyncio.Event() + request_closed = asyncio.Event() + server_started = asyncio.Event() + server_released = asyncio.Event() + state = _generation_state() + + async def lines(): + yield "data: " + json.dumps( + { + "text": None, + "meta_info": { + "output_token_logprobs": [[-0.1, 11, None]], + "output_token_logprobs_length": 1, + }, + } + ) + first_chunk_seen.set() + await asyncio.Event().wait() + + _patch_streaming(monkeypatch, state, _http_stream(lines=lines, on_close=request_closed.set)) + args = _streaming_args() + args.custom_generate_function_path = None + + async def server_generate(args, sample, sampling_params): + server_started.set() + await server_released.wait() + sample.status = Sample.Status.ABORTED + return sample + + async def get_workers(_url): + return {"workers": [{"url": "http://worker:30000"}]} + + async def abort_servers(urls): + assert urls == ["http://worker:30000"] + server_released.set() + + monkeypatch.setattr(sglang_rollout, "generate", server_generate) + monkeypatch.setattr(sglang_rollout, "get", get_workers) + monkeypatch.setattr(sglang_rollout, "abort_servers_until_idle", abort_servers) + + async def exercise(): + request_task = asyncio.create_task( + sglang_rollout.generate_and_rm( + args, + Sample(prompt="hello", generate_function_path="streaming"), + {"max_new_tokens": 8}, + ) + ) + server_task = asyncio.create_task(sglang_rollout.generate_and_rm(args, Sample(), {})) + await asyncio.gather(first_chunk_seen.wait(), server_started.wait()) + await sglang_rollout.abort(args, rollout_id=0) + return await asyncio.gather(request_task, server_task) + + result, server_result = asyncio.run(exercise()) + + assert result.status == Sample.Status.ABORTED + assert server_result.status == Sample.Status.ABORTED + assert (result.tokens, result.rollout_log_probs, result.response) == ([1, 2, 11], [-0.1], "<11>") + assert request_closed.is_set() + assert not state.cancellable_tasks + assert state.active_server_generations == 0 + + +def test_unrelated_stream_cancellation_propagates(monkeypatch): + request_started = asyncio.Event() + state = _generation_state() + + async def lines(): + request_started.set() + await asyncio.Event().wait() + yield "unreachable" + + _patch_streaming(monkeypatch, state, _http_stream(lines=lines)) + args = _streaming_args() + + async def exercise(): + task = asyncio.create_task( + sglang_rollout.generate_and_rm( + args, + Sample(prompt="hello"), + {"max_new_tokens": 8}, + ) + ) + await request_started.wait() + state.aborted = True + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + asyncio.run(exercise()) + assert not state.cancellable_tasks + + +def test_partial_abort_buffers_and_resumes_only_aborted_siblings(monkeypatch): + partial = Sample( + tokens=[1, 11], + response="x", + response_length=1, + rollout_log_probs=[-0.1], + loss_mask=[1], + status=Sample.Status.ABORTED, + ) + terminal = Sample( + tokens=[1, 21, 22], + response="done", + response_length=2, + rollout_log_probs=[-0.3, -0.4], + loss_mask=[1, 1], + reward=1.0, + status=Sample.Status.COMPLETED, + ) + empty = Sample(response="", response_length=0, status=Sample.Status.ABORTED) + mixed_group = [terminal, partial] + + async def exercise(): + state = SimpleNamespace( + aborted=False, + cancellable_tasks=set(), + active_server_generations=0, + pendings={ + asyncio.create_task(asyncio.sleep(0, result=group)) for group in (mixed_group, [terminal], [empty]) + }, + ) + monkeypatch.setattr(sglang_rollout, "GenerateState", lambda _args: state) + args = SimpleNamespace( + partial_rollout=True, + mask_offpolicy_in_partial_rollout=True, + group_rm=False, + custom_generate_function_path="test.resume_generate", + rollout_sample_hook_path=None, + sglang_enable_deterministic_inference=False, + sglang_speculative_algorithm=False, + ) + + buffered_groups = await sglang_rollout.abort(args, rollout_id=7) + assert buffered_groups == [mixed_group] + + async def resume_generate(args, sample, sampling_params): + assert sample is partial + sample.append_response_tokens( + args, + tokens=[12], + log_probs=[-0.2], + meta_info={"finish_reason": {"type": "stop"}}, + text="y", + ) + return sample + + async def reward(_args, sample): + assert sample is partial + return 2.0 + + monkeypatch.setattr(sglang_rollout, "load_function", lambda _path: resume_generate) + monkeypatch.setattr(sglang_rollout, "async_rm", reward) + state.aborted = False + state.semaphore = asyncio.Semaphore(2) + state.dp_rank_context = lambda: nullcontext(None) + + resumed_group = await sglang_rollout.generate_and_rm_group(args, mixed_group, {"max_new_tokens": 1}) + return buffered_groups, resumed_group + + buffered_groups, resumed_group = asyncio.run(exercise()) + + assert buffered_groups == [mixed_group] + assert resumed_group == [terminal, partial] + assert partial.metadata["start_rollout_id"] == 7 + assert terminal.metadata["start_rollout_id"] == 7 + assert terminal.status == Sample.Status.COMPLETED + assert terminal.reward == 1.0 + assert partial.status == Sample.Status.COMPLETED + assert partial.tokens == [1, 11, 12] + assert partial.response == "xy" + assert partial.response_length == 2 + assert partial.rollout_log_probs == [-0.1, -0.2] + assert partial.loss_mask == [0, 1] + assert partial.reward == 2.0 + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__]))