wip
This commit is contained in:
@@ -0,0 +1,214 @@
|
||||
"""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}
|
||||
Reference in New Issue
Block a user