114 lines
3.7 KiB
Python
114 lines
3.7 KiB
Python
# retoor <retoor@molodetz.nl>
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
from typing import Any, Awaitable, Callable
|
|
|
|
from .controller import compute_followup
|
|
from .schedule import now_utc, to_iso
|
|
from .store import TaskStore
|
|
|
|
logger = logging.getLogger("devii.tasks.scheduler")
|
|
|
|
|
|
def _audit_task_execute(row: dict[str, Any], result: str, error: str | None) -> None:
|
|
from devplacepy.services.audit import record as audit
|
|
|
|
owner_kind = row.get("owner_kind") or "user"
|
|
owner_id = row.get("owner_id") or ""
|
|
audit.record_system(
|
|
"devii.task.execute",
|
|
actor_kind="user" if owner_kind == "user" else owner_kind,
|
|
actor_uid=owner_id if owner_kind == "user" else None,
|
|
actor_role="user" if owner_kind == "user" else owner_kind,
|
|
origin="scheduler",
|
|
via_agent=1,
|
|
result="failure" if error else "success",
|
|
target_type="task",
|
|
target_uid=row.get("uid"),
|
|
summary=f"Devii task {row.get('uid')} for {owner_id or owner_kind} executed",
|
|
metadata={"run_count": row.get("run_count"), "error": error, "result": (result or "")[:140]},
|
|
links=[audit.task(row.get("uid"))],
|
|
)
|
|
|
|
PromptExecutor = Callable[[str], Awaitable[str]]
|
|
EventCallback = Callable[[str, dict[str, Any], str], None]
|
|
|
|
DEFAULT_TICK_SECONDS = 1.0
|
|
|
|
|
|
class Scheduler:
|
|
def __init__(
|
|
self,
|
|
store: TaskStore,
|
|
executor: PromptExecutor,
|
|
on_event: EventCallback,
|
|
tick_seconds: float = DEFAULT_TICK_SECONDS,
|
|
) -> None:
|
|
self._store = store
|
|
self._executor = executor
|
|
self._on_event = on_event
|
|
self._tick_seconds = tick_seconds
|
|
self._loop_task: asyncio.Task[None] | None = None
|
|
|
|
def start(self) -> None:
|
|
if self._loop_task is None:
|
|
self._store.recover_running()
|
|
self._loop_task = asyncio.create_task(self._loop())
|
|
logger.info("Scheduler started")
|
|
|
|
async def stop(self) -> None:
|
|
if self._loop_task is None:
|
|
return
|
|
self._loop_task.cancel()
|
|
try:
|
|
await self._loop_task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
self._loop_task = None
|
|
logger.info("Scheduler stopped")
|
|
|
|
async def _loop(self) -> None:
|
|
while True:
|
|
try:
|
|
await self._tick()
|
|
except Exception: # noqa: BLE001 - the scheduler must never die
|
|
logger.exception("Scheduler tick failed")
|
|
await asyncio.sleep(self._tick_seconds)
|
|
|
|
async def _tick(self) -> None:
|
|
now = now_utc()
|
|
for row in self._store.due(to_iso(now)):
|
|
await self._execute(row)
|
|
|
|
async def _execute(self, row: dict[str, Any]) -> None:
|
|
uid = row["uid"]
|
|
self._store.update(uid, {"status": "running"})
|
|
self._on_event("start", row, "")
|
|
logger.info("Executing task uid=%s", uid)
|
|
|
|
try:
|
|
result = await self._executor(row["prompt"])
|
|
changes, finished = compute_followup(row, now_utc())
|
|
changes["last_result"] = result
|
|
changes["last_error"] = None
|
|
except Exception as exc: # noqa: BLE001 - surfaced into the task record
|
|
logger.exception("Task uid=%s crashed", uid)
|
|
self._store.update(
|
|
uid,
|
|
{
|
|
"status": "error",
|
|
"last_error": str(exc),
|
|
"last_run_at": to_iso(now_utc()),
|
|
},
|
|
)
|
|
self._on_event("error", row, str(exc))
|
|
_audit_task_execute(row, "", str(exc))
|
|
return
|
|
|
|
self._store.update(uid, changes)
|
|
self._on_event("done" if not finished else "finished", row, result)
|
|
_audit_task_execute(row, result, None)
|