feat: add OpenVINO NPU prototype services
This commit is contained in:
William Valentin
@@ -0,0 +1,159 @@
|
||||
# OpenVINO NPU document/image triage prototype
|
||||
|
||||
Local-only prototype for triaging screenshots, photos/scans, and PDF page images.
|
||||
It returns structured JSON metadata and explicitly reports CPU vs NPU stages.
|
||||
|
||||
Location: `/home/will/lab/swarm/openvino-doc-image-triage-npu/`
|
||||
|
||||
## Privacy and safety
|
||||
|
||||
- No external uploads.
|
||||
- The only network call is optional localhost-only embeddings at `127.0.0.1:18817`.
|
||||
- Raw OCR/sidecar text is redacted by default and is not logged.
|
||||
- Full source paths are omitted by default; responses include basename and SHA-256.
|
||||
- Allowed roots are enforced for CLI/server requests.
|
||||
- This prototype does not mutate Obsidian, RAG, Chroma, vector collections, routing, or gateway services.
|
||||
|
||||
## CPU vs NPU stages
|
||||
|
||||
CPU:
|
||||
- file intake, allowed-root checks, size checks, hashing
|
||||
- image/PDF decoding/rendering and normalization
|
||||
- optional local text extraction from sidecars or PDF text libraries
|
||||
- regex metadata extraction and rule-based category fallback
|
||||
- final needs-attention rules
|
||||
|
||||
NPU:
|
||||
- needs-attention semantic embedding, via existing local OpenVINO embeddings service on `:18817`
|
||||
- verified with `/sys/class/accel/accel0/device/npu_busy_time_us` before/after each embedding call
|
||||
|
||||
Not configured in v1:
|
||||
- image category classifier on NPU. The JSON reports this as `CPU rule fallback (NPU model not configured in prototype v1)`. A future task can add a static-shape MobileNet/EfficientNet/ResNet OpenVINO IR model.
|
||||
- OCR on NPU. OCR remains CPU/local plumbing in v1.
|
||||
|
||||
## Files
|
||||
|
||||
- `triage.py` — core library and CLI.
|
||||
- `server.py` — stdlib HTTP server with `/healthz`, `/models`, `/triage`, `/triage/batch`.
|
||||
- `make_samples.py` — creates synthetic non-private image/PDF samples.
|
||||
- `tests/smoke_test.py` — end-to-end smoke test, including NPU busy-time verification when `:18817` is reachable.
|
||||
- `samples/` — generated synthetic fixtures.
|
||||
|
||||
## Requirements
|
||||
|
||||
Use the existing NPU venv when available:
|
||||
|
||||
```bash
|
||||
cd /home/will/lab/swarm/openvino-doc-image-triage-npu
|
||||
/home/will/.venvs/npu/bin/python -m pip install pillow
|
||||
```
|
||||
|
||||
`pillow` is already present in the discovered `/home/will/.venvs/npu`. Optional local PDF text/rendering improves PDF support:
|
||||
|
||||
```bash
|
||||
/home/will/.venvs/npu/bin/python -m pip install pypdf pypdfium2
|
||||
```
|
||||
|
||||
The smoke tests do not require external services except the existing localhost `:18817` embeddings service for positive NPU verification.
|
||||
|
||||
## CLI usage
|
||||
|
||||
Generate synthetic samples:
|
||||
|
||||
```bash
|
||||
cd /home/will/lab/swarm/openvino-doc-image-triage-npu
|
||||
/home/will/.venvs/npu/bin/python make_samples.py
|
||||
```
|
||||
|
||||
Triage local files:
|
||||
|
||||
```bash
|
||||
/home/will/.venvs/npu/bin/python triage.py \
|
||||
--allowed-root /home/will/lab/swarm/openvino-doc-image-triage-npu \
|
||||
--pretty \
|
||||
samples/synthetic_invoice.png samples/synthetic_invoice.pdf
|
||||
```
|
||||
|
||||
Disable the local NPU embeddings call if needed:
|
||||
|
||||
```bash
|
||||
/home/will/.venvs/npu/bin/python triage.py --no-embeddings --allowed-root "$PWD" samples/synthetic_receipt.png
|
||||
```
|
||||
|
||||
Include OCR/sidecar text in a single response only when explicitly requested:
|
||||
|
||||
```bash
|
||||
/home/will/.venvs/npu/bin/python triage.py --include-ocr-text --allowed-root "$PWD" samples/synthetic_invoice.png
|
||||
```
|
||||
|
||||
## HTTP usage
|
||||
|
||||
Check that port 18820 is free first:
|
||||
|
||||
```bash
|
||||
ss -ltnp | grep ':18820\b' || true
|
||||
```
|
||||
|
||||
Start local-only server:
|
||||
|
||||
```bash
|
||||
cd /home/will/lab/swarm/openvino-doc-image-triage-npu
|
||||
/home/will/.venvs/npu/bin/python server.py --host 127.0.0.1 --port 18820 --allowed-root "$PWD"
|
||||
```
|
||||
|
||||
Call it:
|
||||
|
||||
```bash
|
||||
curl -sS http://127.0.0.1:18820/healthz | jq
|
||||
curl -sS http://127.0.0.1:18820/models | jq
|
||||
curl -sS -X POST http://127.0.0.1:18820/triage \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{"path":"/home/will/lab/swarm/openvino-doc-image-triage-npu/samples/synthetic_invoice.png","options":{"allowed_roots":["/home/will/lab/swarm/openvino-doc-image-triage-npu"]}}' | jq
|
||||
```
|
||||
|
||||
## Smoke test
|
||||
|
||||
```bash
|
||||
cd /home/will/lab/swarm/openvino-doc-image-triage-npu
|
||||
/home/will/.venvs/npu/bin/python tests/smoke_test.py
|
||||
```
|
||||
|
||||
Expected: JSON ending with `"ok": true`. If the embeddings service is up, the result should show positive NPU busy-time delta and each embedded page should report `verified_npu: true`.
|
||||
|
||||
## Example output shape
|
||||
|
||||
```json
|
||||
{
|
||||
"file_id": "sha256:...",
|
||||
"source_path_basename": "synthetic_invoice.png",
|
||||
"media_type": "image",
|
||||
"page_count": 1,
|
||||
"pages": [
|
||||
{
|
||||
"page_index": 0,
|
||||
"classification": {
|
||||
"label": "bill_or_invoice",
|
||||
"confidence": 0.71,
|
||||
"device": "CPU",
|
||||
"method": "rule_based_fallback"
|
||||
},
|
||||
"needs_attention": {
|
||||
"value": true,
|
||||
"device": "NPU+CPU",
|
||||
"reasons": ["amount_due", "due_date_present"],
|
||||
"embedding": {"verified_npu": true, "npu_busy_delta_us": 12345}
|
||||
},
|
||||
"metadata": {"dates_count": 1, "amounts_count": 1, "raw_values_redacted": true},
|
||||
"ocr": {"available": true, "device": "CPU"}
|
||||
}
|
||||
],
|
||||
"processing_device_summary": {
|
||||
"file_intake": "CPU",
|
||||
"image_category_classification": "CPU rule fallback (NPU model not configured in prototype v1)",
|
||||
"needs_attention_embedding": "NPU via local :18817",
|
||||
"metadata_extraction": "CPU",
|
||||
"npu_verified": true
|
||||
},
|
||||
"privacy": {"external_uploads": false, "raw_text_logged": false}
|
||||
}
|
||||
```
|
||||
@@ -0,0 +1,69 @@
|
||||
#!/usr/bin/env python3
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from PIL import Image, ImageDraw, ImageFilter
|
||||
|
||||
ROOT = Path(__file__).resolve().parent
|
||||
SAMPLES = ROOT / "samples"
|
||||
|
||||
|
||||
def make_doc(path: Path, lines: list[str], size=(900, 1200), rotate: int = 0, blur: bool = False) -> None:
|
||||
img = Image.new("RGB", size, "white")
|
||||
draw = ImageDraw.Draw(img)
|
||||
y = 70
|
||||
for line in lines:
|
||||
draw.text((70, y), line, fill="black")
|
||||
y += 55
|
||||
draw.rectangle((55, 50, size[0] - 55, min(size[1] - 50, y + 30)), outline="gray", width=3)
|
||||
if blur:
|
||||
img = img.filter(ImageFilter.GaussianBlur(2.5))
|
||||
if rotate:
|
||||
img = img.rotate(rotate, expand=True, fillcolor="white")
|
||||
img.save(path)
|
||||
path.with_suffix(path.suffix + ".txt").write_text("\n".join(lines) + "\n")
|
||||
|
||||
|
||||
def main() -> int:
|
||||
SAMPLES.mkdir(exist_ok=True)
|
||||
make_doc(SAMPLES / "synthetic_invoice.png", [
|
||||
"ACME Utilities Invoice",
|
||||
"Invoice No: INV-2026-0604",
|
||||
"Amount Due: $123.45",
|
||||
"Payment due 2026-06-30",
|
||||
"Please submit payment by the due date.",
|
||||
])
|
||||
make_doc(SAMPLES / "synthetic_receipt.png", [
|
||||
"Neighborhood Store Receipt",
|
||||
"Subtotal $14.20",
|
||||
"Tax $1.42",
|
||||
"Total $15.62",
|
||||
"Thank you for shopping",
|
||||
], size=(720, 1100), rotate=3)
|
||||
make_doc(SAMPLES / "synthetic_conversation.png", [
|
||||
"Messages with Alex",
|
||||
"Can you please respond by tomorrow?",
|
||||
"Need signature on the form before Friday.",
|
||||
], size=(1200, 750))
|
||||
make_doc(SAMPLES / "synthetic_sensitive_form.png", [
|
||||
"Sample Government Form - Fake Data",
|
||||
"Applicant: Test Person",
|
||||
"SSN: 123-45-6789",
|
||||
"Signature required",
|
||||
"Submit by Jan 15, 2027",
|
||||
], blur=False)
|
||||
make_doc(SAMPLES / "synthetic_blurry.png", [
|
||||
"Low resolution blurred sample",
|
||||
"No action required",
|
||||
], size=(360, 250), blur=True)
|
||||
# PIL can save a simple local PDF from a synthetic page. This is non-private.
|
||||
pdf_img = Image.open(SAMPLES / "synthetic_invoice.png").convert("RGB")
|
||||
pdf_img.save(SAMPLES / "synthetic_invoice.pdf", "PDF")
|
||||
(SAMPLES / "synthetic_invoice.pdf.txt").write_text((SAMPLES / "synthetic_invoice.png.txt").read_text())
|
||||
print(f"wrote samples under {SAMPLES}")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 4.5 KiB |
@@ -0,0 +1,2 @@
|
||||
Low resolution blurred sample
|
||||
No action required
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 9.1 KiB |
@@ -0,0 +1,3 @@
|
||||
Messages with Alex
|
||||
Can you please respond by tomorrow?
|
||||
Need signature on the form before Friday.
|
||||
Binary file not shown.
@@ -0,0 +1,5 @@
|
||||
ACME Utilities Invoice
|
||||
Invoice No: INV-2026-0604
|
||||
Amount Due: $123.45
|
||||
Payment due 2026-06-30
|
||||
Please submit payment by the due date.
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 13 KiB |
@@ -0,0 +1,5 @@
|
||||
ACME Utilities Invoice
|
||||
Invoice No: INV-2026-0604
|
||||
Amount Due: $123.45
|
||||
Payment due 2026-06-30
|
||||
Please submit payment by the due date.
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 12 KiB |
@@ -0,0 +1,5 @@
|
||||
Neighborhood Store Receipt
|
||||
Subtotal $14.20
|
||||
Tax $1.42
|
||||
Total $15.62
|
||||
Thank you for shopping
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 12 KiB |
@@ -0,0 +1,5 @@
|
||||
Sample Government Form - Fake Data
|
||||
Applicant: Test Person
|
||||
SSN: 123-45-6789
|
||||
Signature required
|
||||
Submit by Jan 15, 2027
|
||||
@@ -0,0 +1,178 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Stdlib localhost HTTP wrapper for the triage prototype.
|
||||
|
||||
Endpoints:
|
||||
- GET /healthz
|
||||
- GET /models
|
||||
- POST /triage JSON: {"path":"/local/file", "options": {...}}
|
||||
- POST /triage/batch JSON: {"paths":["/local/file"], "options": {...}}
|
||||
|
||||
The server binds to 127.0.0.1 by default and accepts only local file paths under
|
||||
configured allowed roots. It never uploads document/image contents externally.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from triage import DEFAULT_EMBED_URL, TriageOptions, read_npu_busy, triage_batch, triage_file
|
||||
|
||||
|
||||
def _roots_within_configured(requested_roots: list[Any], configured_roots: list[Path]) -> list[Path]:
|
||||
"""Return request roots only when they narrow the startup allowlist."""
|
||||
narrowed: list[Path] = []
|
||||
configured = [root.expanduser().resolve() for root in configured_roots]
|
||||
for raw in requested_roots:
|
||||
candidate = Path(str(raw)).expanduser().resolve()
|
||||
if any(candidate == root or candidate.is_relative_to(root) for root in configured):
|
||||
narrowed.append(candidate)
|
||||
else:
|
||||
raise ValueError("requested allowed_roots must be within configured allowed roots")
|
||||
return narrowed
|
||||
|
||||
|
||||
def _validated_embedding_url(raw_url: Any) -> str:
|
||||
"""Allow only the configured local loopback embeddings service."""
|
||||
url = str(raw_url)
|
||||
parsed = urlparse(url)
|
||||
host = parsed.hostname or ""
|
||||
if (
|
||||
parsed.scheme == "http"
|
||||
and host in {"127.0.0.1", "localhost", "::1"}
|
||||
and (parsed.port or 80) == 18817
|
||||
and parsed.path == "/v1/embeddings"
|
||||
and not parsed.username
|
||||
and not parsed.password
|
||||
):
|
||||
return url
|
||||
raise ValueError("embedding_url override must target the configured local loopback embeddings service")
|
||||
|
||||
|
||||
def make_options(payload: dict[str, Any], default_roots: list[Path]) -> TriageOptions:
|
||||
opts = payload.get("options") or {}
|
||||
requested_roots = opts.get("allowed_roots", [])
|
||||
if requested_roots:
|
||||
if not isinstance(requested_roots, list):
|
||||
raise ValueError("allowed_roots must be a list")
|
||||
roots = _roots_within_configured(requested_roots, default_roots)
|
||||
else:
|
||||
roots = default_roots
|
||||
embedding_url = DEFAULT_EMBED_URL
|
||||
if "embedding_url" in opts:
|
||||
embedding_url = _validated_embedding_url(opts["embedding_url"])
|
||||
return TriageOptions(
|
||||
max_pages=int(opts.get("max_pages", 3)),
|
||||
include_ocr_text=bool(opts.get("include_ocr_text", False)),
|
||||
dry_run=bool(opts.get("dry_run", False)),
|
||||
use_embeddings=bool(opts.get("use_embeddings", True)),
|
||||
embedding_url=embedding_url,
|
||||
allowed_roots=roots,
|
||||
include_full_path=bool(opts.get("include_full_path", False)),
|
||||
)
|
||||
|
||||
|
||||
class Handler(BaseHTTPRequestHandler):
|
||||
server_version = "openvino-doc-image-triage-npu/0.1"
|
||||
|
||||
def _json(self, status: int, body: dict[str, Any]) -> None:
|
||||
data = json.dumps(body, sort_keys=True).encode()
|
||||
self.send_response(status)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(data)))
|
||||
self.end_headers()
|
||||
self.wfile.write(data)
|
||||
|
||||
def log_message(self, format: str, *args: Any) -> None:
|
||||
# Do not log request bodies, OCR text, or file paths.
|
||||
return
|
||||
|
||||
@property
|
||||
def allowed_roots(self) -> list[Path]:
|
||||
return self.server.allowed_roots # type: ignore[attr-defined]
|
||||
|
||||
def do_GET(self) -> None: # noqa: N802
|
||||
if self.path in ("/", "/healthz", "/health"):
|
||||
self._json(200, {
|
||||
"ok": True,
|
||||
"service": "openvino-doc-image-triage-npu",
|
||||
"bind_policy": "localhost-default",
|
||||
"npu_busy_time_us": read_npu_busy(),
|
||||
"npu_busy_check_enabled": True,
|
||||
"allowed_roots": [str(p) for p in self.allowed_roots],
|
||||
"privacy": {"external_uploads": False, "raw_text_logged": False},
|
||||
})
|
||||
return
|
||||
if self.path == "/models":
|
||||
self._json(200, {
|
||||
"models": [
|
||||
{
|
||||
"stage": "needs_attention_embedding",
|
||||
"model": "bge-base-en-v1.5-int8-ov via local :18817",
|
||||
"target_device": "NPU",
|
||||
"verification": "sysfs npu_busy_time_us before/after embedding call",
|
||||
},
|
||||
{
|
||||
"stage": "image_category_classification",
|
||||
"model": "rule-based fallback in prototype v1",
|
||||
"target_device": "CPU",
|
||||
"npu_status": "not configured; future static-shape MobileNet/EfficientNet/ResNet OV IR",
|
||||
},
|
||||
{"stage": "ocr_text_extraction", "model": "optional local sidecar/PDF text", "target_device": "CPU"},
|
||||
]
|
||||
})
|
||||
return
|
||||
self._json(404, {"ok": False, "error": "not_found"})
|
||||
|
||||
def _read_payload(self) -> dict[str, Any]:
|
||||
length = int(self.headers.get("Content-Length", "0"))
|
||||
if length > 512 * 1024:
|
||||
raise ValueError("request JSON too large")
|
||||
raw = self.rfile.read(length)
|
||||
if not raw:
|
||||
return {}
|
||||
return json.loads(raw.decode())
|
||||
|
||||
def do_POST(self) -> None: # noqa: N802
|
||||
try:
|
||||
payload = self._read_payload()
|
||||
options = make_options(payload, self.allowed_roots)
|
||||
if self.path == "/triage":
|
||||
path = payload.get("path")
|
||||
if not path:
|
||||
self._json(400, {"ok": False, "error": "missing_path"})
|
||||
return
|
||||
self._json(200, {"ok": True, "result": triage_file(path, options)})
|
||||
return
|
||||
if self.path == "/triage/batch":
|
||||
paths = payload.get("paths") or []
|
||||
if not isinstance(paths, list) or not paths:
|
||||
self._json(400, {"ok": False, "error": "missing_paths"})
|
||||
return
|
||||
self._json(200, triage_batch([str(p) for p in paths], options))
|
||||
return
|
||||
self._json(404, {"ok": False, "error": "not_found"})
|
||||
except Exception as exc:
|
||||
self._json(400, {"ok": False, "error": type(exc).__name__, "message": str(exc)})
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description="Local-only doc/image triage HTTP server")
|
||||
parser.add_argument("--host", default=os.environ.get("DOC_IMAGE_TRIAGE_HOST", "127.0.0.1"))
|
||||
parser.add_argument("--port", type=int, default=int(os.environ.get("DOC_IMAGE_TRIAGE_PORT", "18820")))
|
||||
parser.add_argument("--allowed-root", action="append", default=[], help="allowed local root; may repeat")
|
||||
args = parser.parse_args()
|
||||
roots = [Path(p).expanduser().resolve() for p in args.allowed_root] or [Path.cwd().resolve()]
|
||||
httpd = ThreadingHTTPServer((args.host, args.port), Handler)
|
||||
httpd.allowed_roots = roots # type: ignore[attr-defined]
|
||||
print(json.dumps({"service": "openvino-doc-image-triage-npu", "host": args.host, "port": args.port, "allowed_roots": [str(p) for p in roots]}), flush=True)
|
||||
httpd.serve_forever()
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,127 @@
|
||||
#!/usr/bin/env python3
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
SAMPLES = ROOT / "samples"
|
||||
BUSY = Path("/sys/class/accel/accel0/device/npu_busy_time_us")
|
||||
|
||||
|
||||
def run(cmd: list[str]) -> None:
|
||||
print("+", " ".join(cmd))
|
||||
subprocess.run(cmd, cwd=ROOT, check=True)
|
||||
|
||||
|
||||
def post_json(url: str, payload: dict) -> dict:
|
||||
req = urllib.request.Request(url, data=json.dumps(payload).encode(), headers={"Content-Type": "application/json"})
|
||||
with urllib.request.urlopen(req, timeout=10) as resp:
|
||||
return json.loads(resp.read().decode())
|
||||
|
||||
|
||||
def post_json_status(url: str, payload: dict) -> tuple[int, dict]:
|
||||
req = urllib.request.Request(url, data=json.dumps(payload).encode(), headers={"Content-Type": "application/json"})
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=10) as resp:
|
||||
return resp.status, json.loads(resp.read().decode())
|
||||
except urllib.error.HTTPError as exc:
|
||||
return exc.code, json.loads(exc.read().decode())
|
||||
|
||||
|
||||
def busy() -> int | None:
|
||||
try:
|
||||
return int(BUSY.read_text().strip())
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def main() -> int:
|
||||
run([sys.executable, "make_samples.py"])
|
||||
invoice = SAMPLES / "synthetic_invoice.png"
|
||||
pdf = SAMPLES / "synthetic_invoice.pdf"
|
||||
|
||||
before = busy()
|
||||
raw = subprocess.check_output([
|
||||
sys.executable, "triage.py", "--allowed-root", str(ROOT), "--pretty", str(invoice), str(pdf)
|
||||
], cwd=ROOT, text=True)
|
||||
data = json.loads(raw)
|
||||
assert data["ok"], data
|
||||
first = data["files"][0]["result"]
|
||||
assert first["privacy"]["external_uploads"] is False
|
||||
assert first["pages"][0]["classification"]["label"] == "bill_or_invoice"
|
||||
assert first["pages"][0]["needs_attention"]["value"] is True
|
||||
assert "amount_due" in first["pages"][0]["needs_attention"]["reasons"]
|
||||
assert first["processing_device_summary"]["file_intake"] == "CPU"
|
||||
assert "NPU" in first["processing_device_summary"]["needs_attention_embedding"] or first["pages"][0]["needs_attention"]["device"] == "CPU"
|
||||
after = busy()
|
||||
if before is not None and after is not None:
|
||||
# If :18817 is reachable and text was embedded, NPU delta must be positive.
|
||||
emb = first["pages"][0]["needs_attention"]["embedding"]
|
||||
if emb.get("used"):
|
||||
assert emb.get("verified_npu") is True, emb
|
||||
assert (emb.get("npu_busy_delta_us") or 0) > 0, emb
|
||||
assert after > before, {"before": before, "after": after, "embedding": emb}
|
||||
|
||||
# HTTP smoke on an ephemeral localhost port so we do not collide with 18820 during tests.
|
||||
proc = subprocess.Popen([sys.executable, "server.py", "--host", "127.0.0.1", "--port", "18828", "--allowed-root", str(ROOT)], cwd=ROOT, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
|
||||
try:
|
||||
deadline = time.time() + 5
|
||||
while time.time() < deadline:
|
||||
try:
|
||||
health = urllib.request.urlopen("http://127.0.0.1:18828/healthz", timeout=1).read()
|
||||
assert b"openvino-doc-image-triage-npu" in health
|
||||
break
|
||||
except Exception:
|
||||
time.sleep(0.1)
|
||||
else:
|
||||
raise AssertionError("server did not become ready")
|
||||
resp = post_json("http://127.0.0.1:18828/triage", {"path": str(invoice), "options": {"allowed_roots": [str(ROOT)]}})
|
||||
assert resp["ok"] is True, resp
|
||||
assert resp["result"]["source_path_basename"] == "synthetic_invoice.png"
|
||||
assert "source_path" not in resp["result"]
|
||||
|
||||
# Request bodies may narrow but must not widen the startup --allowed-root policy.
|
||||
with tempfile.NamedTemporaryFile(suffix=".txt") as outside:
|
||||
outside.write(b"sensitive text outside configured artifact root")
|
||||
outside.flush()
|
||||
status, blocked = post_json_status(
|
||||
"http://127.0.0.1:18828/triage",
|
||||
{"path": outside.name, "options": {"allowed_roots": ["/tmp"], "dry_run": True, "use_embeddings": False}},
|
||||
)
|
||||
assert status == 400, blocked
|
||||
assert blocked["ok"] is False, blocked
|
||||
assert "allowed_roots" in blocked.get("message", ""), blocked
|
||||
|
||||
# Request bodies must not redirect extracted text to caller-supplied endpoints.
|
||||
status, blocked = post_json_status(
|
||||
"http://127.0.0.1:18828/triage",
|
||||
{"path": str(invoice), "options": {"embedding_url": "http://198.51.100.1:9/v1/embeddings"}},
|
||||
)
|
||||
assert status == 400, blocked
|
||||
assert blocked["ok"] is False, blocked
|
||||
assert "embedding_url" in blocked.get("message", ""), blocked
|
||||
finally:
|
||||
proc.terminate()
|
||||
proc.wait(timeout=5)
|
||||
|
||||
print(json.dumps({
|
||||
"ok": True,
|
||||
"samples": len(list(SAMPLES.glob("synthetic_*"))),
|
||||
"npu_busy_before": before,
|
||||
"npu_busy_after": after,
|
||||
"npu_delta_observed": None if before is None or after is None else after - before,
|
||||
"triage_label": first["pages"][0]["classification"]["label"],
|
||||
"needs_attention": first["pages"][0]["needs_attention"]["value"],
|
||||
}, indent=2))
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,459 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Local-only document/image triage prototype.
|
||||
|
||||
CPU stages:
|
||||
- local file intake, hashing, MIME/extension checks
|
||||
- image/PDF-page decoding and normalization
|
||||
- optional sidecar/native-text extraction
|
||||
- regex metadata extraction and rule-based category fallback
|
||||
|
||||
NPU stages:
|
||||
- needs-attention semantic embedding via the existing local OpenVINO NPU
|
||||
embeddings service on 127.0.0.1:18817, verified by sysfs busy-time delta.
|
||||
|
||||
No external uploads are performed. The only network call is localhost to the
|
||||
embedding service when enabled.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import dataclasses
|
||||
import datetime as dt
|
||||
import hashlib
|
||||
import io
|
||||
import json
|
||||
import mimetypes
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
try:
|
||||
from PIL import Image, ImageOps
|
||||
except Exception as exc: # pragma: no cover - caught in CLI smoke
|
||||
raise SystemExit("Pillow is required: install pillow in the active Python env") from exc
|
||||
|
||||
NPU_BUSY_PATH = Path("/sys/class/accel/accel0/device/npu_busy_time_us")
|
||||
DEFAULT_EMBED_URL = "http://127.0.0.1:18817/v1/embeddings"
|
||||
DEFAULT_ALLOWED_ROOTS = [Path.cwd()]
|
||||
MAX_FILE_BYTES = 25 * 1024 * 1024
|
||||
CATEGORY_LABELS = [
|
||||
"receipt",
|
||||
"bill_or_invoice",
|
||||
"tax_or_financial",
|
||||
"medical_or_insurance",
|
||||
"legal_or_government",
|
||||
"form_or_application",
|
||||
"travel_or_ticket",
|
||||
"screenshot_conversation",
|
||||
"screenshot_web_or_app",
|
||||
"identity_or_sensitive",
|
||||
"photo_misc",
|
||||
"unknown_or_low_confidence",
|
||||
]
|
||||
|
||||
DATE_PATTERNS = [
|
||||
re.compile(r"\b(20\d{2}[-/](?:0?[1-9]|1[0-2])[-/](?:0?[1-9]|[12]\d|3[01]))\b"),
|
||||
re.compile(r"\b((?:0?[1-9]|1[0-2])[-/](?:0?[1-9]|[12]\d|3[01])[-/](?:20)?\d{2})\b"),
|
||||
re.compile(r"\b((?:Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)[a-z]*\s+\d{1,2},?\s+20\d{2})\b", re.I),
|
||||
]
|
||||
AMOUNT_RE = re.compile(r"(?<!\w)(?:USD\s*)?\$\s?\d{1,4}(?:,\d{3})*(?:\.\d{2})?\b", re.I)
|
||||
EMAIL_RE = re.compile(r"\b[\w.+-]+@[\w.-]+\.[A-Za-z]{2,}\b")
|
||||
PHONE_RE = re.compile(r"\b(?:\+?1[-.\s]?)?(?:\(?\d{3}\)?[-.\s]?){2}\d{4}\b")
|
||||
ACCOUNT_RE = re.compile(r"\b(?:account|acct|policy|invoice|member|case|claim)\s*(?:#|no\.?|id)?\s*[:\-]?\s*[A-Z0-9-]{4,}\b", re.I)
|
||||
SSN_LIKE_RE = re.compile(r"\b\d{3}-\d{2}-\d{4}\b")
|
||||
|
||||
ATTENTION_KEYWORDS = {
|
||||
"due_date_present": ["due", "payment due", "pay by", "deadline"],
|
||||
"amount_due": ["amount due", "balance due", "total due", "$"],
|
||||
"action_required_language": ["action required", "please respond", "complete", "submit", "renew", "verify"],
|
||||
"signature_required": ["signature", "sign and return", "signed"],
|
||||
"appointment_or_deadline": ["appointment", "scheduled", "reservation", "hearing"],
|
||||
"account_security": ["security", "password", "unauthorized", "fraud", "verify your account"],
|
||||
"medical_followup": ["follow up", "lab result", "referral", "insurance"],
|
||||
"tax_deadline": ["irs", "tax", "1099", "w-2", "deadline"],
|
||||
}
|
||||
|
||||
CATEGORY_KEYWORDS = {
|
||||
"receipt": ["receipt", "subtotal", "cashier", "change", "store"],
|
||||
"bill_or_invoice": ["invoice", "amount due", "balance due", "statement", "payment due"],
|
||||
"tax_or_financial": ["tax", "irs", "1099", "w-2", "bank", "routing"],
|
||||
"medical_or_insurance": ["medical", "insurance", "clinic", "patient", "claim"],
|
||||
"legal_or_government": ["court", "government", "department", "notice", "license"],
|
||||
"form_or_application": ["application", "form", "signature", "submit"],
|
||||
"travel_or_ticket": ["boarding", "ticket", "itinerary", "reservation", "gate"],
|
||||
"screenshot_conversation": ["message", "chat", "reply", "conversation"],
|
||||
"screenshot_web_or_app": ["login", "browser", "app", "settings", "dashboard"],
|
||||
"identity_or_sensitive": ["ssn", "passport", "driver license", "social security"],
|
||||
}
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class TriageOptions:
|
||||
max_pages: int = 3
|
||||
include_ocr_text: bool = False
|
||||
dry_run: bool = False
|
||||
use_embeddings: bool = True
|
||||
embedding_url: str = DEFAULT_EMBED_URL
|
||||
allowed_roots: list[Path] = dataclasses.field(default_factory=lambda: DEFAULT_ALLOWED_ROOTS.copy())
|
||||
include_full_path: bool = False
|
||||
timeout_seconds: float = 10.0
|
||||
|
||||
|
||||
def read_npu_busy() -> int | None:
|
||||
try:
|
||||
return int(NPU_BUSY_PATH.read_text().strip())
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def sha256_file(path: Path) -> str:
|
||||
h = hashlib.sha256()
|
||||
with path.open("rb") as f:
|
||||
for chunk in iter(lambda: f.read(1024 * 1024), b""):
|
||||
h.update(chunk)
|
||||
return h.hexdigest()
|
||||
|
||||
|
||||
def under_allowed_root(path: Path, roots: list[Path]) -> bool:
|
||||
resolved = path.resolve()
|
||||
for root in roots:
|
||||
try:
|
||||
resolved.relative_to(root.resolve())
|
||||
return True
|
||||
except ValueError:
|
||||
continue
|
||||
return False
|
||||
|
||||
|
||||
def sidecar_text(path: Path) -> tuple[str, str | None]:
|
||||
for suffix in (path.suffix + ".txt", ".txt"):
|
||||
candidate = path.with_suffix(suffix) if suffix.startswith(path.suffix) else path.with_suffix(suffix)
|
||||
if candidate.exists() and candidate.is_file():
|
||||
try:
|
||||
return candidate.read_text(errors="replace")[:12000], f"sidecar:{candidate.name}"
|
||||
except Exception:
|
||||
return "", "sidecar_unreadable"
|
||||
return "", None
|
||||
|
||||
|
||||
def extract_pdf_text(path: Path, max_pages: int) -> tuple[str, str | None]:
|
||||
# Optional dependency; tests do not require it. Keeps PDF support local-only when installed.
|
||||
try:
|
||||
import pypdf # type: ignore
|
||||
except Exception:
|
||||
return "", "pypdf_not_installed"
|
||||
try:
|
||||
reader = pypdf.PdfReader(str(path))
|
||||
if getattr(reader, "is_encrypted", False):
|
||||
return "", "pdf_encrypted"
|
||||
chunks = []
|
||||
for page in reader.pages[:max_pages]:
|
||||
chunks.append(page.extract_text() or "")
|
||||
return "\n".join(chunks)[:12000], "pypdf_cpu"
|
||||
except Exception as exc:
|
||||
return "", f"pdf_text_error:{type(exc).__name__}"
|
||||
|
||||
|
||||
def load_image_pages(path: Path, max_pages: int) -> tuple[list[Image.Image], str | None]:
|
||||
ext = path.suffix.lower()
|
||||
if ext == ".pdf":
|
||||
try:
|
||||
import pypdfium2 as pdfium # type: ignore
|
||||
except Exception:
|
||||
return [], "pypdfium2_not_installed"
|
||||
try:
|
||||
pdf = pdfium.PdfDocument(str(path))
|
||||
pages = []
|
||||
for i in range(min(len(pdf), max_pages)):
|
||||
bitmap = pdf[i].render(scale=1.5)
|
||||
pages.append(bitmap.to_pil().convert("RGB"))
|
||||
return pages, None
|
||||
except Exception as exc:
|
||||
return [], f"pdf_render_error:{type(exc).__name__}"
|
||||
try:
|
||||
img = Image.open(path)
|
||||
img = ImageOps.exif_transpose(img).convert("RGB")
|
||||
return [img], None
|
||||
except Exception as exc:
|
||||
return [], f"image_decode_error:{type(exc).__name__}"
|
||||
|
||||
|
||||
def normalize_for_hash_features(img: Image.Image) -> dict[str, Any]:
|
||||
small = ImageOps.contain(img.copy(), (224, 224))
|
||||
gray = small.convert("L")
|
||||
hist = gray.histogram()
|
||||
pixels = max(1, gray.width * gray.height)
|
||||
mean = sum(i * c for i, c in enumerate(hist)) / pixels
|
||||
variance = sum(((i - mean) ** 2) * c for i, c in enumerate(hist)) / pixels
|
||||
return {
|
||||
"mean_luma": round(mean, 2),
|
||||
"contrast": round(variance ** 0.5, 2),
|
||||
"aspect_ratio": round(img.width / max(1, img.height), 3),
|
||||
}
|
||||
|
||||
|
||||
def classify_rule(text: str, image_features: dict[str, Any]) -> dict[str, Any]:
|
||||
t = text.lower()
|
||||
best_label = "unknown_or_low_confidence"
|
||||
best_score = 0
|
||||
for label, words in CATEGORY_KEYWORDS.items():
|
||||
score = sum(1 for word in words if word in t)
|
||||
if score > best_score:
|
||||
best_label, best_score = label, score
|
||||
if best_score == 0:
|
||||
ar = image_features.get("aspect_ratio", 1.0)
|
||||
if ar > 1.3:
|
||||
best_label, best_score = "screenshot_web_or_app", 1
|
||||
else:
|
||||
best_label, best_score = "unknown_or_low_confidence", 0
|
||||
confidence = min(0.35 + 0.18 * best_score, 0.92) if best_score else 0.2
|
||||
if confidence < 0.45:
|
||||
best_label = "unknown_or_low_confidence"
|
||||
return {
|
||||
"label": best_label,
|
||||
"confidence": round(confidence, 3),
|
||||
"device": "CPU",
|
||||
"stage": "category_classification",
|
||||
"method": "rule_based_fallback",
|
||||
"npu_status": "not_configured_for_prototype_v1",
|
||||
"candidate_labels": CATEGORY_LABELS,
|
||||
}
|
||||
|
||||
|
||||
def extract_metadata(text: str) -> dict[str, Any]:
|
||||
dates = []
|
||||
for pat in DATE_PATTERNS:
|
||||
dates.extend(m.group(1) for m in pat.finditer(text))
|
||||
amounts = AMOUNT_RE.findall(text)
|
||||
flags = {
|
||||
"org_present": bool(re.search(r"\b(?:inc|llc|clinic|department|bank|insurance|store)\b", text, re.I)),
|
||||
"address_present": bool(re.search(r"\b\d{2,5}\s+[A-Za-z0-9 .]+\s+(?:st|street|ave|avenue|rd|road|blvd|drive|dr)\b", text, re.I)),
|
||||
"phone_present": bool(PHONE_RE.search(text)),
|
||||
"email_present": bool(EMAIL_RE.search(text)),
|
||||
"policy_or_account_id_present": bool(ACCOUNT_RE.search(text)),
|
||||
"identity_number_like_present": bool(SSN_LIKE_RE.search(text)),
|
||||
}
|
||||
return {
|
||||
"dates_count": len(set(dates)),
|
||||
"amounts_count": len(set(amounts)),
|
||||
"detected_entities": flags,
|
||||
"raw_values_redacted": True,
|
||||
}
|
||||
|
||||
|
||||
def call_embeddings(text: str, url: str, timeout: float) -> dict[str, Any]:
|
||||
if not text.strip():
|
||||
return {"used": False, "device": "NPU", "status": "skipped_no_text", "npu_busy_delta_us": 0}
|
||||
before = read_npu_busy()
|
||||
payload = json.dumps({"input": text[:2048], "purpose": "document"}).encode()
|
||||
req = urllib.request.Request(url, data=payload, headers={"Content-Type": "application/json"})
|
||||
t0 = time.perf_counter()
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=timeout) as resp:
|
||||
body = resp.read(1024 * 1024)
|
||||
status = resp.status
|
||||
parsed = json.loads(body.decode())
|
||||
dim = None
|
||||
if isinstance(parsed, dict) and parsed.get("data"):
|
||||
emb = parsed["data"][0].get("embedding", [])
|
||||
dim = len(emb) if isinstance(emb, list) else None
|
||||
after = read_npu_busy()
|
||||
delta = (after - before) if before is not None and after is not None else None
|
||||
return {
|
||||
"used": True,
|
||||
"device": "NPU",
|
||||
"status": "ok" if status == 200 else f"http_{status}",
|
||||
"embedding_dim": dim,
|
||||
"wall_ms": round((time.perf_counter() - t0) * 1000, 2),
|
||||
"npu_busy_delta_us": delta,
|
||||
"verified_npu": bool(delta and delta > 0),
|
||||
"endpoint": "127.0.0.1:18817",
|
||||
}
|
||||
except (urllib.error.URLError, TimeoutError, json.JSONDecodeError) as exc:
|
||||
after = read_npu_busy()
|
||||
delta = (after - before) if before is not None and after is not None else None
|
||||
return {
|
||||
"used": False,
|
||||
"device": "NPU",
|
||||
"status": f"embedding_service_error:{type(exc).__name__}",
|
||||
"npu_busy_delta_us": delta,
|
||||
"verified_npu": False,
|
||||
"endpoint": "127.0.0.1:18817",
|
||||
}
|
||||
|
||||
|
||||
def needs_attention(text: str, embedding_result: dict[str, Any]) -> dict[str, Any]:
|
||||
t = text.lower()
|
||||
reasons = []
|
||||
for reason, words in ATTENTION_KEYWORDS.items():
|
||||
if any(word in t for word in words):
|
||||
reasons.append(reason)
|
||||
meta = extract_metadata(text)
|
||||
if meta["amounts_count"]:
|
||||
reasons.append("amount_due")
|
||||
if meta["dates_count"]:
|
||||
reasons.append("due_date_present")
|
||||
reasons = sorted(set(reasons))
|
||||
value = bool(reasons)
|
||||
confidence = min(0.45 + 0.1 * len(reasons), 0.9) if value else 0.35
|
||||
if embedding_result.get("verified_npu"):
|
||||
confidence = min(confidence + 0.05, 0.95)
|
||||
return {
|
||||
"value": value,
|
||||
"confidence": round(confidence, 3),
|
||||
"reasons": reasons or (["low_confidence"] if not text.strip() else []),
|
||||
"device": "NPU+CPU" if embedding_result.get("used") else "CPU",
|
||||
"stage": "needs_attention",
|
||||
"method": "NPU embedding verification + CPU rules" if embedding_result.get("used") else "CPU rules fallback",
|
||||
"embedding": embedding_result,
|
||||
}
|
||||
|
||||
|
||||
def infer_media_type(path: Path, is_pdf_page: bool = False) -> str:
|
||||
if is_pdf_page:
|
||||
return "pdf_page"
|
||||
mt, _ = mimetypes.guess_type(path.name)
|
||||
if path.suffix.lower() == ".pdf":
|
||||
return "pdf"
|
||||
if mt and mt.startswith("image/"):
|
||||
return "image"
|
||||
return "unknown"
|
||||
|
||||
|
||||
def triage_file(path_like: str | Path, options: TriageOptions | None = None) -> dict[str, Any]:
|
||||
options = options or TriageOptions()
|
||||
path = Path(path_like).expanduser()
|
||||
resolved = path.resolve()
|
||||
if not under_allowed_root(resolved, options.allowed_roots):
|
||||
raise ValueError(f"path is outside allowed roots: {path}")
|
||||
if not resolved.exists() or not resolved.is_file():
|
||||
raise FileNotFoundError(str(path))
|
||||
size = resolved.stat().st_size
|
||||
if size > MAX_FILE_BYTES:
|
||||
raise ValueError(f"file too large for prototype limit: {size} bytes")
|
||||
|
||||
file_hash = sha256_file(resolved)
|
||||
text, text_source = sidecar_text(resolved)
|
||||
pdf_text_status = None
|
||||
if resolved.suffix.lower() == ".pdf" and not text:
|
||||
text, pdf_text_status = extract_pdf_text(resolved, options.max_pages)
|
||||
text_source = pdf_text_status
|
||||
|
||||
pages: list[dict[str, Any]] = []
|
||||
render_error = None
|
||||
if not options.dry_run:
|
||||
images, render_error = load_image_pages(resolved, options.max_pages)
|
||||
else:
|
||||
images = []
|
||||
|
||||
if not images and options.dry_run:
|
||||
images = []
|
||||
elif not images:
|
||||
# Return a file-level record even if PDF rendering is unavailable.
|
||||
images = []
|
||||
|
||||
embedding_result = call_embeddings(text, options.embedding_url, options.timeout_seconds) if options.use_embeddings else {"used": False, "device": "NPU", "status": "disabled", "npu_busy_delta_us": 0, "verified_npu": False}
|
||||
attn = needs_attention(text, embedding_result)
|
||||
meta = extract_metadata(text)
|
||||
|
||||
if images:
|
||||
for idx, img in enumerate(images):
|
||||
features = normalize_for_hash_features(img)
|
||||
classification = classify_rule(text, features)
|
||||
pages.append({
|
||||
"page_index": idx,
|
||||
"media_type": infer_media_type(resolved, resolved.suffix.lower() == ".pdf"),
|
||||
"image": {"width": img.width, "height": img.height, "orientation": "portrait" if img.height >= img.width else "landscape", **features},
|
||||
"classification": classification,
|
||||
"needs_attention": attn,
|
||||
"metadata": meta,
|
||||
"ocr": {"available": bool(text), "quality": 0.7 if text else 0.0, "device": "CPU", "text_source": text_source},
|
||||
})
|
||||
else:
|
||||
classification = classify_rule(text, {"aspect_ratio": 1.0})
|
||||
pages.append({
|
||||
"page_index": 0,
|
||||
"media_type": infer_media_type(resolved, resolved.suffix.lower() == ".pdf"),
|
||||
"image": {"width": None, "height": None, "orientation": None, "render_error": render_error},
|
||||
"classification": classification,
|
||||
"needs_attention": attn,
|
||||
"metadata": meta,
|
||||
"ocr": {"available": bool(text), "quality": 0.7 if text else 0.0, "device": "CPU", "text_source": text_source},
|
||||
})
|
||||
|
||||
result: dict[str, Any] = {
|
||||
"file_id": f"sha256:{file_hash}",
|
||||
"source_path_basename": resolved.name,
|
||||
"media_type": infer_media_type(resolved),
|
||||
"file_size_bytes": size,
|
||||
"page_count": len(pages),
|
||||
"pages": pages,
|
||||
"processing_device_summary": {
|
||||
"file_intake": "CPU",
|
||||
"pdf_rendering": "CPU" if resolved.suffix.lower() == ".pdf" else "not_applicable",
|
||||
"image_category_classification": "CPU rule fallback (NPU model not configured in prototype v1)",
|
||||
"ocr_text_extraction": "CPU/local sidecar or optional local PDF text extractor",
|
||||
"needs_attention_embedding": "NPU via local :18817" if embedding_result.get("used") else "CPU fallback/no text",
|
||||
"metadata_extraction": "CPU",
|
||||
"npu_verified": bool(embedding_result.get("verified_npu")),
|
||||
"npu_busy_delta_us": embedding_result.get("npu_busy_delta_us"),
|
||||
},
|
||||
"privacy": {
|
||||
"external_uploads": False,
|
||||
"localhost_only_embedding_call": bool(options.use_embeddings),
|
||||
"raw_text_logged": False,
|
||||
"raw_values_redacted": True,
|
||||
"full_path_included": options.include_full_path,
|
||||
},
|
||||
"errors": [e for e in [render_error, pdf_text_status if pdf_text_status and not text else None] if e],
|
||||
}
|
||||
if options.include_full_path:
|
||||
result["source_path"] = str(resolved)
|
||||
if options.include_ocr_text:
|
||||
result["ocr_text"] = text
|
||||
return result
|
||||
|
||||
|
||||
def triage_batch(paths: list[str], options: TriageOptions | None = None) -> dict[str, Any]:
|
||||
items = []
|
||||
for p in paths:
|
||||
try:
|
||||
items.append({"ok": True, "result": triage_file(p, options)})
|
||||
except Exception as exc:
|
||||
items.append({"ok": False, "source_path_basename": Path(p).name, "error": type(exc).__name__, "message": str(exc)})
|
||||
return {"ok": all(item["ok"] for item in items), "files": items, "generated_at": dt.datetime.now(dt.UTC).isoformat()}
|
||||
|
||||
|
||||
def cli() -> int:
|
||||
parser = argparse.ArgumentParser(description="Local document/image triage prototype")
|
||||
parser.add_argument("paths", nargs="+", help="local image/PDF paths")
|
||||
parser.add_argument("--allowed-root", action="append", default=[], help="allowed local root; defaults to cwd")
|
||||
parser.add_argument("--max-pages", type=int, default=3)
|
||||
parser.add_argument("--include-ocr-text", action="store_true")
|
||||
parser.add_argument("--include-full-path", action="store_true")
|
||||
parser.add_argument("--no-embeddings", action="store_true", help="disable local NPU embedding call")
|
||||
parser.add_argument("--dry-run", action="store_true")
|
||||
parser.add_argument("--pretty", action="store_true")
|
||||
args = parser.parse_args()
|
||||
roots = [Path(p) for p in args.allowed_root] if args.allowed_root else [Path.cwd()]
|
||||
options = TriageOptions(
|
||||
max_pages=args.max_pages,
|
||||
include_ocr_text=args.include_ocr_text,
|
||||
dry_run=args.dry_run,
|
||||
use_embeddings=not args.no_embeddings,
|
||||
allowed_roots=roots,
|
||||
include_full_path=args.include_full_path,
|
||||
)
|
||||
out = triage_batch(args.paths, options)
|
||||
print(json.dumps(out, indent=2 if args.pretty else None, sort_keys=True))
|
||||
return 0 if out["ok"] else 2
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(cli())
|
||||
@@ -0,0 +1,111 @@
|
||||
# OpenVINO GenAI NPU worker prototype
|
||||
|
||||
Local-only prototype for cheap bounded background generation on Will's Intel NPU. It is intentionally isolated from primary Atlas/Hermes routing.
|
||||
|
||||
## What it does
|
||||
|
||||
- Model: `OpenVINO/Qwen2.5-1.5B-Instruct-int4-ov`.
|
||||
- Runtime: `/home/will/.venvs/npu` with `openvino-genai==2026.2.0.0`.
|
||||
- Device: OpenVINO GenAI `NPU`.
|
||||
- Default bind: `127.0.0.1:18820`.
|
||||
- Jobs: `title`, `summary`, `notification`, `memory_candidate`.
|
||||
- Prompt/input limits: 6000 chars, `MAX_PROMPT_LEN=1024`, max 256 generated tokens.
|
||||
|
||||
The worker does not write memory, does not restart Atlas/Hermes, does not change primary routing, and does not log raw prompt bodies by default.
|
||||
|
||||
## Files
|
||||
|
||||
- `worker.py` — stdlib HTTP API plus CLI wrapper.
|
||||
- `smoke_llm_npu.py` — direct GenAI smoke test with NPU busy-time verification.
|
||||
- `systemd/openvino-genai-npu-worker.service` — optional user-service template; not installed by this prototype.
|
||||
|
||||
## Model/cache
|
||||
|
||||
Downloaded model path:
|
||||
|
||||
```text
|
||||
/home/will/models/openvino-genai/Qwen2.5-1.5B-Instruct-int4-ov
|
||||
```
|
||||
|
||||
OpenVINO compile cache path:
|
||||
|
||||
```text
|
||||
/home/will/.cache/openvino/genai-npu/qwen2.5-1.5b-int4
|
||||
```
|
||||
|
||||
NPU pipeline config used by the prototype:
|
||||
|
||||
```python
|
||||
CACHE_DIR=/home/will/.cache/openvino/genai-npu/qwen2.5-1.5b-int4
|
||||
MAX_PROMPT_LEN=1024
|
||||
MIN_RESPONSE_LEN=64
|
||||
PREFILL_HINT=DYNAMIC
|
||||
GENERATE_HINT=FAST_COMPILE
|
||||
```
|
||||
|
||||
AOT/blob note: first milestone uses `CACHE_DIR` only. Do not switch to manual `EXPORT_BLOB`/`BLOB_PATH` until compile latency is proven to be the bottleneck. If explicit blobs are used later, record OpenVINO version, NPU compiler version, driver version, model id, quantization flags, and source weights path; invalidate blobs after OpenVINO/NPU driver upgrades.
|
||||
|
||||
## Direct smoke test
|
||||
|
||||
```bash
|
||||
cd /home/will/lab/swarm/openvino-genai-npu-worker
|
||||
/home/will/.venvs/npu/bin/python smoke_llm_npu.py
|
||||
```
|
||||
|
||||
Acceptance requires `npu_busy_delta_us > 0`.
|
||||
|
||||
Observed cold-ish smoke after download/cache setup:
|
||||
|
||||
```json
|
||||
{
|
||||
"text": "\"Atlas Summarizes NPU Worker Options Requested by User\"",
|
||||
"timing_ms": {"load": 10989.08, "generate": 3157.94, "total": 14147.02},
|
||||
"npu_busy_delta_us": 2650724
|
||||
}
|
||||
```
|
||||
|
||||
## CLI usage
|
||||
|
||||
```bash
|
||||
/home/will/.venvs/npu/bin/python worker.py \
|
||||
--job title \
|
||||
--input 'Kanban task asks for a small OpenVINO GenAI NPU worker prototype.'
|
||||
```
|
||||
|
||||
## HTTP usage
|
||||
|
||||
Start locally only:
|
||||
|
||||
```bash
|
||||
cd /home/will/lab/swarm/openvino-genai-npu-worker
|
||||
/home/will/.venvs/npu/bin/python worker.py --host 127.0.0.1 --port 18820
|
||||
```
|
||||
|
||||
Endpoints:
|
||||
|
||||
```text
|
||||
GET /healthz
|
||||
GET /models
|
||||
POST /v1/worker/generate
|
||||
POST /v1/worker/extract-memory-candidates
|
||||
POST /v1/worker/condense-notification
|
||||
```
|
||||
|
||||
Example:
|
||||
|
||||
```bash
|
||||
curl -s http://127.0.0.1:18820/v1/worker/generate \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{"job":"summary","input":"Build a bounded local NPU worker for small generation tasks, no primary routing changes.","max_new_tokens":80}' \
|
||||
| python -m json.tool
|
||||
```
|
||||
|
||||
Response includes `npu_busy_delta_us`; treat zero as failure even if HTTP status is 200.
|
||||
|
||||
## Safety boundaries
|
||||
|
||||
- Binds only to `127.0.0.1` by default; non-local bind is refused in code.
|
||||
- No raw request-body logging.
|
||||
- No private external uploads.
|
||||
- No Atlas/Hermes gateway restarts or primary model routing changes.
|
||||
- NPU access is serialized with a process lock because the NPU is a shared resource with existing services.
|
||||
@@ -0,0 +1,73 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Smoke-test OpenVINO GenAI LLMPipeline on Intel NPU.
|
||||
|
||||
This verifies NPU execution by reading /sys/class/accel/accel0/device/npu_busy_time_us
|
||||
before and after generation. HTTP 200/service success is not considered proof.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import openvino_genai as ov_genai
|
||||
|
||||
DEFAULT_MODEL = "/home/will/models/openvino-genai/Qwen2.5-1.5B-Instruct-int4-ov"
|
||||
DEFAULT_CACHE = "/home/will/.cache/openvino/genai-npu/qwen2.5-1.5b-int4"
|
||||
BUSY_PATH = Path("/sys/class/accel/accel0/device/npu_busy_time_us")
|
||||
|
||||
|
||||
def read_busy() -> int:
|
||||
return int(BUSY_PATH.read_text().strip())
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--model", default=DEFAULT_MODEL)
|
||||
parser.add_argument("--cache-dir", default=DEFAULT_CACHE)
|
||||
parser.add_argument("--prompt", default="Write a concise title for: User asked Atlas to summarize NPU worker options.")
|
||||
parser.add_argument("--max-new-tokens", type=int, default=24)
|
||||
args = parser.parse_args()
|
||||
|
||||
model_path = Path(args.model)
|
||||
cache_dir = Path(args.cache_dir)
|
||||
cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
if not model_path.exists():
|
||||
raise SystemExit(f"model path does not exist: {model_path}")
|
||||
|
||||
config = {
|
||||
"CACHE_DIR": str(cache_dir),
|
||||
"MAX_PROMPT_LEN": 1024,
|
||||
"MIN_RESPONSE_LEN": 64,
|
||||
"PREFILL_HINT": "DYNAMIC",
|
||||
"GENERATE_HINT": "FAST_COMPILE",
|
||||
}
|
||||
|
||||
before = read_busy()
|
||||
load_start = time.monotonic()
|
||||
pipe = ov_genai.LLMPipeline(str(model_path), "NPU", config)
|
||||
load_ms = round((time.monotonic() - load_start) * 1000, 2)
|
||||
|
||||
gen_start = time.monotonic()
|
||||
output = pipe.generate(args.prompt, max_new_tokens=args.max_new_tokens)
|
||||
gen_ms = round((time.monotonic() - gen_start) * 1000, 2)
|
||||
after = read_busy()
|
||||
result = {
|
||||
"model": str(model_path),
|
||||
"device": "NPU",
|
||||
"cache_dir": str(cache_dir),
|
||||
"prompt_chars": len(args.prompt),
|
||||
"max_new_tokens": args.max_new_tokens,
|
||||
"text": str(output).strip(),
|
||||
"timing_ms": {"load": load_ms, "generate": gen_ms, "total": round(load_ms + gen_ms, 2)},
|
||||
"npu_busy_before_us": before,
|
||||
"npu_busy_after_us": after,
|
||||
"npu_busy_delta_us": after - before,
|
||||
}
|
||||
print(json.dumps(result, indent=2))
|
||||
return 0 if after > before else 2
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,16 @@
|
||||
[Unit]
|
||||
Description=OpenVINO GenAI NPU worker prototype
|
||||
After=network-online.target
|
||||
|
||||
[Service]
|
||||
Type=simple
|
||||
WorkingDirectory=/home/will/lab/swarm/openvino-genai-npu-worker
|
||||
Environment=OV_GENAI_NPU_MODEL=/home/will/models/openvino-genai/Qwen2.5-1.5B-Instruct-int4-ov
|
||||
Environment=OV_GENAI_NPU_CACHE=/home/will/.cache/openvino/genai-npu/qwen2.5-1.5b-int4
|
||||
Environment=OV_GENAI_NPU_PORT=18820
|
||||
ExecStart=/home/will/.venvs/npu/bin/python /home/will/lab/swarm/openvino-genai-npu-worker/worker.py --host 127.0.0.1 --port 18820
|
||||
Restart=on-failure
|
||||
RestartSec=5
|
||||
|
||||
[Install]
|
||||
WantedBy=default.target
|
||||
@@ -0,0 +1,251 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Local-only OpenVINO GenAI NPU worker.
|
||||
|
||||
Small bounded LLM worker for cheap background tasks. It intentionally does not
|
||||
wire into Atlas/Hermes routing and does not log raw prompts by default.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from pathlib import Path
|
||||
from typing import Any, cast
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import openvino_genai as ov_genai # type: ignore[import-not-found]
|
||||
|
||||
MODEL_ID = "OpenVINO/Qwen2.5-1.5B-Instruct-int4-ov"
|
||||
DEFAULT_MODEL_PATH = "/home/will/models/openvino-genai/Qwen2.5-1.5B-Instruct-int4-ov"
|
||||
DEFAULT_CACHE_DIR = "/home/will/.cache/openvino/genai-npu/qwen2.5-1.5b-int4"
|
||||
BUSY_PATH = Path("/sys/class/accel/accel0/device/npu_busy_time_us")
|
||||
HOST = "127.0.0.1"
|
||||
PORT = 18820
|
||||
MAX_INPUT_CHARS = 6000
|
||||
DEFAULTS = {
|
||||
"title": 32,
|
||||
"summary": 160,
|
||||
"memory_candidate": 192,
|
||||
"notification": 96,
|
||||
}
|
||||
PROMPTS = {
|
||||
"title": "Write one concise title, 8 words or fewer. Return only the title.\n\nInput:\n{input}",
|
||||
"summary": "Summarize the input in one short paragraph or up to 4 bullets. Be factual and concise.\n\nInput:\n{input}",
|
||||
"memory_candidate": (
|
||||
"Extract durable memory candidates from the conversation excerpt. "
|
||||
"Return strict JSON with keys: candidates (array of objects with fact, confidence, reason), notes. "
|
||||
"Do not write memory; only propose candidates.\n\nInput:\n{input}"
|
||||
),
|
||||
"notification": (
|
||||
"Condense this notification or log excerpt for a human. "
|
||||
"Return JSON with keys: severity (info|warning|error), category, summary, action_needed.\n\nInput:\n{input}"
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def read_busy() -> int:
|
||||
return int(BUSY_PATH.read_text().strip())
|
||||
|
||||
|
||||
def coerce_json(text: str) -> Any | None:
|
||||
text = text.strip()
|
||||
if not text:
|
||||
return None
|
||||
try:
|
||||
return json.loads(text)
|
||||
except json.JSONDecodeError:
|
||||
match = re.search(r"(\{.*\}|\[.*\])", text, re.S)
|
||||
if match:
|
||||
try:
|
||||
return json.loads(match.group(1))
|
||||
except json.JSONDecodeError:
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
@dataclass
|
||||
class GenerationResult:
|
||||
text: str
|
||||
parsed_json: Any | None
|
||||
timing_ms: dict[str, float]
|
||||
npu_busy_delta_us: int
|
||||
npu_busy_before_us: int
|
||||
npu_busy_after_us: int
|
||||
|
||||
|
||||
class NpuWorker:
|
||||
def __init__(self, model_path: str, cache_dir: str):
|
||||
self.model_path = Path(model_path)
|
||||
self.cache_dir = Path(cache_dir)
|
||||
self.cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
self._pipe = None
|
||||
self._load_ms: float | None = None
|
||||
self._lock = threading.Lock()
|
||||
self._loaded_at: float | None = None
|
||||
if not self.model_path.exists():
|
||||
raise FileNotFoundError(f"model path does not exist: {self.model_path}")
|
||||
|
||||
def load(self) -> None:
|
||||
if self._pipe is not None:
|
||||
return
|
||||
start = time.monotonic()
|
||||
# NPU GenAI requires bounded prompt/response shapes; CACHE_DIR enables compiled blob caching.
|
||||
self._pipe = ov_genai.LLMPipeline(
|
||||
str(self.model_path),
|
||||
"NPU",
|
||||
CACHE_DIR=str(self.cache_dir),
|
||||
MAX_PROMPT_LEN=1024,
|
||||
MIN_RESPONSE_LEN=64,
|
||||
PREFILL_HINT="DYNAMIC",
|
||||
GENERATE_HINT="FAST_COMPILE",
|
||||
)
|
||||
self._load_ms = round((time.monotonic() - start) * 1000, 2)
|
||||
self._loaded_at = time.time()
|
||||
|
||||
def generate(self, job: str, user_input: str, max_new_tokens: int | None = None) -> GenerationResult:
|
||||
if job not in PROMPTS:
|
||||
raise ValueError(f"unsupported job: {job}")
|
||||
if not isinstance(user_input, str) or not user_input.strip():
|
||||
raise ValueError("input must be a non-empty string")
|
||||
if len(user_input) > MAX_INPUT_CHARS:
|
||||
raise ValueError(f"input too long: {len(user_input)} chars > {MAX_INPUT_CHARS}")
|
||||
max_new_tokens = int(max_new_tokens or DEFAULTS[job])
|
||||
if max_new_tokens < 1 or max_new_tokens > 256:
|
||||
raise ValueError("max_new_tokens must be between 1 and 256")
|
||||
prompt = PROMPTS[job].format(input=user_input.strip())
|
||||
with self._lock:
|
||||
load_start = time.monotonic()
|
||||
self.load()
|
||||
load_ms = round((time.monotonic() - load_start) * 1000, 2)
|
||||
before = read_busy()
|
||||
gen_start = time.monotonic()
|
||||
pipe = cast(Any, self._pipe)
|
||||
text = str(pipe.generate(prompt, max_new_tokens=max_new_tokens)).strip()
|
||||
generate_ms = round((time.monotonic() - gen_start) * 1000, 2)
|
||||
after = read_busy()
|
||||
parsed = coerce_json(text) if job in {"memory_candidate", "notification"} else None
|
||||
if job == "memory_candidate" and isinstance(parsed, list):
|
||||
parsed = {"candidates": parsed, "notes": "model returned a top-level array; worker wrapped it to preserve the API contract"}
|
||||
return GenerationResult(
|
||||
text=text,
|
||||
parsed_json=parsed,
|
||||
timing_ms={"load": load_ms, "initial_load": self._load_ms or 0.0, "generate": generate_ms, "total": round(load_ms + generate_ms, 2)},
|
||||
npu_busy_delta_us=after - before,
|
||||
npu_busy_before_us=before,
|
||||
npu_busy_after_us=after,
|
||||
)
|
||||
|
||||
def health(self) -> dict[str, Any]:
|
||||
return {
|
||||
"ok": True,
|
||||
"model": MODEL_ID,
|
||||
"model_path": str(self.model_path),
|
||||
"device": "NPU",
|
||||
"cache_dir": str(self.cache_dir),
|
||||
"cache_exists": self.cache_dir.exists(),
|
||||
"loaded": self._pipe is not None,
|
||||
"initial_load_ms": self._load_ms,
|
||||
"loaded_at": self._loaded_at,
|
||||
"busy_time_us": read_busy(),
|
||||
"max_input_chars": MAX_INPUT_CHARS,
|
||||
"jobs": sorted(PROMPTS),
|
||||
"bind": f"{HOST}:{PORT}",
|
||||
}
|
||||
|
||||
|
||||
def response_payload(worker: NpuWorker, job: str, result: GenerationResult) -> dict[str, Any]:
|
||||
return {
|
||||
"model": MODEL_ID,
|
||||
"device": "NPU",
|
||||
"job": job,
|
||||
"text": result.text,
|
||||
"json": result.parsed_json,
|
||||
"timing_ms": result.timing_ms,
|
||||
"npu_busy_delta_us": result.npu_busy_delta_us,
|
||||
"npu_busy_before_us": result.npu_busy_before_us,
|
||||
"npu_busy_after_us": result.npu_busy_after_us,
|
||||
"cache_dir": str(worker.cache_dir),
|
||||
}
|
||||
|
||||
|
||||
def make_handler(worker: NpuWorker):
|
||||
class Handler(BaseHTTPRequestHandler):
|
||||
server_version = "openvino-genai-npu-worker/0.1"
|
||||
|
||||
def log_message(self, format: str, *args: Any) -> None:
|
||||
# Log only method/path/status metadata, not raw request bodies.
|
||||
print(f"{self.client_address[0]} {format % args}")
|
||||
|
||||
def send_json(self, status: int, payload: Any) -> None:
|
||||
body = json.dumps(payload, indent=2).encode("utf-8")
|
||||
self.send_response(status)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(body)))
|
||||
self.end_headers()
|
||||
self.wfile.write(body)
|
||||
|
||||
def do_GET(self) -> None: # noqa: N802
|
||||
path = urlparse(self.path).path
|
||||
if path == "/healthz":
|
||||
self.send_json(200, worker.health())
|
||||
elif path == "/models":
|
||||
self.send_json(200, {"models": [{"id": MODEL_ID, "path": str(worker.model_path), "device": "NPU"}]})
|
||||
else:
|
||||
self.send_json(404, {"error": "not found"})
|
||||
|
||||
def do_POST(self) -> None: # noqa: N802
|
||||
path = urlparse(self.path).path
|
||||
route_job = {
|
||||
"/v1/worker/generate": None,
|
||||
"/v1/worker/extract-memory-candidates": "memory_candidate",
|
||||
"/v1/worker/condense-notification": "notification",
|
||||
}.get(path, "__missing__")
|
||||
if route_job == "__missing__":
|
||||
self.send_json(404, {"error": "not found"})
|
||||
return
|
||||
try:
|
||||
length = int(self.headers.get("Content-Length", "0"))
|
||||
payload = json.loads(self.rfile.read(length) or b"{}")
|
||||
job = route_job or str(payload.get("job", "summary"))
|
||||
if job == "memory":
|
||||
job = "memory_candidate"
|
||||
result = worker.generate(job, str(payload.get("input", "")), payload.get("max_new_tokens"))
|
||||
self.send_json(200, response_payload(worker, job, result))
|
||||
except Exception as exc:
|
||||
self.send_json(400, {"error": str(exc)})
|
||||
|
||||
return Handler
|
||||
|
||||
|
||||
def cli(argv: list[str] | None = None) -> int:
|
||||
parser = argparse.ArgumentParser(description="OpenVINO GenAI NPU worker")
|
||||
parser.add_argument("--model-path", default=os.environ.get("OV_GENAI_NPU_MODEL", DEFAULT_MODEL_PATH))
|
||||
parser.add_argument("--cache-dir", default=os.environ.get("OV_GENAI_NPU_CACHE", DEFAULT_CACHE_DIR))
|
||||
parser.add_argument("--host", default=HOST)
|
||||
parser.add_argument("--port", type=int, default=int(os.environ.get("OV_GENAI_NPU_PORT", PORT)))
|
||||
parser.add_argument("--job", choices=sorted(PROMPTS), help="Run one CLI job instead of serving HTTP")
|
||||
parser.add_argument("--input", help="Input text for --job")
|
||||
parser.add_argument("--max-new-tokens", type=int)
|
||||
args = parser.parse_args(argv)
|
||||
|
||||
worker = NpuWorker(args.model_path, args.cache_dir)
|
||||
if args.job:
|
||||
result = worker.generate(args.job, args.input or "", args.max_new_tokens)
|
||||
print(json.dumps(response_payload(worker, args.job, result), indent=2))
|
||||
return 0 if result.npu_busy_delta_us > 0 else 2
|
||||
|
||||
if args.host != "127.0.0.1":
|
||||
raise SystemExit("Refusing non-local bind without code change/explicit approval")
|
||||
server = ThreadingHTTPServer((args.host, args.port), make_handler(worker))
|
||||
print(f"serving {MODEL_ID} on http://{args.host}:{args.port}; raw prompts are not logged")
|
||||
server.serve_forever()
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(cli())
|
||||
@@ -0,0 +1,138 @@
|
||||
# OpenVINO NPU reranker service
|
||||
|
||||
Local-first cross-encoder reranker prototype for second-stage RAG ranking.
|
||||
|
||||
- Default bind: `127.0.0.1:18818`
|
||||
- Default model: `cross-encoder/ms-marco-MiniLM-L6-v2`
|
||||
- Default device: `NPU`
|
||||
- Model cache: `/home/will/.cache/openvino-models/rerankers/ms-marco-MiniLM-L6-v2-int8-ov/`
|
||||
- NPU proof: `/sys/class/accel/accel0/device/npu_busy_time_us` delta before/after inference
|
||||
|
||||
This service is intentionally not wired into live RAG by default.
|
||||
|
||||
## Files
|
||||
|
||||
- `server.py` — stdlib HTTP OpenVINO Runtime service.
|
||||
- `smoke.py` — non-private API/ranking/NPU busy-time smoke test.
|
||||
- `openvino-reranker.service` — optional user-systemd unit.
|
||||
|
||||
## One-time setup
|
||||
|
||||
Use a separate venv so the existing Whisper/embeddings NPU venv is not perturbed:
|
||||
|
||||
```bash
|
||||
python -m venv /home/will/.venvs/openvino-reranker
|
||||
source /home/will/.venvs/openvino-reranker/bin/activate
|
||||
python -m pip install -U pip
|
||||
python -m pip install "openvino>=2026.2" "optimum-intel[openvino]" transformers tokenizers nncf numpy
|
||||
```
|
||||
|
||||
Export the model:
|
||||
|
||||
```bash
|
||||
source /home/will/.venvs/openvino-reranker/bin/activate
|
||||
optimum-cli export openvino \
|
||||
--model cross-encoder/ms-marco-MiniLM-L6-v2 \
|
||||
--task text-classification \
|
||||
--weight-format int8 \
|
||||
--trust-remote-code false \
|
||||
/home/will/.cache/openvino-models/rerankers/ms-marco-MiniLM-L6-v2-int8-ov
|
||||
```
|
||||
|
||||
If INT8 export or NPU compile fails, export an FP16/FP32 IR to a separate directory and point `OPENVINO_RERANKER_MODEL_DIR` at it while debugging. Do not overwrite existing vector/RAG/Chroma collections.
|
||||
|
||||
## Run in foreground
|
||||
|
||||
Check the port and NPU counter first:
|
||||
|
||||
```bash
|
||||
ss -ltnp | grep ':18818 ' || true
|
||||
cat /sys/class/accel/accel0/device/npu_busy_time_us
|
||||
```
|
||||
|
||||
Start locally:
|
||||
|
||||
```bash
|
||||
source /home/will/.venvs/openvino-reranker/bin/activate
|
||||
OPENVINO_RERANKER_HOST=127.0.0.1 \
|
||||
OPENVINO_RERANKER_PORT=18818 \
|
||||
OPENVINO_RERANKER_DEVICE=NPU \
|
||||
OPENVINO_RERANKER_MODEL_DIR=/home/will/.cache/openvino-models/rerankers/ms-marco-MiniLM-L6-v2-int8-ov \
|
||||
python /home/will/lab/swarm/openvino-reranker-npu/server.py
|
||||
```
|
||||
|
||||
Startup performs a non-private smoke inference and fails closed when `OPENVINO_RERANKER_DEVICE=NPU` but `npu_busy_time_us` does not increase.
|
||||
|
||||
## API
|
||||
|
||||
Health:
|
||||
|
||||
```bash
|
||||
curl -sS http://127.0.0.1:18818/healthz | jq
|
||||
curl -sS http://127.0.0.1:18818/readyz | jq
|
||||
```
|
||||
|
||||
Rerank:
|
||||
|
||||
```bash
|
||||
curl -sS http://127.0.0.1:18818/rerank \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{
|
||||
"query":"how do I verify OpenVINO NPU usage?",
|
||||
"documents":[
|
||||
{"id":"good","text":"Check /sys/class/accel/accel0/device/npu_busy_time_us before and after inference."},
|
||||
{"id":"bad","text":"This note is about making sourdough starter."}
|
||||
],
|
||||
"top_k":2
|
||||
}' | jq
|
||||
```
|
||||
|
||||
Compatibility alias:
|
||||
|
||||
```bash
|
||||
curl -sS http://127.0.0.1:18818/v1/rerank \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{"model":"local-reranker","query":"npu busy time","documents":["OpenVINO NPU busy time proves accelerator use."],"top_n":1}' | jq
|
||||
```
|
||||
|
||||
## Smoke test
|
||||
|
||||
```bash
|
||||
source /home/will/.venvs/openvino-reranker/bin/activate
|
||||
python /home/will/lab/swarm/openvino-reranker-npu/smoke.py --url http://127.0.0.1:18818
|
||||
```
|
||||
|
||||
Expected:
|
||||
|
||||
- `/readyz` is HTTP 200 and reports `device=NPU`.
|
||||
- Each fixture returns `ok=true` and a sorted `results` list.
|
||||
- The top result matches the non-private fixture expectation.
|
||||
- Response and sysfs `npu_busy_delta_us` are positive.
|
||||
|
||||
## Optional systemd user service
|
||||
|
||||
Install the unit only after the foreground command and smoke test pass:
|
||||
|
||||
```bash
|
||||
cp /home/will/lab/swarm/openvino-reranker-npu/openvino-reranker.service /home/will/.config/systemd/user/openvino-reranker.service
|
||||
systemctl --user daemon-reload
|
||||
systemctl --user start openvino-reranker.service
|
||||
systemctl --user status openvino-reranker.service --no-pager
|
||||
journalctl --user -u openvino-reranker.service -n 100 --no-pager
|
||||
```
|
||||
|
||||
Do not enable or integrate it into live RAG without explicit approval.
|
||||
|
||||
## Optional RAG integration plan (disabled by default)
|
||||
|
||||
RAG should keep vector search against `obsidian_bge_npu` unchanged, retrieve a larger candidate set, and call this service as a read-only request-time second stage. Suggested disabled-by-default knobs:
|
||||
|
||||
```text
|
||||
RAG_RERANK_ENABLED=false
|
||||
RAG_RERANK_URL=http://127.0.0.1:18818/rerank
|
||||
RAG_RERANK_INITIAL_K=20
|
||||
RAG_RERANK_TOP_K=5
|
||||
RAG_RERANK_TIMEOUT_MS=3000
|
||||
```
|
||||
|
||||
On reranker timeout/error, fall back to vector order and include metadata such as `rerank_error`; do not mutate or reindex Chroma collections.
|
||||
@@ -0,0 +1,19 @@
|
||||
[Unit]
|
||||
Description=OpenVINO NPU Reranker HTTP Service (port 18818)
|
||||
After=network-online.target
|
||||
|
||||
[Service]
|
||||
Type=simple
|
||||
WorkingDirectory=/home/will/lab/swarm/openvino-reranker-npu
|
||||
Environment=OPENVINO_RERANKER_HOST=127.0.0.1
|
||||
Environment=OPENVINO_RERANKER_PORT=18818
|
||||
Environment=OPENVINO_RERANKER_MODEL=cross-encoder/ms-marco-MiniLM-L6-v2
|
||||
Environment=OPENVINO_RERANKER_MODEL_DIR=/home/will/.cache/openvino-models/rerankers/ms-marco-MiniLM-L6-v2-int8-ov
|
||||
Environment=OPENVINO_RERANKER_DEVICE=NPU
|
||||
Environment=OPENVINO_RERANKER_MAX_LENGTH=512
|
||||
ExecStart=/home/will/.venvs/openvino-reranker/bin/python /home/will/lab/swarm/openvino-reranker-npu/server.py
|
||||
Restart=on-failure
|
||||
RestartSec=5
|
||||
|
||||
[Install]
|
||||
WantedBy=default.target
|
||||
Executable
+369
@@ -0,0 +1,369 @@
|
||||
#!/usr/bin/env python3
|
||||
"""OpenVINO NPU cross-encoder reranker HTTP service.
|
||||
|
||||
Default port: 18818
|
||||
Default model: cross-encoder/ms-marco-MiniLM-L6-v2 exported as OpenVINO IR
|
||||
Default device: NPU
|
||||
|
||||
Endpoints:
|
||||
GET /, /healthz, /readyz
|
||||
POST /rerank
|
||||
POST /v1/rerank
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import openvino as ov
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
DEFAULT_MODEL_ID = "cross-encoder/ms-marco-MiniLM-L6-v2"
|
||||
DEFAULT_MODEL_DIR = Path("/home/will/.cache/openvino-models/rerankers/ms-marco-MiniLM-L6-v2-int8-ov")
|
||||
DEFAULT_PORT = 18818
|
||||
DEFAULT_MAX_LENGTH = 512
|
||||
DEFAULT_MAX_DOCUMENTS = 100
|
||||
DEFAULT_MAX_BODY_BYTES = 5 * 1024 * 1024
|
||||
NPU_BUSY_FILE = Path("/sys/class/accel/accel0/device/npu_busy_time_us")
|
||||
|
||||
|
||||
def npu_busy_time_us() -> int | None:
|
||||
try:
|
||||
return int(NPU_BUSY_FILE.read_text().strip())
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def sigmoid(x: float) -> float:
|
||||
if x >= 0:
|
||||
z = math.exp(-x)
|
||||
return 1.0 / (1.0 + z)
|
||||
z = math.exp(x)
|
||||
return z / (1.0 + z)
|
||||
|
||||
|
||||
def softmax_prob(logits: np.ndarray, index: int = 1) -> float:
|
||||
row = np.asarray(logits, dtype=np.float64).reshape(-1)
|
||||
shifted = row - np.max(row)
|
||||
probs = np.exp(shifted) / np.sum(np.exp(shifted))
|
||||
return float(probs[index])
|
||||
|
||||
|
||||
class RerankerService:
|
||||
def __init__(
|
||||
self,
|
||||
model_dir: Path,
|
||||
model_id: str,
|
||||
device: str,
|
||||
max_length: int,
|
||||
startup_smoke: bool = True,
|
||||
) -> None:
|
||||
self.model_dir = model_dir
|
||||
self.model_id = model_id
|
||||
self.device = device
|
||||
self.max_length = int(max_length)
|
||||
self.loaded_at = time.time()
|
||||
self.lock = threading.Lock()
|
||||
self.last_inference: dict[str, Any] | None = None
|
||||
self.startup_smoke: dict[str, Any] | None = None
|
||||
self.ready = False
|
||||
self.ready_error: str | None = None
|
||||
|
||||
if not self.model_dir.exists():
|
||||
raise FileNotFoundError(f"model directory not found: {self.model_dir}")
|
||||
|
||||
self.core = ov.Core()
|
||||
self.available_devices = list(self.core.available_devices)
|
||||
if self.device not in self.available_devices:
|
||||
raise RuntimeError(f"OpenVINO device {self.device!r} unavailable; available={self.available_devices}")
|
||||
|
||||
xml_path = self.model_dir / "openvino_model.xml"
|
||||
if not xml_path.exists():
|
||||
raise FileNotFoundError(f"OpenVINO IR not found: {xml_path}")
|
||||
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(str(self.model_dir), local_files_only=True)
|
||||
model = self.core.read_model(str(xml_path))
|
||||
self._reshape_static(model)
|
||||
self.compiled = self.core.compile_model(model, self.device)
|
||||
self.input_names = {inp.get_any_name() for inp in self.compiled.inputs}
|
||||
self.output = self.compiled.output(0)
|
||||
|
||||
if startup_smoke:
|
||||
try:
|
||||
smoke = self.rerank(
|
||||
"npu busy time",
|
||||
[{"id": "smoke", "text": "OpenVINO NPU usage is verified by npu_busy_time_us."}],
|
||||
top_k=1,
|
||||
return_documents=False,
|
||||
)
|
||||
self.startup_smoke = {
|
||||
"ok": bool(smoke.get("ok")),
|
||||
"duration_ms": smoke.get("duration_ms"),
|
||||
"npu_busy_delta_us": smoke.get("npu_busy_delta_us"),
|
||||
}
|
||||
if self.device == "NPU" and int(smoke.get("npu_busy_delta_us") or 0) <= 0:
|
||||
raise RuntimeError("startup smoke did not increase npu_busy_time_us")
|
||||
except Exception as exc:
|
||||
self.ready_error = f"startup smoke failed: {type(exc).__name__}: {exc}"
|
||||
raise
|
||||
|
||||
self.ready = True
|
||||
|
||||
def _reshape_static(self, model: ov.Model) -> None:
|
||||
shape_by_name: dict[str, list[int]] = {}
|
||||
for inp in model.inputs:
|
||||
name = inp.get_any_name()
|
||||
if name in {"input_ids", "attention_mask", "token_type_ids"}:
|
||||
shape_by_name[name] = [1, self.max_length]
|
||||
if shape_by_name:
|
||||
model.reshape(shape_by_name)
|
||||
|
||||
def _tokenize(self, query: str, document: str) -> dict[str, np.ndarray]:
|
||||
tokens = self.tokenizer(
|
||||
query,
|
||||
document,
|
||||
max_length=self.max_length,
|
||||
padding="max_length",
|
||||
truncation=True,
|
||||
return_tensors="np",
|
||||
)
|
||||
return {name: np.asarray(value) for name, value in tokens.items() if name in self.input_names}
|
||||
|
||||
def _score_pair(self, query: str, document: str) -> dict[str, float | None]:
|
||||
inputs = self._tokenize(query, document)
|
||||
missing = self.input_names - set(inputs)
|
||||
# Some exported BERT models do not use token_type_ids. input_ids and attention_mask are required.
|
||||
required_missing = missing & {"input_ids", "attention_mask"}
|
||||
if required_missing:
|
||||
raise RuntimeError(f"tokenizer did not produce required inputs: {sorted(required_missing)}")
|
||||
outputs = self.compiled(inputs)
|
||||
logits = np.asarray(outputs[self.output])
|
||||
flat = logits.reshape(-1)
|
||||
if flat.size == 1:
|
||||
raw = float(flat[0])
|
||||
return {"score": raw, "raw_score": raw, "probability": sigmoid(raw)}
|
||||
if flat.size >= 2:
|
||||
raw = float(flat[1])
|
||||
return {"score": raw, "raw_score": raw, "probability": softmax_prob(flat, 1)}
|
||||
raise RuntimeError(f"unexpected empty logits shape: {list(logits.shape)}")
|
||||
|
||||
def rerank(
|
||||
self,
|
||||
query: str,
|
||||
documents: list[dict[str, Any]],
|
||||
*,
|
||||
top_k: int | None,
|
||||
return_documents: bool = True,
|
||||
) -> dict[str, Any]:
|
||||
before = npu_busy_time_us()
|
||||
started = time.perf_counter()
|
||||
results: list[dict[str, Any]] = []
|
||||
with self.lock:
|
||||
for idx, doc in enumerate(documents):
|
||||
scored = self._score_pair(query, str(doc["text"]))
|
||||
item: dict[str, Any] = {
|
||||
"index": idx,
|
||||
"score": scored["score"],
|
||||
"raw_score": scored["raw_score"],
|
||||
"probability": scored["probability"],
|
||||
}
|
||||
if doc.get("id") is not None:
|
||||
item["id"] = doc.get("id")
|
||||
if return_documents:
|
||||
item["text"] = doc["text"]
|
||||
item["metadata"] = doc.get("metadata") if isinstance(doc.get("metadata"), dict) else {}
|
||||
results.append(item)
|
||||
after = npu_busy_time_us()
|
||||
results.sort(key=lambda item: (-float(item["score"]), int(item["index"])))
|
||||
clamped_top_k = len(results) if top_k is None else max(1, min(int(top_k), len(results)))
|
||||
duration_ms = round((time.perf_counter() - started) * 1000, 3)
|
||||
npu_delta = None if before is None or after is None else after - before
|
||||
payload = {
|
||||
"ok": True,
|
||||
"model": self.model_id,
|
||||
"model_dir": str(self.model_dir),
|
||||
"device": self.device,
|
||||
"query": query,
|
||||
"input_count": len(documents),
|
||||
"top_k": clamped_top_k,
|
||||
"duration_ms": duration_ms,
|
||||
"npu_busy_delta_us": npu_delta,
|
||||
"results": results[:clamped_top_k],
|
||||
}
|
||||
self.last_inference = {
|
||||
"duration_ms": duration_ms,
|
||||
"docs": len(documents),
|
||||
"npu_busy_delta_us": npu_delta,
|
||||
}
|
||||
return payload
|
||||
|
||||
def health(self) -> dict[str, Any]:
|
||||
status = "ok" if self.ready else "degraded"
|
||||
return {
|
||||
"status": status,
|
||||
"ok": self.ready,
|
||||
"service": "openvino-reranker",
|
||||
"model": self.model_id,
|
||||
"model_dir": str(self.model_dir),
|
||||
"device": self.device,
|
||||
"available_devices": self.available_devices,
|
||||
"max_length": self.max_length,
|
||||
"input_names": sorted(self.input_names),
|
||||
"uptime_s": round(time.time() - self.loaded_at, 3),
|
||||
"npu_busy_time_us": npu_busy_time_us(),
|
||||
"startup_smoke": self.startup_smoke,
|
||||
"last_inference": self.last_inference,
|
||||
"ready_error": self.ready_error,
|
||||
}
|
||||
|
||||
|
||||
def normalize_documents(value: Any, max_documents: int) -> list[dict[str, Any]]:
|
||||
if not isinstance(value, list) or not value:
|
||||
raise ValueError("documents must be a non-empty list")
|
||||
if len(value) > max_documents:
|
||||
raise ValueError(f"documents exceeds max_documents={max_documents}")
|
||||
docs: list[dict[str, Any]] = []
|
||||
for idx, item in enumerate(value):
|
||||
if isinstance(item, str):
|
||||
text = item
|
||||
doc: dict[str, Any] = {"text": text}
|
||||
elif isinstance(item, dict):
|
||||
text = item.get("text")
|
||||
doc = {
|
||||
"id": item.get("id"),
|
||||
"text": text,
|
||||
"metadata": item.get("metadata") if isinstance(item.get("metadata"), dict) else {},
|
||||
}
|
||||
else:
|
||||
raise ValueError(f"documents[{idx}] must be a string or object")
|
||||
if not isinstance(text, str) or not text.strip():
|
||||
raise ValueError(f"documents[{idx}].text must be a non-empty string")
|
||||
docs.append(doc)
|
||||
return docs
|
||||
|
||||
|
||||
class Handler(BaseHTTPRequestHandler):
|
||||
server_version = "OpenVINOReranker/0.1"
|
||||
|
||||
@property
|
||||
def svc(self) -> RerankerService:
|
||||
return self.server.reranker_service # type: ignore[attr-defined]
|
||||
|
||||
@property
|
||||
def max_body_bytes(self) -> int:
|
||||
return self.server.max_body_bytes # type: ignore[attr-defined]
|
||||
|
||||
@property
|
||||
def max_documents(self) -> int:
|
||||
return self.server.max_documents # type: ignore[attr-defined]
|
||||
|
||||
def do_GET(self) -> None:
|
||||
path = self.path.split("?", 1)[0].rstrip("/") or "/"
|
||||
if path == "/":
|
||||
self.write_json({"ok": True, "service": "openvino-reranker", "endpoints": ["/healthz", "/readyz", "/rerank", "/v1/rerank"]})
|
||||
elif path in {"/healthz", "/health"}:
|
||||
self.write_json(self.svc.health(), status=200)
|
||||
elif path == "/readyz":
|
||||
health = self.svc.health()
|
||||
self.write_json(health, status=200 if health.get("ok") else 503)
|
||||
else:
|
||||
self.write_json({"ok": False, "error": "not found", "results": []}, status=404)
|
||||
|
||||
def do_POST(self) -> None:
|
||||
path = self.path.split("?", 1)[0].rstrip("/") or "/"
|
||||
try:
|
||||
if path not in {"/rerank", "/v1/rerank"}:
|
||||
self.write_json({"ok": False, "error": "not found", "results": []}, status=404)
|
||||
return
|
||||
if not self.svc.ready:
|
||||
self.write_json({"ok": False, "error": self.svc.ready_error or "model not ready", "results": []}, status=503)
|
||||
return
|
||||
payload = self.read_json()
|
||||
query = payload.get("query")
|
||||
if not isinstance(query, str) or not query.strip():
|
||||
raise ValueError("query is required")
|
||||
top_k = payload.get("top_k", payload.get("top_n"))
|
||||
documents = normalize_documents(payload.get("documents"), self.max_documents)
|
||||
return_documents = bool(payload.get("return_documents", True))
|
||||
response = self.svc.rerank(query.strip(), documents, top_k=top_k, return_documents=return_documents)
|
||||
self.write_json(response)
|
||||
except RequestTooLarge as exc:
|
||||
self.write_json({"ok": False, "error": str(exc), "results": []}, status=413)
|
||||
except ValueError as exc:
|
||||
self.write_json({"ok": False, "error": str(exc), "results": []}, status=400)
|
||||
except Exception as exc:
|
||||
self.write_json({"ok": False, "error": f"{type(exc).__name__}: {exc}", "results": []}, status=500)
|
||||
|
||||
def read_json(self) -> dict[str, Any]:
|
||||
length = int(self.headers.get("Content-Length") or 0)
|
||||
if length > self.max_body_bytes:
|
||||
raise RequestTooLarge(f"request body exceeds {self.max_body_bytes} bytes")
|
||||
body = self.rfile.read(length).decode("utf-8", "replace") if length else "{}"
|
||||
payload = json.loads(body or "{}")
|
||||
if not isinstance(payload, dict):
|
||||
raise ValueError("JSON body must be an object")
|
||||
return payload
|
||||
|
||||
def write_json(self, payload: dict[str, Any], status: int = 200) -> None:
|
||||
body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
|
||||
self.send_response(status)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(body)))
|
||||
self.end_headers()
|
||||
self.wfile.write(body)
|
||||
|
||||
def log_message(self, format: str, *args: Any) -> None: # noqa: A002 - stdlib override name
|
||||
print(f"{self.address_string()} - {format % args}", file=sys.stderr, flush=True)
|
||||
|
||||
|
||||
class RequestTooLarge(ValueError):
|
||||
pass
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--host", default=os.environ.get("OPENVINO_RERANKER_HOST", "127.0.0.1"))
|
||||
parser.add_argument("--port", type=int, default=int(os.environ.get("OPENVINO_RERANKER_PORT", DEFAULT_PORT)))
|
||||
parser.add_argument("--model-dir", default=os.environ.get("OPENVINO_RERANKER_MODEL_DIR", str(DEFAULT_MODEL_DIR)))
|
||||
parser.add_argument("--model", default=os.environ.get("OPENVINO_RERANKER_MODEL", DEFAULT_MODEL_ID))
|
||||
parser.add_argument("--device", default=os.environ.get("OPENVINO_RERANKER_DEVICE", "NPU"))
|
||||
parser.add_argument("--max-length", type=int, default=int(os.environ.get("OPENVINO_RERANKER_MAX_LENGTH", str(DEFAULT_MAX_LENGTH))))
|
||||
parser.add_argument("--max-documents", type=int, default=int(os.environ.get("OPENVINO_RERANKER_MAX_DOCUMENTS", str(DEFAULT_MAX_DOCUMENTS))))
|
||||
parser.add_argument("--max-body-bytes", type=int, default=int(os.environ.get("OPENVINO_RERANKER_MAX_BODY_BYTES", str(DEFAULT_MAX_BODY_BYTES))))
|
||||
parser.add_argument("--skip-startup-smoke", action="store_true", default=os.environ.get("OPENVINO_RERANKER_SKIP_STARTUP_SMOKE", "").lower() in {"1", "true", "yes"})
|
||||
args = parser.parse_args()
|
||||
|
||||
service = RerankerService(
|
||||
Path(args.model_dir).expanduser(),
|
||||
args.model,
|
||||
args.device,
|
||||
args.max_length,
|
||||
startup_smoke=not args.skip_startup_smoke,
|
||||
)
|
||||
httpd = ThreadingHTTPServer((args.host, args.port), Handler)
|
||||
httpd.reranker_service = service # type: ignore[attr-defined]
|
||||
httpd.max_body_bytes = args.max_body_bytes # type: ignore[attr-defined]
|
||||
httpd.max_documents = args.max_documents # type: ignore[attr-defined]
|
||||
print(
|
||||
f"openvino-reranker listening on {args.host}:{args.port} model={args.model} "
|
||||
f"model_dir={args.model_dir} device={args.device} max_length={args.max_length}",
|
||||
flush=True,
|
||||
)
|
||||
try:
|
||||
httpd.serve_forever()
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
Executable
+167
@@ -0,0 +1,167 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Smoke/benchmark checks for the OpenVINO reranker service.
|
||||
|
||||
Prints a JSON summary and exits non-zero on schema/ranking/NPU verification failure.
|
||||
Uses only non-private fixture text.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import statistics
|
||||
import sys
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
NPU_BUSY_FILE = Path("/sys/class/accel/accel0/device/npu_busy_time_us")
|
||||
|
||||
FIXTURES = [
|
||||
{
|
||||
"query": "how do I verify OpenVINO NPU usage?",
|
||||
"documents": [
|
||||
{"id": "good", "text": "Check /sys/class/accel/accel0/device/npu_busy_time_us before and after inference."},
|
||||
{"id": "bad", "text": "This note is about making sourdough starter."},
|
||||
],
|
||||
"expected_top_id": "good",
|
||||
},
|
||||
{
|
||||
"query": "what port does the reranker service use?",
|
||||
"documents": [
|
||||
{"id": "unrelated", "text": "Whisper transcription accepts audio uploads."},
|
||||
{"id": "port", "text": "The OpenVINO reranker prototype listens locally on port 18818."},
|
||||
],
|
||||
"expected_top_id": "port",
|
||||
},
|
||||
{
|
||||
"query": "why should reranking not mutate vector collections?",
|
||||
"documents": [
|
||||
{"id": "mutation", "text": "Reranking is a read-only second-stage transformation after vector search."},
|
||||
{"id": "cooking", "text": "Boil pasta in salted water until al dente."},
|
||||
],
|
||||
"expected_top_id": "mutation",
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def npu_busy_time_us() -> int | None:
|
||||
try:
|
||||
return int(NPU_BUSY_FILE.read_text().strip())
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def post_json(url: str, payload: dict[str, Any], timeout: float) -> tuple[int, dict[str, Any]]:
|
||||
data = json.dumps(payload).encode("utf-8")
|
||||
req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"}, method="POST")
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=timeout) as resp:
|
||||
body = resp.read().decode("utf-8", "replace")
|
||||
return resp.status, json.loads(body)
|
||||
except urllib.error.HTTPError as exc:
|
||||
body = exc.read().decode("utf-8", "replace")
|
||||
try:
|
||||
parsed = json.loads(body)
|
||||
except Exception:
|
||||
parsed = {"error": body}
|
||||
return exc.code, parsed
|
||||
|
||||
|
||||
def get_json(url: str, timeout: float) -> tuple[int, dict[str, Any]]:
|
||||
try:
|
||||
with urllib.request.urlopen(url, timeout=timeout) as resp:
|
||||
body = resp.read().decode("utf-8", "replace")
|
||||
return resp.status, json.loads(body)
|
||||
except urllib.error.HTTPError as exc:
|
||||
body = exc.read().decode("utf-8", "replace")
|
||||
try:
|
||||
parsed = json.loads(body)
|
||||
except Exception:
|
||||
parsed = {"error": body}
|
||||
return exc.code, parsed
|
||||
|
||||
|
||||
def percentile(values: list[float], pct: float) -> float | None:
|
||||
if not values:
|
||||
return None
|
||||
ordered = sorted(values)
|
||||
idx = min(len(ordered) - 1, max(0, round((pct / 100.0) * (len(ordered) - 1))))
|
||||
return round(ordered[idx], 3)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--url", default="http://127.0.0.1:18818")
|
||||
parser.add_argument("--timeout", type=float, default=20.0)
|
||||
parser.add_argument("--allow-cpu", action="store_true", help="do not fail when health reports a non-NPU device")
|
||||
args = parser.parse_args()
|
||||
|
||||
base = args.url.rstrip("/")
|
||||
failures: list[str] = []
|
||||
health_status, health = get_json(f"{base}/readyz", args.timeout)
|
||||
if health_status != 200 or not health.get("ok"):
|
||||
failures.append(f"readyz failed status={health_status} error={health.get('ready_error') or health.get('error')}")
|
||||
device = health.get("device")
|
||||
if device != "NPU" and not args.allow_cpu:
|
||||
failures.append(f"device is {device!r}, expected 'NPU'")
|
||||
|
||||
latencies: list[float] = []
|
||||
response_npu_total = 0
|
||||
sysfs_npu_total = 0
|
||||
top1_passed = 0
|
||||
|
||||
for case in FIXTURES:
|
||||
before = npu_busy_time_us()
|
||||
started = time.perf_counter()
|
||||
status, payload = post_json(
|
||||
f"{base}/rerank",
|
||||
{"query": case["query"], "documents": case["documents"], "top_k": len(case["documents"]), "return_documents": False},
|
||||
args.timeout,
|
||||
)
|
||||
wall_ms = (time.perf_counter() - started) * 1000
|
||||
after = npu_busy_time_us()
|
||||
latencies.append(float(payload.get("duration_ms") or wall_ms))
|
||||
response_delta = payload.get("npu_busy_delta_us")
|
||||
sysfs_delta = None if before is None or after is None else after - before
|
||||
if isinstance(response_delta, int):
|
||||
response_npu_total += response_delta
|
||||
if isinstance(sysfs_delta, int):
|
||||
sysfs_npu_total += sysfs_delta
|
||||
results = payload.get("results") if isinstance(payload, dict) else None
|
||||
top_id = results[0].get("id") if isinstance(results, list) and results else None
|
||||
if status != 200 or not payload.get("ok"):
|
||||
failures.append(f"case {case['expected_top_id']} HTTP/status failed: status={status} error={payload.get('error')}")
|
||||
if not isinstance(results, list) or len(results) != len(case["documents"]):
|
||||
failures.append(f"case {case['expected_top_id']} returned invalid results")
|
||||
if top_id == case["expected_top_id"]:
|
||||
top1_passed += 1
|
||||
else:
|
||||
failures.append(f"case {case['expected_top_id']} top_id={top_id!r}")
|
||||
if device == "NPU":
|
||||
if not isinstance(response_delta, int) or response_delta <= 0:
|
||||
failures.append(f"case {case['expected_top_id']} response npu delta not positive: {response_delta}")
|
||||
if not isinstance(sysfs_delta, int) or sysfs_delta <= 0:
|
||||
failures.append(f"case {case['expected_top_id']} sysfs npu delta not positive: {sysfs_delta}")
|
||||
|
||||
summary = {
|
||||
"ok": not failures,
|
||||
"url": base,
|
||||
"model": health.get("model"),
|
||||
"device": device,
|
||||
"cases": len(FIXTURES),
|
||||
"top1_passed": top1_passed,
|
||||
"p50_ms": percentile(latencies, 50),
|
||||
"p95_ms": percentile(latencies, 95),
|
||||
"mean_ms": round(statistics.mean(latencies), 3) if latencies else None,
|
||||
"npu_busy_delta_us_total": sysfs_npu_total,
|
||||
"response_npu_busy_delta_us_total": response_npu_total,
|
||||
"failures": failures,
|
||||
}
|
||||
print(json.dumps(summary, indent=2, sort_keys=True))
|
||||
return 0 if not failures else 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -68,7 +68,7 @@ section "HTTP health"
|
||||
http_json "RAG endpoint" "http://127.0.0.1:18810/healthz" || true
|
||||
http_json "RAG/embedding health wrapper" "http://127.0.0.1:18814/healthz" || true
|
||||
http_json "Whisper NPU" "http://127.0.0.1:18816/health" || true
|
||||
http_json "OpenVINO embeddings" "http://127.0.0.1:18817/health" || true
|
||||
http_json "OpenVINO embeddings" "http://127.0.0.1:18817/healthz" || true
|
||||
# Prototypes are expected to be unavailable until explicitly started/approved.
|
||||
http_json "NPU reranker prototype" "http://127.0.0.1:18818/readyz" || true
|
||||
http_json "NPU router classifier prototype" "http://127.0.0.1:18819/healthz" || true
|
||||
@@ -91,10 +91,10 @@ if [[ -z "$response" ]]; then
|
||||
fi
|
||||
delta=$((after - before))
|
||||
printf 'sysfs_before_us=%s\nsysfs_after_us=%s\nsysfs_delta_us=%s\n' "$before" "$after" "$delta"
|
||||
printf '%s' "$response" | python - <<'PY' || true
|
||||
import json, sys
|
||||
RESPONSE_JSON="$response" python - <<'PY' || true
|
||||
import json, os
|
||||
try:
|
||||
data = json.load(sys.stdin)
|
||||
data = json.loads(os.environ.get('RESPONSE_JSON', ''))
|
||||
except Exception as exc:
|
||||
print(f'response_parse_error={type(exc).__name__}: {exc}')
|
||||
raise SystemExit(0)
|
||||
|
||||
Reference in New Issue
Block a user