"""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 from .schema import FLAT_TASK, RewardConfigError, validate_configuration SCHEMA_VERSION = 2 class StorageConflict(RuntimeError): """Optimistic-concurrency revision mismatch.""" 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 TABLE IF NOT EXISTS session_controls( session_id TEXT PRIMARY KEY REFERENCES sessions(id) ON DELETE CASCADE, run_policy TEXT NOT NULL DEFAULT 'continuous' CHECK(run_policy IN ('continuous','step')), dispatch_tokens INTEGER NOT NULL DEFAULT 0 CHECK(dispatch_tokens >= 0), active_base_trial_id TEXT, revision INTEGER NOT NULL DEFAULT 0 ); CREATE TABLE IF NOT EXISTS session_constraints( session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE, path TEXT NOT NULL, kind TEXT NOT NULL CHECK(kind IN ('range','fixed')), min_value REAL, max_value REAL, fixed_value REAL, updated_at TEXT NOT NULL, PRIMARY KEY(session_id, path) ); 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 session_controls(session_id) SELECT id FROM sessions" ) 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: previous = connection.execute( "SELECT id,state FROM sessions WHERE state IN " "('queued','running','evaluating','paused','awaiting_approval')" ).fetchall() for row in previous: connection.execute( "INSERT INTO audit_events(session_id,event_type,payload_json,created_at) " "VALUES (?,?,?,?)", ( row["id"], "service_restart_interrupted", _json( { "previousState": row["state"], "reason": "service_restart", "requiresExplicitResume": True, } ), at, ), ) connection.execute( "UPDATE trials SET state='interrupted', ended_at=?, " "message='服务重启中断，等待显式恢复' " "WHERE state IN ('queued','training','evaluating')", (at,), ) connection.execute( "UPDATE sessions SET state='interrupted', updated_at=?, " "message='服务重启中断，可从完整 checkpoint 恢复' " "WHERE state IN ('queued','running','evaluating','paused','awaiting_approval')", (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 session_controls(session_id) VALUES (?)", (session_id,)) 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", "mode": "mode", "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 get_control(self, session_id: str) -> dict: row = ( self.connection() .execute("SELECT * FROM session_controls WHERE session_id=?", (session_id,)) .fetchone() ) if row is None: raise KeyError(session_id) constraint_rows = ( self.connection() .execute( "SELECT * FROM session_constraints WHERE session_id=? ORDER BY path", (session_id,) ) .fetchall() ) constraints = {} for constraint in constraint_rows: if constraint["kind"] == "fixed": value = {"kind": "fixed", "value": constraint["fixed_value"]} else: value = { "kind": "range", "min": constraint["min_value"], "max": constraint["max_value"], } constraints[constraint["path"]] = value return { "runPolicy": row["run_policy"], "dispatchTokens": row["dispatch_tokens"], "constraintsRevision": row["revision"], "constraints": constraints, "activeBaseTrialId": row["active_base_trial_id"], } def replace_constraints( self, session_id: str, expected_revision: int, constraints: dict ) -> dict: at = now_iso() with self.transaction() as connection: row = connection.execute( "SELECT revision FROM session_controls WHERE session_id=?", (session_id,) ).fetchone() if row is None: raise KeyError(session_id) if row["revision"] != expected_revision: raise StorageConflict( f"参数护栏 revision 已变化（当前 {row['revision']}，请求 {expected_revision}）" ) connection.execute("DELETE FROM session_constraints WHERE session_id=?", (session_id,)) for path, constraint in constraints.items(): connection.execute( "INSERT INTO session_constraints(" "session_id,path,kind,min_value,max_value,fixed_value,updated_at) " "VALUES (?,?,?,?,?,?,?)", ( session_id, path, constraint["kind"], constraint.get("min"), constraint.get("max"), constraint.get("value"), at, ), ) connection.execute( "UPDATE session_controls SET revision=revision+1 WHERE session_id=?", (session_id,), ) return self.get_control(session_id) def grant_dispatch_token(self, session_id: str) -> dict: """Atomically grant the sole outstanding one-Trial token.""" with self.transaction() as connection: cursor = connection.execute( "UPDATE session_controls SET run_policy='step',dispatch_tokens=1 " "WHERE session_id=? AND dispatch_tokens=0", (session_id,), ) if cursor.rowcount != 1: exists = connection.execute( "SELECT 1 FROM session_controls WHERE session_id=?", (session_id,) ).fetchone() if exists is None: raise KeyError(session_id) raise StorageConflict("已有未消费的单步 Trial 令牌") return self.get_control(session_id) def use_dispatch_token(self, session_id: str) -> bool: """Atomically consume one step token; continuous mode never needs a token.""" with self.transaction() as connection: row = connection.execute( "SELECT run_policy,dispatch_tokens FROM session_controls WHERE session_id=?", (session_id,), ).fetchone() if row is None: raise KeyError(session_id) if row["run_policy"] == "continuous": return True if row["dispatch_tokens"] <= 0: return False connection.execute( "UPDATE session_controls SET dispatch_tokens=dispatch_tokens-1 WHERE session_id=?", (session_id,), ) return True def set_run_policy(self, session_id: str, policy: str) -> dict: if policy not in {"continuous", "step"}: raise ValueError(policy) cursor = self.connection().execute( "UPDATE session_controls SET run_policy=?," "dispatch_tokens=CASE WHEN ?='continuous' THEN 0 ELSE dispatch_tokens END " "WHERE session_id=?", (policy, policy, session_id), ) if cursor.rowcount != 1: raise KeyError(session_id) return self.get_control(session_id) def reset_step_gate(self, session_id: str) -> dict: cursor = self.connection().execute( "UPDATE session_controls SET run_policy='step',dispatch_tokens=0 WHERE session_id=?", (session_id,), ) if cursor.rowcount != 1: raise KeyError(session_id) return self.get_control(session_id) def set_active_base(self, session_id: str, trial_id: str | None) -> dict: cursor = self.connection().execute( "UPDATE session_controls SET active_base_trial_id=? WHERE session_id=?", (trial_id, session_id), ) if cursor.rowcount != 1: raise KeyError(session_id) return self.get_control(session_id) 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, after_step: int | None = None, ) -> list[dict]: parameters: list[Any] = [trial_id] clause = "trial_id=?" if tags: clause += f" AND tag IN ({','.join('?' for _ in tags)})" parameters.extend(tags) if after_step is not None: clause += " AND step>?" parameters.append(after_step) 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 _preset_task(self, session_id: str) -> str: # Identity comes only from the persisted source session, never client/preset labels. try: config = self.get_session(session_id)["config"] except (KeyError, ValueError, TypeError) as error: raise RewardConfigError("奖励 preset 来源 session 丢失或损坏") from error if not isinstance(config, dict): raise RewardConfigError("奖励 preset 来源 config 必须是对象") # Original Flat-only sessions did not require taskId. Validate their preset below. return config.get("taskId", FLAT_TASK) def _preset(self, row) -> dict: task_id = self._preset_task(row["session_id"]) try: reward_config = validate_configuration(_decode(row["reward_config_json"]), task_id) except (ValueError, TypeError) as error: raise RewardConfigError("奖励 preset 配置无效：" + str(error)) from error return { "id": row["id"], "name": row["name"], "sessionId": row["session_id"], "trialId": row["trial_id"], "taskId": task_id, "rewardConfig": reward_config, "createdAt": row["created_at"], } def save_preset(self, name: str, session_id: str, trial_id: str, reward_config: dict) -> dict: reward_config = validate_configuration(reward_config, self._preset_task(session_id)) 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 self.get_preset(preset_id) 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 self._preset(row) def list_presets(self) -> list[dict]: rows = ( self.connection().execute("SELECT * FROM presets ORDER BY created_at DESC").fetchall() ) return [self._preset(row) for row in rows]