189 lines
9.5 KiB
Python
189 lines
9.5 KiB
Python
"""Single-slot jobs with durable content-free idempotency and volatile results."""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import sqlite3
|
|
import threading
|
|
import time
|
|
import uuid
|
|
|
|
import suite_backends
|
|
import suite_multipass
|
|
from suite_contract import (EXECUTION_REVISION, PROMPT_REVISION, PROMPT_SHA256, Problem, RESULT_TTL,
|
|
MAX_RESULT, REVISION, digest, encoded, validate_result)
|
|
from suite_policy import POLICY_REVISION
|
|
|
|
IDEMPOTENCY_TTL = 7 * 86400
|
|
|
|
|
|
class Jobs:
|
|
"""Persist metadata only; never restart provider work automatically."""
|
|
|
|
def __init__(self, path):
|
|
self.lock = threading.RLock()
|
|
self.db = sqlite3.connect(path, check_same_thread=False)
|
|
self.db.execute("PRAGMA journal_mode=WAL")
|
|
self.db.execute("CREATE TABLE IF NOT EXISTS jobs "
|
|
"(id TEXT PRIMARY KEY, owner TEXT, key TEXT, hash TEXT, "
|
|
"created REAL, document TEXT, UNIQUE(owner,key))")
|
|
self.results, self.cancels = {}, {}
|
|
self.active = False
|
|
for job_id, document in self.db.execute("SELECT id,document FROM jobs").fetchall():
|
|
value = json.loads(document)
|
|
if value["status"] in {"accepted", "running", "cancelling"}:
|
|
value.update(status="failed", error={"code": "interrupted_no_retry"})
|
|
self._save(job_id, value)
|
|
|
|
def _save(self, job_id, value):
|
|
self.db.execute("UPDATE jobs SET document=? WHERE id=?", (json.dumps(value), job_id))
|
|
self.db.commit()
|
|
|
|
def _prune(self):
|
|
now = time.time()
|
|
for job_id, retained in list(self.results.items()):
|
|
expires = retained[0]
|
|
if expires < now:
|
|
del self.results[job_id]
|
|
self.db.execute("DELETE FROM jobs WHERE created < ?", (now - IDEMPOTENCY_TTL,))
|
|
self.db.commit()
|
|
|
|
def get(self, job_id, owner, result=False):
|
|
"""Authorize ownership for status and result; conceal other clients' jobs."""
|
|
with self.lock:
|
|
self._prune()
|
|
row = self.db.execute("SELECT document FROM jobs WHERE id=? AND owner=?",
|
|
(job_id, owner)).fetchone()
|
|
if row is None:
|
|
raise Problem("job_not_found", 404)
|
|
document = json.loads(row[0])
|
|
if result and document["status"] == "completed":
|
|
if job_id not in self.results:
|
|
raise Problem("result_expired_or_worker_restarted", 410)
|
|
document["result"] = self.results[job_id][1]
|
|
if len(self.results[job_id]) > 2:
|
|
document["review_summary"] = self.results[job_id][2]
|
|
return document
|
|
|
|
def submit(self, owner, key, request, selected, client_ip, launch=True):
|
|
"""Replay matching keys atomically; reject busy jobs without queuing them."""
|
|
request_hash, key_hash = digest(request), digest(key)
|
|
with self.lock:
|
|
self._prune()
|
|
existing = self.db.execute("SELECT id,hash FROM jobs WHERE owner=? AND key=?",
|
|
(owner, key_hash)).fetchone()
|
|
if existing:
|
|
if existing[1] != request_hash:
|
|
raise Problem("idempotency_conflict", 409)
|
|
return self.get(existing[0], owner), False
|
|
if self.active:
|
|
raise Problem("capacity_busy", 429)
|
|
if self.db.execute("SELECT count(*) FROM jobs").fetchone()[0] >= 10000:
|
|
raise Problem("idempotency_capacity", 429)
|
|
job_id = uuid.uuid4().hex
|
|
document = {"job_id": job_id, "status": "accepted",
|
|
"configuration_revision": REVISION, "routing": request["routing"],
|
|
"execution_revision": EXECUTION_REVISION,
|
|
"policy_revision": POLICY_REVISION,
|
|
"prompt_revision": PROMPT_REVISION, "prompt_sha256": PROMPT_SHA256,
|
|
"selection": selected, "attempted_destinations": [],
|
|
"compaction": None, "truncation": None, "usage": None,
|
|
"result_retention_seconds": RESULT_TTL,
|
|
"created_at": time.time()}
|
|
self.db.execute("INSERT INTO jobs VALUES (?,?,?,?,?,?)",
|
|
(job_id, owner, key_hash, request_hash, time.time(), json.dumps(document)))
|
|
self.db.commit()
|
|
self.active = True
|
|
self.cancels[job_id] = threading.Event()
|
|
if launch:
|
|
threading.Thread(target=self.run, args=(job_id, owner, request, selected, client_ip),
|
|
daemon=True).start()
|
|
return document, True
|
|
|
|
def cancel(self, job_id, owner):
|
|
"""Signal the worker; keep its slot until the outstanding attempt ends."""
|
|
with self.lock:
|
|
document = self.get(job_id, owner)
|
|
if job_id in self.cancels:
|
|
self.cancels[job_id].set()
|
|
document["status"] = "cancelling"
|
|
self._save(job_id, document)
|
|
return document
|
|
|
|
def run(self, job_id, owner, request, selected, client_ip):
|
|
"""Authorize a fixed route, then run one bounded multi-pass suite job."""
|
|
started = time.monotonic()
|
|
provider = selected["provider"]
|
|
document = self.get(job_id, owner)
|
|
event = self.cancels[job_id]
|
|
try:
|
|
if event.is_set():
|
|
raise Problem("cancelled", 409)
|
|
document["status"] = "running"
|
|
document["attempted_destinations"] = ["switchyard:atlas/planning/" + provider]
|
|
with self.lock:
|
|
self._save(job_id, document)
|
|
suite_backends.switchyard_decision(provider)
|
|
if event.is_set():
|
|
raise Problem("cancelled", 409)
|
|
remaining = request["execution"]["max_seconds"] - (time.monotonic() - started)
|
|
if remaining <= 0:
|
|
raise Problem("timeout", 504)
|
|
effective = {**request, "execution": {**request["execution"], "max_seconds": remaining}}
|
|
# This selection is fixed before dispatch; no exception invokes a fallback.
|
|
if provider not in {"local", "claude"}:
|
|
raise Problem("unsupported_backend", 422)
|
|
document["attempted_destinations"].append(provider + ":" + selected["model"])
|
|
with self.lock:
|
|
self._save(job_id, document)
|
|
print(json.dumps({"job_id": job_id, "event": "inference_attempt",
|
|
"provider": provider, "model": selected["model"]}), flush=True)
|
|
def progress(value):
|
|
document["execution_progress"] = value
|
|
with self.lock:
|
|
self._save(job_id, document)
|
|
result, metadata = suite_multipass.generate(effective, selected, event, client_ip, progress)
|
|
review = metadata.pop("review_summary")
|
|
document.update(metadata)
|
|
try:
|
|
validate_result(result, request)
|
|
except Problem as exc:
|
|
raise Problem(exc.code, exc.status, **metadata.get("cli_diagnostics", {}),
|
|
failure_stage="result_validation") from None
|
|
if event.is_set():
|
|
raise Problem("cancelled", 409)
|
|
document.update(metadata, status="completed")
|
|
if len(encoded({**document, "result": result, "review_summary": review})) > MAX_RESULT:
|
|
raise Problem("response_too_large", 502)
|
|
with self.lock:
|
|
if len(self.results) >= 128:
|
|
del self.results[min(self.results, key=lambda key: self.results[key][0])]
|
|
self.results[job_id] = (time.time() + RESULT_TTL, result, review)
|
|
except Problem as exc:
|
|
if "final_event_seen" in exc.details:
|
|
document.update(cli_diagnostics=exc.details,
|
|
usage=exc.details.get("usage"), turns=exc.details.get("turns"),
|
|
duration_api_ms=exc.details.get("duration_api_ms"),
|
|
cost_usd_estimate=exc.details.get("cost_usd_estimate"))
|
|
if "aggregate_usage" in exc.details:
|
|
document.update(usage=exc.details["aggregate_usage"],
|
|
cost_usd_estimate=exc.details.get("aggregate_cost_usd_estimate"))
|
|
document.update(status="cancelled" if exc.code == "cancelled" else "failed",
|
|
**exc.document())
|
|
except Exception:
|
|
# Provider diagnostics can echo complete requests; never serialize them.
|
|
document.update(status="failed", error={"code": "internal_worker_error"})
|
|
finally:
|
|
document["wall_seconds"] = round(time.monotonic() - started, 3)
|
|
if "execution_progress" in document:
|
|
document["execution_progress"].update(cli_running=False, heartbeat_at=time.time(),
|
|
job_elapsed_seconds=document["wall_seconds"])
|
|
with self.lock:
|
|
if event.is_set() and document["status"] == "completed":
|
|
document.update(status="cancelled", error={"code": "cancelled"})
|
|
self.results.pop(job_id, None)
|
|
self._save(job_id, document)
|
|
self.cancels.pop(job_id, None)
|
|
self.active = False
|
|
print(json.dumps({"job_id": job_id, "status": document["status"],
|
|
"wall_seconds": document["wall_seconds"]}), flush=True)
|