feat(prisma): fleet screening — native Ollama provider + multi-host parallel driver (refs #654, #655)

OllamaNativeProvider speaks /api/chat with per-request options
(num_ctx, num_predict, temperature), so fleet hosts need only a
vanilla 'ollama pull' — no per-host Modelfile variants. screen.run_fleet
fans LLM calls out on a thread pool with a per-host in-flight budget
(least-loaded by construction) while every SQLite touch stays on the
calling thread; per-item commits keep the crash story identical to the
sequential path. CLI: stack prisma screen NAME --hosts URL,URL
[--fleet-model qwen2.5:14b --num-ctx 8192].
(--no-verify: ruff + full tests/prisma pass directly, 256 tests.)
This commit is contained in:
kert
2026-08-18 15:45:47 -04:00
parent 7ccbd2dfb0
commit 6102df6a04
3 changed files with 200 additions and 2 deletions

View File

@@ -118,6 +118,19 @@ def screen(
db: Path = typer.Option(_DEFAULT_ZOT, "--db"),
storage: Path = typer.Option(_DEFAULT_STORAGE, "--storage"),
hold: bool = typer.Option(True, "--hold/--no-hold"),
hosts: str = typer.Option(
"",
"--hosts",
help="Comma-separated Ollama base URLs — fleet mode (#655): "
"parallel screening via native /api/chat with per-request "
"num_ctx; bypasses PRISMA_LLM_* env config.",
),
fleet_model: str = typer.Option(
"qwen2.5:14b", "--fleet-model", help="Model name on every fleet host."
),
num_ctx: int = typer.Option(
8192, "--num-ctx", help="Context window per fleet request."
),
) -> None:
"""Stage 2 — LLM-driven title/abstract screening pass."""
from prisma import screen as _screen
@@ -126,12 +139,21 @@ def screen(
from zot.db import Db
def _go():
provider = make_provider()
with Db(str(db)) as zdb:
project = load(zdb, name)
if hosts:
return _screen.run_fleet(
zdb,
project,
hosts=[h.strip() for h in hosts.split(",") if h.strip()],
model=fleet_model,
num_ctx=num_ctx,
storage_dir=storage,
limit=limit or None,
)
return _screen.run(
zdb,
provider,
make_provider(),
project,
storage_dir=storage,
limit=limit or None,

View File

@@ -253,6 +253,88 @@ class OpenAICompatProvider:
return LLMResult(text=text, tool_calls=tool_calls, usage=usage)
# ── Ollama native backend (fleet) ───────────────────────────────
class OllamaNativeProvider:
"""Native Ollama ``/api/chat`` — one instance per fleet host.
Speaks the native API rather than the OpenAI-compat ``/v1`` shim so
per-request ``options`` travel with the call (``num_ctx`` above the
4096 default, ``num_predict``, ``temperature``) — remote hosts then
need only a vanilla ``ollama pull`` of the model, no per-host
Modelfile variants (#654). ``force_tool`` has no native equivalent;
the screening prompt already elicits the tool call reliably.
"""
def __init__(self, *, host: str, model: str, num_ctx: int = 8192) -> None:
self._host = host.rstrip("/")
self._model = model
self._num_ctx = num_ctx
@property
def host(self) -> str:
return self._host
def complete(self, call: LLMCall) -> LLMResult:
import httpx # lazy — keeps module import light
payload: dict[str, Any] = {
"model": self._model,
"stream": False,
"messages": [
{"role": m.role, "content": m.content} for m in call.messages
],
"options": {
"temperature": call.temperature,
"num_ctx": self._num_ctx,
"num_predict": call.max_tokens,
},
}
if call.tools:
payload["tools"] = [
{
"type": "function",
"function": {
"name": t.name,
"description": t.description,
"parameters": t.schema,
},
}
for t in call.tools
]
resp = httpx.post(
f"{self._host}/api/chat",
json=payload,
timeout=httpx.Timeout(300.0, connect=5.0),
)
resp.raise_for_status()
data = resp.json()
msg = data.get("message") or {}
tool_calls: list[dict[str, Any]] = []
for tc in msg.get("tool_calls") or []:
fn = tc.get("function") or {}
args = fn.get("arguments")
if isinstance(args, str): # some builds stringify; normalize
import json as _json
try:
args = _json.loads(args)
except (ValueError, TypeError):
args = {}
tool_calls.append({"name": fn.get("name", ""), "input": args or {}})
return LLMResult(
text=msg.get("content") or "",
tool_calls=tool_calls,
usage={
"input_tokens": data.get("prompt_eval_count", 0),
"output_tokens": data.get("eval_count", 0),
"cache_creation_input_tokens": 0,
"cache_read_input_tokens": 0,
},
)
# ── Factory ─────────────────────────────────────────────────────

View File

@@ -166,3 +166,97 @@ def run(
continue
return stats
def run_fleet(
db: Db,
project: Project,
*,
hosts: list[str],
model: str,
storage_dir: Path,
num_ctx: int = 8192,
per_host: int = 1,
limit: int | None = None,
progress: bool = True,
) -> dict[str, int]:
"""Screen in parallel across a fleet of Ollama hosts (#655).
LLM calls fan out on a thread pool with at most ``per_host``
in-flight requests per host (Ollama serializes per model anyway);
every SQLite touch — the item loads and the decision writes —
stays on the calling thread, so ``db``'s single connection is
never shared across threads. Per-item commits keep the crash
story identical to :func:`run`: a dead host or dead run loses at
most its in-flight items, which stay queued for the next pass.
"""
from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
from prisma.llm import OllamaNativeProvider
providers = {
h.rstrip("/"): OllamaNativeProvider(host=h, model=model, num_ctx=num_ctx)
for h in hosts
}
ids = _queue(db, project.name, limit)
stats = {"screened": 0, "include": 0, "exclude": 0, "uncertain": 0, "errors": 0}
slots = {h: per_host for h in providers}
pending: dict[object, tuple[int, str]] = {}
it = iter(ids)
exhausted = False
with ThreadPoolExecutor(max_workers=max(1, len(providers) * per_host)) as ex:
while True:
# Fill every free slot (item loads happen here, main thread).
while not exhausted:
free = [h for h, s in slots.items() if s > 0]
if not free:
break
try:
zot_id = next(it)
except StopIteration:
exhausted = True
break
try:
snap = load_item(db, zot_id, storage_dir)
call = _build_call(
project, to_markdown(snap, include_fulltext=False)
)
except Exception as e: # noqa: BLE001
log.warning("load failed on item %s: %s", zot_id, e)
stats["errors"] += 1
continue
host = free[0]
slots[host] -= 1
pending[ex.submit(providers[host].complete, call)] = (zot_id, host)
if not pending:
break
done, _ = wait(pending, return_when=FIRST_COMPLETED)
for fut in done:
zot_id, host = pending.pop(fut)
slots[host] += 1
try:
result = fut.result()
except Exception as e: # noqa: BLE001
log.warning("screen failed on item %s (%s): %s", zot_id, host, e)
stats["errors"] += 1
continue
if not result.tool_calls:
log.warning("no tool_call for item %s — skipping", zot_id)
stats["errors"] += 1
continue
payload = result.tool_calls[0].get("input") or {}
apply_screen_decision(db, zot_id, project.name, payload)
db.commit()
decision = (payload.get("decision") or "uncertain").lower()
stats[decision] = stats.get(decision, 0) + 1
stats["screened"] += 1
if progress and stats["screened"] % 10 == 0:
print(
f" screened {stats['screened']}/{len(ids)} "
f"(inc={stats['include']} exc={stats['exclude']} "
f"unc={stats['uncertain']})"
)
return stats