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:
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user