Add SURYA_STREAM_MODE streaming (off/on/abort) for chat completions

- _generate_one_stream: streams responses, optional live repeat-loop
  detection (abort) that closes the stream early and flags the partial
  result as repeat=True so the existing retry loop re-runs with higher
  temperature
- _should_retry honors result.repeat; stream_mode threaded through both
  first attempt and retry calls in chat_completions_batch
- GenerationResult.repeat field; SURYA_STREAM_MODE setting (default off)
- diagnostic retry logs (reason=repeat|error|detected, temp)
- docker-compose: SURYA_STREAM_MODE=off, 4090 GPU profile, logs volume
- logger.yaml: TimedRotatingFileHandler for error handler (fixes startup
  crash without maxBytes)
This commit is contained in:
2026-08-06 16:12:56 +03:00
parent 41fc6db9eb
commit 86c579be84
7 changed files with 95 additions and 3 deletions
+74
View File
@@ -99,6 +99,7 @@ def _generate_one(
top_p: float,
timeout: float,
request_logprobs_default: bool,
stream_mode: str = "off",
) -> GenerationResult:
prompt = item.prompt or PROMPT_MAPPING[item.prompt_type]
with timing_span("openai_scale_image", prompt_type=item.prompt_type, image_size=item.image.size):
@@ -134,6 +135,14 @@ def _generate_one(
if item.guided_regex is not None:
kwargs.setdefault("extra_body", {})["guided_regex"] = item.guided_regex
if stream_mode == "off":
return _generate_one_nonstream(
item, client, kwargs, max_tokens, request_logprobs, stream_mode
)
return _generate_one_stream(item, client, kwargs, max_tokens, stream_mode)
def _generate_one_nonstream(item, client, kwargs, max_tokens, request_logprobs, stream_mode):
try:
with timing_span(
"openai_chat_completion",
@@ -173,6 +182,61 @@ def _generate_one(
return GenerationResult(raw="", token_count=0, error=True)
def _generate_one_stream(item, client, kwargs, max_tokens, stream_mode):
"""Streaming variant of _generate_one.
stream_mode="on" — stream only (same retry semantics as non-streaming).
stream_mode="abort" — additionally detect token repetition loops in the
live stream; on detection close the connection (vLLM aborts the
request and frees the slot) and return a partial result flagged
repeat=True so the existing retry loop restarts with higher temperature.
"""
kwargs["stream"] = True
kwargs["stream_options"] = {"include_usage": True}
try:
with timing_span(
"openai_chat_completion_stream",
prompt_type=item.prompt_type,
max_tokens=max_tokens,
):
stream = client.chat.completions.create(**kwargs)
raw_parts = []
usage = None
repeat = False
checked_len = 0
for chunk in stream:
if not chunk.choices:
# Final chunk carrying usage (stream_options include_usage)
usage = getattr(chunk, "usage", None)
continue
delta = chunk.choices[0].delta.content or ""
if not delta:
continue
raw_parts.append(delta)
if stream_mode == "abort":
text_len = sum(len(p) for p in raw_parts)
# Check at most every 30 new chars, only once text is long enough
if text_len - checked_len >= 30 and text_len > 150:
text = "".join(raw_parts)
if detect_repeat_token(text[-1024:]):
logger.info(
f"stream abort: repeat detected at {text_len} chars "
f"(max_tokens={max_tokens}, prompt_type={item.prompt_type})"
)
repeat = True
break
checked_len = text_len
close = getattr(stream, "close", None)
if close:
close()
raw = "".join(raw_parts)
token_count = usage.completion_tokens if usage else 0
return GenerationResult(raw=raw, token_count=token_count, error=False, repeat=repeat)
except Exception as e:
logger.warning(f"Inference error (stream): {e}")
return GenerationResult(raw="", token_count=0, error=True)
def _should_retry(
result: GenerationResult,
retries: int,
@@ -182,6 +246,8 @@ def _should_retry(
return False
if result.error:
return True
if result.repeat:
return True
has_repeat = detect_repeat_token(result.raw) or (
len(result.raw) > 50 and detect_repeat_token(result.raw, cut_from_end=50)
)
@@ -199,6 +265,7 @@ def chat_completions_batch(
max_workers: Optional[int] = None,
max_retries: int = 3,
request_logprobs_default: bool = True,
stream_mode: str = "off",
) -> List[BatchOutputItem]:
"""Run a batch of items through the chat completions endpoint with concurrent workers."""
if not batch:
@@ -219,6 +286,7 @@ def chat_completions_batch(
top_p=top_p,
timeout=timeout,
request_logprobs_default=request_logprobs_default,
stream_mode=stream_mode,
)
retries = 0
while _should_retry(result, retries, max_retries):
@@ -227,6 +295,11 @@ def chat_completions_batch(
time.sleep(backoff)
retry_temp = min(temperature + 0.2 * (retries + 1), 0.8)
retry_top_p = 0.95 if not result.error else top_p
logger.info(
f"retry {retries + 1}/{max_retries} item={item.metadata} "
f"reason={'repeat' if result.repeat else ('error' if result.error else 'detected')} "
f"temp={retry_temp:.1f}"
)
result = _generate_one(
item,
client=client,
@@ -236,6 +309,7 @@ def chat_completions_batch(
top_p=retry_top_p,
timeout=timeout,
request_logprobs_default=request_logprobs_default,
stream_mode=stream_mode,
)
retries += 1
return BatchOutputItem(
+1
View File
@@ -221,4 +221,5 @@ class VllmBackend(Backend):
),
max_retries=settings.SURYA_INFERENCE_MAX_RETRIES,
request_logprobs_default=settings.SURYA_INFERENCE_LOGPROBS,
stream_mode=settings.SURYA_STREAM_MODE,
)