282 lines
11 KiB
Python
282 lines
11 KiB
Python
"""Worker-Prozess: führt jeden Run aus, auch einfache Chat-Antworten (→ Konzept 3, 9.3).
|
|
|
|
Ablauf je Run: beanspruchen → Handler ausführen, währenddessen Lease per
|
|
Heartbeat verlängern → Ergebnis-Status setzen. Wartende Runs belegen keinen
|
|
Worker; nach einer Entscheidung wird der Run neu eingereiht.
|
|
"""
|
|
|
|
import logging
|
|
import os
|
|
import socket
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from collections.abc import Callable
|
|
from contextlib import suppress
|
|
from dataclasses import dataclass, field
|
|
from datetime import timedelta
|
|
from typing import Any
|
|
|
|
from sqlalchemy import func
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.approvals import service as approvals
|
|
from app.changes import service as changes
|
|
from app.core import logs, metrics
|
|
from app.core.config import Settings
|
|
from app.core.database import Database, now
|
|
from app.documents import service as documents
|
|
from app.hugin import service as hugin
|
|
from app.notify import service as notify
|
|
from app.tasks import engine, monitor, states
|
|
from app.tasks.events import EventWriter
|
|
from app.workspaces import removal
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
class Cancelled(Exception):
|
|
pass
|
|
|
|
|
|
class LeaseLost(Exception):
|
|
"""Ein anderer Worker hat den Run übernommen; dieser darf nichts mehr schreiben."""
|
|
|
|
|
|
class RunStopped(Exception):
|
|
"""Vom Handler gezielt beendet, z. B. fehlende Konfiguration oder Guardrail.
|
|
|
|
user_message ist für Nutzer bestimmt; retryable=True versucht es später erneut
|
|
(solange Versuche übrig sind), sonst endet der Run sofort mit status.
|
|
"""
|
|
|
|
def __init__(self, status: str, user_message: str, *, retryable: bool = False, detail: str = ""):
|
|
super().__init__(detail or user_message)
|
|
self.status, self.user_message, self.retryable = status, user_message, retryable
|
|
|
|
|
|
@dataclass
|
|
class RunContext:
|
|
db: Database
|
|
settings: Settings
|
|
run_id: uuid.UUID
|
|
workspace_id: int
|
|
chat_session_id: int | None
|
|
created_by: int
|
|
input: dict[str, Any]
|
|
events: EventWriter
|
|
worker_id: str
|
|
lease_lost: threading.Event = field(default_factory=threading.Event)
|
|
cancel: threading.Event = field(default_factory=threading.Event)
|
|
|
|
def check(self) -> None:
|
|
"""Vor jedem Schritt aufrufen, der Zeit kostet oder etwas schreibt."""
|
|
if self.lease_lost.is_set():
|
|
raise LeaseLost
|
|
if self.cancel.is_set():
|
|
raise Cancelled
|
|
|
|
def guard(self, session: Session) -> None:
|
|
"""In der Transaktion, die Ergebnisse speichert: sperrt den Run und prüft die Lease."""
|
|
if engine.owned(session, self.run_id, self.worker_id) is None:
|
|
raise LeaseLost
|
|
|
|
|
|
# Handler je Run-Art; gibt den Endstatus zurück (Standard: completed)
|
|
Handler = Callable[[RunContext], str | None]
|
|
HANDLERS: dict[str, Handler] = {}
|
|
|
|
|
|
def handler(kind: str) -> Callable[[Handler], Handler]:
|
|
def register(function: Handler) -> Handler:
|
|
HANDLERS[kind] = function
|
|
return function
|
|
return register
|
|
|
|
|
|
class Heartbeat(threading.Thread):
|
|
def __init__(self, context: RunContext, worker_id: str):
|
|
super().__init__(daemon=True, name=f"heartbeat-{context.run_id}")
|
|
self.context, self.worker_id = context, worker_id
|
|
self.stopped = threading.Event()
|
|
|
|
def run(self) -> None:
|
|
interval = max(self.context.settings.run_lease_seconds / 3, 0.2)
|
|
while not self.stopped.wait(interval):
|
|
try:
|
|
beat = engine.heartbeat(self.context.db, self.context.run_id, self.worker_id, self.context.settings)
|
|
except Exception:
|
|
log.warning("Heartbeat fehlgeschlagen für %s", self.context.run_id, exc_info=True)
|
|
continue
|
|
if not beat.alive:
|
|
self.context.lease_lost.set()
|
|
return
|
|
if beat.cancel:
|
|
self.context.cancel.set()
|
|
|
|
def stop(self) -> None:
|
|
self.stopped.set()
|
|
|
|
|
|
# Papierkorb und andere seltene Aufgaben (Sekunden)
|
|
HOUSEKEEPING_SECONDS = 600
|
|
# Wie oft nach fälligen Pull Requests geschaut wird (Sekunden)
|
|
PR_POLL_CHECK_SECONDS = 60
|
|
# Wie oft der Worker nachsieht, ob ein Workspace mit Hugin abzugleichen ist (fällig nach AIW_HUGIN_CHECK_HOURS)
|
|
HUGIN_CHECK_SECONDS = 300
|
|
# Herzschlag in der Datenbank (Sekunden)
|
|
BEAT_SECONDS = 15
|
|
|
|
|
|
class Worker:
|
|
def __init__(self, db: Database, settings: Settings, worker_id: str | None = None):
|
|
self.db, self.settings = db, settings
|
|
self.worker_id = worker_id or f"{socket.gethostname()}-{os.getpid()}-{uuid.uuid4().hex[:6]}"
|
|
self.next_housekeeping = self.next_pr_poll = self.next_hugin_check = 0.0
|
|
self.started, self.runs, self.next_beat = now(), 0, 0.0
|
|
self.probe = monitor.Probe()
|
|
|
|
def housekeeping(self) -> None:
|
|
"""Seltene Aufgaben, höchstens alle HOUSEKEEPING_SECONDS: abgelaufene Workspaces im Papierkorb löschen."""
|
|
if time.monotonic() < self.next_housekeeping:
|
|
return
|
|
self.next_housekeeping = time.monotonic() + HOUSEKEEPING_SECONDS
|
|
removal.purge_expired(self.db, self.settings)
|
|
monitor.prune(self.db)
|
|
|
|
def pull_requests(self) -> None:
|
|
"""Offene Pull Requests der Chat-Branches abfragen; fällig ist jeder alle AIW_PR_POLL_SECONDS."""
|
|
if time.monotonic() < self.next_pr_poll:
|
|
return
|
|
self.next_pr_poll = time.monotonic() + PR_POLL_CHECK_SECONDS
|
|
changes.poll(self.db, self.settings)
|
|
|
|
def hugin_check(self) -> None:
|
|
"""Workspaces mit Hugin abgleichen (nur Befund und Benachrichtigung, nie übernehmen)."""
|
|
if time.monotonic() < self.next_hugin_check:
|
|
return
|
|
self.next_hugin_check = time.monotonic() + HUGIN_CHECK_SECONDS
|
|
hugin.check_due(self.db, self.settings)
|
|
|
|
def heartbeat(self) -> None:
|
|
"""Herzschlag in der Datenbank, höchstens alle BEAT_SECONDS; für /metrics und den Überblick."""
|
|
self.probe.tick()
|
|
if time.monotonic() < self.next_beat:
|
|
return
|
|
self.next_beat = time.monotonic() + BEAT_SECONDS
|
|
monitor.beat(self.db, self.worker_id, self.started, self.runs)
|
|
|
|
def run_once(self) -> bool:
|
|
"""Einen fälligen Run bearbeiten. False, wenn nichts zu tun war."""
|
|
logs.clear()
|
|
logs.bind(worker_id=self.worker_id)
|
|
self.heartbeat()
|
|
engine.recover_expired(self.db, self.settings)
|
|
approvals.expire_due(self.db)
|
|
notify.deliver(self.db, self.settings)
|
|
self.housekeeping()
|
|
self.pull_requests()
|
|
self.hugin_check()
|
|
documents.extract_pending(self.db, self.settings)
|
|
with self.db.session() as session:
|
|
run = engine.claim(session, self.worker_id, self.settings)
|
|
if run is None:
|
|
return False
|
|
context = RunContext(db=self.db, settings=self.settings, run_id=run.id, workspace_id=run.workspace_id,
|
|
chat_session_id=run.chat_session_id, created_by=run.created_by,
|
|
input=dict(run.input), worker_id=self.worker_id,
|
|
events=EventWriter(self.db, run.id, self.settings.event_flush_ms))
|
|
kind = run.kind
|
|
self.process(context, kind)
|
|
return True
|
|
|
|
def process(self, context: RunContext, kind: str) -> None:
|
|
with logs.scope(run_id=str(context.run_id), workspace_id=context.workspace_id, user_id=context.created_by):
|
|
started = time.monotonic()
|
|
log.info("Run gestartet", extra={"event": "run_started", "kind": kind})
|
|
outcome = self.attempt(context, kind)
|
|
seconds = time.monotonic() - started
|
|
self.runs += 1
|
|
metrics.runs_finished.labels(kind, outcome).inc()
|
|
metrics.run_duration.labels(kind).observe(seconds)
|
|
log.info("Run beendet: %s", outcome, extra={"event": "run_finished", "kind": kind, "outcome": outcome,
|
|
"duration_ms": round(seconds * 1000)})
|
|
|
|
def attempt(self, context: RunContext, kind: str) -> str:
|
|
"""Run-Versuch ausführen; liefert das Ergebnis für Metriken und Log."""
|
|
heartbeat = Heartbeat(context, self.worker_id)
|
|
heartbeat.start()
|
|
outcome = "failed"
|
|
try:
|
|
run_handler = HANDLERS.get(kind)
|
|
if run_handler is None:
|
|
raise RuntimeError(f"Kein Handler für Run-Art {kind!r}")
|
|
status = run_handler(context) or states.COMPLETED
|
|
context.check()
|
|
self.finish(context, status)
|
|
outcome = status
|
|
except Cancelled:
|
|
self.finish(context, states.CANCELLED)
|
|
outcome = states.CANCELLED
|
|
except RunStopped as e:
|
|
log.warning("Run %s beendet: %s", context.run_id, e)
|
|
if e.retryable:
|
|
self.failed(context, e, e.user_message)
|
|
outcome = "retry"
|
|
else:
|
|
self.finish(context, e.status, message=e.user_message, error=str(e))
|
|
outcome = e.status
|
|
except LeaseLost:
|
|
log.warning("Lease für %s verloren; Ergebnis wird verworfen", context.run_id)
|
|
outcome = "lease_lost"
|
|
except Exception as e:
|
|
log.exception("Run %s fehlgeschlagen", context.run_id)
|
|
self.failed(context, e)
|
|
finally:
|
|
heartbeat.stop()
|
|
return outcome
|
|
|
|
def finish(self, context: RunContext, status: str, message: str = "", error: str = "") -> None:
|
|
context.events.flush()
|
|
with self.db.session() as session:
|
|
run = engine.owned(session, context.run_id, self.worker_id)
|
|
if run is not None:
|
|
if message:
|
|
run.result = {**run.result, "message": message}
|
|
run.error = error[:2000]
|
|
engine.set_status(session, run, status, **({"message": message} if message else {}))
|
|
|
|
def failed(self, context: RunContext, error: Exception, message: str = "Der Run ist fehlgeschlagen.") -> None:
|
|
"""Erneut versuchen, solange Versuche übrig sind; sonst endgültig fehlgeschlagen.
|
|
|
|
Dem Nutzer wird keine interne Fehlermeldung gezeigt, nur ins Log geschrieben.
|
|
"""
|
|
context.events.flush()
|
|
with self.db.session() as session:
|
|
run = engine.owned(session, context.run_id, self.worker_id)
|
|
if run is None:
|
|
return
|
|
run.error = f"{type(error).__name__}: {error}"[:2000]
|
|
if run.attempts < self.settings.run_max_attempts:
|
|
engine.set_status(session, run, states.QUEUED, reason="retry")
|
|
run.available_at = func.now() + timedelta(seconds=2 ** run.attempts)
|
|
else:
|
|
run.result = {**run.result, "message": message}
|
|
engine.set_status(session, run, states.FAILED, message=message)
|
|
|
|
def run_forever(self, stop: threading.Event) -> None:
|
|
log.info("Worker %s gestartet", self.worker_id)
|
|
server = monitor.serve(self.db, self.settings, self.probe)
|
|
while not stop.is_set():
|
|
try:
|
|
busy = self.run_once()
|
|
except Exception:
|
|
log.exception("Fehler in der Worker-Schleife")
|
|
busy = False
|
|
if not busy:
|
|
stop.wait(self.settings.worker_poll_seconds)
|
|
if server is not None:
|
|
server.shutdown()
|
|
with suppress(Exception):
|
|
monitor.forget(self.db, self.worker_id)
|
|
log.info("Worker %s beendet", self.worker_id)
|