feat(orchestrator): runner async submit + on_stage 回调
This commit is contained in:
+118
-22
@@ -2,6 +2,8 @@
|
||||
Orchestrator for task coordination and execution
|
||||
Manages backtesting tasks with lazy imports
|
||||
"""
|
||||
import asyncio
|
||||
from concurrent.futures import Future
|
||||
from .pool import TaskPool
|
||||
from .task import TaskState
|
||||
|
||||
@@ -14,14 +16,25 @@ class Orchestrator:
|
||||
self.db_path = db_path
|
||||
self.file_dir = file_dir
|
||||
self.pool = TaskPool(max_workers=max_workers)
|
||||
self._pending = {}
|
||||
self._pending: dict[str, dict] = {}
|
||||
self._on_stage = None # async callback(task_id, stage)
|
||||
|
||||
def submit_cta(self, strategy_class, symbol: str, params: dict,
|
||||
def set_on_stage(self, cb):
|
||||
"""Set callback for stage updates (async callable)"""
|
||||
self._on_stage = cb
|
||||
|
||||
async def _notify_stage(self, task_id: str, stage: str):
|
||||
"""Update task stage and fire callback if set"""
|
||||
self.pool.update_stage(task_id, stage)
|
||||
if self._on_stage:
|
||||
await self._on_stage(task_id, stage)
|
||||
|
||||
async def submit_cta(self, strategy_class, symbol: str, params: dict,
|
||||
start: str, end: str, cfg) -> str:
|
||||
"""Submit a CTA backtesting task"""
|
||||
"""Submit a CTA backtesting task asynchronously"""
|
||||
task_id = f"cta_{symbol}_{id(params)}"
|
||||
self.pool.submit(task_id, "cta")
|
||||
self._pending = dict(
|
||||
self._pending[task_id] = dict(
|
||||
strategy_class=strategy_class,
|
||||
symbol=symbol,
|
||||
params=params,
|
||||
@@ -29,31 +42,95 @@ class Orchestrator:
|
||||
end=end,
|
||||
cfg=cfg
|
||||
)
|
||||
return task_id
|
||||
await self._notify_stage(task_id, "排队中")
|
||||
|
||||
def _run_sync(self, task_id: str):
|
||||
"""Execute a task synchronously (lazy import)"""
|
||||
# Lazy import to avoid vnpy dependency issues
|
||||
from sanguo_backtest.cta_engine import run_cta_backtest
|
||||
spec = self._pending[task_id]
|
||||
fut: Future = self.pool.submit_work(
|
||||
task_id, _cta_worker, spec["strategy_class"], spec["symbol"],
|
||||
spec["params"], spec["start"], spec["end"], spec["cfg"], self.db_path
|
||||
)
|
||||
|
||||
task = self.pool.get_task(task_id)
|
||||
task.start()
|
||||
await self._notify_stage(task_id, "回测中")
|
||||
asyncio.ensure_future(self._wait_future(task_id, fut))
|
||||
return task_id
|
||||
|
||||
try:
|
||||
result = run_cta_backtest(
|
||||
self._pending["strategy_class"],
|
||||
self._pending["symbol"],
|
||||
self._pending["params"],
|
||||
self._pending["start"],
|
||||
self._pending["end"],
|
||||
self._pending["cfg"],
|
||||
self.db_path
|
||||
async def submit_optimize(self, strategy_class, symbol: str, grid: dict,
|
||||
start: str, end: str, cfg) -> str:
|
||||
"""Submit a CTA optimization task asynchronously"""
|
||||
task_id = f"opt_{symbol}_{id(grid)}"
|
||||
self.pool.submit(task_id, "optimize")
|
||||
self._pending[task_id] = dict(
|
||||
strategy_class=strategy_class,
|
||||
symbol=symbol,
|
||||
grid=grid,
|
||||
start=start,
|
||||
end=end,
|
||||
cfg=cfg
|
||||
)
|
||||
task.complete(result_id=id(result))
|
||||
except Exception as e:
|
||||
task.fail(f"{type(e).__name__}: {e}")
|
||||
await self._notify_stage(task_id, "参数优化中")
|
||||
|
||||
return task
|
||||
spec = self._pending[task_id]
|
||||
fut: Future = self.pool.submit_work(
|
||||
task_id, _opt_worker, spec["strategy_class"], spec["symbol"],
|
||||
spec["grid"], spec["start"], spec["end"], spec["cfg"], self.db_path
|
||||
)
|
||||
|
||||
task = self.pool.get_task(task_id)
|
||||
task.start()
|
||||
await self._notify_stage(task_id, "优化中")
|
||||
asyncio.ensure_future(self._wait_future(task_id, fut))
|
||||
return task_id
|
||||
|
||||
async def submit_factor(self, symbols: list, factor_names: list,
|
||||
start: str, end: str, cfg, output_dir: str) -> str:
|
||||
"""Submit a factor analysis task asynchronously"""
|
||||
task_id = f"factor_{id(factor_names)}"
|
||||
self.pool.submit(task_id, "factor")
|
||||
self._pending[task_id] = dict(
|
||||
symbols=symbols,
|
||||
factor_names=factor_names,
|
||||
start=start,
|
||||
end=end,
|
||||
cfg=cfg,
|
||||
output_dir=output_dir
|
||||
)
|
||||
await self._notify_stage(task_id, "因子分析中")
|
||||
|
||||
spec = self._pending[task_id]
|
||||
fut: Future = self.pool.submit_work(
|
||||
task_id, _factor_worker, spec["symbols"], spec["factor_names"],
|
||||
spec["start"], spec["end"], spec["cfg"], spec["output_dir"]
|
||||
)
|
||||
|
||||
task = self.pool.get_task(task_id)
|
||||
task.start()
|
||||
await self._notify_stage(task_id, "分析中")
|
||||
asyncio.ensure_future(self._wait_future(task_id, fut))
|
||||
return task_id
|
||||
|
||||
async def _wait_future(self, task_id: str, fut: Future):
|
||||
"""Wait for Future to complete and handle result/exception"""
|
||||
try:
|
||||
result = await asyncio.wrap_future(fut)
|
||||
await self._on_done(task_id, result)
|
||||
except Exception as e:
|
||||
task = self.pool.get_task(task_id)
|
||||
if task:
|
||||
task.fail(f"{type(e).__name__}: {e}")
|
||||
await self._notify_stage(task_id, "失败")
|
||||
|
||||
async def _on_done(self, task_id: str, result):
|
||||
"""Handle task completion (with None-guard for unknown tasks)"""
|
||||
task = self.pool.get_task(task_id)
|
||||
if task is None:
|
||||
# Unknown task - fire callback but don't crash
|
||||
await self._notify_stage(task_id, "完成")
|
||||
return
|
||||
|
||||
task.complete(result_id=id(result))
|
||||
await self._notify_stage(task_id, "完成")
|
||||
|
||||
def get_status(self, task_id: str) -> TaskState | None:
|
||||
"""Get task status by ID"""
|
||||
@@ -67,3 +144,22 @@ class Orchestrator:
|
||||
from sanguo_backtest.result_store import load_result
|
||||
return load_result(task.result_id, self.db_path)
|
||||
return None
|
||||
|
||||
|
||||
# Module-level worker functions (must be top-level for ProcessPoolExecutor pickle)
|
||||
def _cta_worker(strategy_class, symbol: str, params: dict, start: str, end: str, cfg, db_path: str):
|
||||
"""Worker for CTA backtest (lazy import, spawn-friendly)"""
|
||||
from sanguo_backtest.cta_engine import run_cta_backtest
|
||||
return run_cta_backtest(strategy_class, symbol, params, start, end, cfg, db_path)
|
||||
|
||||
|
||||
def _opt_worker(strategy_class, symbol: str, grid: dict, start: str, end: str, cfg, db_path: str):
|
||||
"""Worker for CTA optimization (lazy import, spawn-friendly)"""
|
||||
from sanguo_backtest.cta_optimizer import run_cta_optimization
|
||||
return run_cta_optimization(strategy_class, symbol, grid, start, end, cfg, db_path)
|
||||
|
||||
|
||||
def _factor_worker(symbols: list, factor_names: list, start: str, end: str, cfg, output_dir: str):
|
||||
"""Worker for factor analysis (lazy import, spawn-friendly)"""
|
||||
from sanguo_factor.analyzer import run_factor_analysis
|
||||
return run_factor_analysis(symbols, factor_names, start, end, cfg, output_dir)
|
||||
|
||||
@@ -3,7 +3,7 @@ Tests for sanguo_orchestrator.runner module
|
||||
Tests Orchestrator task coordination
|
||||
"""
|
||||
import pytest
|
||||
from unittest.mock import Mock, patch
|
||||
from unittest.mock import Mock, AsyncMock, patch
|
||||
from sanguo_orchestrator.runner import Orchestrator
|
||||
from sanguo_orchestrator.task import TaskState
|
||||
|
||||
@@ -18,12 +18,14 @@ class TestOrchestrator:
|
||||
assert orchestrator.pool.max_workers == 2
|
||||
assert orchestrator._pending == {}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch('sanguo_orchestrator.runner.TaskPool')
|
||||
def test_submit_cta_creates_task(self, mock_pool_class):
|
||||
"""Test submit_cta() creates task and stores spec"""
|
||||
async def test_submit_cta_creates_task(self, mock_pool_class):
|
||||
"""Test async submit_cta() creates task and stores spec per task_id"""
|
||||
mock_pool = Mock()
|
||||
mock_pool_class.return_value = mock_pool
|
||||
mock_pool.submit.return_value = Mock(task_id="test_1")
|
||||
mock_pool.submit_work.return_value = Mock()
|
||||
|
||||
orchestrator = Orchestrator(db_path="test.db", max_workers=2)
|
||||
strategy_class = Mock
|
||||
@@ -33,7 +35,7 @@ class TestOrchestrator:
|
||||
end = "2024-12-31"
|
||||
cfg = Mock()
|
||||
|
||||
task_id = orchestrator.submit_cta(strategy_class, symbol, params, start, end, cfg)
|
||||
task_id = await orchestrator.submit_cta(strategy_class, symbol, params, start, end, cfg)
|
||||
|
||||
# Verify task was submitted to pool
|
||||
mock_pool.submit.assert_called_once()
|
||||
@@ -41,13 +43,14 @@ class TestOrchestrator:
|
||||
assert call_args[0][0] == task_id # task_id
|
||||
assert call_args[0][1] == "cta" # task_type
|
||||
|
||||
# Verify pending spec was stored
|
||||
assert orchestrator._pending["strategy_class"] == strategy_class
|
||||
assert orchestrator._pending["symbol"] == symbol
|
||||
assert orchestrator._pending["params"] == params
|
||||
assert orchestrator._pending["start"] == start
|
||||
assert orchestrator._pending["end"] == end
|
||||
assert orchestrator._pending["cfg"] == cfg
|
||||
# Verify pending spec was stored per task_id (not global)
|
||||
assert task_id in orchestrator._pending
|
||||
assert orchestrator._pending[task_id]["strategy_class"] == strategy_class
|
||||
assert orchestrator._pending[task_id]["symbol"] == symbol
|
||||
assert orchestrator._pending[task_id]["params"] == params
|
||||
assert orchestrator._pending[task_id]["start"] == start
|
||||
assert orchestrator._pending[task_id]["end"] == end
|
||||
assert orchestrator._pending[task_id]["cfg"] == cfg
|
||||
|
||||
assert task_id.startswith("cta_AAPL_")
|
||||
|
||||
@@ -123,3 +126,43 @@ class TestOrchestrator:
|
||||
result = orchestrator.get_result("nonexistent")
|
||||
|
||||
assert result is None
|
||||
|
||||
|
||||
class TestOrchestratorAsync:
|
||||
"""Test async orchestrator submit and on_stage callback"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_submit_cta_returns_task_id_and_submits(self):
|
||||
"""Test async submit_cta() returns task_id and submits to pool"""
|
||||
orchestrator = Orchestrator(db_path="/tmp/t.db")
|
||||
orchestrator.pool.executor = Mock()
|
||||
mock_future = Mock()
|
||||
orchestrator.pool.executor.submit.return_value = mock_future
|
||||
|
||||
with patch("sanguo_backtest.cta_engine.run_cta_backtest"):
|
||||
task_id = await orchestrator.submit_cta(
|
||||
Mock(), "600000", {}, "2024-01-01", "2024-06-30", cfg=Mock()
|
||||
)
|
||||
|
||||
assert task_id.startswith("cta_600000")
|
||||
orchestrator.pool.executor.submit.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_on_stage_callback_called(self):
|
||||
"""Test set_on_stage() callback fires on task completion"""
|
||||
from sanguo_orchestrator.runner import Orchestrator
|
||||
from sanguo_orchestrator.task import TaskState
|
||||
|
||||
orchestrator = Orchestrator(db_path="/tmp/t.db")
|
||||
orchestrator.pool.executor = Mock()
|
||||
mock_future = Mock()
|
||||
orchestrator.pool.executor.submit.return_value = mock_future
|
||||
|
||||
cb = AsyncMock()
|
||||
orchestrator.set_on_stage(cb)
|
||||
|
||||
# Simulate task completion callback (with None-guard for unknown task)
|
||||
await orchestrator._on_done("t1", {"statistics": {}})
|
||||
|
||||
# Callback should fire even though task t1 was never submitted
|
||||
cb.assert_called_once()
|
||||
Reference in New Issue
Block a user