feat(harness): inference.py — async run_inference + DI
This commit is contained in:
@@ -0,0 +1,441 @@
|
||||
"""async 推理编排 — 训练循环的 forward()。
|
||||
|
||||
从 TRM4 core/harness/inference.py (~560 行) 迁移，重大重构:
|
||||
- 同步 ThreadPoolExecutor → asyncio.Semaphore + asyncio.gather
|
||||
- LLMClient.from_env() 每题构造 → llm: LLMProvider 注入共享
|
||||
- SentenceTransformer/OCR 内部构造 → 调用方通过 tool_dispatch_fn 注入
|
||||
- run_id 必传，空串 → ValueError
|
||||
- _aggregate_results 从内存 results 聚合（非 DB 回读）
|
||||
- record_run 由调用方（Runner）负责
|
||||
- prompt 构建由调用方注入 prompt_builder
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from collections import defaultdict
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from loguru import logger
|
||||
|
||||
from core.agent.loop import AgentLoop
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
from app.harness.log import HarnessLog
|
||||
from core.agent.types import LoopResult
|
||||
from core.protocols import LLMProvider
|
||||
from core.types import GeneratedQuestion
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class InferenceResult:
|
||||
"""推理聚合结果。
|
||||
|
||||
属性:
|
||||
run_id: 运行标识。
|
||||
accuracy: 总正确率。
|
||||
total: 总题数。
|
||||
correct: 正确题数。
|
||||
per_task_type: 按题型分组的指标 {task_type: {accuracy, total, correct}}。
|
||||
steps_mean: 平均步数。
|
||||
token_usage: token 总用量 {prompt_tokens, completion_tokens}。
|
||||
stop_reason_counts: 终止原因计数 {reason: count}。
|
||||
"""
|
||||
|
||||
run_id: str
|
||||
accuracy: float
|
||||
total: int
|
||||
correct: int
|
||||
per_task_type: dict[str, dict]
|
||||
steps_mean: float
|
||||
token_usage: dict[str, int]
|
||||
stop_reason_counts: dict[str, int]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 表 Schema 定义（5 张表，保留 TRM4 全部 schema）
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
PREDICTIONS_SCHEMA: dict[str, str] = {
|
||||
"video_id": "TEXT",
|
||||
"question_id": "TEXT",
|
||||
"task_type": "TEXT",
|
||||
"prediction": "TEXT",
|
||||
"answer": "TEXT",
|
||||
"evidence": "TEXT",
|
||||
"reasoning": "TEXT",
|
||||
"steps_used": "INTEGER",
|
||||
"prompt_tokens": "INTEGER",
|
||||
"completion_tokens": "INTEGER",
|
||||
"stop_reason": "TEXT",
|
||||
"steps_json": "JSON",
|
||||
}
|
||||
|
||||
TRACES_SCHEMA: dict[str, str] = {
|
||||
"video_id": "TEXT",
|
||||
"question_id": "TEXT",
|
||||
"step": "INTEGER",
|
||||
"tool_name": "TEXT",
|
||||
"tool_args": "JSON",
|
||||
"tool_output": "TEXT",
|
||||
"thought": "TEXT",
|
||||
}
|
||||
|
||||
VALIDATION_FLAGS_SCHEMA: dict[str, str] = {
|
||||
"video_id": "TEXT",
|
||||
"question_id": "TEXT",
|
||||
"has_l3_visit": "INTEGER",
|
||||
"l1_count": "INTEGER",
|
||||
"l2_count": "INTEGER",
|
||||
"l3_count": "INTEGER",
|
||||
}
|
||||
|
||||
ANCHOR_CHECK_SCHEMA: dict[str, str] = {
|
||||
"video_id": "TEXT",
|
||||
"question_id": "TEXT",
|
||||
"step": "INTEGER",
|
||||
"n_assertions": "INTEGER",
|
||||
"n_anchored": "INTEGER",
|
||||
"n_illegal": "INTEGER",
|
||||
"n_expanded": "INTEGER",
|
||||
"n_trunc": "INTEGER",
|
||||
"output_chars": "INTEGER",
|
||||
}
|
||||
|
||||
OF_HEALTH_SCHEMA: dict[str, str] = {
|
||||
"video_id": "TEXT",
|
||||
"question_id": "TEXT",
|
||||
"step": "INTEGER",
|
||||
"ocr_injected": "INTEGER",
|
||||
"ocr_chars": "INTEGER",
|
||||
"ocr_failed": "INTEGER",
|
||||
"discrepancy": "INTEGER",
|
||||
"abstain": "INTEGER",
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 内部工具
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _DispatcherAdapter:
|
||||
"""将裸 async callable 包装为 ToolDispatcher Protocol 实例。
|
||||
|
||||
AgentLoop 要求 ToolDispatcher（有 dispatch 方法），而 run_inference
|
||||
接收的 tool_dispatch_fn 是裸 async callable。此适配器桥接两者。
|
||||
|
||||
参数:
|
||||
fn: async def (tool_name, args, *, context) -> str。
|
||||
"""
|
||||
|
||||
def __init__(self, fn: Callable[..., Any]) -> None:
|
||||
self._fn = fn
|
||||
|
||||
async def dispatch(
|
||||
self, tool_name: str, args: dict[str, Any], *, context: dict[str, Any]
|
||||
) -> str:
|
||||
"""转发工具调用给被包装的 callable。"""
|
||||
return await self._fn(tool_name, args, context=context)
|
||||
|
||||
|
||||
def _to_text_field(value: Any) -> str:
|
||||
"""把 prediction 的 evidence/reasoning 归一为可入库的文本。
|
||||
|
||||
LLM 有时把这些字段返回成 list 或 dict（而非字符串）。sqlite 无法绑定
|
||||
非标量类型，直接入库会抛 ProgrammingError 致该题丢失预测行、进而触发
|
||||
rollout 完整性护栏中止整轮。凡非 str 一律 JSON 序列化为文本。
|
||||
|
||||
参数:
|
||||
value: evidence/reasoning 原始值（可能是 str/list/dict）。
|
||||
|
||||
返回:
|
||||
可直接入库的字符串。
|
||||
"""
|
||||
if isinstance(value, str):
|
||||
return value
|
||||
return json.dumps(value, ensure_ascii=False)
|
||||
|
||||
|
||||
def _zero_result(run_id: str) -> InferenceResult:
|
||||
"""空记录时的零值 InferenceResult。
|
||||
|
||||
参数:
|
||||
run_id: 运行标识。
|
||||
|
||||
返回:
|
||||
全零的 InferenceResult。
|
||||
"""
|
||||
return InferenceResult(
|
||||
run_id=run_id,
|
||||
accuracy=0.0,
|
||||
total=0,
|
||||
correct=0,
|
||||
per_task_type={},
|
||||
steps_mean=0.0,
|
||||
token_usage={"prompt_tokens": 0, "completion_tokens": 0},
|
||||
stop_reason_counts={},
|
||||
)
|
||||
|
||||
|
||||
def _group_by_task_type(records: list[dict[str, Any]]) -> dict[str, dict[str, Any]]:
|
||||
"""按 task_type 分组聚合正确率指标。
|
||||
|
||||
参数:
|
||||
records: 预测记录列表。
|
||||
|
||||
返回:
|
||||
{task_type: {accuracy, total, correct}} 映射。
|
||||
"""
|
||||
task_groups: dict[str, list[dict[str, Any]]] = defaultdict(list)
|
||||
for r in records:
|
||||
task_groups[r["task_type"]].append(r)
|
||||
|
||||
per_task_type: dict[str, dict[str, Any]] = {}
|
||||
for task_type, group in task_groups.items():
|
||||
t_total = len(group)
|
||||
t_correct = sum(1 for r in group if r["prediction"] == r["answer"])
|
||||
per_task_type[task_type] = {
|
||||
"accuracy": t_correct / t_total,
|
||||
"total": t_total,
|
||||
"correct": t_correct,
|
||||
}
|
||||
return per_task_type
|
||||
|
||||
|
||||
def _aggregate_results(records: list[dict[str, Any]], run_id: str) -> InferenceResult:
|
||||
"""从内存 records 聚合推理指标。
|
||||
|
||||
TRM4 从 DB 回读 predictions 表聚合；TRM5 改为从内存直接聚合，
|
||||
避免 DB 回读的同步开销和额外依赖。
|
||||
|
||||
参数:
|
||||
records: _run_single_question 返回的 record 列表。
|
||||
run_id: 当前运行标识。
|
||||
|
||||
返回:
|
||||
InferenceResult 冻结实例。
|
||||
"""
|
||||
total = len(records)
|
||||
if total == 0:
|
||||
return _zero_result(run_id)
|
||||
|
||||
correct = sum(1 for r in records if r["prediction"] == r["answer"])
|
||||
stop_counts: dict[str, int] = defaultdict(int)
|
||||
for r in records:
|
||||
stop_counts[r["stop_reason"]] += 1
|
||||
|
||||
return InferenceResult(
|
||||
run_id=run_id,
|
||||
accuracy=correct / total,
|
||||
total=total,
|
||||
correct=correct,
|
||||
per_task_type=_group_by_task_type(records),
|
||||
steps_mean=sum(r["steps_used"] for r in records) / total,
|
||||
token_usage={
|
||||
"prompt_tokens": sum(r["prompt_tokens"] for r in records),
|
||||
"completion_tokens": sum(r["completion_tokens"] for r in records),
|
||||
},
|
||||
stop_reason_counts=dict(stop_counts),
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 单题推理
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def _run_single_question(
|
||||
qa: GeneratedQuestion,
|
||||
*,
|
||||
llm: LLMProvider,
|
||||
tool_dispatch_fn: Callable[..., Any],
|
||||
prompt_builder: Callable[[GeneratedQuestion], tuple[str, str]],
|
||||
log: HarnessLog,
|
||||
max_steps: int,
|
||||
plugins: list[object],
|
||||
) -> dict[str, Any]:
|
||||
"""执行单道题目的 Agent 推理。
|
||||
|
||||
悲观默认值：record 初始 stop_reason="error"，成功后覆盖。
|
||||
prediction 必落库：log.insert 在 try/except 之后（无论成败）。
|
||||
|
||||
参数:
|
||||
qa: 待推理的题目。
|
||||
llm: LLMProvider 共享实例。
|
||||
tool_dispatch_fn: async 工具调度函数 (tool_name, args, *, context) -> str。
|
||||
prompt_builder: (GeneratedQuestion) -> (system_prompt, user_prompt)。
|
||||
log: HarnessLog 实例（线程安全）。
|
||||
max_steps: AgentLoop 最大步数。
|
||||
plugins: pluggy 插件列表。
|
||||
|
||||
返回:
|
||||
预测结果字典（含 video_id, question_id, prediction, answer 等）。
|
||||
"""
|
||||
record: dict[str, Any] = {
|
||||
"video_id": qa.video_id,
|
||||
"question_id": qa.question_id,
|
||||
"task_type": qa.task_type,
|
||||
"prediction": None,
|
||||
"answer": qa.answer,
|
||||
"evidence": "",
|
||||
"reasoning": "",
|
||||
"steps_used": 0,
|
||||
"prompt_tokens": 0,
|
||||
"completion_tokens": 0,
|
||||
"stop_reason": "error", # 悲观默认
|
||||
"steps_json": "[]",
|
||||
}
|
||||
|
||||
try:
|
||||
system_prompt, user_prompt = prompt_builder(qa)
|
||||
dispatcher = _DispatcherAdapter(tool_dispatch_fn)
|
||||
loop = AgentLoop(llm, max_steps=max_steps)
|
||||
loop_result: LoopResult = await loop.run(
|
||||
system_prompt,
|
||||
user_prompt,
|
||||
dispatcher,
|
||||
plugins=plugins,
|
||||
session_id=qa.question_id,
|
||||
)
|
||||
|
||||
result_dict = loop_result.result if isinstance(loop_result.result, dict) else {}
|
||||
evidence = _to_text_field(result_dict.get("evidence", ""))
|
||||
reasoning = _to_text_field(result_dict.get("reasoning", ""))
|
||||
record.update(
|
||||
{
|
||||
"prediction": result_dict.get("answer"),
|
||||
"evidence": evidence,
|
||||
"reasoning": reasoning,
|
||||
"steps_used": loop_result.steps_used,
|
||||
"prompt_tokens": loop_result.token_usage["prompt_tokens"],
|
||||
"completion_tokens": loop_result.token_usage["completion_tokens"],
|
||||
"stop_reason": loop_result.stop_reason,
|
||||
"steps_json": json.dumps(
|
||||
[
|
||||
{
|
||||
"thought": s.thought,
|
||||
"tool_call": s.tool_call,
|
||||
"tool_output": s.tool_output,
|
||||
}
|
||||
for s in loop_result.steps
|
||||
],
|
||||
ensure_ascii=False,
|
||||
),
|
||||
}
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("[{}] QA {} 执行异常", qa.video_id, qa.question_id)
|
||||
|
||||
# prediction 必落库（try 外，无论成败）
|
||||
await asyncio.to_thread(log.insert, "predictions", record)
|
||||
return record
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 建表
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _ensure_tables(log: HarnessLog) -> None:
|
||||
"""创建推理所需的 5 张表。
|
||||
|
||||
参数:
|
||||
log: HarnessLog 实例。
|
||||
"""
|
||||
log.create_table("predictions", PREDICTIONS_SCHEMA)
|
||||
log.create_table("traces", TRACES_SCHEMA)
|
||||
log.create_table("validation_flags", VALIDATION_FLAGS_SCHEMA)
|
||||
log.create_table("anchor_check", ANCHOR_CHECK_SCHEMA)
|
||||
log.create_table("observe_frame_health", OF_HEALTH_SCHEMA)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 公共入口
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def run_inference(
|
||||
questions: list[GeneratedQuestion],
|
||||
*,
|
||||
llm: LLMProvider,
|
||||
tool_dispatch_fn: Callable[..., Any],
|
||||
prompt_builder: Callable[[GeneratedQuestion], tuple[str, str]],
|
||||
log: HarnessLog,
|
||||
run_id: str,
|
||||
concurrency: int,
|
||||
max_steps: int,
|
||||
skill_mode: str,
|
||||
plugins_factory: Callable[[str, str], list[object]] | None = None,
|
||||
) -> InferenceResult:
|
||||
"""在视频树上执行 Agent 推理，对应训练循环的 forward()。
|
||||
|
||||
参数:
|
||||
questions: 待推理的题目列表。
|
||||
llm: LLMProvider 共享实例（依赖注入）。
|
||||
tool_dispatch_fn: async 工具调度函数 (tool_name, args, *, context) -> str。
|
||||
prompt_builder: prompt 构建函数 (GeneratedQuestion) -> (system_prompt, user_prompt)。
|
||||
log: HarnessLog 实例（由调用方管理生命周期）。
|
||||
run_id: 运行标识（必传，空串 → ValueError）。
|
||||
concurrency: 最大并发数（asyncio.Semaphore 控制）。
|
||||
max_steps: AgentLoop 单题最大步数。
|
||||
skill_mode: "auto" / "manual" / "none"（传递给调用方的 prompt/plugin 构建逻辑）。
|
||||
plugins_factory: 可选的插件工厂 (video_id, question_id) -> plugins 列表。
|
||||
|
||||
返回:
|
||||
InferenceResult（含 accuracy、per_task_type 等聚合指标）。
|
||||
|
||||
异常:
|
||||
ValueError: run_id 为空串或纯空白。
|
||||
"""
|
||||
if not run_id or not run_id.strip():
|
||||
raise ValueError("run_id 不得为空串或纯空白")
|
||||
|
||||
_ensure_tables(log)
|
||||
|
||||
if not questions:
|
||||
logger.info("题目列表为空，返回零值 InferenceResult")
|
||||
return _aggregate_results([], run_id)
|
||||
|
||||
sem = asyncio.Semaphore(concurrency)
|
||||
total_count = len(questions)
|
||||
|
||||
async def _bounded(index: int, qa: GeneratedQuestion) -> dict[str, Any]:
|
||||
"""信号量限流的单题推理包装。"""
|
||||
async with sem:
|
||||
plugins = (
|
||||
plugins_factory(qa.video_id, qa.question_id) if plugins_factory is not None else []
|
||||
)
|
||||
result = await _run_single_question(
|
||||
qa,
|
||||
llm=llm,
|
||||
tool_dispatch_fn=tool_dispatch_fn,
|
||||
prompt_builder=prompt_builder,
|
||||
log=log,
|
||||
max_steps=max_steps,
|
||||
plugins=plugins,
|
||||
)
|
||||
logger.info(
|
||||
"[{}/{}] {} QA {} 完成 (stop={})",
|
||||
index + 1,
|
||||
total_count,
|
||||
qa.video_id,
|
||||
qa.question_id,
|
||||
result["stop_reason"],
|
||||
)
|
||||
return result
|
||||
|
||||
results = await asyncio.gather(*[_bounded(i, qa) for i, qa in enumerate(questions)])
|
||||
|
||||
inference_result = _aggregate_results(list(results), run_id)
|
||||
logger.info(
|
||||
"推理完成: accuracy={:.2%} ({}/{})",
|
||||
inference_result.accuracy,
|
||||
inference_result.correct,
|
||||
inference_result.total,
|
||||
)
|
||||
return inference_result
|
||||
Reference in New Issue
Block a user