201 lines
7.2 KiB
Python
201 lines
7.2 KiB
Python
"""M7 定时任务调度入口。
|
|
|
|
支持两种模式:
|
|
--once 立刻执行一次全链路
|
|
(默认) 启动 APScheduler,按 .env 中 SCHEDULE_TIMES 定时执行
|
|
|
|
用法:
|
|
uv run python -m scripts.run_scheduler # 启动定时服务
|
|
uv run python -m scripts.run_scheduler --once # 立即执行一次
|
|
uv run python -m scripts.run_scheduler --once --date 20260616
|
|
uv run python -m scripts.run_scheduler --once --steps crawler,extractor
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import signal
|
|
import sys
|
|
from datetime import date, datetime
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from dotenv import load_dotenv
|
|
from loguru import logger
|
|
|
|
from scheduler import STEP_COMMANDS, run_pipeline
|
|
from scheduler.stock_reporter import generate_all_stock_reports
|
|
|
|
|
|
def _setup_logger(level: str) -> None:
|
|
logger.remove()
|
|
logger.add(
|
|
sys.stderr,
|
|
level=level,
|
|
format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {name} | {message}",
|
|
)
|
|
log_path = Path("logs") / "scheduler.log"
|
|
log_path.parent.mkdir(parents=True, exist_ok=True)
|
|
logger.add(log_path, level="DEBUG", rotation="10 MB", retention=5, encoding="utf-8")
|
|
|
|
|
|
def _parse_schedule_times(raw: str) -> list[tuple[int, int]]:
|
|
"""解析 SCHEDULE_TIMES 环境变量。
|
|
|
|
格式: "07:00,12:00,18:00,22:00"
|
|
返回: [(7,0), (12,0), (18,0), (22,0)]
|
|
"""
|
|
out: list[tuple[int, int]] = []
|
|
for part in raw.split(","):
|
|
part = part.strip()
|
|
if not part:
|
|
continue
|
|
try:
|
|
h, m = part.split(":")
|
|
out.append((int(h), int(m)))
|
|
except ValueError:
|
|
logger.warning("SCHEDULE_TIMES 格式错误: {!r},跳过", part)
|
|
return out
|
|
|
|
|
|
def _once(args: argparse.Namespace) -> int:
|
|
"""单次执行模式。"""
|
|
steps = None
|
|
if args.steps:
|
|
steps = [s.strip() for s in args.steps.split(",")]
|
|
run_pipeline(args.date, steps=steps)
|
|
return 0
|
|
|
|
|
|
def _daemon(args: argparse.Namespace) -> int:
|
|
"""守护进程模式(APScheduler)。"""
|
|
import os
|
|
|
|
from apscheduler.schedulers.background import BackgroundScheduler # noqa: E402
|
|
from apscheduler.triggers.cron import CronTrigger # noqa: E402
|
|
|
|
times_raw = os.environ.get("SCHEDULE_TIMES", "07:00,12:00,18:00,22:00")
|
|
times = _parse_schedule_times(times_raw)
|
|
if not times:
|
|
logger.error("SCHEDULE_TIMES 为空或全部非法,无法启动定时任务")
|
|
return 2
|
|
|
|
# 找出最早的时间(当天首次运行),仅该次追加日报步骤
|
|
sorted_times = sorted(times)
|
|
first_hour, first_minute = sorted_times[0] if sorted_times else (0, 0)
|
|
|
|
scheduler = BackgroundScheduler()
|
|
|
|
# 包装函数:每次触发时重新计算日期,避免 date.today() 在注册时冻结。
|
|
def _scheduled_pipeline(steps: list[str] | None = None) -> None:
|
|
run_pipeline(date.today().strftime("%Y%m%d"), steps=steps)
|
|
|
|
for hour, minute in times:
|
|
trigger = CronTrigger(hour=hour, minute=minute, timezone="Asia/Shanghai")
|
|
is_first = (hour == first_hour and minute == first_minute)
|
|
job_kwargs: dict | None = None
|
|
if is_first:
|
|
job_kwargs = {
|
|
"steps": [k for k in STEP_COMMANDS if k not in ("report", "cninfo_crawl", "cninfo_extract", "cninfo_pdf")] + ["report"]
|
|
}
|
|
scheduler.add_job(
|
|
_scheduled_pipeline,
|
|
trigger=trigger,
|
|
kwargs=job_kwargs,
|
|
id=f"pipeline_{hour:02d}{minute:02d}",
|
|
name=f"全链路 {'+日报' if is_first else ''} {hour:02d}:{minute:02d}",
|
|
)
|
|
logger.info("已注册定时任务: {}每天 {:02d}:{:02d}{}", trigger, hour, minute,
|
|
" (含日报)" if is_first else "")
|
|
|
|
# cninfo 公告管道(可配置,默认 06:30)
|
|
cninfo_raw = os.environ.get("CNINFO_SCHEDULE_TIME", "06:30")
|
|
cninfo_parts = cninfo_raw.split(":")
|
|
cninfo_h, cninfo_m = int(cninfo_parts[0]), int(cninfo_parts[1]) if len(cninfo_parts) > 1 else 0
|
|
cninfo_trigger = CronTrigger(hour=cninfo_h, minute=cninfo_m, timezone="Asia/Shanghai")
|
|
cninfo_steps = ["cninfo_crawl", "cninfo_extract", "cninfo_pdf",
|
|
"dedup", "llm", "embedding", "qdrant"]
|
|
scheduler.add_job(
|
|
_scheduled_pipeline,
|
|
trigger=cninfo_trigger,
|
|
kwargs={"steps": cninfo_steps},
|
|
id="pipeline_cninfo",
|
|
name=f"cninfo 公告管道 {cninfo_h:02d}:{cninfo_m:02d}",
|
|
)
|
|
logger.info("已注册定时任务: cninfo 公告管道 每天 {:02d}:{:02d}", cninfo_h, cninfo_m)
|
|
|
|
# 个股日报(可配置,默认 07:30, 设为空可禁用)
|
|
stock_raw = os.environ.get("STOCK_REPORT_TIME", "07:30")
|
|
if stock_raw:
|
|
stock_parts = stock_raw.split(":")
|
|
stock_h, stock_m = int(stock_parts[0]), int(stock_parts[1]) if len(stock_parts) > 1 else 0
|
|
stock_trigger = CronTrigger(hour=stock_h, minute=stock_m, timezone="Asia/Shanghai")
|
|
scheduler.add_job(
|
|
generate_all_stock_reports,
|
|
trigger=stock_trigger,
|
|
id="stock_report",
|
|
name=f"个股日报 {stock_h:02d}:{stock_m:02d}",
|
|
)
|
|
logger.info("已注册定时任务: 个股日报 每天 {:02d}:{:02d}", stock_h, stock_m)
|
|
else:
|
|
logger.info("STOCK_REPORT_TIME 为空, 已禁用个股日报")
|
|
|
|
# 优雅退出
|
|
def _shutdown(signum: int, frame: Any) -> None:
|
|
logger.info("收到信号 {}, 关闭调度器...", signum)
|
|
scheduler.shutdown(wait=False)
|
|
raise SystemExit(0)
|
|
|
|
signal.signal(signal.SIGINT, _shutdown)
|
|
signal.signal(signal.SIGTERM, _shutdown)
|
|
|
|
scheduler.start()
|
|
logger.info("调度器已启动,等待触发... (按 Ctrl+C 退出)")
|
|
|
|
# 启动时检查是否有因重启/宕机错过的定时任务,30 分钟内补跑
|
|
now = datetime.now()
|
|
for hour, minute in times:
|
|
scheduled = now.replace(hour=hour, minute=minute, second=0, microsecond=0)
|
|
missed_minutes = (now - scheduled).total_seconds() / 60
|
|
if 0 < missed_minutes < 30:
|
|
logger.warning(
|
|
"检测到错过的定时任务 {:02d}:{:02d} ({} 分钟前),立即补跑一次",
|
|
hour, minute, int(missed_minutes),
|
|
)
|
|
steps = [k for k in STEP_COMMANDS if k != "report"]
|
|
if (hour, minute) == sorted_times[0]:
|
|
steps.append("report")
|
|
run_pipeline(date.today().strftime("%Y%m%d"), steps=steps)
|
|
|
|
import contextlib
|
|
|
|
with contextlib.suppress(SystemExit, KeyboardInterrupt):
|
|
# 保持主线程存活,直到收到退出信号
|
|
signal.pause()
|
|
|
|
return 0
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description="A 股新闻定时任务 (M7)")
|
|
parser.add_argument("--once", action="store_true", help="立即执行一次全链路")
|
|
parser.add_argument(
|
|
"--date", default=date.today().strftime("%Y%m%d"),
|
|
help="日期 YYYYMMDD (仅 --once 模式)",
|
|
)
|
|
parser.add_argument("--steps", default=None,
|
|
help="仅执行指定步骤,逗号分隔 (如 crawler,extractor)")
|
|
parser.add_argument("--log-level", default="INFO")
|
|
args = parser.parse_args()
|
|
|
|
_setup_logger(args.log_level)
|
|
load_dotenv()
|
|
|
|
if args.once:
|
|
return _once(args)
|
|
return _daemon(args)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|