214 lines
6.8 KiB
Python
214 lines
6.8 KiB
Python
"""app/context.py —— 生信服务共享实例容器（db/events/kernel/catalog/runner/semaphore/planner）。
|
||
|
||
- 生产：app 启动时 init_context() 加载真实依赖（SQLite、stage_image_map.json、
|
||
vstore FlowStore 本地实例（存在时）、真 docker 执行器）。
|
||
- 测试：build_context() 可注入 fake db/executor/flow_store（mock docker/LLM/vstore）。
|
||
- vstore 不存在/不可用：FlowStore 为 None，planner 自动降级 llm_only（FR-BIO-010 降级）。
|
||
- 并发信号量（BIOFLOW_CONCURRENCY，FR-BIO-031）每 context 一个。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import logging
|
||
from dataclasses import dataclass, field
|
||
from typing import Any
|
||
|
||
from .config import get_settings
|
||
from .db import Database, db as default_db
|
||
from .services.bioflow.catalog import StageCatalog, build_catalog
|
||
from .services.bioflow.events import EventBus
|
||
from .services.bioflow.kernel import Kernel, KernelConfig, load_kernel
|
||
from .services.bioflow.planner import Planner
|
||
from .services.bioflow.runner import Runner
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
@dataclass
|
||
class Context:
|
||
db: Any = None
|
||
events: EventBus | None = None
|
||
kernel: Kernel | None = None
|
||
catalog: StageCatalog | None = None
|
||
runner: Runner | None = None
|
||
semaphore: asyncio.Semaphore | None = None
|
||
planner: Planner | None = None
|
||
flow_store: Any = None
|
||
executor: Any = None
|
||
records: dict[str, Any] = field(default_factory=dict) # job_id → JobRecord（内存）
|
||
_tasks: list[asyncio.Task] = field(default_factory=list)
|
||
|
||
def spawn(self, coro: Any) -> asyncio.Task:
|
||
task = asyncio.create_task(coro)
|
||
self._tasks.append(task)
|
||
task.add_done_callback(lambda t: self._tasks.remove(t) if t in self._tasks else None)
|
||
return task
|
||
|
||
async def close(self) -> None:
|
||
for task in self._tasks:
|
||
task.cancel()
|
||
if self._tasks:
|
||
await asyncio.gather(*self._tasks, return_exceptions=True)
|
||
self._tasks.clear()
|
||
|
||
|
||
# 全局 context（测试可整体替换）
|
||
context: Context = Context()
|
||
|
||
|
||
def init_events(jobs_dir: str) -> EventBus:
|
||
from pathlib import Path
|
||
|
||
return EventBus(Path(jobs_dir) / "events")
|
||
|
||
|
||
def init_flow_store(flow_docs_dir: str) -> Any | None:
|
||
"""vstore FlowStore 本地实例（存在则加载；失败/缺失返回 None → planner 降级）。"""
|
||
try:
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
repo_root = Path(__file__).resolve().parent.parent.parent
|
||
if str(repo_root) not in sys.path:
|
||
sys.path.insert(0, str(repo_root))
|
||
from vstore.flow import FlowStore # P6 A4 交付；缺失则降级
|
||
|
||
return FlowStore(flow_docs_dir)
|
||
except Exception as exc: # noqa: BLE001 - 缺失/损坏降级 llm_only，不阻断服务
|
||
logger.warning("context: vstore FlowStore 不可用（planner 降级 llm_only）: %s", exc)
|
||
return None
|
||
|
||
|
||
def build_context(
|
||
*,
|
||
db: Database | None = None,
|
||
events: EventBus | None = None,
|
||
kernel: Kernel | None = None,
|
||
catalog: StageCatalog | None = None,
|
||
executor: Any = None,
|
||
semaphore: asyncio.Semaphore | None = None,
|
||
flow_store: Any = None,
|
||
runner_kwargs: dict[str, Any] | None = None,
|
||
) -> Context:
|
||
"""构造 context（生产与测试共用；各依赖可注入）。
|
||
|
||
executor=None 时用 kernel.run_stage（真 docker）。
|
||
"""
|
||
settings = get_settings()
|
||
database = db or default_db
|
||
bus = events or init_events(settings.jobs_dir_abs)
|
||
kern = kernel or load_kernel()
|
||
cat = catalog or build_catalog(kern.map)
|
||
sem = semaphore or asyncio.Semaphore(settings.concurrency)
|
||
exec_fn = executor or kern.run_stage
|
||
runner = Runner(
|
||
database,
|
||
bus,
|
||
exec_fn,
|
||
cat,
|
||
sem,
|
||
**(runner_kwargs or {}),
|
||
)
|
||
planner = Planner(
|
||
cat,
|
||
flow_store,
|
||
llm_plan_timeout=settings.llm_plan_timeout,
|
||
plan_candidates=settings.plan_candidates,
|
||
plan_max_rounds=settings.plan_max_rounds,
|
||
archive_finder=_make_archive_finder(database),
|
||
)
|
||
return Context(
|
||
db=database,
|
||
events=bus,
|
||
kernel=kern,
|
||
catalog=cat,
|
||
runner=runner,
|
||
semaphore=sem,
|
||
planner=planner,
|
||
flow_store=flow_store,
|
||
executor=exec_fn,
|
||
)
|
||
|
||
|
||
def _make_archive_finder(database: Database):
|
||
"""planner 复用归档查找回调（async；store 查询）。"""
|
||
from .services.bioflow import store
|
||
|
||
async def find(name: str) -> dict[str, Any] | None:
|
||
try:
|
||
return await store.find_archive_by_name(database, name)
|
||
except Exception: # noqa: BLE001
|
||
return None
|
||
|
||
return find
|
||
|
||
|
||
async def init_context(db_path: str | None = None) -> Context:
|
||
"""应用启动：初始化全局 context（幂等；db_path 供测试替换）。"""
|
||
global context
|
||
from .db import init_db
|
||
|
||
database = context.db
|
||
if db_path is not None:
|
||
database = await init_db(db_path)
|
||
elif context.db is None:
|
||
database = await init_db()
|
||
context = build_context(db=database)
|
||
return context
|
||
|
||
|
||
async def shutdown_context() -> None:
|
||
await context.close()
|
||
|
||
|
||
# ---- JobRecord 注册与恢复（路由层操作任务的内存记录）----
|
||
|
||
|
||
def register_record(record: Any) -> None:
|
||
context.records[record.id] = record
|
||
|
||
|
||
def _record_from_job(job: dict[str, Any], stages: list[dict[str, Any]]) -> Any:
|
||
"""DB 行 → JobRecord（无运行中内存记录时恢复，跨重启安全）。"""
|
||
from .services.bioflow.runner import JobRecord, params_dict
|
||
|
||
record = JobRecord(
|
||
id=job["id"],
|
||
session_id=job["session_id"],
|
||
title=job["title"],
|
||
status=job["status"],
|
||
stage_chain=list(job["stage_chain"] or []),
|
||
params=params_dict(job.get("params")),
|
||
auto_continue=bool(job.get("auto_continue")),
|
||
archive_id=job.get("archive_id"),
|
||
)
|
||
for s in stages:
|
||
record._stage_states[s["stage"]] = {
|
||
"status": s["status"],
|
||
"params": params_dict(s.get("params")),
|
||
"attempt": int(s.get("attempt") or 0),
|
||
"manifest_path": s.get("manifest_path"),
|
||
"log_tail": s.get("log_tail") or "",
|
||
"error": s.get("error"),
|
||
"started_at": s.get("started_at"),
|
||
"finished_at": s.get("finished_at"),
|
||
}
|
||
return record
|
||
|
||
|
||
async def load_record(job_id: str) -> Any | None:
|
||
"""取任务 JobRecord：内存优先，否则 DB 恢复（并注册回内存）。"""
|
||
if job_id in context.records:
|
||
return context.records[job_id]
|
||
if context.db is None:
|
||
return None
|
||
from .services.bioflow import store
|
||
|
||
job = await store.get_job(context.db, job_id)
|
||
if job is None:
|
||
return None
|
||
stages = await store.list_stages(context.db, job_id)
|
||
record = _record_from_job(job, stages)
|
||
context.records[job_id] = record
|
||
return record
|