215 lines
7.3 KiB
Python
215 lines
7.3 KiB
Python
"""Huey task definitions and shared batch state for the admin panel."""
|
|
from __future__ import annotations
|
|
|
|
import threading
|
|
import time
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
from typing import Any
|
|
|
|
from huey import MemoryHuey
|
|
|
|
huey = MemoryHuey("seed", immediate=False)
|
|
|
|
# ── Shared state (all accessed under _lock) ───────────────────────────────────
|
|
_lock = threading.Lock()
|
|
|
|
_current_batch_id: str | None = None
|
|
_current_result: Any = None # Huey Result — kept for revocation
|
|
|
|
_batch_meta: dict[str, Any] = {} # batch_id → aggregate metadata
|
|
_worker_slots: dict[str, dict] = {} # f"{batch_id}:{slot:04d}" → worker info
|
|
_stop_events: dict[str, threading.Event] = {}
|
|
|
|
|
|
# ── Public helpers ─────────────────────────────────────────────────────────────
|
|
|
|
def get_current_batch_id() -> str | None:
|
|
with _lock:
|
|
return _current_batch_id
|
|
|
|
|
|
def set_current(batch_id: str, result: Any) -> None:
|
|
global _current_batch_id, _current_result
|
|
with _lock:
|
|
_current_batch_id = batch_id
|
|
_current_result = result
|
|
|
|
|
|
def stop_current() -> bool:
|
|
"""Signal the running batch to stop. Returns True if there was something to stop."""
|
|
with _lock:
|
|
bid = _current_batch_id
|
|
result = _current_result
|
|
|
|
if bid is None:
|
|
return False
|
|
|
|
# Revoke task if it hasn't started yet (still pending in Huey queue)
|
|
if result is not None:
|
|
try:
|
|
result.revoke(revoke_once=True)
|
|
except Exception:
|
|
pass
|
|
|
|
# Signal running workers to stop after their current step
|
|
ev = _stop_events.get(bid)
|
|
if ev:
|
|
ev.set()
|
|
|
|
with _lock:
|
|
if bid in _batch_meta:
|
|
_batch_meta[bid]["stopping"] = True
|
|
|
|
return True
|
|
|
|
|
|
def batch_status() -> dict[str, Any]:
|
|
with _lock:
|
|
bid = _current_batch_id
|
|
if bid is None:
|
|
return {"batch_id": None, "running": False}
|
|
|
|
meta = dict(_batch_meta.get(bid, {}))
|
|
prefix = f"{bid}:"
|
|
workers = [dict(v) for k, v in _worker_slots.items() if k.startswith(prefix)]
|
|
|
|
meta["batch_id"] = bid
|
|
meta["workers_detail"] = sorted(workers, key=lambda w: w.get("slot", 0))
|
|
meta["active_workers"] = sum(1 for w in workers if w.get("status") == "running")
|
|
return meta
|
|
|
|
|
|
# ── Huey task ─────────────────────────────────────────────────────────────────
|
|
|
|
@huey.task(name="run_batch")
|
|
def run_batch_task(
|
|
batch_id: str,
|
|
cfg: dict[str, Any],
|
|
workers: int,
|
|
total_runs: int,
|
|
headless: bool,
|
|
stagger_ms: int,
|
|
) -> dict[str, Any]:
|
|
import asyncio
|
|
|
|
from crawler.driver import driver_session
|
|
from crawler.dynamic_flow import DynamicFlow
|
|
from crawler.scenarios.fresh_session import FreshSessionScenario
|
|
from db.repositories.runs import runs_repo
|
|
from logger import get_logger
|
|
|
|
log = get_logger(__name__)
|
|
|
|
stop_ev = threading.Event()
|
|
with _lock:
|
|
_stop_events[batch_id] = stop_ev
|
|
_batch_meta[batch_id] = {
|
|
"running": True,
|
|
"stopping": False,
|
|
"total": total_runs,
|
|
"completed": 0,
|
|
"succeeded": 0,
|
|
"failed": 0,
|
|
"configured_workers": workers,
|
|
"started_at": time.time(),
|
|
"finished_at": None,
|
|
"error": None,
|
|
}
|
|
|
|
def _run_one(slot: int) -> bool:
|
|
slot_key = f"{batch_id}:{slot:04d}"
|
|
|
|
if stop_ev.is_set():
|
|
with _lock:
|
|
_worker_slots[slot_key] = {
|
|
"slot": slot, "status": "skipped",
|
|
"started_at": None, "finished_at": None, "duration_ms": None, "error": "stopped",
|
|
}
|
|
return False
|
|
|
|
with _lock:
|
|
_worker_slots[slot_key] = {
|
|
"slot": slot, "status": "running",
|
|
"started_at": time.time(), "finished_at": None, "duration_ms": None, "error": None,
|
|
}
|
|
|
|
t0 = time.monotonic()
|
|
identifier = f"batch-{batch_id[:8]}-{slot:04d}"
|
|
log.info("[batch:%s] slot %d starting", batch_id[:8], slot)
|
|
|
|
try:
|
|
target_url: str = cfg.get("target_url") or ""
|
|
with driver_session(headless=headless) as driver:
|
|
FreshSessionScenario(driver, target_url).run()
|
|
result = DynamicFlow(driver, cfg).run() if cfg else None
|
|
|
|
total_ms = int((time.monotonic() - t0) * 1000)
|
|
success = result.success if result else True
|
|
reason = result.failure_reason if result else None
|
|
steps = [
|
|
{"name": s.name, "success": s.success, "error": s.error, "duration_ms": s.duration_ms}
|
|
for s in (result.steps if result else [])
|
|
]
|
|
asyncio.run(runs_repo.insert("batch", identifier, success, reason, steps, total_ms))
|
|
|
|
status = "ok" if success else "failed"
|
|
with _lock:
|
|
_worker_slots[slot_key].update({
|
|
"status": status, "finished_at": time.time(), "duration_ms": total_ms,
|
|
})
|
|
log.info("[batch:%s] slot %d %s — %d ms", batch_id[:8], slot, status, total_ms)
|
|
return success
|
|
|
|
except Exception as exc:
|
|
total_ms = int((time.monotonic() - t0) * 1000)
|
|
asyncio.run(runs_repo.insert("batch", identifier, False, str(exc), [], total_ms))
|
|
with _lock:
|
|
_worker_slots[slot_key].update({
|
|
"status": "error", "error": str(exc),
|
|
"finished_at": time.time(), "duration_ms": total_ms,
|
|
})
|
|
log.error("[batch:%s] slot %d error — %s", batch_id[:8], slot, exc)
|
|
return False
|
|
|
|
succeeded = 0
|
|
failed = 0
|
|
|
|
try:
|
|
with ThreadPoolExecutor(
|
|
max_workers=workers,
|
|
thread_name_prefix=f"batch-{batch_id[:8]}",
|
|
) as pool:
|
|
futures = []
|
|
for i in range(total_runs):
|
|
if stop_ev.is_set():
|
|
break
|
|
futures.append(pool.submit(_run_one, i))
|
|
if stagger_ms > 0 and i < total_runs - 1:
|
|
if stop_ev.wait(timeout=stagger_ms / 1000):
|
|
break # stop was signalled during the stagger delay
|
|
|
|
for fut in as_completed(futures):
|
|
try:
|
|
ok = fut.result()
|
|
except Exception:
|
|
ok = False
|
|
if ok:
|
|
succeeded += 1
|
|
else:
|
|
failed += 1
|
|
with _lock:
|
|
_batch_meta[batch_id]["completed"] = succeeded + failed
|
|
_batch_meta[batch_id]["succeeded"] = succeeded
|
|
_batch_meta[batch_id]["failed"] = failed
|
|
|
|
except Exception as exc:
|
|
with _lock:
|
|
_batch_meta[batch_id]["error"] = str(exc)
|
|
|
|
finally:
|
|
with _lock:
|
|
_batch_meta[batch_id]["running"] = False
|
|
_batch_meta[batch_id]["finished_at"] = time.time()
|
|
|
|
return {"batch_id": batch_id, "succeeded": succeeded, "failed": failed}
|