Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
84 changes: 53 additions & 31 deletions packages/sie_server/src/sie_server/adapters/sglang/generation.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,14 +34,16 @@
import uuid
from collections.abc import AsyncIterator
from pathlib import Path
from typing import Any, cast
from typing import Any, Literal, cast

import httpx

from sie_server.adapters._generation_base import (
FinishReason,
GenerationAdapter,
GenerationChunk,
GenerationError,
GenerationInvalidRequestError,
GenerationResult,
)
from sie_server.adapters._spec import AdapterSpec
Expand Down Expand Up @@ -156,17 +158,29 @@ def _encode_image_data(images: list[ImageInput] | None) -> list[str] | None:
return encoded


def _raise_for_sglang_event_error(event: Any) -> None:
"""Propagate an error carried inside SGLang's HTTP-200 SSE stream."""
if not isinstance(event, dict) or "error" not in event:
return
error = event["error"]
message = error.get("message") if isinstance(error, dict) else error
if not isinstance(message, str) or not message.strip():
message = "unknown in-band SGLang error"
message = message.strip()[:500]
logger.error("SGLang /generate in-band error: %s", message)
raise RuntimeError(f"SGLang /generate error: {message}")
def _raise_for_sglang_event_error(event: Any, *, grammar: GrammarSpec | None = None, terminal: bool = False) -> None:
"""Reject upstream failures before consuming their text or usage."""
if not isinstance(event, dict):
raise GenerationError("SGLang /generate returned an invalid event")
if "error" in event:
raise GenerationError("SGLang /generate returned an in-band error")
meta = event.get("meta_info")
finish = meta.get("finish_reason") if isinstance(meta, dict) else None
kind = finish.get("type") if isinstance(finish, dict) else finish
if finish is not None and kind is None:
raise GenerationError("SGLang /generate returned a malformed finish reason")
if kind == "abort":
if isinstance(finish, dict) and finish.get("status_code") == 400:
if grammar is not None:
raise GenerationInvalidRequestError(
"grammar", f"The generation backend rejected the requested {grammar.kind} grammar"
)
raise GenerationInvalidRequestError("prompt", "The generation backend rejected the request")
raise GenerationError("SGLang /generate aborted the request")
if kind is not None and kind not in ("stop", "length"):
raise GenerationError("SGLang /generate returned an unsupported finish reason")
if (terminal or event.get("finished")) and kind is None:
raise GenerationError("SGLang /generate completed without a finish reason")


def _tail_file(path: str, *, max_lines: int = 200) -> str:
Expand Down Expand Up @@ -1233,6 +1247,7 @@ async def generate(
sbody["top_logprobs_num"] = top_logprobs
sclient = await self._get_or_create_http_client()
last_text: dict[int, str] = {}
completed_candidates: set[int] = set()
# Per-candidate logprob watermark: SGLang's
# ``meta_info.output_token_logprobs`` is a per-candidate cumulative
# list growing across events for that index. Slicing
Expand All @@ -1257,8 +1272,12 @@ async def generate(
event = json.loads(line)
except json.JSONDecodeError:
continue
_raise_for_sglang_event_error(event)
idx = int(event.get("index", 0))
_raise_for_sglang_event_error(event, grammar=grammar)
Comment thread
coderabbitai[bot] marked this conversation as resolved.
idx = event.get("index")
if not isinstance(idx, int) or isinstance(idx, bool) or not 0 <= idx < return_count:
raise GenerationError("SGLang /generate returned an invalid candidate index")
if idx in completed_candidates:
raise GenerationError("SGLang /generate returned an event after the candidate terminal")
cumulative = event.get("text", "")
if not isinstance(cumulative, str):
cumulative = last_text.get(idx, "")
Expand All @@ -1270,6 +1289,8 @@ async def generate(
fr = meta.get("finish_reason")
fr_type = fr.get("type") if isinstance(fr, dict) else fr
candidate_done = fr_type is not None
if candidate_done:
completed_candidates.add(idx)
if candidate_done and isinstance(meta.get("completion_tokens"), int):
total_completion += meta["completion_tokens"]
# Per-candidate logprob slice — same shape conversion as the
Expand Down Expand Up @@ -1332,6 +1353,8 @@ async def generate(
choice_index=idx,
logprobs=chunk_logprobs,
)
if len(completed_candidates) != return_count:
raise GenerationError("SGLang /generate ended before all candidates completed")
# Single global terminal closes the multi-candidate stream (carries
# aggregate usage). Each candidate already received its own
# ``finish_reason`` on the per-choice completion chunk above; this
Expand Down Expand Up @@ -1378,6 +1401,12 @@ async def generate(
# tolerate a single dict defensively.
if isinstance(results, dict):
results = [results]
if not isinstance(results, list):
raise GenerationError("SGLang /generate returned an invalid candidate list")
for result in results:
_raise_for_sglang_event_error(result, grammar=grammar, terminal=True)
if len(results) != gen_count:
raise GenerationError("SGLang /generate returned an incorrect candidate count")
if rank:
# Highest cumulative token-logprob first; keep the top return_count.
results = sorted(results, key=_cumulative_logprob, reverse=True)[:return_count]
Expand Down Expand Up @@ -1438,7 +1467,7 @@ async def generate(
candidates.append(
{
"text": r.get("text", "") if isinstance(r, dict) else "",
"finish_reason": fr_type if isinstance(fr_type, str) else "stop",
"finish_reason": fr_type,
"logprobs": cand_logprobs,
}
)
Expand Down Expand Up @@ -1506,13 +1535,8 @@ async def generate(
try:
async with client.stream("POST", f"{self._server_url}/generate", json=body) as response:
if response.status_code != 200:
# Drain a bit of the body for diagnostics, then raise.
body_preview = await response.aread()
logger.error(
"SGLang /generate stream error %d: %s",
response.status_code,
body_preview[:500],
)
await response.aread()
logger.error("SGLang /generate stream error %d", response.status_code)
response.raise_for_status()

last_cumulative_text = ""
Expand Down Expand Up @@ -1550,10 +1574,10 @@ async def generate(
try:
event = json.loads(line)
except json.JSONDecodeError:
logger.warning("SGLang stream: skipping non-JSON line: %s", line[:200])
logger.warning("SGLang stream: skipping non-JSON line")
continue

_raise_for_sglang_event_error(event)
_raise_for_sglang_event_error(event, grammar=grammar)
chunk = _chunk_from_sglang_event(
event,
previous_cumulative_text=last_cumulative_text,
Expand Down Expand Up @@ -1815,6 +1839,7 @@ def _chunk_from_sglang_event(
OpenAI ``ChatCompletionTokenLogprob`` shape for the tokens
introduced by *this* event.
"""
_raise_for_sglang_event_error(event)
if not isinstance(event, dict):
return None
cumulative = event.get("text", "")
Expand Down Expand Up @@ -1897,11 +1922,7 @@ def _chunk_from_sglang_event(
chunk_logprobs = tuple(built)

if is_terminal:
finish_reason: FinishReason
if raw_finish in ("stop", "length", "cancelled", "error"):
finish_reason = raw_finish # type: ignore[assignment]
else:
finish_reason = "stop"
finish_reason = cast("Literal['stop', 'length']", raw_finish)
prompt_tokens = meta.get("prompt_tokens") if isinstance(meta, dict) else None
completion_tokens = meta.get("completion_tokens") if isinstance(meta, dict) else None
return GenerationChunk(
Expand Down Expand Up @@ -1976,6 +1997,7 @@ def _parse_sglang_generate_response(result: Any) -> GenerationResult:
msg = f"SGLang /generate returned unexpected shape: {type(result).__name__}"
raise RuntimeError(msg)

_raise_for_sglang_event_error(result, terminal=True)
text = result.get("text", "")
if not isinstance(text, str):
msg = "SGLang /generate response missing 'text'"
Expand All @@ -1988,11 +2010,11 @@ def _parse_sglang_generate_response(result: Any) -> GenerationResult:
raw_finish = meta.get("finish_reason")
if isinstance(raw_finish, dict):
raw_finish = raw_finish.get("type")
finish_reason = raw_finish if raw_finish in ("stop", "length") else "stop"
finish_reason = cast("Literal['stop', 'length']", raw_finish)

return GenerationResult(
text=text,
finish_reason=finish_reason, # type: ignore[arg-type]
finish_reason=finish_reason,
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
)
Loading
Loading