334 lines
11 KiB
Python
334 lines
11 KiB
Python
"""HarnessLog：SQLite 薄包装 + RunLogImpl 只读查询端口。
|
||
|
||
HarnessLog 提供统一的结构化日志接口，从 TRM4 直搬，保留全部线程安全与幂等语义。
|
||
RunLogImpl 实现 core/evolution/protocols.py::RunLog Protocol，用独立连接做只读 SELECT，
|
||
不经 HarnessLog 生命周期（不触发 _runs INSERT OR IGNORE），避免污染运行状态。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
import sqlite3
|
||
import subprocess
|
||
import threading
|
||
from datetime import UTC, datetime
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
|
||
def _get_git_sha() -> str | None:
|
||
"""获取当前 git commit SHA。"""
|
||
try:
|
||
result = subprocess.run(
|
||
["git", "rev-parse", "HEAD"],
|
||
capture_output=True,
|
||
text=True,
|
||
check=True,
|
||
)
|
||
return result.stdout.strip()
|
||
except (subprocess.CalledProcessError, FileNotFoundError):
|
||
return None
|
||
|
||
|
||
def _now_iso() -> str:
|
||
"""返回当前 UTC 时间的 ISO 格式字符串。"""
|
||
return datetime.now(UTC).isoformat()
|
||
|
||
|
||
class HarnessLog:
|
||
"""SQLite 薄包装，为科研项目提供统一的结构化日志接口。
|
||
|
||
关键设计：
|
||
- WAL 模式 + threading.Lock 保证共享连接下并发安全。
|
||
- INSERT OR IGNORE INTO _runs 保证幂等（同 run_id 多次创建不报错）。
|
||
- query 也持锁：共享连接（check_same_thread=False）下并发 SELECT + INSERT
|
||
在同一连接上 execute 会损坏游标状态，故读也须串行化。
|
||
- context manager 语义：正常退出 completed，异常退出 failed。
|
||
|
||
参数:
|
||
db_path: SQLite 数据库文件路径。
|
||
run_id: 本次运行的唯一标识。
|
||
git_sha: 代码版本，默认自动获取。
|
||
config_snapshot: 本次运行的配置快照。
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
db_path: str,
|
||
run_id: str,
|
||
git_sha: str | None = None,
|
||
config_snapshot: dict[str, Any] | None = None,
|
||
) -> None:
|
||
self._run_id = run_id
|
||
Path(db_path).parent.mkdir(parents=True, exist_ok=True)
|
||
self._conn = sqlite3.connect(db_path, check_same_thread=False)
|
||
self._lock = threading.Lock()
|
||
self._conn.row_factory = sqlite3.Row
|
||
self._conn.execute("PRAGMA journal_mode=WAL")
|
||
self._init_fixed_tables()
|
||
resolved_sha = git_sha or _get_git_sha()
|
||
config_json = json.dumps(config_snapshot, ensure_ascii=False) if config_snapshot else None
|
||
self._conn.execute(
|
||
"INSERT OR IGNORE INTO _runs"
|
||
" (run_id, git_sha, started_at, config, status)"
|
||
" VALUES (?, ?, ?, ?, ?)",
|
||
(run_id, resolved_sha, _now_iso(), config_json, "running"),
|
||
)
|
||
self._conn.commit()
|
||
|
||
def _init_fixed_tables(self) -> None:
|
||
"""创建 _runs 和 _events 固定表。"""
|
||
self._conn.execute("""
|
||
CREATE TABLE IF NOT EXISTS _runs (
|
||
run_id TEXT PRIMARY KEY,
|
||
git_sha TEXT,
|
||
started_at TEXT,
|
||
finished_at TEXT,
|
||
config JSON,
|
||
status TEXT DEFAULT 'running',
|
||
skills_version TEXT,
|
||
prompts_version TEXT,
|
||
questions_ref TEXT
|
||
)
|
||
""")
|
||
self._conn.execute("""
|
||
CREATE TABLE IF NOT EXISTS _events (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
run_id TEXT,
|
||
timestamp TEXT,
|
||
event_type TEXT,
|
||
payload JSON
|
||
)
|
||
""")
|
||
self._conn.commit()
|
||
|
||
def create_table(
|
||
self,
|
||
name: str,
|
||
columns: dict[str, str],
|
||
primary_key: str | None = None,
|
||
) -> None:
|
||
"""创建自定义表，自动追加 run_id 和 timestamp 列。
|
||
|
||
参数:
|
||
name: 表名。
|
||
columns: 列定义，如 {"epoch": "INTEGER", "loss": "REAL"}。
|
||
primary_key: 主键列名。
|
||
"""
|
||
all_columns = {"run_id": "TEXT", "timestamp": "TEXT"}
|
||
all_columns.update(columns)
|
||
col_defs = []
|
||
for col_name, col_type in all_columns.items():
|
||
pk_suffix = " PRIMARY KEY" if col_name == primary_key else ""
|
||
col_defs.append(f"{col_name} {col_type}{pk_suffix}")
|
||
sql = f"CREATE TABLE IF NOT EXISTS {name} ({', '.join(col_defs)})"
|
||
self._conn.execute(sql)
|
||
self._conn.commit()
|
||
|
||
def insert(self, table: str, record: dict[str, Any], mode: str = "append") -> None:
|
||
"""插入一条记录，自动填充 run_id 和 timestamp。
|
||
|
||
参数:
|
||
table: 目标表名。
|
||
record: 要插入的数据。
|
||
mode: "append" 或 "upsert"。
|
||
"""
|
||
enriched = {"run_id": self._run_id, "timestamp": _now_iso()}
|
||
enriched.update(record)
|
||
cols = list(enriched.keys())
|
||
placeholders = ", ".join(["?"] * len(cols))
|
||
col_names = ", ".join(cols)
|
||
values = [enriched[c] for c in cols]
|
||
if mode == "upsert":
|
||
sql = f"INSERT OR REPLACE INTO {table} ({col_names}) VALUES ({placeholders})"
|
||
else:
|
||
sql = f"INSERT INTO {table} ({col_names}) VALUES ({placeholders})"
|
||
with self._lock:
|
||
self._conn.execute(sql, values)
|
||
self._conn.commit()
|
||
|
||
def insert_many(self, table: str, records: list[dict[str, Any]], mode: str = "append") -> None:
|
||
"""批量插入多条记录。
|
||
|
||
参数:
|
||
table: 目标表名。
|
||
records: 要插入的数据列表。
|
||
mode: "append" 或 "upsert"。
|
||
"""
|
||
for record in records:
|
||
self.insert(table, record, mode=mode)
|
||
|
||
def execute(self, sql: str, params: tuple[Any, ...] = ()) -> None:
|
||
"""执行原生 SQL 写操作。
|
||
|
||
参数:
|
||
sql: SQL 语句。
|
||
params: 参数元组。
|
||
"""
|
||
with self._lock:
|
||
self._conn.execute(sql, params)
|
||
self._conn.commit()
|
||
|
||
def query(self, sql: str, params: tuple[Any, ...] = ()) -> list[dict[str, Any]]:
|
||
"""执行原生 SQL 查询，返回 list[dict]。
|
||
|
||
与所有写方法同持 self._lock：共享连接（check_same_thread=False）下，
|
||
并发 SELECT 与 INSERT 在同一连接上 execute 会损坏游标状态，故读也须串行化。
|
||
|
||
参数:
|
||
sql: SQL 查询语句。
|
||
params: 参数元组。
|
||
|
||
返回:
|
||
查询结果列表，每行为一个字典。
|
||
"""
|
||
with self._lock:
|
||
cursor = self._conn.execute(sql, params)
|
||
columns = [desc[0] for desc in cursor.description]
|
||
return [dict(zip(columns, row, strict=True)) for row in cursor.fetchall()]
|
||
|
||
def log_event(self, event_type: str, payload: dict[str, Any]) -> None:
|
||
"""向 _events 表写入一条事件。
|
||
|
||
参数:
|
||
event_type: 事件类型标识。
|
||
payload: 事件数据。
|
||
"""
|
||
with self._lock:
|
||
self._conn.execute(
|
||
"INSERT INTO _events (run_id, timestamp, event_type, payload) VALUES (?, ?, ?, ?)",
|
||
(
|
||
self._run_id,
|
||
_now_iso(),
|
||
event_type,
|
||
json.dumps(payload, ensure_ascii=False),
|
||
),
|
||
)
|
||
self._conn.commit()
|
||
|
||
def close(self, status: str = "completed") -> None:
|
||
"""更新运行状态并关闭连接。
|
||
|
||
参数:
|
||
status: 最终状态，"completed" 或 "failed"。
|
||
"""
|
||
with self._lock:
|
||
self._conn.execute(
|
||
"UPDATE _runs SET finished_at = ?, status = ? WHERE run_id = ?",
|
||
(_now_iso(), status, self._run_id),
|
||
)
|
||
self._conn.commit()
|
||
self._conn.close()
|
||
|
||
def __enter__(self) -> HarnessLog:
|
||
return self
|
||
|
||
def __exit__(
|
||
self,
|
||
exc_type: type[BaseException] | None,
|
||
exc_val: BaseException | None,
|
||
exc_tb: Any,
|
||
) -> None:
|
||
status = "failed" if exc_type is not None else "completed"
|
||
self.close(status=status)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# RunLogImpl — core/evolution/protocols.py::RunLog 的只读实现
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _read_table(
|
||
db_path: str,
|
||
table: str,
|
||
run_id: str,
|
||
*,
|
||
question_ids: list[str] | None = None,
|
||
) -> list[dict[str, Any]]:
|
||
"""纯读某表指定 run 的行——不经 HarnessLog 生命周期，避免回读污染 _runs 运行状态。
|
||
|
||
HarnessLog.__enter__/__exit__ 会对 run_id 做 INSERT OR IGNORE 并在退出时标 completed；
|
||
回读指标绝不应改运行状态，故走独立只读连接（仅 SELECT）。
|
||
|
||
参数:
|
||
db_path: SQLite 路径。
|
||
table: 表名（内部固定常量，非外部输入，无注入风险）。
|
||
run_id: 过滤的 run ID。
|
||
question_ids: 可选的 question_id 过滤列表。
|
||
|
||
返回:
|
||
行 dict 列表；表尚未建（没写过）视为无数据返 []。
|
||
"""
|
||
conn = sqlite3.connect(db_path)
|
||
conn.row_factory = sqlite3.Row
|
||
try:
|
||
exists = conn.execute(
|
||
"SELECT name FROM sqlite_master WHERE type='table' AND name=?", (table,)
|
||
).fetchone()
|
||
if exists is None:
|
||
return []
|
||
|
||
if question_ids is not None:
|
||
placeholders = ", ".join(["?"] * len(question_ids))
|
||
sql = f"SELECT * FROM {table} WHERE run_id = ? AND question_id IN ({placeholders})"
|
||
rows = conn.execute(sql, (run_id, *question_ids)).fetchall()
|
||
else:
|
||
rows = conn.execute(f"SELECT * FROM {table} WHERE run_id = ?", (run_id,)).fetchall()
|
||
|
||
return [dict(r) for r in rows]
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
class RunLogImpl:
|
||
"""RunLog Protocol 的只读实现。
|
||
|
||
用独立 sqlite3.connect 做 SELECT，不经 HarnessLog 生命周期（不触发 _runs INSERT），
|
||
asyncio.to_thread 包装同步 SQL 查询，避免引入 aiosqlite 新依赖。
|
||
|
||
参数:
|
||
db_path: SQLite 数据库文件路径。
|
||
"""
|
||
|
||
def __init__(self, db_path: str) -> None:
|
||
self._db_path = db_path
|
||
|
||
async def get_predictions(
|
||
self,
|
||
run_id: str,
|
||
*,
|
||
question_ids: list[str] | None = None,
|
||
) -> list[dict[str, Any]]:
|
||
"""查询指定 run 的预测记录。
|
||
|
||
参数:
|
||
run_id: 运行标识。
|
||
question_ids: 可选的题目 ID 过滤列表。
|
||
|
||
返回:
|
||
预测记录字典列表。
|
||
"""
|
||
return await asyncio.to_thread(
|
||
_read_table, self._db_path, "predictions", run_id, question_ids=question_ids
|
||
)
|
||
|
||
async def get_traces(
|
||
self,
|
||
run_id: str,
|
||
*,
|
||
question_ids: list[str] | None = None,
|
||
) -> list[dict[str, Any]]:
|
||
"""查询指定 run 的推理轨迹。
|
||
|
||
参数:
|
||
run_id: 运行标识。
|
||
question_ids: 可选的题目 ID 过滤列表。
|
||
|
||
返回:
|
||
轨迹记录字典列表。
|
||
"""
|
||
return await asyncio.to_thread(
|
||
_read_table, self._db_path, "traces", run_id, question_ids=question_ids
|
||
)
|