From 86c579be8411a3503684345aa62688e91a422260 Mon Sep 17 00:00:00 2001 From: nkozobrod Date: Thu, 6 Aug 2026 16:12:56 +0300 Subject: [PATCH] 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) --- docker-compose.yml | 8 ++- logger.yaml | 6 +- surya/inference/backends/openai_client.py | 74 +++++++++++++++++++++++ surya/inference/backends/vllm.py | 1 + surya/inference/schema.py | 3 + surya/settings.py | 5 ++ vllm_tools.py | 1 + 7 files changed, 95 insertions(+), 3 deletions(-) diff --git a/docker-compose.yml b/docker-compose.yml index be1e8d6..ba3f5a0 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -5,12 +5,18 @@ services: runtime: nvidia environment: - VLLM_DTYPE=bfloat16 - - VLLM_GPU_TYPE=a100 + - VLLM_GPU_TYPE=4090 + - VLLM_MAX_NUM_SEQS=32 + - VLLM_MAX_BATCHED_TOKENS=8192 + - SURYA_INFERENCE_MAX_INFLIGHT=32 + # "off" | "on" | "abort" — streaming mode for chat completions + - SURYA_STREAM_MODE=off - NVIDIA_VISIBLE_DEVICES=all ports: - "5002:5002" volumes: - ~/.cache/huggingface:/root/.cache/huggingface + - ./logs:/opt/suya-ocr/logs deploy: resources: reservations: diff --git a/logger.yaml b/logger.yaml index df3929d..10c17d6 100644 --- a/logger.yaml +++ b/logger.yaml @@ -18,12 +18,14 @@ handlers: interval: 1 backupCount: 30 error_file_handler: - class: logging.handlers.RotatingFileHandler + class: logging.handlers.TimedRotatingFileHandler level: ERROR formatter: simple filename: logs/errors.log - backupCount: 20 encoding: utf8 + when: d + interval: 1 + backupCount: 20 root: level: INFO handlers: [console, info_file_handler, error_file_handler] \ No newline at end of file diff --git a/surya/inference/backends/openai_client.py b/surya/inference/backends/openai_client.py index 5ae167d..baf6922 100644 --- a/surya/inference/backends/openai_client.py +++ b/surya/inference/backends/openai_client.py @@ -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( diff --git a/surya/inference/backends/vllm.py b/surya/inference/backends/vllm.py index a51420c..facd579 100644 --- a/surya/inference/backends/vllm.py +++ b/surya/inference/backends/vllm.py @@ -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, ) diff --git a/surya/inference/schema.py b/surya/inference/schema.py index f171da5..458c606 100644 --- a/surya/inference/schema.py +++ b/surya/inference/schema.py @@ -29,6 +29,9 @@ class GenerationResult: raw: str token_count: int error: bool = False + # True when streaming mode aborted generation early due to a detected + # token repetition loop (see SURYA_STREAM_MODE=abort) + repeat: bool = False # Mean of exp(logprob) across response tokens, if logprobs requested mean_token_prob: Optional[float] = None # Per-token logprobs (raw OpenAI-style content list), if requested - phase 2 use diff --git a/surya/settings.py b/surya/settings.py index 81edfda..e0fb45c 100644 --- a/surya/settings.py +++ b/surya/settings.py @@ -73,6 +73,11 @@ class Settings(BaseSettings): SURYA_INFERENCE_STARTUP_TIMEOUT: float = 600.0 SURYA_INFERENCE_LOGPROBS: bool = True SURYA_INFERENCE_MAX_RETRIES: int = 1 + # "off" = non-streaming chat completions (default, unchanged behaviour) + # "on" = stream the response (same retry semantics, measures streaming overhead) + # "abort" = stream + abort early on token repetition loop, then retry with + # higher temperature (see openai_client._generate_one_stream) + SURYA_STREAM_MODE: str = "off" # Force layout/table_rec output through a JSON schema via guided decoding. # Eliminates malformed-JSON failures at small decode-throughput cost. SURYA_GUIDED_LAYOUT: bool = True diff --git a/vllm_tools.py b/vllm_tools.py index dd12c2c..90adda4 100644 --- a/vllm_tools.py +++ b/vllm_tools.py @@ -14,6 +14,7 @@ os.environ.setdefault("SURYA_INFERENCE_AUTOSTART", "false") os.environ.setdefault("SURYA_INFERENCE_KEEP_ALIVE", "false") os.environ.setdefault("SURYA_INFERENCE_LOGPROBS", "false") os.environ.setdefault("SURYA_INFERENCE_MAX_RETRIES", "1") +os.environ.setdefault("SURYA_STREAM_MODE", "off") os.environ.setdefault("SURYA_INFERENCE_PARALLEL", "8") os.environ.setdefault("SURYA_MAX_TOKENS_FULL_PAGE", "6144") os.environ.setdefault("SURYA_MAX_BLOCKS_PER_PAGE", "80")