85 lines
2.3 KiB
Python
85 lines
2.3 KiB
Python
|
"""知习 RAG Worker — 文档导入主进程"""
|
|||
|
|
|
|||
|
|
import asyncio
|
|||
|
|
import signal
|
|||
|
|
import sys
|
|||
|
|
from config import WORKER_ID, POLL_INTERVAL, HEARTBEAT_INTERVAL
|
|||
|
|
from api_client import get_next_job, claim_job, heartbeat, update_job_status
|
|||
|
|
from pipelines.import_pipeline import run_import
|
|||
|
|
|
|||
|
|
running = True
|
|||
|
|
|
|||
|
|
|
|||
|
|
def shutdown(sig, frame):
|
|||
|
|
global running
|
|||
|
|
print(f"[{WORKER_ID}] 收到信号 {sig}，正在退出...")
|
|||
|
|
running = False
|
|||
|
|
|
|||
|
|
|
|||
|
|
signal.signal(signal.SIGINT, shutdown)
|
|||
|
|
signal.signal(signal.SIGTERM, shutdown)
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def heartbeat_loop():
|
|||
|
|
"""心跳循环（所有活跃任务）"""
|
|||
|
|
# 简化实现：worker 级心跳，后续可扩展到 per-job 心跳
|
|||
|
|
while running:
|
|||
|
|
await asyncio.sleep(HEARTBEAT_INTERVAL)
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def work_loop():
|
|||
|
|
"""主工作循环：轮询 → 认领 → 执行"""
|
|||
|
|
print(f"[{WORKER_ID}] RAG Worker 已启动")
|
|||
|
|
|
|||
|
|
while running:
|
|||
|
|
try:
|
|||
|
|
job = await get_next_job()
|
|||
|
|
if not job:
|
|||
|
|
await asyncio.sleep(POLL_INTERVAL)
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
job_id = job.get("id") or job.get("jobId")
|
|||
|
|
if not job_id:
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
# 认领任务
|
|||
|
|
claimed = await claim_job(job_id)
|
|||
|
|
if not claimed:
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
print(f"[{WORKER_ID}] 开始处理任务 {job_id}")
|
|||
|
|
|
|||
|
|
# 启动心跳（后台任务）
|
|||
|
|
hb_task = asyncio.create_task(_per_job_heartbeat(job_id))
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
await run_import(job)
|
|||
|
|
print(f"[{WORKER_ID}] 任务 {job_id} 完成")
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"[{WORKER_ID}] 任务 {job_id} 失败: {e}")
|
|||
|
|
await update_job_status(job_id, "FAILED_RETRYABLE", {
|
|||
|
|
"errorMessage": str(e)[:500],
|
|||
|
|
})
|
|||
|
|
finally:
|
|||
|
|
hb_task.cancel()
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"[{WORKER_ID}] 轮询异常: {e}")
|
|||
|
|
await asyncio.sleep(POLL_INTERVAL)
|
|||
|
|
|
|||
|
|
print(f"[{WORKER_ID}] Worker 已停止")
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def _per_job_heartbeat(job_id: str):
|
|||
|
|
"""单个任务的心跳上报"""
|
|||
|
|
while running:
|
|||
|
|
try:
|
|||
|
|
await heartbeat(job_id)
|
|||
|
|
except Exception:
|
|||
|
|
pass
|
|||
|
|
await asyncio.sleep(HEARTBEAT_INTERVAL)
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
asyncio.run(work_loop())
|