313 lines
10 KiB
Python
313 lines
10 KiB
Python
"""Backtest result storage using SQLite + parquet files."""
|
|
import os
|
|
import sqlite3
|
|
import json
|
|
import logging
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
import pandas as pd
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _json_default(obj):
|
|
"""json.dumps default: vnpy calculate_result 的 statistics 含 numpy.int64/float64,
|
|
标准 json 不能序列化（optimize 的统计直接来自 vnpy，未像 cta 那样被 compute_metrics
|
|
覆盖成 python float）。用 duck typing 把 numpy 标量/数组转原生，不引入 numpy 硬依赖。"""
|
|
import datetime as _dt
|
|
if isinstance(obj, (_dt.date, _dt.datetime)):
|
|
return obj.isoformat()
|
|
if hasattr(obj, "item") and callable(getattr(obj, "item")):
|
|
try:
|
|
return obj.item()
|
|
except Exception:
|
|
pass
|
|
if hasattr(obj, "tolist") and callable(getattr(obj, "tolist")):
|
|
try:
|
|
return obj.tolist()
|
|
except Exception:
|
|
pass
|
|
raise TypeError(f"Object of type {type(obj).__name__} is not JSON serializable")
|
|
|
|
|
|
@dataclass
|
|
class BacktestResult:
|
|
"""Backtest result data structure."""
|
|
task_id: str
|
|
type: str
|
|
status: str
|
|
strategy: str
|
|
symbol: str
|
|
params: dict
|
|
start: str
|
|
end: str
|
|
statistics: dict
|
|
equity_curve: Optional[pd.DataFrame] = None
|
|
trades: Optional[pd.DataFrame] = None
|
|
error_msg: Optional[str] = None
|
|
id: Optional[int] = None
|
|
created_at: Optional[str] = None # 任务创建时间(DB TIMESTAMP)
|
|
|
|
|
|
# SQLite schema for backtest stats
|
|
_SCHEMA = """CREATE TABLE IF NOT EXISTS backtest_stats (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
task_id TEXT,
|
|
type TEXT,
|
|
status TEXT,
|
|
strategy TEXT,
|
|
symbol TEXT,
|
|
params TEXT,
|
|
start TEXT,
|
|
end TEXT,
|
|
statistics TEXT,
|
|
equity_path TEXT,
|
|
trades_path TEXT,
|
|
error_msg TEXT,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
|
);
|
|
|
|
-- 提交时间表:任务提交即写一行(backtest_stats.created_at 是存结果时=完成时间,
|
|
-- 不是真创建时间)。完成后两表相减即真实耗时。
|
|
CREATE TABLE IF NOT EXISTS task_submissions (
|
|
task_id TEXT PRIMARY KEY,
|
|
submitted_at TIMESTAMP
|
|
);"""
|
|
|
|
|
|
def _connect(db_path: str) -> sqlite3.Connection:
|
|
"""Create database connection and initialize schema."""
|
|
conn = sqlite3.connect(db_path)
|
|
conn.executescript(_SCHEMA)
|
|
return conn
|
|
|
|
|
|
def save_result(result: BacktestResult, db_path: str, file_dir: Optional[str] = None) -> int:
|
|
"""
|
|
Save backtest result to database and optionally save DataFrame fields to parquet files.
|
|
|
|
Args:
|
|
result: BacktestResult object to save
|
|
db_path: Path to SQLite database file
|
|
file_dir: Optional directory to save parquet files (equity_curve and trades)
|
|
|
|
Returns:
|
|
int: The ID of the inserted record
|
|
"""
|
|
conn = _connect(db_path)
|
|
try:
|
|
equity_path = None
|
|
trades_path = None
|
|
|
|
# Save DataFrames to parquet if file_dir is provided
|
|
if file_dir:
|
|
fdir = Path(file_dir)
|
|
fdir.mkdir(parents=True, exist_ok=True)
|
|
|
|
if result.equity_curve is not None and not result.equity_curve.empty:
|
|
equity_path = str(fdir / f"{result.task_id}_equity.json")
|
|
result.equity_curve.to_json(equity_path, orient="records", date_format="iso", force_ascii=False)
|
|
|
|
if result.trades is not None and not result.trades.empty:
|
|
trades_path = str(fdir / f"{result.task_id}_trades.json")
|
|
result.trades.to_json(trades_path, orient="records", date_format="iso", force_ascii=False)
|
|
|
|
# Insert record into database
|
|
cur = conn.execute(
|
|
"""INSERT INTO backtest_stats
|
|
(task_id, type, status, strategy, symbol, params, start, end,
|
|
statistics, equity_path, trades_path, error_msg)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
|
|
(
|
|
result.task_id,
|
|
result.type,
|
|
result.status,
|
|
result.strategy,
|
|
result.symbol,
|
|
json.dumps(result.params, default=_json_default),
|
|
result.start,
|
|
result.end,
|
|
json.dumps(result.statistics, default=_json_default),
|
|
equity_path,
|
|
trades_path,
|
|
result.error_msg
|
|
)
|
|
)
|
|
conn.commit()
|
|
result.id = cur.lastrowid
|
|
return cur.lastrowid
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def record_submission(task_id: str, db_path: str) -> str:
|
|
"""Record task submission time (UTC, same format as CURRENT_TIMESTAMP).
|
|
|
|
提交即写:真创建时间。失败只 warning 不抛——时间戳缺失只影响展示,
|
|
不应影响任务提交本身。
|
|
"""
|
|
from datetime import datetime, timezone
|
|
ts = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S")
|
|
try:
|
|
conn = _connect(db_path)
|
|
try:
|
|
conn.execute(
|
|
"INSERT OR REPLACE INTO task_submissions (task_id, submitted_at) VALUES (?, ?)",
|
|
(task_id, ts),
|
|
)
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
except Exception as e:
|
|
logger.warning("record_submission %s failed: %s", task_id, e)
|
|
return ts
|
|
|
|
|
|
def load_submissions(db_path: str) -> dict[str, str]:
|
|
"""Return {task_id: submitted_at(UTC)} for all recorded submissions."""
|
|
conn = _connect(db_path)
|
|
try:
|
|
return {
|
|
row[0]: row[1]
|
|
for row in conn.execute("SELECT task_id, submitted_at FROM task_submissions")
|
|
}
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def delete_result_by_task_id(task_id: str, db_path: str) -> int:
|
|
"""Delete all persisted traces of a task: backtest_stats rows +
|
|
equity/trades JSON files + task_submissions row.
|
|
|
|
Returns deleted backtest_stats row count. Missing files tolerated
|
|
(historical records may reference stale paths after host migration).
|
|
"""
|
|
conn = _connect(db_path)
|
|
try:
|
|
rows = conn.execute(
|
|
"SELECT equity_path, trades_path FROM backtest_stats WHERE task_id=?",
|
|
(task_id,),
|
|
).fetchall()
|
|
for equity_path, trades_path in rows:
|
|
for path in (equity_path, trades_path):
|
|
if not path:
|
|
continue
|
|
try:
|
|
os.remove(path)
|
|
except OSError:
|
|
pass
|
|
n = conn.execute(
|
|
"DELETE FROM backtest_stats WHERE task_id=?", (task_id,)
|
|
).rowcount
|
|
conn.execute("DELETE FROM task_submissions WHERE task_id=?", (task_id,))
|
|
conn.commit()
|
|
return n
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def _safe_read_json(path: Optional[str]) -> Optional[pd.DataFrame]:
|
|
"""Read a JSON equity/trades file; return None if missing or unreadable.
|
|
|
|
Historical records may reference paths from a previous host (e.g. NAS
|
|
absolute paths after migration to VPS). Swallow those failures so the
|
|
record stays listable instead of crashing list_results.
|
|
"""
|
|
if not path:
|
|
return None
|
|
try:
|
|
return pd.read_json(path, orient="records")
|
|
except (FileNotFoundError, ValueError, OSError) as e:
|
|
logger.warning("result_store: skipping unreadable JSON file %s: %s", path, e)
|
|
return None
|
|
|
|
|
|
def load_result(rid: int, db_path: str) -> BacktestResult:
|
|
"""
|
|
Load backtest result by ID from database.
|
|
|
|
Args:
|
|
rid: Result ID to load
|
|
db_path: Path to SQLite database file
|
|
|
|
Returns:
|
|
BacktestResult: Loaded result object
|
|
|
|
Raises:
|
|
KeyError: If result ID not found
|
|
"""
|
|
conn = _connect(db_path)
|
|
try:
|
|
row = conn.execute("SELECT * FROM backtest_stats WHERE id=?", (rid,)).fetchone()
|
|
if not row:
|
|
raise KeyError(f"result {rid} not found")
|
|
|
|
# Get column names from table description
|
|
cols = [d[0] for d in conn.execute("SELECT * FROM backtest_stats LIMIT 0").description]
|
|
d = dict(zip(cols, row))
|
|
|
|
# Load JSON files if paths exist (equity_curve/trades persisted as JSON).
|
|
# Tolerate stale paths (e.g. NAS absolute paths left after VPS migration):
|
|
# missing/unreadable file -> None, so record still appears in list_results.
|
|
equity = _safe_read_json(d.get("equity_path"))
|
|
trades = _safe_read_json(d.get("trades_path"))
|
|
|
|
return BacktestResult(
|
|
task_id=d["task_id"],
|
|
type=d["type"],
|
|
status=d["status"],
|
|
strategy=d["strategy"],
|
|
symbol=d["symbol"],
|
|
params=json.loads(d["params"]),
|
|
start=d["start"],
|
|
end=d["end"],
|
|
statistics=json.loads(d["statistics"]) if d["statistics"] else {},
|
|
equity_curve=equity,
|
|
trades=trades,
|
|
error_msg=d.get("error_msg"),
|
|
created_at=d.get("created_at"),
|
|
)
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def load_result_by_task_id(task_id: str, db_path: str) -> BacktestResult | None:
|
|
"""Load the most recent result for a task_id (historical lookup after restart)."""
|
|
conn = _connect(db_path)
|
|
try:
|
|
row = conn.execute(
|
|
"SELECT id FROM backtest_stats WHERE task_id=? ORDER BY id DESC LIMIT 1",
|
|
(task_id,),
|
|
).fetchone()
|
|
if not row:
|
|
return None
|
|
return load_result(row[0], db_path)
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def list_results(type_filter: Optional[str] = None, db_path: str = "") -> list[BacktestResult]:
|
|
"""
|
|
List all backtest results, optionally filtered by type.
|
|
|
|
Args:
|
|
type_filter: Optional filter for result type (e.g., 'cta', 'factor')
|
|
db_path: Path to SQLite database file
|
|
|
|
Returns:
|
|
list[BacktestResult]: List of loaded result objects
|
|
"""
|
|
conn = _connect(db_path)
|
|
try:
|
|
query = "SELECT id FROM backtest_stats"
|
|
args = ()
|
|
|
|
if type_filter:
|
|
query = query + " WHERE type=?"
|
|
args = (type_filter,)
|
|
|
|
return [load_result(row[0], db_path) for row in conn.execute(query, args).fetchall()]
|
|
finally:
|
|
conn.close()
|