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"),
|
||||
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,
|
||||
|
||||
@@ -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 ─────────────────────────────────────────────────────
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user