181 lines
7.4 KiB
Python
181 lines
7.4 KiB
Python
"""Deep-research job dispatcher.
|
|
|
|
Mirrors :class:`RoundtableJobDispatcher` exactly in shape — same wake-mechanism
|
|
(``asyncio.Event`` + poll), same claim flow (``list_claimable`` → ``claim_job``
|
|
CAS → ``executor.start_job``), same startup-scan, same cancel-reaper — minus the
|
|
roundtable outbox compensator (deep research has no versioned draft to project
|
|
back into; the session row IS the projection target, written directly by the
|
|
executor).
|
|
|
|
All workers run the same dispatcher; whichever ``claim_job`` returns a row
|
|
(rowcount==1) executes — consistency comes from the DB, not in-process locks.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
from contextlib import suppress
|
|
from datetime import UTC, datetime, timedelta
|
|
|
|
from app.gateway.deep_research_job_executor import (
|
|
_LEASE_TTL_SECONDS,
|
|
WORKER_ID,
|
|
DeepResearchJobExecutor,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_DISPATCH_POLL_INTERVAL = 3.0
|
|
_CLAIM_BATCH = 20
|
|
|
|
|
|
class DeepResearchJobDispatcher:
|
|
"""Cross-worker scheduler for deep-research jobs (one instance per worker)."""
|
|
|
|
def __init__(self, store, executor: DeepResearchJobExecutor, *, worker_id: str | None = None) -> None:
|
|
self._store = store
|
|
self._executor = executor
|
|
self._worker_id = worker_id or WORKER_ID
|
|
self._wake = asyncio.Event()
|
|
self._task: asyncio.Task | None = None
|
|
self._stopping = False
|
|
|
|
# ── lifecycle ───────────────────────────────────────────────────────────
|
|
|
|
def start(self) -> None:
|
|
"""Start the background loop (called from lifespan startup). Idempotent."""
|
|
if self._task is not None and not self._task.done():
|
|
return
|
|
self._stopping = False
|
|
self._task = asyncio.create_task(self._loop(), name="deep-research-job-dispatcher")
|
|
logger.info("[deep-research] dispatcher started (worker=%s)", self._worker_id)
|
|
|
|
async def stop(self) -> None:
|
|
"""Stop the background loop (called from lifespan shutdown)."""
|
|
self._stopping = True
|
|
self._wake.set()
|
|
if self._task is not None:
|
|
self._task.cancel()
|
|
with suppress(asyncio.CancelledError, Exception):
|
|
await self._task
|
|
self._task = None
|
|
|
|
def cancel_running(self, job_id: str) -> bool:
|
|
"""Stop an in-process executor task. Safe if the job is not on this worker."""
|
|
return self._executor.cancel(job_id)
|
|
|
|
def nudge(self) -> None:
|
|
"""Wake the loop immediately (called after start/resume to cut latency)."""
|
|
self._wake.set()
|
|
|
|
# ── main loop ───────────────────────────────────────────────────────────
|
|
|
|
async def _loop(self) -> None:
|
|
# Startup scan: pick up historically queued / expired-lease jobs.
|
|
await self._safe_dispatch()
|
|
while not self._stopping:
|
|
try:
|
|
await asyncio.wait_for(self._wake.wait(), timeout=_DISPATCH_POLL_INTERVAL)
|
|
except TimeoutError:
|
|
pass
|
|
self._wake.clear()
|
|
if self._stopping:
|
|
break
|
|
await self._safe_dispatch()
|
|
|
|
async def _safe_dispatch(self) -> None:
|
|
try:
|
|
await self.dispatch_once()
|
|
except Exception: # noqa: BLE001
|
|
logger.exception("[deep-research] dispatch round failed")
|
|
|
|
# ── claim ───────────────────────────────────────────────────────────────
|
|
|
|
async def dispatch_once(self) -> int:
|
|
"""Scan for claimable jobs and try to claim each. Returns started count."""
|
|
now = datetime.now(UTC)
|
|
await self._reap_cancel_requested(now)
|
|
try:
|
|
candidate_ids = await self._store.list_claimable(now=now, limit=_CLAIM_BATCH)
|
|
except Exception: # noqa: BLE001
|
|
logger.warning("[deep-research] list_claimable failed", exc_info=True)
|
|
return 0
|
|
if not candidate_ids:
|
|
return 0
|
|
|
|
started = 0
|
|
lease_until = now + timedelta(seconds=_LEASE_TTL_SECONDS)
|
|
for job_id in candidate_ids:
|
|
try:
|
|
peek = await self._store.get_unscoped(job_id)
|
|
except Exception: # noqa: BLE001
|
|
peek = None
|
|
# Expired-lease ``running`` rows used to be reclaimed and resumed.
|
|
# After a full Gateway restart that re-occupied the per-user 429
|
|
# latch. Drop them instead of auto-resuming; queued jobs still start.
|
|
if peek is not None and peek.get("status") == "running":
|
|
owner = str(peek.get("user_id") or "")
|
|
with suppress(Exception):
|
|
if owner:
|
|
await self._store.request_cancel(job_id, user_id=owner)
|
|
else:
|
|
await self._store.request_cancel(job_id)
|
|
with suppress(Exception):
|
|
await self._store.finalize_cancel(job_id)
|
|
logger.info("[deep-research] dropped stale running job %s instead of reclaiming", job_id)
|
|
continue
|
|
try:
|
|
row = await self._store.claim_job(
|
|
job_id, lease_owner=self._worker_id, lease_until=lease_until
|
|
)
|
|
except Exception: # noqa: BLE001
|
|
logger.warning("[deep-research] claim_job failed for %s", job_id, exc_info=True)
|
|
continue
|
|
if row is None:
|
|
continue # another worker won / status changed
|
|
|
|
snapshot = row.get("input_snapshot")
|
|
if not snapshot or not isinstance(snapshot, dict):
|
|
logger.error("[deep-research] job %s claimed but snapshot missing; marking error", job_id)
|
|
with suppress(Exception):
|
|
await self._store.update_progress(
|
|
job_id,
|
|
lease_owner=self._worker_id,
|
|
status="failed",
|
|
error_code="SNAPSHOT_MISSING",
|
|
error_message="input_snapshot 缺失或损坏，无法在后台重建研究",
|
|
)
|
|
continue
|
|
|
|
self._executor.start_job(snapshot, job_id=job_id, lease_owner=self._worker_id)
|
|
started += 1
|
|
logger.info(
|
|
"[deep-research] worker=%s claimed job %s (attempt=%s)",
|
|
self._worker_id,
|
|
job_id,
|
|
row.get("attempt"),
|
|
)
|
|
return started
|
|
|
|
async def _reap_cancel_requested(self, now: datetime) -> int:
|
|
"""Finalize orphaned cancellations whose lease expired (worker crash)."""
|
|
try:
|
|
candidate_ids = await self._store.list_cancel_reapable(now=now, limit=_CLAIM_BATCH)
|
|
except Exception: # noqa: BLE001
|
|
return 0
|
|
|
|
finalized = 0
|
|
for job_id in candidate_ids:
|
|
try:
|
|
row = await self._store.finalize_cancel(job_id)
|
|
except Exception: # noqa: BLE001
|
|
continue
|
|
if row is not None:
|
|
finalized += 1
|
|
logger.info("[deep-research] reaped orphaned cancellation for job %s", job_id)
|
|
return finalized
|
|
|
|
|
|
__all__ = ["DeepResearchJobDispatcher"]
|