From 66aa27e807ce3c3493bba17070ba20128abfed65 Mon Sep 17 00:00:00 2001
From: claude_dev <claude-dev@sanguo.local>
Date: Mon, 6 Jul 2026 18:19:24 +0800
Subject: [PATCH] =?UTF-8?q?feat(orchestrator):=20pool=20=E5=BC=82=E6=AD=A5?=
 =?UTF-8?q?=E5=8C=96=EF=BC=88ProcessPoolExecutor=20spawn=20+=20stage=20?=
 =?UTF-8?q?=E8=BF=BD=E8=B8=AA=EF=BC=89?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 sanguo_orchestrator/pool.py     | 20 +++++++++++++++++++
 sanguo_orchestrator/task.py     |  1 +
 tests/orchestrator/test_pool.py | 34 ++++++++++++++++++++++++++++++++-
 3 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/sanguo_orchestrator/pool.py b/sanguo_orchestrator/pool.py
index 9141573..09adef7 100644
--- a/sanguo_orchestrator/pool.py
+++ b/sanguo_orchestrator/pool.py
@@ -2,6 +2,8 @@
 Task pool for managing multiple tasks in memory
 Provides task storage and status tracking (not actual multiprocessing)
 """
+from concurrent.futures import ProcessPoolExecutor, Future
+from multiprocessing import get_context
 from .task import Task, TaskState
 
 
@@ -12,6 +14,9 @@ class TaskPool:
         """Initialize task pool with maximum workers"""
         self.max_workers = max_workers
         self._tasks: dict[str, Task] = {}
+        self.executor = ProcessPoolExecutor(
+            max_workers=max_workers, mp_context=get_context("spawn")
+        )
 
     def submit(self, task_id: str, task_type: str) -> Task:
         """Submit a new task to the pool"""
@@ -19,11 +24,26 @@ class TaskPool:
         self._tasks[task_id] = task
         return task
 
+    def submit_work(self, task_id: str, func, *args) -> Future:
+        """Submit work to the process pool executor"""
+        return self.executor.submit(func, *args)
+
+    def update_stage(self, task_id: str, stage: str):
+        """Update the stage of a task"""
+        t = self._tasks.get(task_id)
+        if t:
+            t.stage = stage
+
     def get_status(self, task_id: str) -> TaskState | None:
         """Get task status by ID"""
         task = self._tasks.get(task_id)
         return task.status if task else None
 
+    def get_stage(self, task_id: str) -> str | None:
+        """Get task stage by ID"""
+        t = self._tasks.get(task_id)
+        return t.stage if t else None
+
     def get_task(self, task_id: str) -> Task | None:
         """Get task object by ID"""
         return self._tasks.get(task_id)
\ No newline at end of file
diff --git a/sanguo_orchestrator/task.py b/sanguo_orchestrator/task.py
index 6d0fcbf..83f67da 100644
--- a/sanguo_orchestrator/task.py
+++ b/sanguo_orchestrator/task.py
@@ -22,6 +22,7 @@ class Task:
     status: TaskState = TaskState.PENDING
     result_id: int | None = None
     error_msg: str | None = None
+    stage: str = ""  # Current stage (数据加载/算因子/回测中...)
 
     def start(self):
         """Transition from PENDING to RUNNING"""
diff --git a/tests/orchestrator/test_pool.py b/tests/orchestrator/test_pool.py
index 17beda2..a9af08d 100644
--- a/tests/orchestrator/test_pool.py
+++ b/tests/orchestrator/test_pool.py
@@ -59,4 +59,36 @@ class TestTaskPool:
         task2 = pool.submit(task_id="test_2", task_type="test")
         assert len(pool._tasks) == 2
         assert pool.get_status("test_1") == TaskState.PENDING
-        assert pool.get_status("test_2") == TaskState.PENDING
\ No newline at end of file
+        assert pool.get_status("test_2") == TaskState.PENDING
+
+
+class TestTaskPoolStageAndAsync:
+    """Test TaskPool stage tracking and async execution (Task 3)"""
+
+    def test_task_has_stage_field(self):
+        """Test Task has stage field with default empty string"""
+        from sanguo_orchestrator.task import Task
+        t = Task(task_id="t1", task_type="cta")
+        assert t.stage == ""
+
+    def test_pool_submit_work_returns_future(self):
+        """Test submit_work() returns Future from executor"""
+        from unittest.mock import MagicMock
+        from sanguo_orchestrator.pool import TaskPool
+        pool = TaskPool(max_workers=2)
+        pool.executor = MagicMock()  # mock executor to avoid spawning real processes
+        mock_future = MagicMock()
+        pool.executor.submit.return_value = mock_future
+        fut = pool.submit_work("t1", func=lambda: 1)
+        assert fut is mock_future
+        pool.executor.submit.assert_called_once()
+
+    def test_pool_update_and_get_stage(self):
+        """Test update_stage() and get_stage() methods"""
+        from sanguo_orchestrator.pool import TaskPool
+        from sanguo_orchestrator.task import Task
+        pool = TaskPool(max_workers=2)
+        pool.submit("t1", "cta")
+        pool.update_stage("t1", "回测中")
+        assert pool.get_stage("t1") == "回测中"
+        assert pool.get_task("t1").stage == "回测中"
\ No newline at end of file