updates :)
This commit is contained in:
+299
-113
@@ -1,137 +1,306 @@
|
||||
"""Huey task definitions and shared batch state for the admin panel."""
|
||||
"""Batch task execution via RQ (Redis Queue) — fully decoupled from FastAPI's event loop."""
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
import time
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
import uuid
|
||||
from typing import Any
|
||||
|
||||
from huey import MemoryHuey
|
||||
import redis
|
||||
from rq import Queue
|
||||
|
||||
huey = MemoryHuey("seed", immediate=False)
|
||||
from db.connection import DB_PATH
|
||||
|
||||
# ── Shared state (all accessed under _lock) ───────────────────────────────────
|
||||
_lock = threading.Lock()
|
||||
REDIS_URL = os.environ.get("REDIS_URL", "redis://localhost:6379/0")
|
||||
|
||||
_current_batch_id: str | None = None
|
||||
_current_result: Any = None # Huey Result — kept for revocation
|
||||
# Redis key helpers
|
||||
_K_CURRENT = "seed:current" # string: current batch_id (or empty)
|
||||
_K_META = "seed:meta:{}" # JSON string: batch aggregate meta
|
||||
_K_SLOTS = "seed:slots:{}" # hash: slot_key → JSON worker info
|
||||
_K_STOP = "seed:stop:{}" # string: "1" if stop requested
|
||||
_K_QUEUE = "seed:pending_queue" # list: JSON-encoded pending items
|
||||
|
||||
_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] = {}
|
||||
|
||||
def _r() -> redis.Redis: # type: ignore[type-arg]
|
||||
return redis.from_url(REDIS_URL, decode_responses=True)
|
||||
|
||||
|
||||
def _q() -> Queue:
|
||||
return Queue("batch", connection=redis.from_url(REDIS_URL))
|
||||
|
||||
|
||||
# ── Sync DB helpers ────────────────────────────────────────────────────────────
|
||||
|
||||
def _db_insert_run(
|
||||
scenario: str,
|
||||
identifier: str,
|
||||
success: bool,
|
||||
failure_reason: str | None,
|
||||
steps: list[dict[str, Any]],
|
||||
total_ms: int,
|
||||
) -> None:
|
||||
conn = sqlite3.connect(DB_PATH)
|
||||
try:
|
||||
conn.execute(
|
||||
"INSERT INTO runs (scenario, identifier, success, failure_reason, steps_json, total_ms)"
|
||||
" VALUES (?,?,?,?,?,?)",
|
||||
(scenario, identifier, int(success), failure_reason, json.dumps(steps), total_ms),
|
||||
)
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def _db_insert_batch_run(
|
||||
batch_id: str,
|
||||
config_id: int | None,
|
||||
config_name: str,
|
||||
workers: int,
|
||||
total_runs: int,
|
||||
stagger_ms: int,
|
||||
headless: bool,
|
||||
) -> None:
|
||||
conn = sqlite3.connect(DB_PATH)
|
||||
try:
|
||||
conn.execute(
|
||||
"INSERT INTO batch_runs"
|
||||
" (batch_id, config_id, config_name, workers, total_runs, stagger_ms, headless)"
|
||||
" VALUES (?,?,?,?,?,?,?)",
|
||||
(batch_id, config_id, config_name, workers, total_runs, stagger_ms, int(headless)),
|
||||
)
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def _db_finish_batch_run(
|
||||
batch_id: str,
|
||||
succeeded: int,
|
||||
failed: int,
|
||||
stopped: bool,
|
||||
error: str | None,
|
||||
) -> None:
|
||||
conn = sqlite3.connect(DB_PATH)
|
||||
try:
|
||||
conn.execute(
|
||||
"UPDATE batch_runs"
|
||||
" SET finished_at=strftime('%Y-%m-%dT%H:%M:%SZ','now'),"
|
||||
" succeeded=?, failed=?, stopped=?, error=?"
|
||||
" WHERE batch_id=?",
|
||||
(succeeded, failed, int(stopped), error, batch_id),
|
||||
)
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
# ── Public helpers ─────────────────────────────────────────────────────────────
|
||||
|
||||
def get_current_batch_id() -> str | None:
|
||||
with _lock:
|
||||
return _current_batch_id
|
||||
def get_queue() -> list[dict[str, Any]]:
|
||||
r = _r()
|
||||
raw = r.lrange(_K_QUEUE, 0, -1)
|
||||
items = []
|
||||
for s in raw:
|
||||
try:
|
||||
item = json.loads(s)
|
||||
items.append({k: v for k, v in item.items() if k != "cfg"})
|
||||
except Exception:
|
||||
pass
|
||||
return items
|
||||
|
||||
|
||||
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 enqueue(
|
||||
cfg: dict[str, Any],
|
||||
workers: int,
|
||||
total_runs: int,
|
||||
stagger_ms: int,
|
||||
headless: bool,
|
||||
) -> dict[str, Any]:
|
||||
queue_id = uuid.uuid4().hex
|
||||
|
||||
item: dict[str, Any] = {
|
||||
"queue_id": queue_id,
|
||||
"cfg": cfg,
|
||||
"config_id": cfg.get("id"),
|
||||
"config_name": cfg.get("name") or "",
|
||||
"workers": workers,
|
||||
"total_runs": total_runs,
|
||||
"stagger_ms": stagger_ms,
|
||||
"headless": headless,
|
||||
"added_at": time.time(),
|
||||
}
|
||||
|
||||
r = _r()
|
||||
current = r.get(_K_CURRENT)
|
||||
if current:
|
||||
meta_raw = r.get(_K_META.format(current))
|
||||
if meta_raw:
|
||||
meta = json.loads(meta_raw)
|
||||
if meta.get("running"):
|
||||
r.rpush(_K_QUEUE, json.dumps(item))
|
||||
return {"queue_id": queue_id, "status": "queued"}
|
||||
|
||||
_launch(item)
|
||||
return {"queue_id": queue_id, "status": "started"}
|
||||
|
||||
|
||||
def dequeue(queue_id: str) -> bool:
|
||||
r = _r()
|
||||
raw = r.lrange(_K_QUEUE, 0, -1)
|
||||
for s in raw:
|
||||
try:
|
||||
item = json.loads(s)
|
||||
if item.get("queue_id") == queue_id:
|
||||
r.lrem(_K_QUEUE, 1, s)
|
||||
return True
|
||||
except Exception:
|
||||
pass
|
||||
return False
|
||||
|
||||
|
||||
def clear_queue() -> int:
|
||||
r = _r()
|
||||
raw = r.lrange(_K_QUEUE, 0, -1)
|
||||
count = len(raw)
|
||||
r.delete(_K_QUEUE)
|
||||
return count
|
||||
|
||||
|
||||
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:
|
||||
r = _r()
|
||||
current = r.get(_K_CURRENT)
|
||||
if not current:
|
||||
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
|
||||
|
||||
meta_raw = r.get(_K_META.format(current))
|
||||
if not meta_raw:
|
||||
return False
|
||||
meta = json.loads(meta_raw)
|
||||
if not meta.get("running"):
|
||||
return False
|
||||
r.set(_K_STOP.format(current), "1")
|
||||
meta["stopping"] = True
|
||||
r.set(_K_META.format(current), json.dumps(meta))
|
||||
return True
|
||||
|
||||
|
||||
def batch_status() -> dict[str, Any]:
|
||||
with _lock:
|
||||
bid = _current_batch_id
|
||||
if bid is None:
|
||||
return {"batch_id": None, "running": False}
|
||||
r = _r()
|
||||
bid = r.get(_K_CURRENT)
|
||||
if not bid:
|
||||
return {"batch_id": None, "running": False, "queue": get_queue()}
|
||||
|
||||
meta = dict(_batch_meta.get(bid, {}))
|
||||
prefix = f"{bid}:"
|
||||
workers = [dict(v) for k, v in _worker_slots.items() if k.startswith(prefix)]
|
||||
meta_raw = r.get(_K_META.format(bid))
|
||||
meta = json.loads(meta_raw) if meta_raw else {}
|
||||
|
||||
slot_data = r.hgetall(_K_SLOTS.format(bid))
|
||||
workers = []
|
||||
for v in slot_data.values():
|
||||
try:
|
||||
workers.append(json.loads(v))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
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")
|
||||
meta["queue"] = get_queue()
|
||||
return meta
|
||||
|
||||
|
||||
# ── Huey task ─────────────────────────────────────────────────────────────────
|
||||
# ── Internal ───────────────────────────────────────────────────────────────────
|
||||
|
||||
@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
|
||||
def _launch(item: dict[str, Any]) -> None:
|
||||
batch_id = uuid.uuid4().hex
|
||||
r = _r()
|
||||
r.set(_K_CURRENT, batch_id)
|
||||
|
||||
q = _q()
|
||||
q.enqueue(
|
||||
run_batch_job,
|
||||
batch_id,
|
||||
item,
|
||||
job_id=batch_id,
|
||||
job_timeout=-1,
|
||||
result_ttl=3600,
|
||||
)
|
||||
|
||||
|
||||
def _start_next_in_queue() -> None:
|
||||
r = _r()
|
||||
raw = r.lpop(_K_QUEUE)
|
||||
if not raw:
|
||||
return
|
||||
try:
|
||||
item = json.loads(raw)
|
||||
except Exception:
|
||||
return
|
||||
_launch(item)
|
||||
|
||||
|
||||
# ── RQ job (runs in worker process, no asyncio) ────────────────────────────────
|
||||
|
||||
def run_batch_job(batch_id: str, item: dict[str, Any]) -> None:
|
||||
"""Top-level RQ job function. Runs in a forked worker process."""
|
||||
import threading
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
|
||||
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,
|
||||
}
|
||||
r = redis.from_url(REDIS_URL, decode_responses=True)
|
||||
|
||||
cfg = item["cfg"]
|
||||
workers = item["workers"]
|
||||
total_runs = item["total_runs"]
|
||||
stagger_ms = item["stagger_ms"]
|
||||
headless = item["headless"]
|
||||
config_id: int | None = int(cfg["id"]) if cfg and cfg.get("id") is not None else None
|
||||
config_name: str = str(cfg.get("name") or "") if cfg else ""
|
||||
|
||||
_db_insert_batch_run(batch_id, config_id, config_name, workers, total_runs, stagger_ms, headless)
|
||||
|
||||
meta: dict[str, Any] = {
|
||||
"running": True,
|
||||
"stopping": False,
|
||||
"total": total_runs,
|
||||
"completed": 0,
|
||||
"succeeded": 0,
|
||||
"failed": 0,
|
||||
"configured_workers": workers,
|
||||
"config_id": config_id,
|
||||
"config_name": config_name,
|
||||
"started_at": time.time(),
|
||||
"finished_at": None,
|
||||
"error": None,
|
||||
}
|
||||
r.set(_K_META.format(batch_id), json.dumps(meta))
|
||||
|
||||
slots_key = _K_SLOTS.format(batch_id)
|
||||
stop_key = _K_STOP.format(batch_id)
|
||||
meta_key = _K_META.format(batch_id)
|
||||
|
||||
_slots_lock = threading.Lock()
|
||||
|
||||
def _run_one(slot: int) -> bool:
|
||||
slot_key = f"{batch_id}:{slot:04d}"
|
||||
slot_key = f"{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",
|
||||
}
|
||||
if r.exists(stop_key):
|
||||
r.hset(slots_key, slot_key, json.dumps({
|
||||
"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,
|
||||
}
|
||||
r.hset(slots_key, slot_key, json.dumps({
|
||||
"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}"
|
||||
@@ -150,24 +319,23 @@ def run_batch_task(
|
||||
{"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))
|
||||
_db_insert_run("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,
|
||||
})
|
||||
r.hset(slots_key, slot_key, json.dumps({
|
||||
"slot": slot, "status": status,
|
||||
"finished_at": time.time(), "duration_ms": total_ms, "error": None,
|
||||
}))
|
||||
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,
|
||||
})
|
||||
_db_insert_run("batch", identifier, False, str(exc), [], total_ms)
|
||||
r.hset(slots_key, slot_key, json.dumps({
|
||||
"slot": slot, "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
|
||||
|
||||
@@ -181,12 +349,17 @@ def run_batch_task(
|
||||
) as pool:
|
||||
futures = []
|
||||
for i in range(total_runs):
|
||||
if stop_ev.is_set():
|
||||
if r.exists(stop_key):
|
||||
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
|
||||
deadline = time.monotonic() + stagger_ms / 1000
|
||||
while time.monotonic() < deadline:
|
||||
if r.exists(stop_key):
|
||||
break
|
||||
time.sleep(0.1)
|
||||
if r.exists(stop_key):
|
||||
break
|
||||
|
||||
for fut in as_completed(futures):
|
||||
try:
|
||||
@@ -197,18 +370,31 @@ def run_batch_task(
|
||||
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
|
||||
|
||||
current_meta_raw = r.get(meta_key)
|
||||
if current_meta_raw:
|
||||
current_meta = json.loads(current_meta_raw)
|
||||
current_meta["completed"] = succeeded + failed
|
||||
current_meta["succeeded"] = succeeded
|
||||
current_meta["failed"] = failed
|
||||
r.set(meta_key, json.dumps(current_meta))
|
||||
|
||||
except Exception as exc:
|
||||
with _lock:
|
||||
_batch_meta[batch_id]["error"] = str(exc)
|
||||
current_meta_raw = r.get(meta_key)
|
||||
if current_meta_raw:
|
||||
current_meta = json.loads(current_meta_raw)
|
||||
current_meta["error"] = str(exc)
|
||||
r.set(meta_key, json.dumps(current_meta))
|
||||
|
||||
finally:
|
||||
with _lock:
|
||||
_batch_meta[batch_id]["running"] = False
|
||||
_batch_meta[batch_id]["finished_at"] = time.time()
|
||||
final_meta_raw = r.get(meta_key)
|
||||
error_msg: str | None = None
|
||||
if final_meta_raw:
|
||||
final_meta = json.loads(final_meta_raw)
|
||||
error_msg = final_meta.get("error")
|
||||
final_meta["running"] = False
|
||||
final_meta["finished_at"] = time.time()
|
||||
r.set(meta_key, json.dumps(final_meta))
|
||||
|
||||
return {"batch_id": batch_id, "succeeded": succeeded, "failed": failed}
|
||||
_db_finish_batch_run(batch_id, succeeded, failed, bool(r.exists(stop_key)), error_msg)
|
||||
_start_next_in_queue()
|
||||
|
||||
Reference in New Issue
Block a user