feat(orchestrator): pool 异步化(ProcessPoolExecutor spawn + stage 追踪)

This commit is contained in:
2026-07-06 18:19:24 +08:00
parent 48a9058cb2
commit 66aa27e807
3 changed files with 54 additions and 1 deletions
+20
View File
@@ -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)
+1
View File
@@ -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"""