diff --git a/sanguo_orchestrator/runner.py b/sanguo_orchestrator/runner.py index b32988b..e0ad7ad 100644 --- a/sanguo_orchestrator/runner.py +++ b/sanguo_orchestrator/runner.py @@ -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, - start: str, end: str, cfg) -> str: - """Submit a CTA backtesting task""" + 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 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 + 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 + ) + await self._notify_stage(task_id, "参数优化中") + + 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 = 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 - ) - task.complete(result_id=id(result)) + result = await asyncio.wrap_future(fut) + await self._on_done(task_id, result) except Exception as e: - task.fail(f"{type(e).__name__}: {e}") + task = self.pool.get_task(task_id) + if task: + task.fail(f"{type(e).__name__}: {e}") + await self._notify_stage(task_id, "失败") - return task + 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""" @@ -66,4 +143,23 @@ class Orchestrator: # Lazy import to avoid vnpy dependency issues from sanguo_backtest.result_store import load_result return load_result(task.result_id, self.db_path) - return None \ No newline at end of file + 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) diff --git a/tests/orchestrator/test_runner.py b/tests/orchestrator/test_runner.py index 9278f14..ca1e1f7 100644 --- a/tests/orchestrator/test_runner.py +++ b/tests/orchestrator/test_runner.py @@ -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_") @@ -122,4 +125,44 @@ class TestOrchestrator: orchestrator = Orchestrator(db_path="test.db", max_workers=2) result = orchestrator.get_result("nonexistent") - assert result is None \ No newline at end of file + 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() \ No newline at end of file