Enforce the analysis code timeout per call
The VFS bridge blocks Monty's worker thread for the duration of a read, and Monty checks its duration budget between interpreter steps, so it cannot check while a read is outstanding. Code that reads document after document overran a 60s budget by minutes: measured 20ms per read over a 2789-document corpus, so a full scan spends ~55s in reads alone. Check the deadline before starting each read. Monty also spends `max_duration_secs` across the session rather than per call, and the session is reused so variables persist. Budget it for `code_timeout * max_executions`; at the old per-call value the first slow execution starved every later one, which is what failed 5 of 6 executions in a local analysis run.
This commit is contained in:
parent
044da7ae99
commit
a553f518bf
3 changed files with 117 additions and 1 deletions
|
|
@ -7,12 +7,14 @@
|
||||||
- `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`.
|
- `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`.
|
- `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.
|
- 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.
|
||||||
|
- The analysis sandbox gives Monty a duration budget of `analysis.code_timeout * analysis.max_executions` for the session, was `analysis.code_timeout`.
|
||||||
|
|
||||||
### Fixed
|
### Fixed
|
||||||
|
|
||||||
- `Store.set_haiku_version` stamps the store's own config into a recreated settings row instead of the process-global `Config`.
|
- `Store.set_haiku_version` stamps the store's own config into a recreated settings row instead of the process-global `Config`.
|
||||||
- `check_source_accessible` returns `False` for a URI it cannot resolve (unparseable host, unreadable path) instead of raising and aborting a full rebuild.
|
- `check_source_accessible` returns `False` for a URI it cannot resolve (unparseable host, unreadable path) instead of raising and aborting a full rebuild.
|
||||||
- `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.
|
- `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.
|
||||||
|
- `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
|
### Removed
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -142,6 +142,7 @@ class Sandbox:
|
||||||
_repl: MontyRepl | None
|
_repl: MontyRepl | None
|
||||||
_vfs: OSAccess | None
|
_vfs: OSAccess | None
|
||||||
_loop: asyncio.AbstractEventLoop | None
|
_loop: asyncio.AbstractEventLoop | None
|
||||||
|
_deadline: float | None
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
|
|
@ -164,6 +165,7 @@ class Sandbox:
|
||||||
self._repl = None
|
self._repl = None
|
||||||
self._vfs = None
|
self._vfs = None
|
||||||
self._loop = None
|
self._loop = None
|
||||||
|
self._deadline = None
|
||||||
|
|
||||||
@asynccontextmanager
|
@asynccontextmanager
|
||||||
async def _connection(self) -> "AsyncIterator[HaikuRAG]":
|
async def _connection(self) -> "AsyncIterator[HaikuRAG]":
|
||||||
|
|
@ -187,10 +189,21 @@ class Sandbox:
|
||||||
|
|
||||||
Called from Monty's worker thread while ``feed_run_async`` leaves the
|
Called from Monty's worker thread while ``feed_run_async`` leaves the
|
||||||
loop free, so scheduling onto it and blocking for the result is safe.
|
loop free, so scheduling onto it and blocking for the result is safe.
|
||||||
|
|
||||||
|
Blocking here parks the worker thread, so Monty cannot check its own
|
||||||
|
duration budget until the read returns. Enforce the budget before
|
||||||
|
starting another one, or code that reads in a loop overruns it by
|
||||||
|
however long the outstanding reads take.
|
||||||
"""
|
"""
|
||||||
assert self._loop is not None, (
|
assert self._loop is not None, (
|
||||||
"VFS reads happen during execute(); the loop must be captured first."
|
"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()
|
return asyncio.run_coroutine_threadsafe(coro, self._loop).result()
|
||||||
|
|
||||||
def close(self) -> None:
|
def close(self) -> None:
|
||||||
|
|
@ -417,9 +430,18 @@ class Sandbox:
|
||||||
"""Initialize the REPL session and VFS on first use."""
|
"""Initialize the REPL session and VFS on first use."""
|
||||||
if self._repl is None:
|
if self._repl is None:
|
||||||
self._vfs = await self._build_vfs()
|
self._vfs = await self._build_vfs()
|
||||||
|
# Monty spends `max_duration_secs` across the REPL's whole life, and
|
||||||
|
# the REPL 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 it all.
|
||||||
self._repl = MontyRepl(
|
self._repl = MontyRepl(
|
||||||
limits={
|
limits={
|
||||||
"max_duration_secs": self._config.analysis.code_timeout,
|
"max_duration_secs": (
|
||||||
|
self._config.analysis.code_timeout
|
||||||
|
* self._config.analysis.max_executions
|
||||||
|
),
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
assert self._repl is not None and self._vfs is not None
|
assert self._repl is not None and self._vfs is not None
|
||||||
|
|
@ -432,6 +454,7 @@ class Sandbox:
|
||||||
"""
|
"""
|
||||||
# Monty's synchronous file callbacks bridge DB reads back to this loop.
|
# Monty's synchronous file callbacks bridge DB reads back to this loop.
|
||||||
self._loop = asyncio.get_running_loop()
|
self._loop = asyncio.get_running_loop()
|
||||||
|
self._deadline = self._loop.time() + self._config.analysis.code_timeout
|
||||||
repl, vfs = await self._ensure_initialized()
|
repl, vfs = await self._ensure_initialized()
|
||||||
external_fns = self._build_external_functions()
|
external_fns = self._build_external_functions()
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -7,6 +7,7 @@ import pytest
|
||||||
from haiku.rag.client import HaikuRAG
|
from haiku.rag.client import HaikuRAG
|
||||||
from haiku.rag.config.models import AppConfig
|
from haiku.rag.config.models import AppConfig
|
||||||
from haiku.rag.sandbox import AnalysisContext, Sandbox, SandboxResult
|
from haiku.rag.sandbox import AnalysisContext, Sandbox, SandboxResult
|
||||||
|
from haiku.rag.store.models.chunk import Chunk
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture(scope="module")
|
@pytest.fixture(scope="module")
|
||||||
|
|
@ -618,3 +619,93 @@ class TestSandboxHeldConnection:
|
||||||
assert not any(t.name == "sandbox-vfs" for t in threading.enumerate())
|
assert not any(t.name == "sandbox-vfs" for t in threading.enumerate())
|
||||||
sb.close()
|
sb.close()
|
||||||
sb.close()
|
sb.close()
|
||||||
|
|
||||||
|
|
||||||
|
class TestSandboxReadDeadline:
|
||||||
|
"""The VFS bridge blocks Monty's worker thread, so its duration budget
|
||||||
|
cannot be checked while a read 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
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_session_budget_covers_every_permitted_execution(
|
||||||
|
self, temp_db_path, monkeypatch
|
||||||
|
):
|
||||||
|
"""Monty spends its duration budget across the REPL's whole life, so a
|
||||||
|
per-call value would let the first call starve the rest."""
|
||||||
|
import haiku.rag.sandbox.sandbox as sandbox_mod
|
||||||
|
|
||||||
|
captured: dict[str, float] = {}
|
||||||
|
real_repl = sandbox_mod.MontyRepl
|
||||||
|
|
||||||
|
def _capture(*args, **kwargs):
|
||||||
|
captured.update(kwargs.get("limits", {}))
|
||||||
|
return real_repl(*args, **kwargs)
|
||||||
|
|
||||||
|
monkeypatch.setattr(sandbox_mod, "MontyRepl", _capture)
|
||||||
|
|
||||||
|
config = AppConfig()
|
||||||
|
config.analysis.code_timeout = 5.0
|
||||||
|
config.analysis.max_executions = 3
|
||||||
|
|
||||||
|
async with HaikuRAG(temp_db_path, create=True):
|
||||||
|
pass
|
||||||
|
sb = Sandbox(db_path=temp_db_path, config=config, context=AnalysisContext())
|
||||||
|
result = await sb.execute("x = 1")
|
||||||
|
|
||||||
|
assert result.success
|
||||||
|
assert captured["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())
|
||||||
|
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
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue