522 lines
20 KiB
Python
522 lines
20 KiB
Python
"""SQLite persistence for tuning sessions, trials, proposals and scalar data."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import sqlite3
|
|
import threading
|
|
import uuid
|
|
from collections.abc import Iterator
|
|
from contextlib import contextmanager
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
SCHEMA_VERSION = 1
|
|
|
|
|
|
def now_iso() -> str:
|
|
return datetime.now(UTC).isoformat()
|
|
|
|
|
|
def _json(value: Any) -> str:
|
|
return json.dumps(value, ensure_ascii=False, separators=(",", ":"), allow_nan=False)
|
|
|
|
|
|
def _decode(value: str | None) -> Any:
|
|
return json.loads(value) if value else None
|
|
|
|
|
|
def _lttb(points: list[dict], threshold: int) -> list[dict]:
|
|
"""Largest-Triangle-Three-Buckets downsampling preserving peaks and endpoints."""
|
|
if threshold >= len(points) or threshold < 3:
|
|
return points[:threshold]
|
|
sampled = [points[0]]
|
|
bucket_width = (len(points) - 2) / (threshold - 2)
|
|
anchor_index = 0
|
|
for bucket in range(threshold - 2):
|
|
average_start = int((bucket + 1) * bucket_width) + 1
|
|
average_end = min(int((bucket + 2) * bucket_width) + 1, len(points))
|
|
average_bucket = points[average_start:average_end] or [points[-1]]
|
|
average_x = sum(point["step"] for point in average_bucket) / len(average_bucket)
|
|
average_y = sum(point["value"] for point in average_bucket) / len(average_bucket)
|
|
range_start = int(bucket * bucket_width) + 1
|
|
range_end = min(int((bucket + 1) * bucket_width) + 1, len(points) - 1)
|
|
anchor = points[anchor_index]
|
|
selected_index = range_start
|
|
maximum_area = -1.0
|
|
for index in range(range_start, max(range_start + 1, range_end)):
|
|
point = points[index]
|
|
area = abs(
|
|
(anchor["step"] - average_x) * (point["value"] - anchor["value"])
|
|
- (anchor["step"] - point["step"]) * (average_y - anchor["value"])
|
|
)
|
|
if area > maximum_area:
|
|
maximum_area = area
|
|
selected_index = index
|
|
sampled.append(points[selected_index])
|
|
anchor_index = selected_index
|
|
sampled.append(points[-1])
|
|
return sampled
|
|
|
|
|
|
class TuningStorage:
|
|
def __init__(self, path: Path):
|
|
self.path = path.expanduser().resolve()
|
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
|
self.local = threading.local()
|
|
self._migrate()
|
|
|
|
def connection(self) -> sqlite3.Connection:
|
|
connection = getattr(self.local, "connection", None)
|
|
if connection is None:
|
|
connection = sqlite3.connect(self.path, timeout=10, isolation_level=None)
|
|
connection.row_factory = sqlite3.Row
|
|
connection.execute("PRAGMA foreign_keys=ON")
|
|
connection.execute("PRAGMA journal_mode=WAL")
|
|
connection.execute("PRAGMA busy_timeout=10000")
|
|
self.local.connection = connection
|
|
return connection
|
|
|
|
@contextmanager
|
|
def transaction(self) -> Iterator[sqlite3.Connection]:
|
|
connection = self.connection()
|
|
connection.execute("BEGIN IMMEDIATE")
|
|
try:
|
|
yield connection
|
|
connection.execute("COMMIT")
|
|
except Exception:
|
|
connection.execute("ROLLBACK")
|
|
raise
|
|
|
|
def _migrate(self) -> None:
|
|
connection = self.connection()
|
|
connection.executescript(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS schema_migrations(version INTEGER PRIMARY KEY);
|
|
CREATE TABLE IF NOT EXISTS sessions(
|
|
id TEXT PRIMARY KEY, state TEXT NOT NULL, mode TEXT NOT NULL,
|
|
created_at TEXT NOT NULL, updated_at TEXT NOT NULL,
|
|
config_json TEXT NOT NULL, objective_json TEXT NOT NULL,
|
|
message TEXT NOT NULL, current_trial_id TEXT, best_trial_id TEXT,
|
|
consecutive_no_improve INTEGER NOT NULL DEFAULT 0,
|
|
fallback_enabled INTEGER NOT NULL DEFAULT 0
|
|
);
|
|
CREATE TABLE IF NOT EXISTS trials(
|
|
id TEXT PRIMARY KEY,
|
|
session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
|
|
number INTEGER NOT NULL, state TEXT NOT NULL, rung INTEGER NOT NULL,
|
|
target_iterations INTEGER NOT NULL, reward_config_json TEXT NOT NULL,
|
|
proposal_id TEXT, run_dir TEXT NOT NULL, checkpoint_path TEXT,
|
|
policy_path TEXT, evaluation_json TEXT, score REAL, eligible INTEGER,
|
|
created_at TEXT NOT NULL, started_at TEXT, ended_at TEXT, message TEXT NOT NULL,
|
|
UNIQUE(session_id, number, rung)
|
|
);
|
|
CREATE TABLE IF NOT EXISTS proposals(
|
|
id TEXT PRIMARY KEY,
|
|
session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
|
|
base_trial_id TEXT, state TEXT NOT NULL, source TEXT NOT NULL,
|
|
patch_json TEXT NOT NULL, rationale TEXT NOT NULL,
|
|
expected_json TEXT, confidence REAL NOT NULL,
|
|
created_at TEXT NOT NULL, decided_at TEXT, feedback TEXT
|
|
);
|
|
CREATE TABLE IF NOT EXISTS metric_points(
|
|
trial_id TEXT NOT NULL REFERENCES trials(id) ON DELETE CASCADE,
|
|
tag TEXT NOT NULL, step INTEGER NOT NULL,
|
|
wall_time REAL NOT NULL, value REAL NOT NULL,
|
|
PRIMARY KEY(trial_id, tag, step)
|
|
);
|
|
CREATE TABLE IF NOT EXISTS audit_events(
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
|
|
event_type TEXT NOT NULL, payload_json TEXT NOT NULL, created_at TEXT NOT NULL
|
|
);
|
|
CREATE TABLE IF NOT EXISTS presets(
|
|
id TEXT PRIMARY KEY, name TEXT NOT NULL UNIQUE, session_id TEXT NOT NULL,
|
|
trial_id TEXT NOT NULL, reward_config_json TEXT NOT NULL, created_at TEXT NOT NULL
|
|
);
|
|
CREATE INDEX IF NOT EXISTS idx_trials_session ON trials(session_id, number, rung);
|
|
CREATE INDEX IF NOT EXISTS idx_proposals_session ON proposals(session_id, created_at);
|
|
CREATE INDEX IF NOT EXISTS idx_metrics_trial_tag ON metric_points(trial_id, tag, step);
|
|
"""
|
|
)
|
|
connection.execute(
|
|
"INSERT OR IGNORE INTO schema_migrations(version) VALUES (?)", (SCHEMA_VERSION,)
|
|
)
|
|
|
|
def recover_interrupted(self) -> None:
|
|
at = now_iso()
|
|
with self.transaction() as connection:
|
|
connection.execute(
|
|
"UPDATE trials SET state='interrupted', ended_at=?, "
|
|
"message='服务重启中断，等待显式恢复' "
|
|
"WHERE state IN ('training','evaluating')",
|
|
(at,),
|
|
)
|
|
connection.execute(
|
|
"UPDATE sessions SET state='interrupted', updated_at=?, "
|
|
"message='服务重启中断，可从完整 checkpoint 恢复' "
|
|
"WHERE state IN ('running','evaluating')",
|
|
(at,),
|
|
)
|
|
|
|
def create_session(self, mode: str, config: dict, objective: dict, fallback: bool) -> dict:
|
|
session_id, at = uuid.uuid4().hex, now_iso()
|
|
with self.transaction() as connection:
|
|
connection.execute(
|
|
"INSERT INTO sessions("
|
|
"id,state,mode,created_at,updated_at,config_json,objective_json,"
|
|
"message,fallback_enabled) "
|
|
"VALUES (?, 'queued', ?, ?, ?, ?, ?, '等待基线训练', ?)",
|
|
(session_id, mode, at, at, _json(config), _json(objective), int(fallback)),
|
|
)
|
|
connection.execute(
|
|
"INSERT INTO audit_events(session_id,event_type,payload_json,created_at) "
|
|
"VALUES (?,?,?,?)",
|
|
(session_id, "session_created", _json({"mode": mode}), at),
|
|
)
|
|
return self.get_session(session_id)
|
|
|
|
def _session(self, row: sqlite3.Row) -> dict:
|
|
return {
|
|
"id": row["id"],
|
|
"state": row["state"],
|
|
"mode": row["mode"],
|
|
"createdAt": row["created_at"],
|
|
"updatedAt": row["updated_at"],
|
|
"config": _decode(row["config_json"]),
|
|
"objectiveWeights": _decode(row["objective_json"]),
|
|
"message": row["message"],
|
|
"currentTrialId": row["current_trial_id"],
|
|
"bestTrialId": row["best_trial_id"],
|
|
"consecutiveNoImprove": row["consecutive_no_improve"],
|
|
"fallbackEnabled": bool(row["fallback_enabled"]),
|
|
}
|
|
|
|
def get_session(self, session_id: str) -> dict:
|
|
row = (
|
|
self.connection().execute("SELECT * FROM sessions WHERE id=?", (session_id,)).fetchone()
|
|
)
|
|
if row is None:
|
|
raise KeyError(session_id)
|
|
return self._session(row)
|
|
|
|
def list_sessions(self, limit: int = 50) -> list[dict]:
|
|
rows = (
|
|
self.connection()
|
|
.execute("SELECT * FROM sessions ORDER BY created_at DESC LIMIT ?", (limit,))
|
|
.fetchall()
|
|
)
|
|
return [self._session(row) for row in rows]
|
|
|
|
def update_session(self, session_id: str, **changes: Any) -> bool:
|
|
columns = {
|
|
"state": "state",
|
|
"message": "message",
|
|
"current_trial_id": "current_trial_id",
|
|
"best_trial_id": "best_trial_id",
|
|
"consecutive_no_improve": "consecutive_no_improve",
|
|
}
|
|
values, assignments = [], []
|
|
for key, value in changes.items():
|
|
if key not in columns:
|
|
raise ValueError(key)
|
|
assignments.append(f"{columns[key]}=?")
|
|
values.append(value)
|
|
assignments.append("updated_at=?")
|
|
values.extend((now_iso(), session_id))
|
|
cursor = self.connection().execute(
|
|
f"UPDATE sessions SET {', '.join(assignments)} WHERE id=?", values
|
|
)
|
|
return cursor.rowcount == 1
|
|
|
|
def create_trial(
|
|
self,
|
|
session_id: str,
|
|
number: int,
|
|
rung: int,
|
|
target: int,
|
|
reward_config: dict,
|
|
proposal_id: str | None,
|
|
run_dir: str,
|
|
) -> dict:
|
|
trial_id, at = uuid.uuid4().hex, now_iso()
|
|
with self.transaction() as connection:
|
|
connection.execute(
|
|
"INSERT INTO trials("
|
|
"id,session_id,number,state,rung,target_iterations,reward_config_json,"
|
|
"proposal_id,run_dir,created_at,message) "
|
|
"VALUES (?,?,?,'queued',?,?,?,?,?,?,'等待训练')",
|
|
(
|
|
trial_id,
|
|
session_id,
|
|
number,
|
|
rung,
|
|
target,
|
|
_json(reward_config),
|
|
proposal_id,
|
|
run_dir,
|
|
at,
|
|
),
|
|
)
|
|
connection.execute(
|
|
"UPDATE sessions SET current_trial_id=?,updated_at=? WHERE id=?",
|
|
(trial_id, at, session_id),
|
|
)
|
|
return self.get_trial(trial_id)
|
|
|
|
def _trial(self, row: sqlite3.Row) -> dict:
|
|
return {
|
|
"id": row["id"],
|
|
"sessionId": row["session_id"],
|
|
"number": row["number"],
|
|
"state": row["state"],
|
|
"rung": row["rung"],
|
|
"targetIterations": row["target_iterations"],
|
|
"rewardConfig": _decode(row["reward_config_json"]),
|
|
"proposalId": row["proposal_id"],
|
|
"runDir": row["run_dir"],
|
|
"checkpointPath": row["checkpoint_path"],
|
|
"policyPath": row["policy_path"],
|
|
"evaluation": _decode(row["evaluation_json"]),
|
|
"score": row["score"],
|
|
"eligible": None if row["eligible"] is None else bool(row["eligible"]),
|
|
"createdAt": row["created_at"],
|
|
"startedAt": row["started_at"],
|
|
"endedAt": row["ended_at"],
|
|
"message": row["message"],
|
|
}
|
|
|
|
def get_trial(self, trial_id: str) -> dict:
|
|
row = self.connection().execute("SELECT * FROM trials WHERE id=?", (trial_id,)).fetchone()
|
|
if row is None:
|
|
raise KeyError(trial_id)
|
|
return self._trial(row)
|
|
|
|
def list_trials(self, session_id: str) -> list[dict]:
|
|
rows = (
|
|
self.connection()
|
|
.execute("SELECT * FROM trials WHERE session_id=? ORDER BY number,rung", (session_id,))
|
|
.fetchall()
|
|
)
|
|
return [self._trial(row) for row in rows]
|
|
|
|
def delete_trial(self, trial_id: str) -> bool:
|
|
cursor = self.connection().execute(
|
|
"DELETE FROM trials WHERE id=? AND state='interrupted'", (trial_id,)
|
|
)
|
|
return cursor.rowcount == 1
|
|
|
|
def update_trial(self, trial_id: str, **changes: Any) -> bool:
|
|
columns = {
|
|
"state": "state",
|
|
"message": "message",
|
|
"checkpoint_path": "checkpoint_path",
|
|
"policy_path": "policy_path",
|
|
"score": "score",
|
|
"eligible": "eligible",
|
|
"started_at": "started_at",
|
|
"ended_at": "ended_at",
|
|
"evaluation": "evaluation_json",
|
|
}
|
|
values, assignments = [], []
|
|
for key, value in changes.items():
|
|
if key not in columns:
|
|
raise ValueError(key)
|
|
if key == "evaluation":
|
|
value = _json(value)
|
|
if key == "eligible":
|
|
value = int(value)
|
|
assignments.append(f"{columns[key]}=?")
|
|
values.append(value)
|
|
values.append(trial_id)
|
|
cursor = self.connection().execute(
|
|
f"UPDATE trials SET {', '.join(assignments)} WHERE id=?", values
|
|
)
|
|
return cursor.rowcount == 1
|
|
|
|
def create_proposal(
|
|
self,
|
|
session_id: str,
|
|
base_trial_id: str | None,
|
|
patch: dict,
|
|
rationale: str,
|
|
expected: Any,
|
|
confidence: float,
|
|
source: str = "agent",
|
|
) -> dict:
|
|
proposal_id, at = uuid.uuid4().hex, now_iso()
|
|
self.connection().execute(
|
|
"INSERT INTO proposals("
|
|
"id,session_id,base_trial_id,state,source,patch_json,rationale,"
|
|
"expected_json,confidence,created_at) "
|
|
"VALUES (?,?,?,'pending',?,?,?,?,?,?)",
|
|
(
|
|
proposal_id,
|
|
session_id,
|
|
base_trial_id,
|
|
source,
|
|
_json(patch),
|
|
rationale,
|
|
_json(expected),
|
|
confidence,
|
|
at,
|
|
),
|
|
)
|
|
return self.get_proposal(proposal_id)
|
|
|
|
def _proposal(self, row: sqlite3.Row) -> dict:
|
|
return {
|
|
"id": row["id"],
|
|
"sessionId": row["session_id"],
|
|
"baseTrialId": row["base_trial_id"],
|
|
"state": row["state"],
|
|
"source": row["source"],
|
|
"patch": _decode(row["patch_json"]),
|
|
"rationale": row["rationale"],
|
|
"expectedImpact": _decode(row["expected_json"]),
|
|
"confidence": row["confidence"],
|
|
"createdAt": row["created_at"],
|
|
"decidedAt": row["decided_at"],
|
|
"feedback": row["feedback"],
|
|
}
|
|
|
|
def get_proposal(self, proposal_id: str) -> dict:
|
|
row = (
|
|
self.connection()
|
|
.execute("SELECT * FROM proposals WHERE id=?", (proposal_id,))
|
|
.fetchone()
|
|
)
|
|
if row is None:
|
|
raise KeyError(proposal_id)
|
|
return self._proposal(row)
|
|
|
|
def list_proposals(self, session_id: str) -> list[dict]:
|
|
rows = (
|
|
self.connection()
|
|
.execute(
|
|
"SELECT * FROM proposals WHERE session_id=? ORDER BY created_at", (session_id,)
|
|
)
|
|
.fetchall()
|
|
)
|
|
return [self._proposal(row) for row in rows]
|
|
|
|
def decide_proposal(
|
|
self, proposal_id: str, state: str, feedback: str | None, patch: dict | None = None
|
|
) -> bool:
|
|
at = now_iso()
|
|
assignments, values = ["state=?", "feedback=?", "decided_at=?"], [state, feedback, at]
|
|
if patch is not None:
|
|
assignments.append("patch_json=?")
|
|
values.append(_json(patch))
|
|
values.extend((proposal_id,))
|
|
cursor = self.connection().execute(
|
|
f"UPDATE proposals SET {', '.join(assignments)} WHERE id=? AND state='pending'", values
|
|
)
|
|
return cursor.rowcount == 1
|
|
|
|
def insert_metrics(self, trial_id: str, points: list[tuple[str, int, float, float]]) -> None:
|
|
self.connection().executemany(
|
|
"INSERT INTO metric_points(trial_id,tag,step,wall_time,value) "
|
|
"VALUES (?,?,?,?,?) ON CONFLICT(trial_id,tag,step) DO UPDATE SET "
|
|
"wall_time=excluded.wall_time,value=excluded.value",
|
|
[(trial_id, *point) for point in points],
|
|
)
|
|
|
|
def metrics(
|
|
self, trial_id: str, tags: list[str] | None = None, max_points: int = 1000
|
|
) -> list[dict]:
|
|
parameters: list[Any] = [trial_id]
|
|
clause = "trial_id=?"
|
|
if tags:
|
|
clause += f" AND tag IN ({','.join('?' for _ in tags)})"
|
|
parameters.extend(tags)
|
|
rows = (
|
|
self.connection()
|
|
.execute(
|
|
f"SELECT tag,step,wall_time,value FROM metric_points "
|
|
f"WHERE {clause} ORDER BY tag,step",
|
|
parameters,
|
|
)
|
|
.fetchall()
|
|
)
|
|
grouped: dict[str, list[dict]] = {}
|
|
for row in rows:
|
|
grouped.setdefault(row["tag"], []).append(
|
|
{"step": row["step"], "wallTime": row["wall_time"], "value": row["value"]}
|
|
)
|
|
series = []
|
|
for tag, values in grouped.items():
|
|
if len(values) > max_points:
|
|
values = _lttb(values, max_points)
|
|
series.append({"tag": tag, "points": values})
|
|
return series
|
|
|
|
def audit(self, session_id: str, event_type: str, payload: Any) -> None:
|
|
self.connection().execute(
|
|
"INSERT INTO audit_events(session_id,event_type,payload_json,created_at) "
|
|
"VALUES (?,?,?,?)",
|
|
(session_id, event_type, _json(payload), now_iso()),
|
|
)
|
|
|
|
def audit_events(self, session_id: str) -> list[dict]:
|
|
rows = (
|
|
self.connection()
|
|
.execute("SELECT * FROM audit_events WHERE session_id=? ORDER BY id", (session_id,))
|
|
.fetchall()
|
|
)
|
|
return [
|
|
{
|
|
"id": row["id"],
|
|
"type": row["event_type"],
|
|
"payload": _decode(row["payload_json"]),
|
|
"createdAt": row["created_at"],
|
|
}
|
|
for row in rows
|
|
]
|
|
|
|
def save_preset(self, name: str, session_id: str, trial_id: str, reward_config: dict) -> dict:
|
|
preset_id, at = uuid.uuid4().hex, now_iso()
|
|
self.connection().execute(
|
|
"INSERT INTO presets(id,name,session_id,trial_id,reward_config_json,created_at) "
|
|
"VALUES (?,?,?,?,?,?)",
|
|
(preset_id, name, session_id, trial_id, _json(reward_config), at),
|
|
)
|
|
return {
|
|
"id": preset_id,
|
|
"name": name,
|
|
"sessionId": session_id,
|
|
"trialId": trial_id,
|
|
"rewardConfig": reward_config,
|
|
"createdAt": at,
|
|
}
|
|
|
|
def get_preset(self, preset_id: str) -> dict:
|
|
row = self.connection().execute("SELECT * FROM presets WHERE id=?", (preset_id,)).fetchone()
|
|
if row is None:
|
|
raise KeyError(preset_id)
|
|
return {
|
|
"id": row["id"],
|
|
"name": row["name"],
|
|
"sessionId": row["session_id"],
|
|
"trialId": row["trial_id"],
|
|
"rewardConfig": _decode(row["reward_config_json"]),
|
|
"createdAt": row["created_at"],
|
|
}
|
|
|
|
def list_presets(self) -> list[dict]:
|
|
rows = (
|
|
self.connection().execute("SELECT * FROM presets ORDER BY created_at DESC").fetchall()
|
|
)
|
|
return [
|
|
{
|
|
"id": row["id"],
|
|
"name": row["name"],
|
|
"sessionId": row["session_id"],
|
|
"trialId": row["trial_id"],
|
|
"rewardConfig": _decode(row["reward_config_json"]),
|
|
"createdAt": row["created_at"],
|
|
}
|
|
for row in rows
|
|
]
|