Simon
861a4051ca
- Job kind=selection：executor 分支用 SelectionService（spec=SelectionQuery），结果 落 result_json、自动归档 Experiment（summary：as_of + 选出 N）；/api/jobs 与 /api/experiments 的 result 解码支持 SelectionResult - POST /api/selections/jobs：提交 SelectionQuery 为异步 Job（BackgroundTasks/subprocess） - Web 选股页新增「异步（全市场）」按钮：提交 Job → waitJob 轮询结果（解决同步 60s+） - tests/test_selection_job.py：executor 执行归档（SelectionResult/Experiment kind）、 API 提交→轮询→结果与实验列表；全量 pytest + tsc 通过
127 lines
4.4 KiB
Python
127 lines
4.4 KiB
Python
"""异步研究 Job API（Phase 4）：提交 / 查询 / SSE 进度。
|
||
|
||
POST /api/jobs 创建 Job（BackgroundTasks 后台执行），立即返回 job_id
|
||
GET /api/jobs/{id} 状态 + 结果（成功时内嵌 result）
|
||
GET /api/jobs/{id}/events SSE 进度（queued→running→success|failed）
|
||
|
||
执行模式见 config job.mode：subprocess 时研究任务在独立子进程跑（内存隔离），
|
||
API worker 不被重任务拖垮（内存优化专项）。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
from datetime import datetime
|
||
from typing import Annotated
|
||
|
||
from fastapi import APIRouter, BackgroundTasks, HTTPException, Query
|
||
from fastapi.responses import StreamingResponse
|
||
|
||
from app.api.deps import DbSession, JobRepoDep
|
||
from app.application.services.job_executor import new_id, run_job_background, terminate_active
|
||
from app.domain.entities.research import (
|
||
BacktestResult,
|
||
FactorTestReport,
|
||
JobRecord,
|
||
JobStatus,
|
||
ResearchSpec,
|
||
)
|
||
from app.infrastructure.persistence.sqlalchemy.repositories.jobs_impl import (
|
||
SqlAlchemyJobRepository,
|
||
)
|
||
from app.infrastructure.persistence.sqlalchemy.session import SessionLocal
|
||
|
||
router = APIRouter(prefix="/jobs", tags=["jobs"])
|
||
|
||
|
||
def _decode_result(kind: str, result_json: str | None):
|
||
from app.domain.entities.selection import SelectionResult
|
||
|
||
if result_json is None:
|
||
return None
|
||
if kind == "selection":
|
||
return SelectionResult.model_validate_json(result_json)
|
||
model = BacktestResult if kind == "backtest" else FactorTestReport
|
||
return model.model_validate_json(result_json)
|
||
|
||
|
||
def _job_view(job: JobRecord) -> dict:
|
||
view = job.model_dump()
|
||
view["result"] = _decode_result(job.kind, job.result_json)
|
||
view.pop("result_json", None)
|
||
view["spec"] = json.loads(job.spec_json)
|
||
return view
|
||
|
||
|
||
@router.post("", summary="创建异步研究 Job")
|
||
def create_job(
|
||
spec: ResearchSpec,
|
||
background: BackgroundTasks,
|
||
session: DbSession,
|
||
job_repo: JobRepoDep,
|
||
) -> dict:
|
||
job = JobRecord(
|
||
id=new_id("JOB"),
|
||
kind=spec.type,
|
||
spec_json=spec.model_dump_json(),
|
||
status=JobStatus.QUEUED,
|
||
created_at=datetime.now(),
|
||
)
|
||
job_repo.create(job)
|
||
session.commit()
|
||
background.add_task(run_job_background, job.id)
|
||
return {"job_id": job.id, "status": job.status}
|
||
|
||
|
||
@router.get("", summary="Job 列表")
|
||
def list_jobs(
|
||
job_repo: JobRepoDep,
|
||
kind: Annotated[str | None, Query(description="按类型过滤（backtest/factor_test）")] = None,
|
||
limit: Annotated[int, Query(ge=1, le=200)] = 20,
|
||
) -> list[dict]:
|
||
return [_job_view(j) for j in job_repo.list_recent(kind=kind, limit=limit)]
|
||
|
||
|
||
@router.post("/{job_id}/cancel", summary="取消 Job（queued/running）")
|
||
def cancel_job(job_id: str, session: DbSession, job_repo: JobRepoDep) -> dict:
|
||
job = job_repo.get(job_id)
|
||
if job is None:
|
||
raise HTTPException(status_code=404, detail=f"Job {job_id} 不存在")
|
||
if job.status not in (JobStatus.QUEUED, JobStatus.RUNNING):
|
||
return {"job_id": job_id, "status": job.status, "cancelled": False}
|
||
job.status = JobStatus.CANCELLED
|
||
job.stage = None
|
||
job_repo.update(job)
|
||
session.commit()
|
||
terminate_active(job_id) # 终止研究子进程（若有）；父进程兜底已跳过 CANCELLED
|
||
return {"job_id": job_id, "status": JobStatus.CANCELLED, "cancelled": True}
|
||
|
||
|
||
@router.get("/{job_id}", summary="查询 Job 状态与结果")
|
||
def get_job(job_id: str, job_repo: JobRepoDep) -> dict:
|
||
job = job_repo.get(job_id)
|
||
if job is None:
|
||
raise HTTPException(status_code=404, detail=f"Job {job_id} 不存在")
|
||
return _job_view(job)
|
||
|
||
|
||
@router.get("/{job_id}/events", summary="Job 进度 SSE")
|
||
async def job_events(job_id: str) -> StreamingResponse:
|
||
async def gen():
|
||
while True:
|
||
with SessionLocal() as session:
|
||
job = SqlAlchemyJobRepository(session).get(job_id)
|
||
if job is None:
|
||
yield "event: error\ndata: job not found\n\n"
|
||
return
|
||
payload = json.dumps(
|
||
{"job_id": job.id, "status": job.status, "stage": job.stage}, ensure_ascii=False
|
||
)
|
||
yield f"data: {payload}\n\n"
|
||
if job.status in (JobStatus.SUCCESS, JobStatus.FAILED, JobStatus.CANCELLED):
|
||
return
|
||
await asyncio.sleep(0.4)
|
||
|
||
return StreamingResponse(gen(), media_type="text/event-stream")
|