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:
@@ -118,6 +118,19 @@ def screen(
|
|||||||
db: Path = typer.Option(_DEFAULT_ZOT, "--db"),
|
db: Path = typer.Option(_DEFAULT_ZOT, "--db"),
|
||||||
storage: Path = typer.Option(_DEFAULT_STORAGE, "--storage"),
|
storage: Path = typer.Option(_DEFAULT_STORAGE, "--storage"),
|
||||||
hold: bool = typer.Option(True, "--hold/--no-hold"),
|
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:
|
) -> None:
|
||||||
"""Stage 2 — LLM-driven title/abstract screening pass."""
|
"""Stage 2 — LLM-driven title/abstract screening pass."""
|
||||||
from prisma import screen as _screen
|
from prisma import screen as _screen
|
||||||
@@ -126,12 +139,21 @@ def screen(
|
|||||||
from zot.db import Db
|
from zot.db import Db
|
||||||
|
|
||||||
def _go():
|
def _go():
|
||||||
provider = make_provider()
|
|
||||||
with Db(str(db)) as zdb:
|
with Db(str(db)) as zdb:
|
||||||
project = load(zdb, name)
|
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(
|
return _screen.run(
|
||||||
zdb,
|
zdb,
|
||||||
provider,
|
make_provider(),
|
||||||
project,
|
project,
|
||||||
storage_dir=storage,
|
storage_dir=storage,
|
||||||
limit=limit or None,
|
limit=limit or None,
|
||||||
|
|||||||
@@ -253,6 +253,88 @@ class OpenAICompatProvider:
|
|||||||
return LLMResult(text=text, tool_calls=tool_calls, usage=usage)
|
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 ─────────────────────────────────────────────────────
|
# ── Factory ─────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -166,3 +166,97 @@ def run(
|
|||||||
continue
|
continue
|
||||||
|
|
||||||
return stats
|
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
|
||||||
|
|||||||
Reference in New Issue
Block a user