717 lines
28 KiB
Python
717 lines
28 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
|
||
|
||
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]
|