danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Check the code timeout before each document read

Browse source 
The VFS bridge suspends the Monty worker for the length of a read. Monty checks
its duration budget between interpreter steps, so it cannot check while a read
is in flight. Code that reads in a loop overran a 60s budget by minutes. A read
takes about 20ms on a 2789-document corpus, so a full scan spends about 55s in
reads alone.

Check the deadline before each read. Raising from inside the callback answers the
worker's suspension, which keeps the session usable.

Monty also spends max_duration_secs across the session rather than per call, and
the sandbox reuses the session so that variables persist. Budget it for
code_timeout * max_executions. At the old per-call value the first slow call
starved every later one.

Do not wrap feed_run in asyncio.wait_for. Cancelling during pure compute is
clean, but cancelling while a read waits for an answer wedges the session with a
protocol RuntimeError that escapes execute(). A call that computes without
reading stays bounded by the session budget alone.
This commit is contained in:
Yiorgis Gozadinos 2026-07-28 15:55:44 +03:00
parent feaca386d3
commit 522959d9b4
No known key found for this signature in database
4 changed files with 110 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
CHANGELOG.md
View file

@ -9,6 +9,7 @@
- Require `pydantic-ai-slim>=2.18,<3`.
- `pydantic-monty` bumped to `>=0.0.19`; sandbox code now runs in a subprocess worker pool.
- The analysis sandbox gives Monty a duration budget of `analysis.code_timeout * analysis.max_executions` for the session, was `analysis.code_timeout`.
- `enable_thinking` maps onto Pydantic AI's unified `thinking` setting for the `anthropic`, `gemini`, `groq` and `bedrock` providers. The Anthropic thinking budget is now Pydantic AI's default of 10000 tokens, was 4096, and `max_tokens` must exceed it on budget-based Claude models; `enable_thinking: false` disables Gemini thinking rather than only hiding thoughts; Groq maps reasoning effort rather than `groq_reasoning_format`; Bedrock Qwen with `enable_thinking: false` no longer sends `reasoning_config`, while Bedrock-served Claude keeps an explicit `thinking: disabled`. The `openai` and `ollama` providers still map to `openai_reasoning_effort`.
- `provider: bedrock` with a proprietary OpenAI model such as `openai.o3-mini-v1:0` raises `UserError`: Bedrock Converse serves only the `gpt-oss` family. Use `provider: bedrock-mantle`.
- Tool failures raise `pydantic_ai.ToolFailed` instead of returning failure text: search and code-execution limits, sandbox execution errors, and `get_document`/`summarize_document` misses.
@ -20,6 +21,7 @@
- `evaluations run` opens the database read-only outside the population phase, so an embedder identity differing from the stored one warns instead of aborting the run.
- `docs/installation.md` documents the `jina`, `s3` and `ingester` extras, names the extras the full package actually pulls, drops the removed MixedBread AI reranker, and no longer lists Anthropic as a built-in provider.
- `docs/tuning.md` no longer points at the removed `claim_timeout_s` setting.
- `analysis.code_timeout` is checked before each document read; code that reads in a loop no longer overruns it by the duration of the outstanding reads.
### Removed

4
docs/configuration/qa.md
View file

@ -50,13 +50,13 @@ analysis:
provider: anthropic
name: claude-sonnet-4-20250514
temperature: 0.0 # Default: 0.0 (deterministic for code generation)
code_timeout: 60.0 # Max seconds for code execution
code_timeout: 60.0 # Max seconds a call may spend reading documents
max_output_chars: 50000 # Truncate output after this many chars
max_executions: 15 # Max execute_code calls per question
```
- **model**: LLM configuration (see [Providers](providers.md#model-settings)). When unset, falls back to `qa.model`.
- **code_timeout**: Maximum seconds for each code execution (default: 60)
- **code_timeout**: Seconds a single `execute_code` call may spend reading documents (default: 60). The sandbox refuses further reads past this point. Code that computes without reading is bounded instead by the session budget of `code_timeout * max_executions`.
- **max_output_chars**: Truncate code output after this many characters (default: 50000)
- **max_executions**: Maximum `execute_code` calls per question before the capability is told to answer from what it has (default: 15)

35
haiku_rag_slim/haiku/rag/sandbox/sandbox.py
View file

@ -14,6 +14,7 @@ from pydantic_monty import (
CallbackFile,
MemoryFile,
OSAccess,
ResourceLimits,
)
from haiku.rag.config.models import AppConfig
@ -151,6 +152,7 @@ class Sandbox:
_session: AsyncMontySession | None
_vfs: OSAccess | None
_loop: asyncio.AbstractEventLoop | None
_deadline: float | None
def __init__(
self,
@ -174,6 +176,7 @@ class Sandbox:
self._session = None
self._vfs = None
self._loop = None
self._deadline = None
@asynccontextmanager
async def _connection(self) -> "AsyncIterator[HaikuRAG]":
@ -197,10 +200,24 @@ class Sandbox:
Called off the loop while ``feed_run`` is awaited, so scheduling onto it
and blocking for the result is safe.
Blocking here suspends the worker, and Monty checks its duration budget
between interpreter steps, so it cannot check while a read is in flight.
Enforce the budget before starting another read, or code that reads in a
loop overruns it by however long the outstanding reads take. Raising from
inside the callback answers the worker's suspension, which keeps the
session usable — cancelling ``feed_run`` from outside does not, and wedges
the protocol.
"""
assert self._loop is not None, (
"VFS reads happen during execute(); the loop must be captured first."
)
if self._deadline is not None and self._loop.time() > self._deadline:
coro.close()
raise TimeoutError(
"time limit exceeded: no further document reads after "
f"{self._config.analysis.code_timeout}s"
)
return asyncio.run_coroutine_threadsafe(coro, self._loop).result()
async def close(self) -> None:
@ -428,15 +445,26 @@ class Sandbox:
return OSAccess(files)
def _session_limits(self) -> ResourceLimits:
"""Resource limits for the worker session.
Monty spends ``max_duration_secs`` across the session's whole life, and
the session is reused so variables persist between calls. Budget it for
the run rather than for one call, or the first slow call starves every
later one. ``code_timeout`` is enforced per call by the read deadline in
``_run_on_loop``. This is the backstop for code that computes without
reading, and one such call can spend all of it.
"""
analysis = self._config.analysis
return {"max_duration_secs": analysis.code_timeout * analysis.max_executions}
async def _ensure_initialized(self) -> tuple[AsyncMontySession, OSAccess]:
"""Check out a worker session and build the VFS on first use."""
if self._session is None:
self._vfs = await self._build_vfs()
self._pool = AsyncMonty()
await self._pool.__aenter__()
session = self._pool.checkout(
limits={"max_duration_secs": self._config.analysis.code_timeout},
)
session = self._pool.checkout(limits=self._session_limits())
await session.__aenter__()
self._session = session
assert self._session is not None and self._vfs is not None
@ -449,6 +477,7 @@ class Sandbox:
"""
# Monty's synchronous file callbacks bridge DB reads back to this loop.
self._loop = asyncio.get_running_loop()
self._deadline = self._loop.time() + self._config.analysis.code_timeout
session, vfs = await self._ensure_initialized()
external_fns = self._build_external_functions()

74
tests/sandbox/test_sandbox.py
View file

@ -733,3 +733,77 @@ class TestSandboxHeldConnection:
assert not any(t.name == "sandbox-vfs" for t in threading.enumerate())
await sb.close()
await sb.close()
class TestSandboxReadDeadline:
"""The VFS bridge suspends the worker for the length of a read, so Monty
cannot check its duration budget while one is in flight. The sandbox
enforces the budget itself, before each read."""
@pytest.mark.asyncio
async def test_read_after_deadline_raises_without_scheduling(self, sandbox):
"""A read attempted past the deadline fails instead of querying."""
scheduled = False
async def _never_runs():
nonlocal scheduled
scheduled = True
sandbox._loop = asyncio.get_running_loop()
sandbox._deadline = sandbox._loop.time() - 1.0
coro = _never_runs()
with pytest.raises(TimeoutError, match="time limit exceeded"):
sandbox._run_on_loop(coro)
coro.close()
assert scheduled is False
def test_session_budget_covers_every_permitted_execution(self, temp_db_path):
"""Monty spends its duration budget across the session's whole life, so a
per-call value would let the first call starve the rest."""
config = AppConfig()
config.analysis.code_timeout = 5.0
config.analysis.max_executions = 3
sb = Sandbox(db_path=temp_db_path, config=config, context=AnalysisContext())
assert sb._session_limits() == {"max_duration_secs": 15.0}
@pytest.mark.asyncio
async def test_document_read_past_the_deadline_fails_the_execution(
self, temp_db_path
):
"""The overrun surfaces as a failed result, not a raised exception."""
from docling_core.types.doc.document import DoclingDocument
from docling_core.types.doc.labels import DocItemLabel
config = AppConfig()
config.analysis.code_timeout = 0.0
docling = DoclingDocument(name="d")
docling.add_text(label=DocItemLabel.TEXT, text="Foxes and dogs.")
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.import_document(
docling,
[
Chunk(
content="Foxes and dogs.",
embedding=[0.1] * config.embeddings.model.vector_dim,
order=0,
)
],
uri="test://deadline",
)
sb = Sandbox(db_path=temp_db_path, config=config, context=AnalysisContext())
try:
result = await sb.execute(
"from pathlib import Path\n"
f"print(Path('/documents/{doc.id}/content.txt').read_text())"
)
assert result.success is False
assert "no further document reads" in result.stderr
finally:
await sb.close()