"""Spawn-isolated strategy backtest and optimizer task runner.""" from __future__ import annotations import json import multiprocessing as mp import os import queue import threading import time import traceback from collections.abc import Callable from contextlib import suppress from dataclasses import asdict from datetime import date from pathlib import Path from typing import Any import psutil class BacktestWorkerError(RuntimeError): """Raised when a spawned worker fails before returning a task result.""" _CANCEL_GRACE_SECONDS = 5.0 class _PeakRssSampler: """Track whole-task and resettable phase RSS peaks with one sampling thread.""" def __init__(self, interval_seconds: float = 0.05) -> None: if interval_seconds <= 0: raise ValueError("RSS sample interval must be positive") self._process = psutil.Process(os.getpid()) self._interval_seconds = float(interval_seconds) self._stop = threading.Event() self._thread = threading.Thread(target=self._sample, daemon=True) self._lock = threading.Lock() self._started = False current = int(self._process.memory_info().rss) self.peak_rss_bytes = current self._phase_peak_rss_bytes = current def start(self) -> None: if self._started: raise RuntimeError("RSS sampler has already started") self._started = True self._thread.start() def stop(self) -> int: if self._started: self._stop.set() self._thread.join(timeout=1.0) self._record_current() return self.peak_rss_bytes def reset_phase(self) -> None: current = int(self._process.memory_info().rss) with self._lock: self._phase_peak_rss_bytes = current def phase_peak_rss_bytes(self) -> int: self._record_current() with self._lock: return self._phase_peak_rss_bytes def _record_current(self) -> None: current = int(self._process.memory_info().rss) with self._lock: self.peak_rss_bytes = max(self.peak_rss_bytes, current) self._phase_peak_rss_bytes = max(self._phase_peak_rss_bytes, current) def _sample(self) -> None: while not self._stop.wait(self._interval_seconds): self._record_current() def _rss_bytes() -> int: return int(psutil.Process(os.getpid()).memory_info().rss) def _strategy_dirs(data_dir: Path) -> list[Path]: app_dir = Path(__file__).resolve().parents[1] return [ app_dir / "strategy" / "builtin", data_dir / "strategies" / "custom", data_dir / "strategies" / "ai", data_dir / "strategies" / "composite", ] def _decode_backtest_config(payload: dict[str, Any]): from app.backtest.strategy import StrategyBacktestConfig values = dict(payload) values["start"] = date.fromisoformat(values["start"]) values["end"] = date.fromisoformat(values["end"]) return StrategyBacktestConfig(**values) def _decode_optimize_config(payload: dict[str, Any]): from app.backtest.optimizer import OptimizeConfig values = dict(payload) values["start"] = date.fromisoformat(values["start"]) values["end"] = date.fromisoformat(values["end"]) return OptimizeConfig(**values) def _decode_walkforward_config(payload: dict[str, Any]): from app.backtest.walkforward import WalkForwardConfig values = dict(payload) values["start"] = date.fromisoformat(values["start"]) values["end"] = date.fromisoformat(values["end"]) return WalkForwardConfig(**values) def encode_backtest_config(config) -> dict[str, Any]: payload = asdict(config) payload["start"] = config.start.isoformat() payload["end"] = config.end.isoformat() return payload def encode_optimize_config(config) -> dict[str, Any]: payload = asdict(config) payload["start"] = config.start.isoformat() payload["end"] = config.end.isoformat() return payload def make_worker_task(kind: str, data_dir: Path, config) -> dict[str, Any]: if kind == "backtest": encoded = encode_backtest_config(config) elif kind == "optimize": encoded = encode_optimize_config(config) elif kind == "walkforward": encoded = asdict(config) encoded["start"] = config.start.isoformat() encoded["end"] = config.end.isoformat() elif kind == "mining": if not isinstance(config, dict): raise TypeError("mining worker config must be a dict") encoded = dict(config) else: raise ValueError(f"unsupported worker task kind: {kind}") return { "kind": kind, "data_dir": str(data_dir.resolve()), "config": encoded, } def _attach_worker_metrics( kind: str, result: dict[str, Any], metrics: dict[str, Any], ) -> None: if kind == "backtest": result.setdefault("stats", {})["worker"] = metrics else: result["worker"] = metrics def _worker_entry(task: dict[str, Any], event_queue, cancel_event) -> None: sampler = _PeakRssSampler() sampler.start() started = time.perf_counter() store = None try: from app.backtest.engine import BacktestEngine from app.backtest.optimizer import StrategyOptimizer from app.backtest.strategy import StrategyBacktestService from app.strategy import config as strategy_config from app.strategy.engine import StrategyEngine from app.tickflow.repository import DataStore, KlineRepository data_dir = Path(task["data_dir"]) store = DataStore(data_dir) repo = KlineRepository(store) strategy_engine = StrategyEngine( strategy_dirs=_strategy_dirs(data_dir), override_loader=lambda sid: strategy_config.load_override(data_dir, sid), ) service = StrategyBacktestService(BacktestEngine(repo), strategy_engine) def _progress(message: dict) -> None: event_queue.put({"type": "progress", "payload": message}) kind = task["kind"] if kind == "backtest": config = _decode_backtest_config(task["config"]) result = asdict(service.run(config, _progress, cancel_event)) elif kind == "optimize": config = _decode_optimize_config(task["config"]) optimizer = StrategyOptimizer(service, strategy_engine) result = optimizer.optimize( config, _progress, cancel_event, rss_sampler=sampler, ) elif kind == "walkforward": from app.backtest.walkforward import WalkForwardService config = _decode_walkforward_config(task["config"]) optimizer = StrategyOptimizer(service, strategy_engine) walkforward = WalkForwardService(optimizer, service, strategy_engine) result = walkforward.run(config, _progress, cancel_event) elif kind == "mining": from app.backtest.mining_runtime import run_mining_runtime result = run_mining_runtime( task["config"], data_dir=data_dir, service=service, strategy_engine=strategy_engine, progress_cb=_progress, cancel_check=cancel_event, rss_sampler=sampler, ) else: raise ValueError(f"unsupported worker task kind: {kind}") serialization_started = time.perf_counter() serialized_bytes = len( json.dumps(result, ensure_ascii=False, default=str).encode("utf-8") ) serialization_ms = round( (time.perf_counter() - serialization_started) * 1000, 1, ) peak_rss = sampler.stop() metrics = { "pid": os.getpid(), "peak_rss_bytes": peak_rss, "final_rss_bytes": _rss_bytes(), "serialization_ms": serialization_ms, "serialized_result_bytes": serialized_bytes, "task_elapsed_ms": round((time.perf_counter() - started) * 1000, 1), } _attach_worker_metrics(kind, result, metrics) event_queue.put({"type": "result", "payload": result}) except BaseException as exc: with suppress(Exception): sampler.stop() event_queue.put({ "type": "error", "message": str(exc), "traceback": traceback.format_exc(), }) finally: if store is not None: with suppress(Exception): store.db.close() # 保证结果消息在进程退出前完整刷入管道: put 只是入队, # 实际写管道的是后台 feeder 线程; 不 join 的话主线程先退出, # feeder 随进程销毁, 消息尾部丢失 → 父进程误判 "exited without result"。 with suppress(Exception): event_queue.close() event_queue.join_thread() def run_worker_task( task: dict[str, Any], progress_cb: Callable[[dict], None] | None = None, cancel_event: threading.Event | None = None, ) -> dict[str, Any]: """Run one complete task in a spawned process and wait for deterministic exit.""" context = mp.get_context("spawn") events = context.Queue() process_cancel = context.Event() process = context.Process( target=_worker_entry, args=(task, events, process_cancel), daemon=False, ) parent_rss_before = _rss_bytes() try: process.start() except BaseException: events.close() events.join_thread() raise result: dict[str, Any] | None = None failure: dict[str, Any] | None = None ipc_started = time.perf_counter() cancel_started: float | None = None try: while result is None and failure is None: if cancel_event is not None and cancel_event.is_set(): process_cancel.set() if cancel_started is None: cancel_started = time.monotonic() elif time.monotonic() - cancel_started >= _CANCEL_GRACE_SECONDS: process.terminate() process.join(timeout=5.0) raise BacktestWorkerError( "backtest worker did not stop within 5 seconds after cancellation" ) try: message = events.get(timeout=0.1) except queue.Empty: if not process.is_alive(): break continue message_type = message.get("type") if message_type == "progress": if progress_cb is not None: progress_cb(message["payload"]) elif message_type == "result": result = message["payload"] elif message_type == "error": failure = message # 子进程退出后, 队列读线程可能尚未把管道尾部的 result/error 搬进本地缓冲 # (0.1s 轮询在系统高负载下会先看到 Empty+进程已死)。join 后做一次兜底排空, # 只要消息完整刷入过管道就一定能取到。 if result is None and failure is None: for _ in range(2): try: message = events.get(timeout=1.0) except queue.Empty: break message_type = message.get("type") if message_type == "result": result = message["payload"] elif message_type == "error": failure = message process.join(timeout=10.0) if process.is_alive(): process.terminate() process.join(timeout=5.0) raise BacktestWorkerError("backtest worker returned but did not exit within 10 seconds") if failure is not None: raise BacktestWorkerError( f"{failure.get('message', 'worker failed')}\n{failure.get('traceback', '')}".rstrip() ) if result is None: raise BacktestWorkerError( f"backtest worker exited without result (exitcode={process.exitcode})" ) parent_metrics = { "ipc_elapsed_ms": round((time.perf_counter() - ipc_started) * 1000, 1), "parent_rss_before_bytes": parent_rss_before, "parent_rss_after_worker_exit_bytes": _rss_bytes(), "worker_exitcode": process.exitcode, } kind = task["kind"] if kind == "backtest": result.setdefault("stats", {}).setdefault("worker", {}).update(parent_metrics) else: result.setdefault("worker", {}).update(parent_metrics) return result finally: if process.is_alive(): process.terminate() process.join(timeout=5.0) events.close() events.join_thread()