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:
+7
-1
@@ -5,12 +5,18 @@ services:
|
|||||||
runtime: nvidia
|
runtime: nvidia
|
||||||
environment:
|
environment:
|
||||||
- VLLM_DTYPE=bfloat16
|
- 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
|
- NVIDIA_VISIBLE_DEVICES=all
|
||||||
ports:
|
ports:
|
||||||
- "5002:5002"
|
- "5002:5002"
|
||||||
volumes:
|
volumes:
|
||||||
- ~/.cache/huggingface:/root/.cache/huggingface
|
- ~/.cache/huggingface:/root/.cache/huggingface
|
||||||
|
- ./logs:/opt/suya-ocr/logs
|
||||||
deploy:
|
deploy:
|
||||||
resources:
|
resources:
|
||||||
reservations:
|
reservations:
|
||||||
|
|||||||
+4
-2
@@ -18,12 +18,14 @@ handlers:
|
|||||||
interval: 1
|
interval: 1
|
||||||
backupCount: 30
|
backupCount: 30
|
||||||
error_file_handler:
|
error_file_handler:
|
||||||
class: logging.handlers.RotatingFileHandler
|
class: logging.handlers.TimedRotatingFileHandler
|
||||||
level: ERROR
|
level: ERROR
|
||||||
formatter: simple
|
formatter: simple
|
||||||
filename: logs/errors.log
|
filename: logs/errors.log
|
||||||
backupCount: 20
|
|
||||||
encoding: utf8
|
encoding: utf8
|
||||||
|
when: d
|
||||||
|
interval: 1
|
||||||
|
backupCount: 20
|
||||||
root:
|
root:
|
||||||
level: INFO
|
level: INFO
|
||||||
handlers: [console, info_file_handler, error_file_handler]
|
handlers: [console, info_file_handler, error_file_handler]
|
||||||
@@ -99,6 +99,7 @@ def _generate_one(
|
|||||||
top_p: float,
|
top_p: float,
|
||||||
timeout: float,
|
timeout: float,
|
||||||
request_logprobs_default: bool,
|
request_logprobs_default: bool,
|
||||||
|
stream_mode: str = "off",
|
||||||
) -> GenerationResult:
|
) -> GenerationResult:
|
||||||
prompt = item.prompt or PROMPT_MAPPING[item.prompt_type]
|
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):
|
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:
|
if item.guided_regex is not None:
|
||||||
kwargs.setdefault("extra_body", {})["guided_regex"] = item.guided_regex
|
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:
|
try:
|
||||||
with timing_span(
|
with timing_span(
|
||||||
"openai_chat_completion",
|
"openai_chat_completion",
|
||||||
@@ -173,6 +182,61 @@ def _generate_one(
|
|||||||
return GenerationResult(raw="", token_count=0, error=True)
|
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(
|
def _should_retry(
|
||||||
result: GenerationResult,
|
result: GenerationResult,
|
||||||
retries: int,
|
retries: int,
|
||||||
@@ -182,6 +246,8 @@ def _should_retry(
|
|||||||
return False
|
return False
|
||||||
if result.error:
|
if result.error:
|
||||||
return True
|
return True
|
||||||
|
if result.repeat:
|
||||||
|
return True
|
||||||
has_repeat = detect_repeat_token(result.raw) or (
|
has_repeat = detect_repeat_token(result.raw) or (
|
||||||
len(result.raw) > 50 and detect_repeat_token(result.raw, cut_from_end=50)
|
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_workers: Optional[int] = None,
|
||||||
max_retries: int = 3,
|
max_retries: int = 3,
|
||||||
request_logprobs_default: bool = True,
|
request_logprobs_default: bool = True,
|
||||||
|
stream_mode: str = "off",
|
||||||
) -> List[BatchOutputItem]:
|
) -> List[BatchOutputItem]:
|
||||||
"""Run a batch of items through the chat completions endpoint with concurrent workers."""
|
"""Run a batch of items through the chat completions endpoint with concurrent workers."""
|
||||||
if not batch:
|
if not batch:
|
||||||
@@ -219,6 +286,7 @@ def chat_completions_batch(
|
|||||||
top_p=top_p,
|
top_p=top_p,
|
||||||
timeout=timeout,
|
timeout=timeout,
|
||||||
request_logprobs_default=request_logprobs_default,
|
request_logprobs_default=request_logprobs_default,
|
||||||
|
stream_mode=stream_mode,
|
||||||
)
|
)
|
||||||
retries = 0
|
retries = 0
|
||||||
while _should_retry(result, retries, max_retries):
|
while _should_retry(result, retries, max_retries):
|
||||||
@@ -227,6 +295,11 @@ def chat_completions_batch(
|
|||||||
time.sleep(backoff)
|
time.sleep(backoff)
|
||||||
retry_temp = min(temperature + 0.2 * (retries + 1), 0.8)
|
retry_temp = min(temperature + 0.2 * (retries + 1), 0.8)
|
||||||
retry_top_p = 0.95 if not result.error else top_p
|
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(
|
result = _generate_one(
|
||||||
item,
|
item,
|
||||||
client=client,
|
client=client,
|
||||||
@@ -236,6 +309,7 @@ def chat_completions_batch(
|
|||||||
top_p=retry_top_p,
|
top_p=retry_top_p,
|
||||||
timeout=timeout,
|
timeout=timeout,
|
||||||
request_logprobs_default=request_logprobs_default,
|
request_logprobs_default=request_logprobs_default,
|
||||||
|
stream_mode=stream_mode,
|
||||||
)
|
)
|
||||||
retries += 1
|
retries += 1
|
||||||
return BatchOutputItem(
|
return BatchOutputItem(
|
||||||
|
|||||||
@@ -221,4 +221,5 @@ class VllmBackend(Backend):
|
|||||||
),
|
),
|
||||||
max_retries=settings.SURYA_INFERENCE_MAX_RETRIES,
|
max_retries=settings.SURYA_INFERENCE_MAX_RETRIES,
|
||||||
request_logprobs_default=settings.SURYA_INFERENCE_LOGPROBS,
|
request_logprobs_default=settings.SURYA_INFERENCE_LOGPROBS,
|
||||||
|
stream_mode=settings.SURYA_STREAM_MODE,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -29,6 +29,9 @@ class GenerationResult:
|
|||||||
raw: str
|
raw: str
|
||||||
token_count: int
|
token_count: int
|
||||||
error: bool = False
|
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 of exp(logprob) across response tokens, if logprobs requested
|
||||||
mean_token_prob: Optional[float] = None
|
mean_token_prob: Optional[float] = None
|
||||||
# Per-token logprobs (raw OpenAI-style content list), if requested - phase 2 use
|
# Per-token logprobs (raw OpenAI-style content list), if requested - phase 2 use
|
||||||
|
|||||||
@@ -73,6 +73,11 @@ class Settings(BaseSettings):
|
|||||||
SURYA_INFERENCE_STARTUP_TIMEOUT: float = 600.0
|
SURYA_INFERENCE_STARTUP_TIMEOUT: float = 600.0
|
||||||
SURYA_INFERENCE_LOGPROBS: bool = True
|
SURYA_INFERENCE_LOGPROBS: bool = True
|
||||||
SURYA_INFERENCE_MAX_RETRIES: int = 1
|
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.
|
# Force layout/table_rec output through a JSON schema via guided decoding.
|
||||||
# Eliminates malformed-JSON failures at small decode-throughput cost.
|
# Eliminates malformed-JSON failures at small decode-throughput cost.
|
||||||
SURYA_GUIDED_LAYOUT: bool = True
|
SURYA_GUIDED_LAYOUT: bool = True
|
||||||
|
|||||||
@@ -14,6 +14,7 @@ os.environ.setdefault("SURYA_INFERENCE_AUTOSTART", "false")
|
|||||||
os.environ.setdefault("SURYA_INFERENCE_KEEP_ALIVE", "false")
|
os.environ.setdefault("SURYA_INFERENCE_KEEP_ALIVE", "false")
|
||||||
os.environ.setdefault("SURYA_INFERENCE_LOGPROBS", "false")
|
os.environ.setdefault("SURYA_INFERENCE_LOGPROBS", "false")
|
||||||
os.environ.setdefault("SURYA_INFERENCE_MAX_RETRIES", "1")
|
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_INFERENCE_PARALLEL", "8")
|
||||||
os.environ.setdefault("SURYA_MAX_TOKENS_FULL_PAGE", "6144")
|
os.environ.setdefault("SURYA_MAX_TOKENS_FULL_PAGE", "6144")
|
||||||
os.environ.setdefault("SURYA_MAX_BLOCKS_PER_PAGE", "80")
|
os.environ.setdefault("SURYA_MAX_BLOCKS_PER_PAGE", "80")
|
||||||
|
|||||||
Reference in New Issue
Block a user