Simon
7c268e43df
- ResearchService.run_backtest/run_factor_test 支持 on_stage 回调（data_loading →
factor_calculation/backtesting → analysis）；job_executor 经独立短会话把 stage 写库
（子进程模式同样走 DB），收尾回读最后阶段避免覆盖
- /api/jobs：GET 列表（kind/limit）、POST /{id}/cancel（queued/running → CANCELLED +
终止子进程 terminate_active；父进程兜底已跳过 CANCELLED）
- SSE /jobs/{id}/events 现会携带 stage
- tests/test_job_stages.py：成功 job 终态 stage=analysis；executor 取消 queued/幂等/
已完成不可取消；API 列表+终态不可取消+404；全量 pytest 通过
438 lines
15 KiB
Python
438 lines
15 KiB
Python
"""Job 执行编排与 Experiment 归档（Phase 4）+ 内存隔离执行（内存优化专项）。
|
||
|
||
execute_job 运行状态机 queued→running→(success|failed)，
|
||
成功时自动把 spec/result 存为 Experiment（含代码版本），实现「研究可复现」（AGENT.md §21）。
|
||
独立 Session 生命周期（不依赖请求 scope），可被 FastAPI BackgroundTasks 或测试直接调用。
|
||
|
||
执行模式（config job.mode）：
|
||
- local —— 本进程内执行（开发 / 单测；无内存隔离）
|
||
- subprocess —— 独立子进程执行（python -m app.cli.run_job <job_id>），
|
||
子进程内施加 RLIMIT_AS 上限（config job.max_memory_gb），研究任务 OOM 只会
|
||
MemoryError 失败归档，不会拖垮 API worker；父进程在子进程异常退出时补记 failed。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import os
|
||
import subprocess
|
||
import sys
|
||
import threading
|
||
import uuid
|
||
from collections.abc import Callable
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
|
||
from app.domain.entities.research import (
|
||
BacktestResult,
|
||
ExperimentRecord,
|
||
FactorTestReport,
|
||
JobRecord,
|
||
JobStatus,
|
||
ResearchSpec,
|
||
)
|
||
from app.quant.service import ResearchService
|
||
|
||
# backend/app/application/services/job_executor.py → parents[4] = 项目根
|
||
PROJECT_ROOT = Path(__file__).resolve().parents[4]
|
||
BACKEND_ROOT = PROJECT_ROOT / "backend"
|
||
|
||
# 本进程内正在运行的子进程 Job（并发上限保护宿主内存）
|
||
_active_jobs: dict[str, subprocess.Popen | None] = {}
|
||
_active_lock = threading.Lock()
|
||
|
||
|
||
def new_id(prefix: str) -> str:
|
||
return f"{prefix}-{uuid.uuid4().hex[:8].upper()}"
|
||
|
||
|
||
def _git_short_rev() -> str | None:
|
||
try:
|
||
out = subprocess.run(
|
||
["git", "rev-parse", "--short", "HEAD"],
|
||
cwd=PROJECT_ROOT,
|
||
capture_output=True,
|
||
text=True,
|
||
timeout=5,
|
||
check=False,
|
||
)
|
||
return out.stdout.strip() or None
|
||
except Exception: # noqa: BLE001
|
||
return None
|
||
|
||
|
||
def _summary_text(kind: str, result: BacktestResult | FactorTestReport) -> str | None:
|
||
if kind == "backtest" and isinstance(result, BacktestResult):
|
||
s = result.summary
|
||
return f"总收益 {s.total_return_pct:.2f}% · 年化 {s.annual_return_pct:.2f}% · 回撤 {s.max_drawdown_pct:.2f}%"
|
||
if isinstance(result, FactorTestReport):
|
||
return f"IC {result.ic_mean:.4f} · RankIC {result.rank_ic_mean:.4f} · 样本 {result.sample_days} 日"
|
||
return None
|
||
|
||
|
||
def _execute_inner(
|
||
job_id: str,
|
||
*,
|
||
session_factory: Callable,
|
||
job_repo_factory: Callable,
|
||
experiment_repo_factory: Callable,
|
||
stock_repo_factory: Callable,
|
||
daily_repo_factory: Callable,
|
||
engine,
|
||
) -> None:
|
||
with session_factory() as session:
|
||
job_repo = job_repo_factory(session)
|
||
experiment_repo = experiment_repo_factory(session)
|
||
job = job_repo.get(job_id)
|
||
if job is None:
|
||
return
|
||
job.status = JobStatus.RUNNING
|
||
job.started_at = datetime.now()
|
||
job_repo.update(job)
|
||
session.commit()
|
||
try:
|
||
spec = ResearchSpec.model_validate_json(job.spec_json)
|
||
service = ResearchService(
|
||
stock_repo_factory(session), daily_repo_factory(session), engine
|
||
)
|
||
|
||
def _set_stage(name: str) -> None:
|
||
"""阶段上报（v3 §23）：独立短会话写 job.stage 并 commit（子进程同样走 DB）。"""
|
||
try:
|
||
with session_factory() as st_sess:
|
||
st = job_repo_factory(st_sess).get(job_id)
|
||
if st is not None and st.status == JobStatus.RUNNING:
|
||
st.stage = name
|
||
job_repo_factory(st_sess).update(st)
|
||
st_sess.commit()
|
||
except Exception: # noqa: BLE001 —— 阶段上报失败不阻断执行
|
||
pass
|
||
|
||
if spec.type == "backtest":
|
||
result = service.run_backtest(spec, on_stage=_set_stage)
|
||
else:
|
||
result = service.run_factor_test(spec, on_stage=_set_stage)
|
||
|
||
result_json = json.dumps(result.model_dump(mode="json"), ensure_ascii=False)
|
||
experiment = ExperimentRecord(
|
||
id=new_id("EXP"),
|
||
kind=spec.type,
|
||
spec_json=job.spec_json,
|
||
result_json=result_json,
|
||
summary_text=_summary_text(spec.type, result),
|
||
code_version=_git_short_rev(),
|
||
job_id=job.id,
|
||
created_at=datetime.now(),
|
||
)
|
||
experiment_repo.save(experiment)
|
||
|
||
job.result_json = result_json
|
||
job.experiment_id = experiment.id
|
||
job.status = JobStatus.SUCCESS
|
||
job.error = None
|
||
except Exception as exc: # noqa: BLE001 —— 统一记为 failed 供前端展示
|
||
job.status = JobStatus.FAILED
|
||
job.error = f"{type(exc).__name__}: {exc}"
|
||
finally:
|
||
job.finished_at = datetime.now()
|
||
# 回读最后一次阶段上报（_set_stage 经独立会话写库），避免被本会话覆盖
|
||
try:
|
||
with session_factory() as last_sess:
|
||
last = job_repo_factory(last_sess).get(job_id)
|
||
if last is not None:
|
||
job.stage = last.stage
|
||
except Exception: # noqa: BLE001
|
||
pass
|
||
job_repo.update(job)
|
||
session.commit()
|
||
|
||
|
||
def execute_job(
|
||
job_id: str,
|
||
*,
|
||
session_factory: Callable,
|
||
job_repo_factory: Callable,
|
||
experiment_repo_factory: Callable,
|
||
stock_repo_factory: Callable,
|
||
daily_repo_factory: Callable,
|
||
engine,
|
||
) -> None:
|
||
"""入口包装：任何未预期异常都将 Job 标记 failed（防止卡在 queued/running）。"""
|
||
try:
|
||
_execute_inner(
|
||
job_id,
|
||
session_factory=session_factory,
|
||
job_repo_factory=job_repo_factory,
|
||
experiment_repo_factory=experiment_repo_factory,
|
||
stock_repo_factory=stock_repo_factory,
|
||
daily_repo_factory=daily_repo_factory,
|
||
engine=engine,
|
||
)
|
||
except Exception as exc: # noqa: BLE001
|
||
try:
|
||
with session_factory() as session:
|
||
repo = job_repo_factory(session)
|
||
job = repo.get(job_id)
|
||
if job is not None:
|
||
job.status = JobStatus.FAILED
|
||
job.error = f"内部错误: {type(exc).__name__}: {exc}"
|
||
job.finished_at = datetime.now()
|
||
repo.update(job)
|
||
session.commit()
|
||
except Exception: # noqa: BLE001
|
||
pass
|
||
|
||
|
||
def default_factories() -> dict:
|
||
"""后台执行所需的独立 Session / Repository / 引擎装配（跨请求生命周期）。"""
|
||
from app.infrastructure.persistence.sqlalchemy.repositories.jobs_impl import (
|
||
SqlAlchemyExperimentRepository,
|
||
SqlAlchemyJobRepository,
|
||
)
|
||
from app.infrastructure.persistence.sqlalchemy.repositories.market_impl import (
|
||
SqlAlchemyDailyBarRepository,
|
||
SqlAlchemyStockRepository,
|
||
)
|
||
from app.infrastructure.persistence.sqlalchemy.session import SessionLocal
|
||
from app.quant.engine import LocalEngine
|
||
|
||
return {
|
||
"session_factory": SessionLocal,
|
||
"job_repo_factory": lambda s: SqlAlchemyJobRepository(s),
|
||
"experiment_repo_factory": lambda s: SqlAlchemyExperimentRepository(s),
|
||
"stock_repo_factory": lambda s: SqlAlchemyStockRepository(s),
|
||
"daily_repo_factory": lambda s: SqlAlchemyDailyBarRepository(s),
|
||
"engine": LocalEngine(),
|
||
}
|
||
|
||
|
||
def submit_and_run(spec: ResearchSpec, *, factories: dict | None = None) -> JobRecord:
|
||
"""创建并执行一个 Job（复用 Job 状态机与 Experiment 归档），返回终态 Job。
|
||
|
||
按 job.mode 调度：subprocess 模式在独立进程内执行（内存隔离），否则本进程。
|
||
"""
|
||
facts = factories or default_factories()
|
||
session_factory = facts["session_factory"]
|
||
job_repo = facts["job_repo_factory"]
|
||
job = JobRecord(
|
||
id=new_id("JOB"),
|
||
kind=spec.type,
|
||
spec_json=json.dumps(spec.model_dump(mode="json"), ensure_ascii=False),
|
||
status=JobStatus.QUEUED,
|
||
created_at=datetime.now(),
|
||
)
|
||
with session_factory() as session:
|
||
job_repo(session).create(job)
|
||
session.commit()
|
||
if _job_mode() == "subprocess":
|
||
_run_in_subprocess(job.id, session_factory=session_factory, job_repo_factory=job_repo)
|
||
else:
|
||
execute_job(
|
||
job.id,
|
||
session_factory=session_factory,
|
||
job_repo_factory=job_repo,
|
||
experiment_repo_factory=facts["experiment_repo_factory"],
|
||
stock_repo_factory=facts["stock_repo_factory"],
|
||
daily_repo_factory=facts["daily_repo_factory"],
|
||
engine=facts["engine"],
|
||
)
|
||
with session_factory() as session:
|
||
done = job_repo(session).get(job.id)
|
||
assert done is not None
|
||
return done
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 执行模式调度（config job.mode：local | subprocess）
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _job_mode() -> str:
|
||
from app.core.config import get_settings
|
||
|
||
return get_settings().job_mode
|
||
|
||
|
||
def _job_memory_limit_gb() -> int:
|
||
from app.core.config import get_settings
|
||
|
||
return get_settings().job_memory_limit_gb
|
||
|
||
|
||
def _job_max_concurrent() -> int:
|
||
from app.core.config import get_settings
|
||
|
||
return get_settings().job_max_concurrent
|
||
|
||
|
||
def _job_worker_cmd(job_id: str) -> list[str]:
|
||
"""研究子进程命令行（独立进程执行，cwd=backend 使 `-m app.cli.run_job` 可导入）。"""
|
||
return [sys.executable, "-m", "app.cli.run_job", job_id]
|
||
|
||
|
||
def run_job_background(
|
||
job_id: str,
|
||
*,
|
||
session_factory: Callable | None = None,
|
||
job_repo_factory: Callable | None = None,
|
||
) -> None:
|
||
"""API 后台任务入口：按 job.mode 调度 Job 执行。"""
|
||
if _job_mode() == "subprocess":
|
||
_run_in_subprocess(
|
||
job_id, session_factory=session_factory, job_repo_factory=job_repo_factory
|
||
)
|
||
else:
|
||
execute_job(job_id, **default_factories())
|
||
|
||
|
||
def _acquire_slot(job_id: str) -> bool:
|
||
with _active_lock:
|
||
if len(_active_jobs) >= max(_job_max_concurrent(), 1):
|
||
return False
|
||
_active_jobs[job_id] = None
|
||
return True
|
||
|
||
|
||
def _release_slot(job_id: str) -> None:
|
||
with _active_lock:
|
||
_active_jobs.pop(job_id, None)
|
||
|
||
|
||
def _mark_failed(
|
||
job_id: str,
|
||
error: str,
|
||
*,
|
||
session_factory: Callable | None = None,
|
||
job_repo_factory: Callable | None = None,
|
||
) -> None:
|
||
"""把非终态 Job 标记 failed（子进程异常退出 / 并发超限时兜底，防永久 running）。"""
|
||
if session_factory is None or job_repo_factory is None:
|
||
facts = default_factories()
|
||
session_factory = session_factory or facts["session_factory"]
|
||
job_repo_factory = job_repo_factory or facts["job_repo_factory"]
|
||
try:
|
||
with session_factory() as session:
|
||
repo = job_repo_factory(session)
|
||
job = repo.get(job_id)
|
||
if job is not None and job.status in (JobStatus.QUEUED, JobStatus.RUNNING):
|
||
job.status = JobStatus.FAILED
|
||
job.error = error
|
||
job.finished_at = datetime.now()
|
||
repo.update(job)
|
||
session.commit()
|
||
except Exception: # noqa: BLE001 —— 兜底标记失败自身异常不向上抛
|
||
pass
|
||
|
||
|
||
def _run_in_subprocess(
|
||
job_id: str,
|
||
*,
|
||
session_factory: Callable | None = None,
|
||
job_repo_factory: Callable | None = None,
|
||
) -> None:
|
||
"""在独立 python 进程执行 Job：子进程 RLIMIT 上限（防 OOM 整机），
|
||
父进程等待；子进程异常退出（如被杀）时把 Job 补记 failed。"""
|
||
if not _acquire_slot(job_id):
|
||
_mark_failed(
|
||
job_id,
|
||
"系统繁忙：并发研究任务已达上限，请稍后重试",
|
||
session_factory=session_factory,
|
||
job_repo_factory=job_repo_factory,
|
||
)
|
||
return
|
||
env = dict(os.environ)
|
||
env["QLIB_JOB_MEM_LIMIT_GB"] = str(_job_memory_limit_gb())
|
||
rc = -1
|
||
already_marked = False
|
||
try:
|
||
proc = subprocess.Popen(
|
||
_job_worker_cmd(job_id),
|
||
cwd=BACKEND_ROOT,
|
||
env=env,
|
||
stdout=subprocess.DEVNULL,
|
||
stderr=subprocess.DEVNULL,
|
||
)
|
||
with _active_lock:
|
||
_active_jobs[job_id] = proc
|
||
rc = proc.wait()
|
||
except Exception as exc: # noqa: BLE001
|
||
_mark_failed(
|
||
job_id,
|
||
f"研究子进程启动失败: {type(exc).__name__}: {exc}",
|
||
session_factory=session_factory,
|
||
job_repo_factory=job_repo_factory,
|
||
)
|
||
already_marked = True
|
||
finally:
|
||
_release_slot(job_id)
|
||
if rc != 0 and not already_marked:
|
||
_mark_failed(
|
||
job_id,
|
||
f"研究子进程异常退出(code={rc})：任务未完成",
|
||
session_factory=session_factory,
|
||
job_repo_factory=job_repo_factory,
|
||
)
|
||
|
||
|
||
def terminate_active(job_id: str) -> bool:
|
||
"""终止该 Job 的活动子进程（如有）。返回是否找到并终止。"""
|
||
with _active_lock:
|
||
proc = _active_jobs.get(job_id)
|
||
if proc is None:
|
||
return False
|
||
import contextlib
|
||
|
||
with contextlib.suppress(Exception):
|
||
proc.terminate()
|
||
return True
|
||
|
||
|
||
def cancel_job(
|
||
job_id: str,
|
||
*,
|
||
session_factory: Callable | None = None,
|
||
job_repo_factory: Callable | None = None,
|
||
) -> bool:
|
||
"""取消 queued/running Job（置 CANCELLED 并终止子进程）；不可取消返回 False。"""
|
||
facts = default_factories()
|
||
sf = session_factory or facts["session_factory"]
|
||
jr = job_repo_factory or facts["job_repo_factory"]
|
||
with sf() as session:
|
||
repo = jr(session)
|
||
job = repo.get(job_id)
|
||
if job is None or job.status not in (JobStatus.QUEUED, JobStatus.RUNNING):
|
||
return False
|
||
job.status = JobStatus.CANCELLED
|
||
job.stage = None
|
||
repo.update(job)
|
||
session.commit()
|
||
terminate_active(job_id)
|
||
return True
|
||
|
||
|
||
def mark_stale_jobs_failed(
|
||
*,
|
||
session_factory: Callable | None = None,
|
||
job_repo_factory: Callable | None = None,
|
||
reason: str | None = None,
|
||
) -> int:
|
||
"""服务启动调用：把上次进程异常退出遗留的 queued/running Job 标记 failed。
|
||
|
||
返回处理数量。防「进程被杀后 Job 永久 running」（AGENT.md §20 状态机闭环）。
|
||
"""
|
||
facts = default_factories()
|
||
sf = session_factory or facts["session_factory"]
|
||
jr = job_repo_factory or facts["job_repo_factory"]
|
||
counted = 0
|
||
with sf() as session:
|
||
repo = jr(session)
|
||
for status in (JobStatus.QUEUED, JobStatus.RUNNING):
|
||
for job in repo.list_by_status(status):
|
||
job.status = JobStatus.FAILED
|
||
job.error = reason or "服务重启：上次未完成任务被中断"
|
||
job.finished_at = datetime.now()
|
||
repo.update(job)
|
||
counted += 1
|
||
session.commit()
|
||
return counted
|