You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
423 lines
15 KiB
423 lines
15 KiB
"""
|
|
GEO 持续采集 Worker v2 — 多进程架构
|
|
|
|
每个平台独立进程池, 主进程拉任务→分发→收结果→提交。
|
|
各子进程独立加载 Client / WASM / cookie, 无 GIL 竞争。
|
|
|
|
用法:
|
|
python main.py
|
|
python main.py --config worker_config.json
|
|
python main.py --once
|
|
python main.py --dry-run
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import re
|
|
import signal
|
|
import sys
|
|
import time
|
|
from datetime import datetime
|
|
from multiprocessing import Process, Queue, Event
|
|
from typing import Any
|
|
|
|
sys.stdout.reconfigure(encoding="utf-8")
|
|
|
|
import urllib3
|
|
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
|
|
|
|
from config import (
|
|
load_config, normalize_platform_name, API_TO_INTERNAL,
|
|
SUPPORTED_API_PLATFORMS, _parse_bool, _default_switches,
|
|
)
|
|
from api_client import GeoApiClient
|
|
from payload import build_submit_payload
|
|
from worker import WORKER_FUNCS
|
|
|
|
_stop_event = Event()
|
|
|
|
|
|
def is_stopping():
|
|
return _stop_event.is_set()
|
|
|
|
|
|
def _handle_signal(sig, frame):
|
|
if _stop_event.is_set():
|
|
print("\n[FORCE] 二次 Ctrl+C, 立即退出")
|
|
os._exit(1)
|
|
_stop_event.set()
|
|
print("\n[stop] 收到停止信号, 完成当前任务后退出...")
|
|
|
|
|
|
# ---------- 平台开关 ----------
|
|
def _merge_switches(base, override):
|
|
if not override or not isinstance(override, dict):
|
|
return dict(base)
|
|
merged = dict(base)
|
|
for k, v in override.items():
|
|
plat = normalize_platform_name(str(k))
|
|
if plat and plat in merged:
|
|
try:
|
|
merged[plat] = _parse_bool(v)
|
|
except ValueError:
|
|
pass
|
|
return merged
|
|
|
|
|
|
class PlatformSwitchBoard:
|
|
def __init__(self, switches, hot_reload=True, config_path=""):
|
|
self._switches = dict(switches)
|
|
self.hot_reload = hot_reload
|
|
self.config_path = config_path
|
|
self._last_mtime = 0.0
|
|
|
|
def is_enabled(self, platform):
|
|
plat = normalize_platform_name(platform)
|
|
return bool(self._switches.get(plat))
|
|
|
|
def enabled_list(self):
|
|
return [p for p in SUPPORTED_API_PLATFORMS if self._switches.get(p)]
|
|
|
|
def query_platforms(self):
|
|
return ",".join(self.enabled_list())
|
|
|
|
def summary(self):
|
|
on = [p for p in SUPPORTED_API_PLATFORMS if self._switches.get(p)]
|
|
off = [p for p in SUPPORTED_API_PLATFORMS if not self._switches.get(p)]
|
|
return f"开启=[{','.join(on) or '-'}] 关闭=[{','.join(off) or '-'}]"
|
|
|
|
def maybe_reload(self) -> bool:
|
|
if not self.hot_reload or not os.path.exists(self.config_path):
|
|
return False
|
|
try:
|
|
mtime = os.path.getmtime(self.config_path)
|
|
if mtime <= self._last_mtime:
|
|
return False
|
|
self._last_mtime = mtime
|
|
with open(self.config_path, "r", encoding="utf-8") as f:
|
|
file_cfg = json.load(f)
|
|
new_switches = _merge_switches(_default_switches(), file_cfg.get("platform_switches"))
|
|
if new_switches != self._switches:
|
|
old = self.summary()
|
|
self._switches = new_switches
|
|
print(f"[hot-reload] 平台开关更新: {old} -> {self.summary()}")
|
|
return True
|
|
except Exception as e:
|
|
print(f"[hot-reload] 失败: {e}")
|
|
return False
|
|
|
|
|
|
# ---------- 工具 ----------
|
|
FAILED_DIR = ""
|
|
|
|
|
|
def _save_failed(stage, task, detail):
|
|
os.makedirs(FAILED_DIR, exist_ok=True)
|
|
ts = datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
did = task.get("dispatchId") or task.get("id") or "unknown"
|
|
safe_did = re.sub(r'[\\/:*?"<>|]', '_', str(did))
|
|
path = os.path.join(FAILED_DIR, f"{ts}_{safe_did}_{stage}.json")
|
|
with open(path, "w", encoding="utf-8") as f:
|
|
json.dump({"stage": stage, "task": task, "detail": detail}, f, ensure_ascii=False, indent=2)
|
|
print(f" [saved] {path}")
|
|
|
|
|
|
# ---------- 主流程 ----------
|
|
def run_worker(config_path=None, overrides=None):
|
|
global FAILED_DIR
|
|
|
|
cfg = load_config(config_path)
|
|
if overrides:
|
|
cfg.update(overrides)
|
|
base_dir = os.path.dirname(cfg["_config_path"])
|
|
FAILED_DIR = os.path.join(base_dir, "worker_failed")
|
|
os.makedirs(FAILED_DIR, exist_ok=True)
|
|
|
|
if not cfg.get("spider_token"):
|
|
print(f"[FATAL] 请在配置文件填写 spider_token:\n {cfg['_config_path']}")
|
|
return
|
|
|
|
switches = _merge_switches(_default_switches(), cfg.get("platform_switches"))
|
|
board = PlatformSwitchBoard(switches, hot_reload=bool(cfg.get("hot_reload", True)), config_path=cfg["_config_path"])
|
|
|
|
once = bool(cfg.get("once"))
|
|
dry_run = bool(cfg.get("dry_run"))
|
|
idle_sleep = float(cfg.get("idle_sleep") or 30)
|
|
enable_commit = bool(cfg.get("enable_task_commit", True))
|
|
location = cfg.get("location") or ""
|
|
conc = cfg.get("concurrency") or {}
|
|
|
|
# 按平台创建队列和进程 (每个平台 1 个进程, 内部按 concurrency 并发)
|
|
platform_queues = {}
|
|
result_queue = Queue()
|
|
workers = []
|
|
|
|
enabled = board.enabled_list()
|
|
conc_str_parts = []
|
|
for plat in ("deepseek", "kimi", "tongyi", "baiduai", "doubao"):
|
|
n = max(1, int(conc.get(plat, 1)))
|
|
conc_str_parts.append(f"{plat}={n}")
|
|
if plat not in enabled:
|
|
continue
|
|
if plat not in WORKER_FUNCS:
|
|
continue
|
|
q = Queue()
|
|
platform_queues[plat] = q
|
|
p = Process(
|
|
target=WORKER_FUNCS[plat],
|
|
args=(0, q, result_queue, _stop_event, cfg),
|
|
daemon=True,
|
|
)
|
|
workers.append((plat, 0, p))
|
|
|
|
print("=" * 60)
|
|
print("GEO Task Worker v2 (多进程)")
|
|
print(f" 配置文件 : {cfg['_config_path']}")
|
|
print(f" API : {cfg['api_base']}")
|
|
print(f" 拉题平台 : {board.query_platforms() or '(无)'}")
|
|
print(f" 启停控制 : {board.summary()}")
|
|
print(f" 并发(进程): {' | '.join(conc_str_parts)}")
|
|
print(f" count : {cfg['batch_count']}")
|
|
print(f" idle : {idle_sleep}s")
|
|
print(f" commit : {enable_commit}")
|
|
print(f" dry_run : {dry_run}")
|
|
print(f" once : {once}")
|
|
print(f" 进程总数 : {len(workers)}")
|
|
print("=" * 60)
|
|
|
|
if not platform_queues:
|
|
print("无可用平台, 请在 worker_config.json 中开启至少一个平台")
|
|
return
|
|
|
|
# 启动进程
|
|
for plat, wid, p in workers:
|
|
p.start()
|
|
print(f" [{plat}-{wid}] started pid={p.pid}")
|
|
|
|
geo = GeoApiClient(cfg["api_base"], cfg["spider_token"], cfg.get("vendor_token") or "")
|
|
stats = {"pulled": 0, "ok": 0, "fail": 0, "skip": 0, "disabled": 0, "submitted": 0}
|
|
|
|
def collect_results(block=False, timeout=0.1):
|
|
"""从 result_queue 中取所有可用结果并提交"""
|
|
while True:
|
|
try:
|
|
r = result_queue.get(block=block, timeout=timeout)
|
|
except Exception:
|
|
break
|
|
task = r["task"]
|
|
did = task.get("dispatchId") or task.get("id") or "?"
|
|
|
|
if not r["ok"]:
|
|
stats["fail"] += 1
|
|
err = r.get("error", "unknown")
|
|
print(f" [main] ✗ did={did} error={err}")
|
|
_save_failed("collect", task, err)
|
|
continue
|
|
|
|
parsed = r["parsed"]
|
|
payload = build_submit_payload(task, parsed, location=location)
|
|
if not payload.get("id"):
|
|
stats["fail"] += 1
|
|
_save_failed("submit", task, "missing_id")
|
|
continue
|
|
|
|
if dry_run:
|
|
print(f" [main] [dry-run] did={did} 跳过提交")
|
|
stats["ok"] += 1
|
|
continue
|
|
|
|
try:
|
|
resp = geo.submit(payload)
|
|
code = resp.get("code")
|
|
print(f" [main] [submit] did={did} code={code} msg={resp.get('msg')}")
|
|
if code in (0, 200, "0"):
|
|
stats["ok"] += 1
|
|
stats["submitted"] += 1
|
|
else:
|
|
stats["fail"] += 1
|
|
_save_failed("submit", task, resp)
|
|
except Exception as e:
|
|
stats["fail"] += 1
|
|
print(f" [main] [submit] did={did} 异常: {e}")
|
|
_save_failed("submit", task, str(e))
|
|
|
|
try:
|
|
while not is_stopping():
|
|
board.maybe_reload()
|
|
pull_platforms = board.query_platforms()
|
|
if not pull_platforms:
|
|
print(f"[idle] 全部平台已关闭, {idle_sleep}s 后重试...")
|
|
if once:
|
|
break
|
|
time.sleep(idle_sleep)
|
|
continue
|
|
|
|
try:
|
|
tasks = []
|
|
for plat in pull_platforms.split(","):
|
|
if is_stopping():
|
|
break
|
|
plat = plat.strip()
|
|
if not plat:
|
|
continue
|
|
plat_tasks = geo.get_tasks(plat, cfg["batch_count"])
|
|
tasks.extend(plat_tasks)
|
|
except PermissionError as e:
|
|
print(f"[FATAL] {e}")
|
|
break
|
|
except Exception as e:
|
|
if is_stopping():
|
|
break
|
|
print(f"[get_task] 异常: {e}, {idle_sleep}s 后重试")
|
|
time.sleep(idle_sleep)
|
|
if once:
|
|
break
|
|
continue
|
|
|
|
if is_stopping():
|
|
break
|
|
|
|
if not tasks:
|
|
collect_results()
|
|
print(f"[idle] 无任务 platforms={pull_platforms}, {idle_sleep}s 后重试...")
|
|
if once:
|
|
break
|
|
time.sleep(idle_sleep)
|
|
continue
|
|
|
|
# 统计分布
|
|
plat_counts = {}
|
|
for t in tasks:
|
|
p = (t.get("platform") or "unknown").strip().lower()
|
|
plat_counts[p] = plat_counts.get(p, 0) + 1
|
|
breakdown = " | ".join(f"{p}={n}" for p, n in sorted(plat_counts.items()))
|
|
print(f"\n[pull] 取得 {len(tasks)} 条任务 (platforms={pull_platforms})")
|
|
print(f" 分布: {breakdown}")
|
|
|
|
# commit + 分发
|
|
dispatched = 0
|
|
for task in tasks:
|
|
if is_stopping():
|
|
break
|
|
stats["pulled"] += 1
|
|
api_plat = (task.get("platform") or "").strip().lower()
|
|
plat = normalize_platform_name(api_plat) or api_plat
|
|
|
|
if API_TO_INTERNAL.get(plat) is None:
|
|
print(f" skip 不支持平台 {api_plat} dispatchId={task.get('dispatchId')}")
|
|
stats["skip"] += 1
|
|
continue
|
|
if not board.is_enabled(plat):
|
|
print(f" skip 平台已停用 {plat}")
|
|
stats["disabled"] += 1
|
|
continue
|
|
if plat not in platform_queues:
|
|
print(f" skip 无 worker: {plat}")
|
|
stats["skip"] += 1
|
|
continue
|
|
|
|
if enable_commit:
|
|
try:
|
|
resp = geo.task_commit(task)
|
|
code = resp.get("code")
|
|
print(f" [commit] did={task.get('dispatchId')} code={code}")
|
|
except Exception as e:
|
|
print(f" [commit] 异常: {e}")
|
|
|
|
platform_queues[plat].put(task)
|
|
dispatched += 1
|
|
|
|
print(f" 已分发 {dispatched} 条任务到子进程")
|
|
|
|
# 等待本轮结果 (非阻塞轮询, 直到所有分发的任务都返回)
|
|
collected = 0
|
|
wait_start = time.time()
|
|
while collected < dispatched and not is_stopping():
|
|
try:
|
|
r = result_queue.get(timeout=5)
|
|
collected += 1
|
|
task = r["task"]
|
|
did = task.get("dispatchId") or task.get("id") or "?"
|
|
|
|
if not r["ok"]:
|
|
stats["fail"] += 1
|
|
print(f" [main] ✗ did={did} error={r.get('error', '?')}")
|
|
_save_failed("collect", task, r.get("error", "unknown"))
|
|
continue
|
|
|
|
parsed = r["parsed"]
|
|
payload = build_submit_payload(task, parsed, location=location)
|
|
if not payload.get("id"):
|
|
stats["fail"] += 1
|
|
_save_failed("submit", task, "missing_id")
|
|
continue
|
|
|
|
if dry_run:
|
|
print(f" [main] [dry-run] did={did}")
|
|
stats["ok"] += 1
|
|
continue
|
|
|
|
try:
|
|
resp = geo.submit(payload)
|
|
code = resp.get("code")
|
|
print(f" [main] [submit] did={did} code={code}")
|
|
if code in (0, 200, "0"):
|
|
stats["ok"] += 1
|
|
stats["submitted"] += 1
|
|
else:
|
|
stats["fail"] += 1
|
|
_save_failed("submit", task, resp)
|
|
except Exception as e:
|
|
stats["fail"] += 1
|
|
_save_failed("submit", task, str(e))
|
|
except Exception:
|
|
if time.time() - wait_start > 600:
|
|
print("[main] 等待结果超时 (600s)")
|
|
break
|
|
|
|
print(f"\n[stats] pulled={stats['pulled']} ok={stats['ok']} "
|
|
f"fail={stats['fail']} skip={stats['skip']} submitted={stats['submitted']}")
|
|
print(f"[switch] {board.summary()}")
|
|
if once or is_stopping():
|
|
break
|
|
|
|
finally:
|
|
print("\n[main] 正在关闭子进程...")
|
|
_stop_event.set()
|
|
for plat, q in platform_queues.items():
|
|
# 多发几个毒丸确保每个 worker 都能收到
|
|
n = max(1, int(conc.get(plat, 1)))
|
|
for _ in range(n):
|
|
try:
|
|
q.put(None)
|
|
except Exception:
|
|
pass
|
|
for plat, wid, p in workers:
|
|
p.join(timeout=30)
|
|
if p.is_alive():
|
|
print(f" [{plat}-{wid}] 强制终止")
|
|
p.terminate()
|
|
print(f"\n[main] 结束 {stats}")
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="GEO 持续采集 Worker v2 (多进程)")
|
|
parser.add_argument("--config", default=None, help="配置文件路径")
|
|
parser.add_argument("--once", action="store_true", help="只跑一轮")
|
|
parser.add_argument("--dry-run", action="store_true", help="不提交结果")
|
|
args = parser.parse_args()
|
|
|
|
signal.signal(signal.SIGINT, _handle_signal)
|
|
signal.signal(signal.SIGTERM, _handle_signal)
|
|
|
|
overrides = {}
|
|
if args.once:
|
|
overrides["once"] = True
|
|
if args.dry_run:
|
|
overrides["dry_run"] = True
|
|
run_worker(args.config, overrides)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|