409 lines
13 KiB
Python
409 lines
13 KiB
Python
"""Batch task execution via RQ (Redis Queue) — fully decoupled from FastAPI's event loop."""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import sqlite3
|
|
import time
|
|
import uuid
|
|
from typing import Any
|
|
|
|
import redis
|
|
from rq import Queue
|
|
|
|
from db.connection import DB_PATH
|
|
|
|
REDIS_URL = os.environ.get("REDIS_URL", "redis://localhost:6379/0")
|
|
|
|
# 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
|
|
|
|
|
|
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_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 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:
|
|
r = _r()
|
|
current = r.get(_K_CURRENT)
|
|
if not current:
|
|
return False
|
|
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]:
|
|
r = _r()
|
|
bid = r.get(_K_CURRENT)
|
|
if not bid:
|
|
return {"batch_id": None, "running": False, "queue": get_queue()}
|
|
|
|
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
|
|
|
|
|
|
# ── Internal ───────────────────────────────────────────────────────────────────
|
|
|
|
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 logger import get_logger
|
|
|
|
log = get_logger(__name__)
|
|
|
|
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"{slot:04d}"
|
|
|
|
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
|
|
|
|
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}"
|
|
log.info("[batch:%s] slot %d starting", batch_id[:8], slot)
|
|
|
|
try:
|
|
target_url: str = cfg.get("target_url") or ""
|
|
scenario: str = cfg.get("scenario") or "dynamic"
|
|
with driver_session(headless=headless, recording_name=identifier) as driver:
|
|
FreshSessionScenario(driver, target_url).run()
|
|
if cfg:
|
|
if scenario == "digipay":
|
|
from crawler.scenarios.digipay import DigiPayFlow
|
|
result = DigiPayFlow(driver, cfg).run()
|
|
else:
|
|
result = DynamicFlow(driver, cfg).run()
|
|
else:
|
|
result = 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 [])
|
|
]
|
|
_db_insert_run("batch", identifier, success, reason, steps, total_ms)
|
|
|
|
status = "ok" if success else "failed"
|
|
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)
|
|
_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
|
|
|
|
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 r.exists(stop_key):
|
|
break
|
|
futures.append(pool.submit(_run_one, i))
|
|
if stagger_ms > 0 and i < total_runs - 1:
|
|
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:
|
|
ok = fut.result()
|
|
except Exception:
|
|
ok = False
|
|
if ok:
|
|
succeeded += 1
|
|
else:
|
|
failed += 1
|
|
|
|
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:
|
|
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:
|
|
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))
|
|
|
|
_db_finish_batch_run(batch_id, succeeded, failed, bool(r.exists(stop_key)), error_msg)
|
|
_start_next_in_queue()
|