160 lines
7.6 KiB
Python
160 lines
7.6 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
|
|
from suite_contract import Problem, RESULT_TTL, REVISION, digest, validate_result
|
|
|
|
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, (expires, _) in list(self.results.items()):
|
|
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]
|
|
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"],
|
|
"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 Switchyard decision, then make one inference attempt."""
|
|
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)
|
|
if provider == "local":
|
|
result, metadata = suite_backends.local_generate(effective, event, client_ip)
|
|
else:
|
|
result, metadata = suite_backends.claude_generate(effective, event)
|
|
validate_result(result, request)
|
|
if event.is_set():
|
|
raise Problem("cancelled", 409)
|
|
document.update(metadata, status="completed")
|
|
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)
|
|
except Problem as exc:
|
|
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)
|
|
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)
|