mirror of
https://ghfast.top/https://github.com/aeroxw/tick-stock-panel.git
synced 2026-09-12 21:24:16 +08:00
118 lines
4.6 KiB
Python
118 lines
4.6 KiB
Python
"""盘后管道 API — 异步触发 + 进度跟踪。"""
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import concurrent.futures as _cf
|
||
import logging
|
||
|
||
from fastapi import APIRouter, HTTPException, Request
|
||
|
||
from app.jobs import daily_pipeline
|
||
from app.services.pipeline_jobs import (
|
||
JobCancelledError,
|
||
job_store,
|
||
release_run_slot,
|
||
run_with_capacity,
|
||
try_acquire_run_slot,
|
||
)
|
||
from app.api.data import invalidate_storage_cache
|
||
|
||
# 长时间任务专用线程池（隔离于 FastAPI 默认线程池，防止阻塞请求处理）
|
||
_long_task_executor = _cf.ThreadPoolExecutor(max_workers=2, thread_name_prefix="long-task")
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
router = APIRouter(prefix="/api/pipeline", tags=["pipeline"])
|
||
|
||
|
||
@router.post("/run")
|
||
async def run_now(request: Request) -> dict:
|
||
"""异步触发盘后管道,立即返回 job_id。客户端轮询 /jobs/{id} 拿进度。
|
||
|
||
若已有任务在跑,**返回该任务 id 而不是开新任务**(防止并发拉数据撞限流)。
|
||
卡死判定按「进度停滞」而非总时长(慢带宽下长任务不会被误杀), 见 reap_stale。
|
||
"""
|
||
repo = request.app.state.repo
|
||
capset = request.app.state.capabilities
|
||
|
||
# 检测卡死的 running job (如 reload 后孤儿 task / 网络读无限阻塞)。
|
||
# reap_stale 会在 /run 和 /jobs/{id} 轮询端点都调用,保证卡死后能自愈。
|
||
job_store.reap_stale()
|
||
|
||
# 单飞: 复用任何活跃 (pending∨running) 任务, is_new=False 时不再调度新任务
|
||
job_id, is_new = job_store.create()
|
||
if not is_new:
|
||
return {"job_id": job_id, "reused": True}
|
||
|
||
# 在 executor 里跑同步任务(pipeline 内部都是阻塞 IO + CPU)
|
||
async def task() -> None:
|
||
# 重任务执行槽: 防僵尸并发(reap 后线程仍活时新任务不得并行写 parquet)
|
||
if not try_acquire_run_slot(job_id):
|
||
job_store.fail(job_id, "已有数据任务在运行(或上一次任务卡死未结束),请稍后再试")
|
||
return
|
||
# 管道运行期间暂停实时行情取数, 防止覆写同一批 parquet 竞态
|
||
qs = getattr(request.app.state, "quote_service", None)
|
||
try:
|
||
loop = asyncio.get_event_loop()
|
||
|
||
def progress(stage: str, pct: int, msg: str, stage_pct: int | None = None,
|
||
skip_log: bool = False) -> None:
|
||
job_store.progress(job_id, stage, pct, msg, stage_pct=stage_pct, skip_log=skip_log)
|
||
|
||
def _run() -> dict:
|
||
try:
|
||
if qs:
|
||
with qs.paused():
|
||
return daily_pipeline.run_now(repo, capset, on_progress=progress)
|
||
return daily_pipeline.run_now(repo, capset, on_progress=progress)
|
||
finally:
|
||
repo.refresh_cache()
|
||
|
||
result = await loop.run_in_executor(_long_task_executor, run_with_capacity, job_id, _run)
|
||
job_store.succeed(job_id, result)
|
||
invalidate_storage_cache()
|
||
except JobCancelledError:
|
||
# 已被 reap/手动取消终止: job 状态已由 terminate() 写为 failed,
|
||
# 拉取线程在分块回调处自行退出, 这里无需(也无法)再写状态。
|
||
logger.warning("pipeline job %s cancelled", job_id)
|
||
except Exception as e: # noqa: BLE001
|
||
logger.exception("pipeline failed")
|
||
job_store.fail(job_id, str(e))
|
||
invalidate_storage_cache()
|
||
finally:
|
||
release_run_slot(job_id)
|
||
|
||
asyncio.create_task(task())
|
||
return {"job_id": job_id, "reused": False}
|
||
|
||
|
||
@router.get("/jobs/{job_id}")
|
||
def get_job(job_id: str) -> dict:
|
||
# 每次轮询都检查卡死 job — 前端持续轮询, 进度停滞超阈值后必定自愈,
|
||
# 无需用户再次手动点「同步」。
|
||
job_store.reap_stale()
|
||
j = job_store.get(job_id)
|
||
if not j:
|
||
raise HTTPException(status_code=404, detail="job not found")
|
||
return j
|
||
|
||
|
||
@router.post("/jobs/{job_id}/cancel")
|
||
def cancel_job(job_id: str) -> dict:
|
||
"""手动取消一个 running 的 job(协作式: 拉取线程在当前分块完成后自行退出)。"""
|
||
j = job_store.get(job_id)
|
||
if not j:
|
||
raise HTTPException(status_code=404, detail="job not found")
|
||
if j["status"] not in ("running", "pending"):
|
||
raise HTTPException(status_code=400, detail=f"job status is {j['status']}, cannot cancel")
|
||
job_store.terminate(job_id, "用户手动取消")
|
||
return {"cancelled": job_id}
|
||
|
||
|
||
@router.get("/jobs")
|
||
def list_jobs(limit: int = 20) -> dict:
|
||
return {
|
||
"active_id": job_store.active_id(),
|
||
"jobs": job_store.list_recent(limit=limit),
|
||
}
|