395 lines
20 KiB
Python
395 lines
20 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""常驻 worker: 从 COS 自动认领任务分片、处理、回写完成标记, 没任务就空闲轮询;
|
||
每天指定时刻把本地 数据.jsonl 备份到 COS。
|
||
|
||
用法:
|
||
python run_worker.py --max-workers 32 --executor process
|
||
|
||
机器ID、取任务地址(jobs_root)、每日备份位置(daily_backup_path)等全部读 config.dispatch。
|
||
"""
|
||
import argparse
|
||
import logging
|
||
from logging.handlers import RotatingFileHandler
|
||
import os
|
||
import random
|
||
import threading
|
||
import time
|
||
from datetime import datetime
|
||
|
||
from config.config_utils import config
|
||
from tool import job_cos
|
||
# 复用 main.py 既有的处理流水线与清理逻辑
|
||
from main import (VideoProcessor, load_processed,
|
||
force_rmtree, JSONL_PATH, get_available_disk_gb)
|
||
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format='%(asctime)s [%(levelname)s] %(message)s',
|
||
handlers=[
|
||
# 常驻worker一跑就是数天/数周, 普通FileHandler会无限增长; 100M/份, 最多留5份(共500M封顶)
|
||
RotatingFileHandler('batch_processor.log', maxBytes=100 * 1024 * 1024, backupCount=5, encoding='utf-8'),
|
||
logging.StreamHandler(),
|
||
],
|
||
)
|
||
|
||
|
||
# ------------------------------------------------------------ 取下一片 ----
|
||
|
||
def _next_shard(machine_id: str, shuffle: bool, claim_match=None):
|
||
"""返回下一片要做的 shard 名; 没有可领返回 None。
|
||
优先级: ①续做我名下未完成的(重启续跑, 不受 claim_match 限制) -> ②原子认领一片新的(受 claim_match 过滤)。"""
|
||
# ① 粘性续做: 我上次没跑完的分片(别的机器不会碰; 已经是我的, 不再过滤)
|
||
for shard in job_cos.my_open_claims(machine_id):
|
||
logging.info(f"[续做] 接回我名下未完成分片: {shard}")
|
||
return shard
|
||
# ② 认领新分片
|
||
todo = job_cos.list_todo()
|
||
# 排除已完成 / 已被锁的
|
||
todo = [s for s in todo if not job_cos.is_done(s)]
|
||
# 只认领名称匹配的(claim_match 留空=全部)
|
||
if claim_match:
|
||
todo = [s for s in todo if job_cos.shard_matches(s, claim_match)]
|
||
if shuffle:
|
||
random.shuffle(todo)
|
||
for shard in todo:
|
||
if job_cos.read_lock(shard) is not None:
|
||
continue # 已被人锁
|
||
if job_cos.claim(shard, machine_id):
|
||
logging.info(f"[认领] 抢到分片: {shard}")
|
||
return shard
|
||
# 409: 刚好被别人抢走, 试下一片
|
||
return None
|
||
|
||
|
||
def _claim_long_batch(machine_id: str, shuffle: bool, min_batch: int):
|
||
"""长视频队列专用: 一次性多领几片, 累计任务数够 min_batch(或没片可领了)才停手 ——
|
||
长视频队列每天上传的还是各机器各自攒的小分片(可能就几十条), 这里只是在【领取】这一侧
|
||
把多个小分片攒成一大批一起处理, 不改上传那边的行为。
|
||
先照 _next_shard() 同样的优先级续做粘性未完成分片(不受 min_batch 限制, 必须先了结),
|
||
没有才开始新领; min_batch<=0 视为1(至少领一片，不做批量)。
|
||
返回 [(shard, tasks), ...]；某一片没片可领时提前结束(不会因为凑不够 min_batch 空等)。"""
|
||
min_batch = max(1, int(min_batch or 1))
|
||
batch = []
|
||
total = 0
|
||
|
||
# ① 粘性续做: 我名下所有未完成分片先原样接回来, 不受 min_batch 限制(必须先了结)
|
||
open_claims = job_cos.my_open_claims(machine_id)
|
||
for shard in open_claims:
|
||
tasks = job_cos.download_shard(shard)
|
||
batch.append((shard, tasks))
|
||
total += len(tasks)
|
||
if batch:
|
||
logging.info(f"[长视频批量] 接回 {len(batch)} 片未完成分片, 共{total}条, 直接处理(不再新领)")
|
||
return batch
|
||
|
||
# ② 新领: 不断认领新分片直到攒够 min_batch 条(或没片可领了)
|
||
todo = job_cos.list_todo()
|
||
todo = [s for s in todo if not job_cos.is_done(s)]
|
||
if shuffle:
|
||
random.shuffle(todo)
|
||
for shard in todo:
|
||
if total >= min_batch:
|
||
break
|
||
if job_cos.read_lock(shard) is not None:
|
||
continue # 已被人锁
|
||
if not job_cos.claim(shard, machine_id):
|
||
continue # 409: 刚好被别人抢走
|
||
tasks = job_cos.download_shard(shard)
|
||
batch.append((shard, tasks))
|
||
total += len(tasks)
|
||
logging.info(f"[长视频批量] 认领分片: {shard} ({len(tasks)}条, 累计{total}/{min_batch})")
|
||
return batch
|
||
|
||
|
||
def _process_shard_batch(batch, processor, machine_id, split_cfg,
|
||
coscli_path, cos_config_path, hb_interval, global_processed=None):
|
||
"""处理一批长视频分片(多片任务清单先合并成一个大列表再统一跑一次 process_batch,
|
||
而不是一片一片分别跑) —— 逻辑和 _process_shard 一致, 只是输入是 [(shard, tasks), ...]。
|
||
完成后把批次里每一片都各自 mark_done + release(全部标同一份聚合统计, 只是省事,
|
||
不是精确到每片的独立统计, 不影响正确性)。"""
|
||
shards = [s for s, _ in batch]
|
||
tasks = [t for _, ts in batch for t in ts]
|
||
processed = load_processed(JSONL_PATH)
|
||
if global_processed:
|
||
processed = processed | global_processed
|
||
todo_tasks = [t for t in tasks if t not in processed]
|
||
skipped = len(tasks) - len(todo_tasks)
|
||
logging.info(f"长视频批次({len(shards)}片): 共 {len(tasks)} 条, 跳过已处理 {skipped}, "
|
||
f"待处理 {len(todo_tasks)}")
|
||
|
||
hb_stop = threading.Event()
|
||
def _hb():
|
||
while not hb_stop.wait(hb_interval):
|
||
for shard in shards:
|
||
try:
|
||
job_cos.heartbeat(shard, machine_id)
|
||
except Exception as e:
|
||
logging.warning(f"心跳刷新失败 {shard}: {e}")
|
||
threading.Thread(target=_hb, daemon=True).start()
|
||
|
||
stats_dict = {"total": len(todo_tasks), "success": 0, "no_subtitle": 0, "error": 0}
|
||
try:
|
||
if todo_tasks:
|
||
stats = processor.process_batch(todo_tasks, coscli_path, cos_config_path, split_cfg)
|
||
stats_dict.update({"success": stats.success,
|
||
"no_subtitle": stats.no_subtitle, "error": stats.error})
|
||
finally:
|
||
hb_stop.set()
|
||
|
||
for shard in shards:
|
||
job_cos.mark_done(shard, machine_id, stats_dict)
|
||
job_cos.release(shard)
|
||
logging.info(f"长视频批次({len(shards)}片)完成并释放: {stats_dict}")
|
||
|
||
for d in ["./info", "./info_error", "./log", "./video"]:
|
||
if os.path.exists(d):
|
||
force_rmtree(d)
|
||
os.makedirs("./video", exist_ok=True)
|
||
|
||
|
||
def _fetch_global_processed(attempts: int = 3, backoff: float = 10.0) -> set:
|
||
"""启动时拉全局已处理集合, 失败重试(COS 抖动/网络问题很常见, 直接放弃太可惜)。
|
||
每次失败等 backoff 秒(线性,不是指数,重试次数少没必要搞复杂)再试; 全部失败返回空集合,
|
||
调用方会据此退化为只按本机 数据.jsonl 去重, 不阻塞 worker 启动。"""
|
||
last_err = None
|
||
for i in range(1, attempts + 1):
|
||
try:
|
||
return job_cos.global_processed_set()
|
||
except Exception as e:
|
||
last_err = e
|
||
logging.warning(f"[全局去重] 第{i}/{attempts}次拉取失败: {e}")
|
||
if i < attempts:
|
||
time.sleep(backoff)
|
||
logging.warning(f"[全局去重] 重试{attempts}次仍失败, 本次运行退化为只按本机 数据.jsonl 去重: {last_err}")
|
||
return set()
|
||
|
||
|
||
def _process_shard(shard, processor, machine_id, split_cfg,
|
||
coscli_path, cos_config_path, hb_interval, global_processed=None):
|
||
"""处理一个分片: 下载清单 -> 跳过(本机已处理 ∪ 全局已处理) -> 跑流水线 -> 写 done -> 释放。
|
||
global_processed: 启动时拉取的全局已处理集合(见 main() 里的 job_cos.global_processed_set()),
|
||
用于跨机器去重 —— 本机 数据.jsonl 只挡得住"这台机器"重复跑同一个视频, 挡不住这个视频被
|
||
错误地灌进了另一台机器认领的分片; 传 None/空集合时退化为只按本机去重(旧行为)。"""
|
||
tasks = job_cos.download_shard(shard)
|
||
processed = load_processed(JSONL_PATH) # 本机已处理(续跑去重)
|
||
if global_processed:
|
||
processed = processed | global_processed # 并上全局已处理(跨机器去重)
|
||
todo_tasks = [t for t in tasks if t not in processed]
|
||
skipped = len(tasks) - len(todo_tasks)
|
||
logging.info(f"分片 {shard}: 共 {len(tasks)} 条, 跳过已处理 {skipped}, 待处理 {len(todo_tasks)}")
|
||
|
||
# 心跳线程: 处理期间周期刷新锁的 hb_ts(留作监控/未来回收用)
|
||
hb_stop = threading.Event()
|
||
def _hb():
|
||
while not hb_stop.wait(hb_interval):
|
||
try:
|
||
job_cos.heartbeat(shard, machine_id)
|
||
except Exception as e:
|
||
logging.warning(f"心跳刷新失败 {shard}: {e}")
|
||
threading.Thread(target=_hb, daemon=True).start()
|
||
|
||
stats_dict = {"total": len(todo_tasks), "success": 0, "no_subtitle": 0, "error": 0}
|
||
try:
|
||
if todo_tasks:
|
||
stats = processor.process_batch(todo_tasks, coscli_path, cos_config_path, split_cfg)
|
||
stats_dict.update({"success": stats.success,
|
||
"no_subtitle": stats.no_subtitle, "error": stats.error})
|
||
finally:
|
||
hb_stop.set()
|
||
|
||
# 完成: 写 done -> 删 todo + lock
|
||
job_cos.mark_done(shard, machine_id, stats_dict)
|
||
job_cos.release(shard)
|
||
logging.info(f"分片 {shard} 完成并释放: {stats_dict}")
|
||
|
||
# 清理共享临时目录(沿用 main 的批次后清理)
|
||
for d in ["./info", "./info_error", "./log", "./video"]:
|
||
if os.path.exists(d):
|
||
force_rmtree(d)
|
||
os.makedirs("./video", exist_ok=True)
|
||
|
||
|
||
def main():
|
||
ap = argparse.ArgumentParser(description="常驻 worker: 从 COS 认领任务并处理")
|
||
ap.add_argument('--max-workers', type=int, default=None,
|
||
help='不填/0 = 自动按本机CPU核数选')
|
||
ap.add_argument('--mem-guard-gb', type=float, default=2.0)
|
||
ap.add_argument('--disk-guard-gb', type=float, default=10.0,
|
||
help='本地磁盘可用空间低于此值(GB)时暂停启动新视频，给数据处理预留空间；设 0 关闭')
|
||
ap.add_argument('--max-tasks-per-child', type=int, default=30)
|
||
ap.add_argument('--executor', choices=['process', 'thread'], default='process')
|
||
args = ap.parse_args()
|
||
|
||
if not args.max_workers:
|
||
cpu = os.cpu_count() or 0
|
||
args.max_workers = job_cos.auto_max_workers(cpu, disk_guard_gb=args.disk_guard_gb)
|
||
logging.info(f"--max-workers 未指定, 自动选为 {args.max_workers}(本机CPU核数={cpu or '未知'}; "
|
||
f"若本机还跑了OCR服务但不在 tencentcloud.ocr.endpoints 里, 自动探测不到, 请手动传小一点的值; "
|
||
f"同时按可用磁盘 ÷ 每视频预估占用(dispatch.est_disk_gb_per_video)做了收紧)")
|
||
# 供子进程(forkserver fork出来的worker)读取: splitframe.ffmpeg_threads=auto 时按此摊分CPU
|
||
os.environ['OCR_MAX_WORKERS'] = str(args.max_workers)
|
||
|
||
dcfg = config.get_dispatch_config() or {}
|
||
if not dcfg.get("enabled", True):
|
||
logging.error("config.dispatch.enabled=false, 退出。如需启用请改 config。")
|
||
return
|
||
|
||
machine_id = job_cos.get_machine_id()
|
||
keep_alive = bool(dcfg.get("keep_alive", True))
|
||
idle_poll = int(dcfg.get("idle_poll_seconds", 60))
|
||
shuffle = bool(dcfg.get("claim_shuffle", True))
|
||
hb_interval = int(dcfg.get("heartbeat_seconds", 300))
|
||
backup_hour = int(dcfg.get("daily_backup_hour", 10))
|
||
claim_match = dcfg.get("claim_match") or None # 只领匹配的分片(留空=全部)
|
||
|
||
run_long = bool(dcfg.get("run_long_video"))
|
||
long_root = dcfg.get("long_list_backup_path") or None
|
||
# 只有"本来跑主队列"的机器(非专职长视频机)、且配了长视频队列地址、且没有被显式关掉时,
|
||
# 才会在主队列空闲时顺手兼跑长视频队列 —— 专职长视频机本来就一直在长视频队列里, 不涉及
|
||
# 这个开关。auto_fallback_long_video 默认 true(维持原行为); 想让某台机器只专心跑主队列、
|
||
# 空闲就等着(不去碰长视频)的话, 在这台机器的 config.json 里配 false 即可, 不影响它自己
|
||
# 超长待重跑清单的每日备份(那个走 job_cos.upload_longlist_backup 各自独立)。
|
||
fallback_enabled = dcfg.get("auto_fallback_long_video", True)
|
||
auto_fallback_long = (not run_long) and bool(long_root) and bool(fallback_enabled)
|
||
active_root = job_cos._active_jobs_root()
|
||
logging.info("=" * 56)
|
||
logging.info(f"机器ID: {machine_id} | {'【专跑长视频】队列' if run_long else '主队列'}: {active_root}")
|
||
if auto_fallback_long:
|
||
logging.info(f"主队列空闲时自动兼跑长视频队列: {long_root}")
|
||
if claim_match:
|
||
logging.info(f"只认领匹配分片: {claim_match}")
|
||
logging.info(f"常驻: {keep_alive} 空闲轮询: {idle_poll}s | 每日备份: 每天{backup_hour}点")
|
||
disk = get_available_disk_gb(".")
|
||
disk_str = f"{disk:.1f}G" if disk is not None else "未知"
|
||
logging.info(f"磁盘保护阈值: {args.disk_guard_gb}G | 当前可用磁盘: {disk_str}")
|
||
if disk is not None and args.disk_guard_gb >= disk:
|
||
logging.warning(f"⚠️ 磁盘保护阈值({args.disk_guard_gb}G) ≥ 当前可用磁盘({disk:.1f}G)，"
|
||
f"几乎所有视频都会被卡住等待，请调小 --disk-guard-gb 或清理磁盘")
|
||
logging.info("=" * 56)
|
||
|
||
# 启动时拉一次全局已处理集合(各机器最新一天的已处理清单, 见 job_cos.global_processed_set),
|
||
# 跟本机 数据.jsonl 一起用于分片去重 —— 挡住"这个视频被错误地灌进了另一台机器分片"的情况。
|
||
# 只在启动时拉一次、缓存到本次运行结束, 不是每领一片分片都重新拉(清单本身也就每天更新一次,
|
||
# 没必要频繁重拉); 失败会重试(见 _fetch_global_processed), 重试仍失败才退化为只按本机去重。
|
||
# global_dedup_on_startup=false 时直接跳过这一步、只按本机 数据.jsonl 去重 —— 每台机器的
|
||
# 已处理清单可能有几MB, 20多台机器挨个拉下来这一步可能耗时较久(且下载期间没有中间日志输出,
|
||
# 看起来像卡住了), 网络不好/机器多的场合可以关掉。
|
||
global_processed = set()
|
||
if dcfg.get("daily_backup_path") and dcfg.get("global_dedup_on_startup", True):
|
||
global_processed = _fetch_global_processed()
|
||
if global_processed:
|
||
logging.info(f"[全局去重] 启动时拉取全局已处理集合: {len(global_processed)} 条")
|
||
else:
|
||
logging.info("[全局去重] 已关闭(global_dedup_on_startup=false), 只按本机 数据.jsonl 去重")
|
||
|
||
split_cfg = config.get_splitframe_config()
|
||
coscli_path = split_cfg.get('coscli_path', './coscli-linux-amd64')
|
||
cos_config_path = split_cfg.get('cosconfig_path', '.cos.yaml')
|
||
processor = VideoProcessor(max_workers=args.max_workers, mem_guard_gb=args.mem_guard_gb,
|
||
executor=args.executor, max_tasks_per_child=args.max_tasks_per_child,
|
||
disk_guard_gb=args.disk_guard_gb)
|
||
|
||
os.makedirs("./video", exist_ok=True)
|
||
|
||
# 每日备份后台线程(数据.jsonl + 超长列表)
|
||
stop = threading.Event()
|
||
threading.Thread(target=job_cos.run_daily_backup_loop,
|
||
args=(machine_id, backup_hour, stop, JSONL_PATH),
|
||
daemon=True).start()
|
||
|
||
# 每小时监控摘要上传线程(常态化监控用, 见 docs/superpowers/specs/2026-07-08-fleet-monitoring-design.md)
|
||
threading.Thread(target=job_cos.run_hourly_digest_loop,
|
||
args=(machine_id, stop, JSONL_PATH),
|
||
daemon=True).start()
|
||
|
||
# 长视频队列每次【领取】攒够这么多条任务才开始处理(领的时候多片一起领, 上传那边还是各机器
|
||
# 各自每天传小分片不受影响) —— 见 _claim_long_batch。<=0 视为1(不做批量, 领一片处理一片)。
|
||
min_long_batch = int(dcfg.get("long_list_min_batch", 1000) or 1)
|
||
|
||
try:
|
||
while True:
|
||
if job_cos.stop_requested():
|
||
logging.info("检测到 COS STOP 标记, 优雅停机。")
|
||
break
|
||
|
||
# 专职长视频机: 全程都在长视频队列, 直接批量领+批量处理, 不走主队列那套单片逻辑
|
||
if run_long:
|
||
try:
|
||
batch = _claim_long_batch(machine_id, shuffle, min_long_batch)
|
||
except Exception as e:
|
||
logging.error(f"✗ 访问长视频队列失败(稍后重试): {e}")
|
||
time.sleep(min(idle_poll, 30))
|
||
continue
|
||
if not batch:
|
||
if keep_alive:
|
||
logging.info(f"长视频队列无可领任务, 空闲等待 {idle_poll}s...")
|
||
time.sleep(idle_poll)
|
||
continue
|
||
else:
|
||
logging.info("长视频队列无可领任务, keep_alive=false, 退出。")
|
||
break
|
||
try:
|
||
_process_shard_batch(batch, processor, machine_id, split_cfg,
|
||
coscli_path, cos_config_path, hb_interval, global_processed)
|
||
except Exception as e:
|
||
logging.error(f"✗ 长视频批次处理出错(锁保留, 重启可续做): {e}")
|
||
time.sleep(5)
|
||
continue
|
||
|
||
# ① 先看主队列(单片领取, 行为不变)
|
||
try:
|
||
shard = _next_shard(machine_id, shuffle, claim_match)
|
||
except Exception as e:
|
||
logging.error(f"✗ 访问 COS 队列失败(稍后重试): {e}")
|
||
time.sleep(min(idle_poll, 30))
|
||
continue
|
||
|
||
# ② 主队列没活 -> 顺手看看长视频队列有没有能领的(批量领, 不是单片)
|
||
used_long = False
|
||
long_batch = None
|
||
if shard is None and auto_fallback_long:
|
||
job_cos.set_active_root_override(long_root)
|
||
try:
|
||
long_batch = _claim_long_batch(machine_id, shuffle, min_long_batch)
|
||
except Exception as e:
|
||
logging.error(f"✗ 访问长视频队列失败(稍后重试): {e}")
|
||
job_cos.set_active_root_override(None)
|
||
time.sleep(min(idle_poll, 30))
|
||
continue
|
||
if not long_batch:
|
||
job_cos.set_active_root_override(None) # 长视频队列也没活, 恢复主队列上下文
|
||
else:
|
||
used_long = True
|
||
logging.info(f"主队列空闲, 转去处理长视频队列批次: {len(long_batch)}片")
|
||
|
||
if shard is None and not long_batch:
|
||
# 两个队列都真·没活了
|
||
if keep_alive:
|
||
logging.info(f"无可领任务, 空闲等待 {idle_poll}s(加新任务到 todo/ 会自动接上)...")
|
||
time.sleep(idle_poll)
|
||
continue
|
||
else:
|
||
logging.info("无可领任务, keep_alive=false, 退出。")
|
||
break
|
||
|
||
try:
|
||
if long_batch:
|
||
_process_shard_batch(long_batch, processor, machine_id, split_cfg,
|
||
coscli_path, cos_config_path, hb_interval, global_processed)
|
||
else:
|
||
_process_shard(shard, processor, machine_id, split_cfg,
|
||
coscli_path, cos_config_path, hb_interval, global_processed)
|
||
except Exception as e:
|
||
logging.error(f"✗ 分片处理出错(锁保留, 重启可续做): {e}")
|
||
time.sleep(5)
|
||
finally:
|
||
if used_long:
|
||
job_cos.set_active_root_override(None) # 这批长视频处理完/出错, 下一轮先回去看主队列
|
||
finally:
|
||
stop.set()
|
||
logging.info("worker 退出。")
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main()
|