diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 75da00451e..80b9739cc1 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -122,15 +122,11 @@ _INTENT_SIGNAL = re.compile( ) _MAX_REPROMPTS = 1 -# Without max_tokens, llama-server defaults n_predict = n_ctx (up to 262144 for -# Qwen3.5), causing many-minute zombie decodes when cancel fails. -# t_max_predict_ms is a wall-clock backstop but per the llama.cpp README only -# fires after a newline, so we keep a token cap as the front-line limiter. -# The cap is the effective context length when known, else this floor. 4096 was -# too low: Qwen3 / gpt-oss reasoning traces and max_tokens-omitting OpenAI-API -# callers (langchain, llama-index, curl) got truncated mid-sentence. +# Default max_tokens to the effective context when known. The floor is high +# enough for reasoning-heavy GGUFs and max_tokens-omitting API clients. _DEFAULT_MAX_TOKENS_FLOOR = 32768 -_DEFAULT_T_MAX_PREDICT_MS = 600_000 # 10 min +_DEFAULT_FIRST_TOKEN_TIMEOUT_S = 1200.0 # 20 min +_DEFAULT_STREAM_STALL_TIMEOUT_S = 120.0 # 2 min _REPROMPT_MAX_CHARS = 2000 _FORCED_REPEAT_PLAN_SIGNAL = re.compile( r"\b(?:i\s+will|i'll|let\s+me|going\s+to|need\s+to|call|use|run|search|fetch|render)\b", @@ -5308,28 +5304,84 @@ class LlamaCppBackend: @staticmethod def _iter_text_cancellable( - response: "httpx.Response", cancel_event: Optional[threading.Event] = None + response: "httpx.Response", + cancel_event: Optional[threading.Event] = None, + stall_timeout_s: float = _DEFAULT_STREAM_STALL_TIMEOUT_S, + first_token_deadline: Optional[float] = None, + post_first_chunk_read_timeout_s: Optional[float] = _DEFAULT_STREAM_STALL_TIMEOUT_S, ) -> Generator[str, None, None]: - """Iterate an httpx streaming response with cancel support. - - Checks cancel_event between chunks and on ReadTimeout; the - _stream_with_retry watcher also closes the response on cancel. - """ + """Iterate a stream while polling cancel and stall timeouts.""" text_iter = response.iter_text() + if first_token_deadline is None: + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S + last_chunk_at: Optional[float] = None while True: if cancel_event is not None and cancel_event.is_set(): response.close() return try: + if last_chunk_at is None: + remaining_s = first_token_deadline - time.monotonic() + if remaining_s <= 0: + raise httpx.ReadTimeout("The model did not produce a first token in time.") + LlamaCppBackend._set_stream_read_timeout(response, remaining_s) chunk = next(text_iter) + if chunk: + if last_chunk_at is None and post_first_chunk_read_timeout_s is not None: + LlamaCppBackend._set_stream_read_timeout( + response, + post_first_chunk_read_timeout_s, + ) + last_chunk_at = time.monotonic() yield chunk except StopIteration: return except httpx.ReadTimeout: - # No data within the timeout window -- loop back and re-check - # cancel_event. + now = time.monotonic() + if last_chunk_at is None: + if now >= first_token_deadline: + raise + elif now - last_chunk_at >= stall_timeout_s: + raise httpx.ReadTimeout("The model stopped producing tokens mid-response.") continue + @staticmethod + def _set_stream_read_timeout(response: "httpx.Response", read_timeout_s: float) -> None: + """Lower only post-header stream reads; keep prefill timeout long.""" + try: + timeout_ext = response.request.extensions.get("timeout") + if isinstance(timeout_ext, dict): + timeout_ext["read"] = read_timeout_s + except Exception: + logger.debug("Could not lower response read timeout", exc_info = True) + + @staticmethod + def _shutdown_active_httpx_sockets(client: "httpx.Client") -> None: + """Best-effort interrupt for a sync httpx request blocked before headers.""" + try: + pool = getattr(getattr(client, "_transport", None), "_pool", None) + connections = list(getattr(pool, "_connections", []) or []) + for connection in connections: + inner = getattr(connection, "_connection", None) + stream = getattr(inner, "_network_stream", None) + sock = getattr(stream, "_sock", None) + if sock is None: + continue + try: + sock.shutdown(socket.SHUT_RDWR) + except OSError: + pass + try: + sock.close() + except OSError: + pass + except Exception: + logger.debug("Could not shutdown active httpx socket", exc_info = True) + try: + client.close() + except Exception: + logger.debug("Could not close httpx client", exc_info = True) + @staticmethod @contextlib.contextmanager def _stream_with_retry( @@ -5338,38 +5390,28 @@ class LlamaCppBackend: payload: dict, cancel_event: Optional[threading.Event] = None, headers: Optional[dict] = None, + first_token_deadline: Optional[float] = None, ): - """Open an httpx streaming POST with cancel support. - - Sends once with a long read timeout (120 s) so prefill finishes without - a retry storm (the old 0.5 s timeout caused duplicate POSTs every half - second). A watcher thread cancels by closing the response. httpx can't - interrupt a blocked read before the response exists, so cancel during - the header wait (1-5 s prefill) is deferred until headers arrive. - """ + """Open one streaming POST and let cancel interrupt prefill or reads.""" if cancel_event is not None and cancel_event.is_set(): raise GeneratorExit - # Background watcher: close the response if cancel is requested. - # Only effective after response headers arrive (httpx limitation). _cancel_closed = threading.Event() _response_ref: list = [None] def _cancel_watcher(): while not _cancel_closed.is_set(): if cancel_event.wait(timeout = 0.3): - # Cancel requested. Poll until the response object exists - # so we can close it, or until the main thread finishes - # (_cancel_closed set in finally). while not _cancel_closed.is_set(): r = _response_ref[0] - if r is not None: - try: + try: + if r is not None: r.close() - return - except Exception as e: - logger.debug(f"Error closing response in cancel watcher: {e}") - # Response not created yet -- wait briefly and retry + else: + LlamaCppBackend._shutdown_active_httpx_sockets(client) + return + except Exception as e: + logger.debug(f"Error closing request in cancel watcher: {e}") _cancel_closed.wait(timeout = 0.1) return @@ -5379,12 +5421,12 @@ class LlamaCppBackend: watcher.start() try: - # Long read timeout so prefill can finish without a retry storm. - # Cancel during prefill and streaming is handled by the watcher - # thread closing the response, unblocking any httpx read. + if first_token_deadline is None: + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S + prefill_read_timeout = max(0.1, first_token_deadline - time.monotonic()) prefill_timeout = httpx.Timeout( connect = 30, - read = 120.0, + read = prefill_read_timeout, write = 10, pool = 10, ) @@ -5400,7 +5442,7 @@ class LlamaCppBackend: raise GeneratorExit yield response return - except (httpx.ReadError, httpx.RemoteProtocolError, httpx.CloseError): + except (httpx.RequestError, RuntimeError): # Response was closed by the cancel watcher if cancel_event is not None and cancel_event.is_set(): raise GeneratorExit @@ -5455,14 +5497,12 @@ class LlamaCppBackend: ) if _reasoning_kw is not None: payload["chat_template_kwargs"] = _reasoning_kw - # Cap to the effective context length when known, else the floor. - # The wall-clock backstop below stops a stuck model regardless. + # Default cap to the model context when known. payload["max_tokens"] = ( max_tokens if max_tokens is not None else (self._effective_context_length or _DEFAULT_MAX_TOKENS_FLOOR) ) - payload["t_max_predict_ms"] = _DEFAULT_T_MAX_PREDICT_MS if stop: payload["stop"] = stop if seed is not None: @@ -5478,20 +5518,20 @@ class LlamaCppBackend: _metadata_finish_reason = None try: - # _stream_with_retry uses a 120 s read timeout so prefill can - # finish. Cancel during streaming is handled by the watcher - # thread (closes the response on cancel_event). + # Prefill can use the long first-token timeout; body reads are lowered after headers. stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10) _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None with httpx.Client( timeout = stream_timeout, limits = httpx.Limits(max_keepalive_connections = 0) ) as client: + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S with self._stream_with_retry( client, url, payload, cancel_event, headers = _auth_headers, + first_token_deadline = first_token_deadline, ) as response: if response.status_code != 200: error_body = response.read().decode() @@ -5502,7 +5542,11 @@ class LlamaCppBackend: buffer = "" has_content_tokens = False reasoning_text = "" - for raw_chunk in self._iter_text_cancellable(response, cancel_event): + for raw_chunk in self._iter_text_cancellable( + response, + cancel_event, + first_token_deadline = first_token_deadline, + ): buffer += raw_chunk while "\n" in buffer: line, buffer = buffer.split("\n", 1) @@ -5732,7 +5776,6 @@ class LlamaCppBackend: if max_tokens is not None else (self._effective_context_length or _DEFAULT_MAX_TOKENS_FLOOR) ) - payload["t_max_predict_ms"] = _DEFAULT_T_MAX_PREDICT_MS if stop: payload["stop"] = stop if seed is not None: @@ -5778,12 +5821,14 @@ class LlamaCppBackend: timeout = stream_timeout, limits = httpx.Limits(max_keepalive_connections = 0), ) as client: + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S with self._stream_with_retry( client, url, payload, cancel_event, headers = _auth_headers, + first_token_deadline = first_token_deadline, ) as response: if response.status_code != 200: error_body = response.read().decode() @@ -5795,6 +5840,7 @@ class LlamaCppBackend: for raw_chunk in self._iter_text_cancellable( response, cancel_event, + first_token_deadline = first_token_deadline, ): raw_buf += raw_chunk while "\n" in raw_buf: @@ -6444,7 +6490,6 @@ class LlamaCppBackend: if max_tokens is not None else (self._effective_context_length or _DEFAULT_MAX_TOKENS_FLOOR) ) - stream_payload["t_max_predict_ms"] = _DEFAULT_T_MAX_PREDICT_MS if stop: stream_payload["stop"] = stop if seed is not None: @@ -6467,12 +6512,14 @@ class LlamaCppBackend: with httpx.Client( timeout = stream_timeout, limits = httpx.Limits(max_keepalive_connections = 0) ) as client: + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S with self._stream_with_retry( client, url, stream_payload, cancel_event, headers = _auth_headers, + first_token_deadline = first_token_deadline, ) as response: if response.status_code != 200: error_body = response.read().decode() @@ -6481,7 +6528,11 @@ class LlamaCppBackend: ) buffer = "" - for raw_chunk in self._iter_text_cancellable(response, cancel_event): + for raw_chunk in self._iter_text_cancellable( + response, + cancel_event, + first_token_deadline = first_token_deadline, + ): buffer += raw_chunk while "\n" in buffer: line, buffer = buffer.split("\n", 1) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index fa708d0e39..e9353cd803 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -121,6 +121,19 @@ def _template_raise_message(error_text: str, chat_template: Optional[str]) -> Op def _friendly_error(exc: Exception) -> str: """Extract a user-friendly message from known llama-server errors.""" + if isinstance(exc, httpx.ReadTimeout): + if "stopped producing tokens" in str(exc).lower(): + return ( + "The model stopped producing tokens before the response " + "completed. Try stopping and retrying, or reduce max tokens." + ) + return ( + "The model is still processing the prompt but did not produce a " + "first token within 20 minutes. Try reducing context length, " + "using more GPU offload, or loading a smaller model." + ) + if isinstance(exc, httpx.TimeoutException): + return "Timed out communicating with the model server. Try again shortly." # httpx transport failures from the async pass-through helpers. Any # RequestError subclass (ConnectError, ReadError, RemoteProtocolError, # WriteError, PoolTimeout, ...) means the llama-server subprocess is @@ -224,7 +237,11 @@ def _openai_stream_error_chunk(exc) -> dict: (a code-less error hides it).""" _cls = _classify_llama_generation_error(exc) if _cls: - return openai_error_body(_friendly_error(exc), status = 400, code = "context_length_exceeded") + return openai_error_body( + _friendly_error(exc), + status = 400, + code = "context_length_exceeded", + ) if _cls is False: return openai_error_body(_friendly_error(exc), status = 400) return openai_error_body(_friendly_error(exc), status = 500) @@ -415,17 +432,15 @@ def _apply_overflow_truncation(body: dict, err_text: str) -> bool: return True -def _anthropic_stream_error_event(exc): - """Anthropic in-band SSE ``error`` event for a mid-stream failure, or ``None`` - to fall through to a normal message_delta finish. Returns an event only for a - classifiable upstream client error (context overflow / 4xx) so a streaming - over-context request surfaces a real error instead of a silent empty - end_turn message.""" - if _classify_llama_generation_error(exc) is None: +def _anthropic_stream_error_event(exc, *, force: bool = False): + """Return an Anthropic in-band stream error event when one is useful.""" + _cls = _classify_llama_generation_error(exc) + if _cls is None and not force: return None + status = 400 if _cls is not None else 500 return build_anthropic_sse_event( "error", - anthropic_error_body(_friendly_error(exc), status = 400), + anthropic_error_body(_friendly_error(exc), status = status), ) @@ -588,8 +603,9 @@ try: from core.inference import get_inference_backend from core.inference.llama_cpp import ( LlamaCppBackend, + _DEFAULT_FIRST_TOKEN_TIMEOUT_S, _DEFAULT_MAX_TOKENS_FLOOR, - _DEFAULT_T_MAX_PREDICT_MS, + _DEFAULT_STREAM_STALL_TIMEOUT_S, _canonicalize_spec_mode, _extra_args_set_spec_type, _hf_offline_if_dns_dead, @@ -621,8 +637,9 @@ except ImportError: from core.inference import get_inference_backend from core.inference.llama_cpp import ( LlamaCppBackend, + _DEFAULT_FIRST_TOKEN_TIMEOUT_S, _DEFAULT_MAX_TOKENS_FLOOR, - _DEFAULT_T_MAX_PREDICT_MS, + _DEFAULT_STREAM_STALL_TIMEOUT_S, _canonicalize_spec_mode, _extra_args_set_spec_type, _hf_offline_if_dns_dead, @@ -648,6 +665,152 @@ except ImportError: verify_native_path_lease, ) + +def _llama_non_streaming_generation_timeout() -> httpx.Timeout: + return httpx.Timeout( + connect = _DEFAULT_FIRST_TOKEN_TIMEOUT_S, + read = _DEFAULT_FIRST_TOKEN_TIMEOUT_S, + write = _DEFAULT_FIRST_TOKEN_TIMEOUT_S, + pool = _DEFAULT_FIRST_TOKEN_TIMEOUT_S, + ) + + +def _llama_streaming_generation_timeout() -> httpx.Timeout: + return httpx.Timeout( + connect = _DEFAULT_FIRST_TOKEN_TIMEOUT_S, + read = _DEFAULT_FIRST_TOKEN_TIMEOUT_S, + write = _DEFAULT_FIRST_TOKEN_TIMEOUT_S, + pool = _DEFAULT_FIRST_TOKEN_TIMEOUT_S, + ) + + +def _set_stream_response_read_timeout( + response: httpx.Response, read_timeout_s: float = _DEFAULT_STREAM_STALL_TIMEOUT_S +) -> None: + try: + timeout_ext = response.request.extensions.get("timeout") + if isinstance(timeout_ext, dict): + timeout_ext["read"] = read_timeout_s + except Exception: + pass + + +async def _preheader_cancelled(cancel_event = None, request: Optional[Request] = None) -> bool: + if cancel_event is not None and cancel_event.is_set(): + return True + if request is not None and await request.is_disconnected(): + if cancel_event is not None: + cancel_event.set() + return True + return False + + +async def _wait_preheader_cancel(cancel_event = None, request: Optional[Request] = None) -> None: + while not await _preheader_cancelled(cancel_event, request): + await asyncio.sleep(0.05) + + +async def _send_stream_with_preheader_cancel( + client: httpx.AsyncClient, + req: httpx.Request, + cancel_event = None, + request: Optional[Request] = None, +) -> Optional[httpx.Response]: + if cancel_event is None and request is None: + return await client.send(req, stream = True) + if await _preheader_cancelled(cancel_event, request): + return None + + send_task = asyncio.create_task(client.send(req, stream = True)) + cancel_task = asyncio.create_task(_wait_preheader_cancel(cancel_event, request)) + + async def _stop_send_task() -> None: + try: + await client.aclose() + except Exception: + pass + send_task.cancel() + try: + await send_task + except (asyncio.CancelledError, Exception): + pass + + try: + done, _pending = await asyncio.wait( + {send_task, cancel_task}, + return_when = asyncio.FIRST_COMPLETED, + ) + if send_task in done: + return await send_task + + await _stop_send_task() + return None + except asyncio.CancelledError: + if cancel_event is not None: + cancel_event.set() + await _stop_send_task() + raise + finally: + cancel_task.cancel() + try: + await cancel_task + except (asyncio.CancelledError, Exception): + pass + + +async def _aiter_llama_stream_items( + async_iter, + *, + cancel_event = None, + request: Optional[Request] = None, + first_token_deadline: Optional[float] = None, + response: Optional[httpx.Response] = None, + post_first_item_read_timeout_s: Optional[float] = _DEFAULT_STREAM_STALL_TIMEOUT_S, +): + if first_token_deadline is None: + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S + last_item_at: Optional[float] = None + while True: + if cancel_event is not None and cancel_event.is_set(): + return + if request is not None and await request.is_disconnected(): + if cancel_event is not None: + cancel_event.set() + return + waiting_first_item = last_item_at is None + try: + if waiting_first_item: + remaining_s = first_token_deadline - time.monotonic() + if remaining_s <= 0: + raise httpx.ReadTimeout("The model did not produce a first token in time.") + if response is not None: + _set_stream_response_read_timeout(response, remaining_s) + item = await asyncio.wait_for(async_iter.__anext__(), timeout = remaining_s) + else: + item = await async_iter.__anext__() + except asyncio.TimeoutError as exc: + if waiting_first_item: + raise httpx.ReadTimeout("The model did not produce a first token in time.") from exc + raise + except StopAsyncIteration: + return + except httpx.ReadTimeout: + now = time.monotonic() + if last_item_at is None: + if now >= first_token_deadline: + raise + continue + raise httpx.ReadTimeout("The model stopped producing tokens mid-response.") + if ( + last_item_at is None + and response is not None + and post_first_item_read_timeout_s is not None + ): + _set_stream_response_read_timeout(response, post_first_item_read_timeout_s) + last_item_at = time.monotonic() + yield item + + from models.inference import ( LoadRequest, UnloadRequest, @@ -4935,6 +5098,8 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge ) body = await request.json() + if body.get("max_tokens") is None: + body["max_tokens"] = llama_backend.context_length or _DEFAULT_MAX_TOKENS_FLOOR target_url = f"{llama_backend.base_url}/v1/completions" is_stream = body.get("stream", False) @@ -4952,15 +5117,27 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge # honor stream_options.include_usage per event, while keeping SSE # framing and token bytes intact. _include_usage = bool((body.get("stream_options") or {}).get("include_usage")) - client = httpx.AsyncClient(timeout = 600) + client = httpx.AsyncClient(timeout = _llama_streaming_generation_timeout()) resp = None bytes_iter = None try: req = client.build_request("POST", target_url, json = body) - resp = await client.send(req, stream = True) + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S + resp = await _send_stream_with_preheader_cancel(client, req, request = request) + if resp is None: + return + if resp.status_code != 200: + err_bytes = await resp.aread() + err_text = err_bytes.decode("utf-8", errors = "replace") + raise RuntimeError(f"llama-server returned {resp.status_code}: {err_text}") bytes_iter = resp.aiter_bytes() buffer = b"" - async for chunk in bytes_iter: + async for chunk in _aiter_llama_stream_items( + bytes_iter, + request = request, + first_token_deadline = first_token_deadline, + response = resp, + ): buffer += chunk while b"\n\n" in buffer: event, buffer = buffer.split(b"\n\n", 1) @@ -4976,6 +5153,9 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge yield out + b"\n\n" except Exception as e: logger.error("openai_completions stream error: %s", e) + error_chunk = _openai_stream_error_chunk(e) + yield f"data: {json.dumps(error_chunk)}\n\n".encode("utf-8") + return finally: if bytes_iter is not None: try: @@ -4995,7 +5175,11 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge return StreamingResponse(_stream(), media_type = "text/event-stream") else: async with httpx.AsyncClient() as client: - resp = await client.post(target_url, json = body, timeout = 600) + resp = await client.post( + target_url, + json = body, + timeout = _llama_non_streaming_generation_timeout(), + ) if resp.status_code != 200: raise _openai_passthrough_error(resp.status_code, resp.text) @@ -5033,7 +5217,7 @@ async def openai_embeddings(request: Request, current_subject: str = Depends(get target_url = f"{llama_backend.base_url}/v1/embeddings" async with httpx.AsyncClient() as client: - resp = await client.post(target_url, json = body, timeout = 600) + resp = await client.post(target_url, json = body, timeout = _DEFAULT_FIRST_TOKEN_TIMEOUT_S) return Response( content = resp.content, status_code = resp.status_code, @@ -5775,6 +5959,28 @@ async def _responses_stream( ) return [item for _, item in sorted(indexed_items, key = lambda pair: pair[0])] + def _failed_response_payload(exc: Exception, status_code: int) -> dict: + return { + "type": "response.failed", + "response": { + "id": resp_id, + "object": "response", + "created_at": created_at, + "status": "failed", + "model": payload.model, + "output": _snapshot_output(), + "usage": { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "total_tokens": input_tokens + output_tokens, + }, + "error": { + "code": status_code, + "message": _friendly_error(exc), + }, + }, + } + # ── Preamble events ── yield _sse( "response.created", @@ -5798,13 +6004,16 @@ async def _responses_stream( # `async with`, explicit aclose of lines_iter BEFORE resp / client so # the innermost httpcore byte stream is finalised in this task (not via # the asyncgen GC in a sibling task). - client = httpx.AsyncClient(timeout = 600) + client = httpx.AsyncClient(timeout = _llama_streaming_generation_timeout()) resp = None lines_iter = None try: req = client.build_request("POST", target_url, json = body) + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S try: - resp = await client.send(req, stream = True) + resp = await _send_stream_with_preheader_cancel(client, req, request = request) + if resp is None: + return except httpx.RequestError as e: logger.error("responses stream: upstream unreachable: %s", e) yield _sse( @@ -5853,9 +6062,12 @@ async def _responses_stream( return lines_iter = resp.aiter_lines() - async for raw_line in lines_iter: - if await request.is_disconnected(): - break + async for raw_line in _aiter_llama_stream_items( + lines_iter, + request = request, + first_token_deadline = first_token_deadline, + response = resp, + ): if not raw_line: continue if not raw_line.startswith("data: "): @@ -5974,6 +6186,12 @@ async def _responses_stream( output_tokens = usage.get("completion_tokens", output_tokens) except Exception as e: logger.error("responses stream error: %s", e) + status_code = 400 if _classify_llama_generation_error(e) is not None else 500 + yield _sse( + "response.failed", + _failed_response_payload(e, status_code), + ) + return finally: if lines_iter is not None: try: @@ -7083,7 +7301,6 @@ def _build_passthrough_payload( body["max_tokens"] = ( max_tokens if max_tokens is not None else (backend_ctx or _DEFAULT_MAX_TOKENS_FLOOR) ) - body["t_max_predict_ms"] = _DEFAULT_T_MAX_PREDICT_MS # Normalize stop the same way the non-passthrough path does (the passthrough # was previously the one path that forwarded an empty stop string verbatim). _stop = _normalize_stop_sequences(stop) @@ -7193,7 +7410,7 @@ async def _anthropic_passthrough_stream( # `try: ... except Exception: pass` so nested anyio cleanup noise can't # bubble out. client = httpx.AsyncClient( - timeout = 600, + timeout = _llama_streaming_generation_timeout(), limits = httpx.Limits(max_keepalive_connections = 0), ) resp = None @@ -7201,7 +7418,12 @@ async def _anthropic_passthrough_stream( cancel_watcher = None try: req = client.build_request("POST", target_url, json = body) - resp = await client.send(req, stream = True) + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S + resp = await _send_stream_with_preheader_cancel( + client, req, cancel_event, request = request + ) + if resp is None: + return # Upstream client error (e.g. over-context 400) arrives before any # SSE. The 200 stream headers are already flushed, so surface it as @@ -7230,12 +7452,13 @@ async def _anthropic_passthrough_stream( # The watcher closes `resp` on cancel, raising in aiter_lines. cancel_watcher = asyncio.create_task(_await_cancel_then_close(cancel_event, resp)) lines_iter = resp.aiter_lines() - async for raw_line in lines_iter: - if cancel_event.is_set(): - break - if await request.is_disconnected(): - cancel_event.set() - break + async for raw_line in _aiter_llama_stream_items( + lines_iter, + cancel_event = cancel_event, + request = request, + first_token_deadline = first_token_deadline, + response = resp, + ): if not raw_line or not raw_line.startswith("data: "): continue data_str = raw_line[6:] @@ -7249,11 +7472,26 @@ async def _anthropic_passthrough_stream( _drop_parallel_tool_call_deltas(chunk) for line in emitter.feed_chunk(chunk): yield line - except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError): + except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError) as e: if not cancel_event.is_set(): - raise + logger.error("anthropic_messages passthrough stream error: %s", e) + event = _anthropic_stream_error_event( + e, + force = True, + ) + if event is not None: + yield event + return except Exception as e: - logger.error("anthropic_messages passthrough stream error: %s", e) + if not cancel_event.is_set(): + logger.error("anthropic_messages passthrough stream error: %s", e) + event = _anthropic_stream_error_event( + e, + force = True, + ) + if event is not None: + yield event + return finally: if cancel_watcher is not None: cancel_watcher.cancel() @@ -7327,7 +7565,11 @@ async def _anthropic_passthrough_non_streaming( ) async with httpx.AsyncClient() as client: - resp = await client.post(target_url, json = body, timeout = 600) + resp = await client.post( + target_url, + json = body, + timeout = _llama_non_streaming_generation_timeout(), + ) if resp.status_code != 200: raise HTTPException( @@ -7681,15 +7923,13 @@ async def _openai_passthrough_stream( _tracker = _TrackedCancel(cancel_event, *_cancel_keys) _tracker.__enter__() - # Outer guard: asyncio.CancelledError at `await client.send(...)` is a - # BaseException that bypasses `except httpx.RequestError`; without this the - # tracker leaks. The generator's finally only runs once iteration starts. + # Keep tracker cleanup paired if pre-header dispatch is cancelled. try: # Dispatch BEFORE returning StreamingResponse so transport errors and # non-200 upstream statuses surface as real HTTP errors -- OpenAI SDKs # rely on status codes to raise APIError/BadRequestError. client = httpx.AsyncClient( - timeout = 600, + timeout = _llama_streaming_generation_timeout(), limits = httpx.Limits(max_keepalive_connections = 0), ) resp = None @@ -7699,7 +7939,10 @@ async def _openai_passthrough_stream( while True: try: req = client.build_request("POST", target_url, json = body) - resp = await client.send(req, stream = True) + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S + resp = await _send_stream_with_preheader_cancel( + client, req, cancel_event, request = request + ) except httpx.RequestError as e: # llama-server subprocess crashed / starting / unreachable. logger.error("openai passthrough stream: upstream unreachable: %s", e) @@ -7716,6 +7959,21 @@ async def _openai_passthrough_stream( status_code = 502, detail = _friendly_error(e), ) + if resp is None: + try: + await client.aclose() + except Exception: + pass + _tracker.__exit__(None, None, None) + return StreamingResponse( + iter(()), + media_type = "text/event-stream", + headers = { + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", + }, + ) if resp.status_code == 200: break @@ -7759,12 +8017,13 @@ async def _openai_passthrough_stream( cancel_watcher = asyncio.create_task(_await_cancel_then_close(cancel_event, resp)) try: lines_iter = resp.aiter_lines() - async for raw_line in lines_iter: - if cancel_event.is_set(): - break - if await request.is_disconnected(): - cancel_event.set() - break + async for raw_line in _aiter_llama_stream_items( + lines_iter, + cancel_event = cancel_event, + request = request, + first_token_deadline = first_token_deadline, + response = resp, + ): if not raw_line: continue if not raw_line.startswith("data: "): @@ -7843,7 +8102,11 @@ async def _openai_passthrough_non_streaming(llama_backend, payload, model_name): while True: try: async with httpx.AsyncClient() as client: - resp = await client.post(target_url, json = body, timeout = 600) + resp = await client.post( + target_url, + json = body, + timeout = _llama_non_streaming_generation_timeout(), + ) except httpx.RequestError as e: # llama-server subprocess crashed / starting / unreachable. Surface the # same friendly message the sync chat path emits so operators don't see diff --git a/studio/backend/tests/test_gguf_route_cursor_reset.py b/studio/backend/tests/test_gguf_route_cursor_reset.py index fdbf31f9de..dc397b2d08 100644 --- a/studio/backend/tests/test_gguf_route_cursor_reset.py +++ b/studio/backend/tests/test_gguf_route_cursor_reset.py @@ -59,11 +59,16 @@ def _make_backend(monkeypatch, streams: list[list[str]], payloads: list[dict]): payload, _cancel_event, headers = None, + first_token_deadline = None, ): payloads.append(copy.deepcopy(payload)) yield type("FakeResponse", (), {"status_code": 200, "chunks": streams.pop(0)})() - def fake_iter_text_cancellable(response, _cancel_event): + def fake_iter_text_cancellable( + response, + _cancel_event, + first_token_deadline = None, + ): yield from response.chunks monkeypatch.setattr(backend, "_stream_with_retry", fake_stream_with_retry) diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py index 3c121c281d..aacb029ff6 100644 --- a/studio/backend/tests/test_llama_cpp_tool_loop.py +++ b/studio/backend/tests/test_llama_cpp_tool_loop.py @@ -52,11 +52,16 @@ def _make_backend(monkeypatch, streams: list[list[str]], payloads: list[dict]): payload, _cancel_event, headers = None, + first_token_deadline = None, ): payloads.append(copy.deepcopy(payload)) yield type("FakeResponse", (), {"status_code": 200, "chunks": streams.pop(0)})() - def fake_iter_text_cancellable(response, _cancel_event): + def fake_iter_text_cancellable( + response, + _cancel_event, + first_token_deadline = None, + ): yield from response.chunks monkeypatch.setattr(backend, "_stream_with_retry", fake_stream_with_retry) diff --git a/studio/backend/tests/test_llama_route_timeouts.py b/studio/backend/tests/test_llama_route_timeouts.py new file mode 100644 index 0000000000..5aee6198ba --- /dev/null +++ b/studio/backend/tests/test_llama_route_timeouts.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +import asyncio +import os +import sys +import time +from types import SimpleNamespace + +_backend = os.path.join(os.path.dirname(__file__), "..") +sys.path.insert(0, _backend) + +import routes.inference as inf_mod # noqa: E402 + + +def test_non_streaming_generation_timeout_has_read_deadline(): + timeout = inf_mod._llama_non_streaming_generation_timeout() + assert timeout.read == inf_mod._DEFAULT_FIRST_TOKEN_TIMEOUT_S + + +def test_stream_first_item_deadline_after_headers(): + async def _run(): + class _Never: + async def __anext__(self): + await asyncio.Future() + + started = time.monotonic() + try: + async for _ in inf_mod._aiter_llama_stream_items( + _Never(), + first_token_deadline = started + 0.02, + ): + pass + except inf_mod.httpx.ReadTimeout: + pass + else: + raise AssertionError("first item deadline did not fire") + assert time.monotonic() - started < 0.5 + + asyncio.run(_run()) + + +def test_preheader_send_cleanup_on_disconnect_and_cancel(): + async def _run(cancel_parent): + state = SimpleNamespace(disconnected = False, closed = False, cancelled = False) + started = asyncio.Event() + + class _Client: + async def send( + self, + req, + stream = False, + ): + started.set() + try: + await asyncio.Future() + except asyncio.CancelledError: + state.cancelled = True + raise + + async def aclose(self): + state.closed = True + + class _Request: + async def is_disconnected(self): + return state.disconnected + + task = asyncio.create_task( + inf_mod._send_stream_with_preheader_cancel(_Client(), object(), request = _Request()) + ) + await started.wait() + if cancel_parent: + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + else: + raise AssertionError("helper cancellation did not propagate") + else: + state.disconnected = True + assert await task is None + assert state.closed + assert state.cancelled + + asyncio.run(_run(False)) + asyncio.run(_run(True)) diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index baa3e50b95..b54f4c130d 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -1,13 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. -"""Tests for the OpenAI /v1/chat/completions client-side tool pass-through. - -Covers ChatMessage tool/assistant roles, ChatCompletionRequest tool fields and -extra="allow", anthropic_tool_choice_to_openai, _build_passthrough_payload -tool_choice propagation, and _friendly_error's httpx-to-"Lost connection" -mapping. No server or GPU required. -""" +"""Tests for the OpenAI /v1/chat/completions client-side tool pass-through.""" import os import sys @@ -914,12 +908,6 @@ class TestOpenAICompatibilityHelpers: class TestFriendlyErrorHttpx: - """When llama-server is down, httpx RequestError strings lack the - "Lost connection to llama-server" substring the sync path keys off, so the - old substring-only `_friendly_error` returned a useless generic message. - These tests pin the new isinstance-based mapping. - """ - def _req(self): return httpx.Request("POST", "http://127.0.0.1:65535/v1/chat/completions") @@ -937,7 +925,7 @@ class TestFriendlyErrorHttpx: def test_read_timeout_mapped(self): exc = httpx.ReadTimeout("timed out", request = self._req()) - assert "Lost connection" in _friendly_error(exc) + assert "first token within 20 minutes" in _friendly_error(exc) def test_non_httpx_unchanged(self): # Non-httpx exceptions still fall through to the substring heuristics diff --git a/tests/studio/test_llama_cpp_wall_clock_cap.py b/tests/studio/test_llama_cpp_wall_clock_cap.py index 671abea823..f173cfbc55 100644 --- a/tests/studio/test_llama_cpp_wall_clock_cap.py +++ b/tests/studio/test_llama_cpp_wall_clock_cap.py @@ -1,22 +1,4 @@ -""" -Tests for the llama-server wall-clock cap (t_max_predict_ms). - -The UI always sends max_tokens = context_length, so gating -t_max_predict_ms on `max_tokens is None` makes the safety net dead -code. The fix applies the wall-clock cap unconditionally on all three -streaming payload sites and raises the default to 10 minutes so slow -CPU / macOS / Windows installs are not cut off mid-generation. - -Verifies: - - t_max_predict_ms is assigned unconditionally at the three - payload-builder sites (not inside an `if max_tokens is None` else - branch). - - _DEFAULT_T_MAX_PREDICT_MS is at least 10 minutes (previously - 120_000). - - The default max_tokens path still applies _DEFAULT_MAX_TOKENS. - - The three payload variable names (payload x2, stream_payload x1) - each get both `max_tokens` and `t_max_predict_ms`. -""" +"""Timeout policy checks for Studio's local llama-server path.""" from __future__ import annotations @@ -36,88 +18,25 @@ SRC = SOURCE_PATH.read_text() TREE = ast.parse(SRC) -def _is_subscript_assign(stmt: ast.stmt, target_name: str, key: str) -> bool: - if not isinstance(stmt, ast.Assign) or len(stmt.targets) != 1: - return False - t = stmt.targets[0] - if not isinstance(t, ast.Subscript): - return False - if not (isinstance(t.value, ast.Name) and t.value.id == target_name): - return False - slc = t.slice - return isinstance(slc, ast.Constant) and slc.value == key - - -def _collect_assignments(tree, target_name, key): - """Return list of (node, stack_of_enclosing_ifs) for each match.""" - hits = [] - - def visit(node, stack): - if _is_subscript_assign(node, target_name, key): - hits.append((node, stack)) - for child in ast.iter_child_nodes(node): - if isinstance(child, ast.If): - for sub in child.body: - visit(sub, stack + [(child, "body")]) - for sub in child.orelse: - visit(sub, stack + [(child, "orelse")]) - else: - visit(child, stack) - - visit(tree, []) - return hits - - -def test_default_t_max_predict_ms_is_at_least_ten_minutes(): +def _module_constant(name: str): for node in TREE.body: if isinstance(node, ast.Assign) and len(node.targets) == 1: t = node.targets[0] - if isinstance(t, ast.Name) and t.id == "_DEFAULT_T_MAX_PREDICT_MS": + if isinstance(t, ast.Name) and t.id == name: value = node.value assert isinstance(value, ast.Constant) - assert value.value >= 600_000, ( - f"_DEFAULT_T_MAX_PREDICT_MS must be >= 10 minutes " - f"(600_000 ms) to avoid cutting off slow-CPU generations; " - f"got {value.value}" - ) - return - raise AssertionError("_DEFAULT_T_MAX_PREDICT_MS constant missing") + return value.value + raise AssertionError(f"{name} constant missing") -def test_t_max_predict_ms_set_unconditionally_at_three_sites(): - hits_payload = _collect_assignments(TREE, "payload", "t_max_predict_ms") - hits_stream = _collect_assignments(TREE, "stream_payload", "t_max_predict_ms") - total = len(hits_payload) + len(hits_stream) - assert total == 3, ( - f"expected 3 total t_max_predict_ms assignments " - f"(payload x2 + stream_payload x1), got {total}" - ) - for node, stack in hits_payload + hits_stream: - for parent_if, branch in stack: - # The assignment must not be gated by a test that checks - # `max_tokens is None` (which would make it dead code for - # the UI path where max_tokens is always set). - test_src = ast.unparse(parent_if.test) - assert "max_tokens" not in test_src, ( - f"t_max_predict_ms at line {node.lineno} is nested under " - f"`if {test_src}:` -- it must be applied unconditionally so " - f"the wall-clock cap is not dead code for callers that set " - f"max_tokens" - ) +def test_first_token_timeout_is_at_least_twenty_minutes(): + value = _module_constant("_DEFAULT_FIRST_TOKEN_TIMEOUT_S") + assert value >= 1200.0 + + +def test_studio_chat_payloads_do_not_set_wall_clock_generation_cap(): + assert "t_max_predict_ms" not in SRC def test_max_tokens_default_cap_still_applied(): - # _DEFAULT_MAX_TOKENS must still kick in when caller passes None. - # We check the conditional expression `max_tokens if max_tokens is not - # None else _DEFAULT_MAX_TOKENS` appears at each site. - matches = 0 - for node in ast.walk(TREE): - if not isinstance(node, ast.IfExp): - continue - src = ast.unparse(node) - if "max_tokens" in src and "_DEFAULT_MAX_TOKENS" in src: - matches += 1 - assert matches >= 3, ( - f"expected >=3 `max_tokens if max_tokens is not None else " - f"_DEFAULT_MAX_TOKENS` expressions; got {matches}" - ) + assert SRC.count("_DEFAULT_MAX_TOKENS_FLOOR") >= 3