celestialflow 3.0.1__py3-none-any.whl

celestialflow/task_nodes.py

@@ -0,0 +1,160 @@
+import json
+import time
+import redis
+
+from .task_manage import TaskManager
+
+
+class RemoteWorkerError(Exception):
+    pass
+
+
+class TaskSplitter(TaskManager):
+    def __init__(self):
+        """
+        初始化 TaskSplitter
+        """
+        super().__init__(
+            func=self._split_task,
+            execution_mode="serial",
+            max_retries=0,
+        )
+
+    def _split_task(self, *task):
+        """
+        实际上这个函数不执行逻辑，仅用于符合 TaskManager 架构
+        """
+        return task
+
+    def get_args(self, task):
+        return task
+
+    def put_split_result(self, result: tuple):
+        split_count = 0
+        for item in result:
+            self.put_result_queues(item)
+            split_count += 1
+
+        self.extra_stats["split_output_count"].value += split_count
+        return split_count
+
+    def process_result(self, task, result):
+        """
+        处理不可迭代的任务结果
+        """
+        if not hasattr(result, "__iter__") or isinstance(result, (str, bytes)):
+            result = (result,)
+        elif isinstance(result, list):
+            result = tuple(result)
+
+        return result
+
+    def process_task_success(self, task, result, start_time):
+        """
+        统一处理成功任务
+
+        :param task: 完成的任务
+        :param result: 任务的结果
+        :param start_time: 任务开始时间
+        """
+        processed_result = self.process_result(task, result)
+
+        if self.enable_result_cache:
+            self.success_dict[task] = processed_result
+
+        # ✅ 清理 retry_time_dict
+        task_id = self.get_task_id(task)
+        self.retry_time_dict.pop(task_id, None)
+
+        split_count = self.put_split_result(result)
+        self.update_success_counter()
+
+        self.task_logger.splitter_success(
+            self.func.__name__,
+            self.get_task_info(task),
+            split_count,
+            time.time() - start_time,
+        )
+
+
+class TaskRedisTransfer(TaskManager):
+    def __init__(
+        self,
+        worker_limit=50,
+        unpack_task_args=False,
+        host="localhost",
+        port=6379,
+        db=0,
+        fetch_timeout=10,
+        result_timeout=10,
+    ):
+        """
+        初始化 TaskRedisTransfer
+        :param worker_limit: 并行工作线程数
+        :param unpack_task_args: 是否将任务参数解包
+        :param host: Redis 主机地址
+        :param port: Redis 端口
+        :param db: Redis 数据库
+        :param fetch_timeout: Redis 任务等待超时时间
+        :param result_timeout: Redis 结果等待超时时间
+        """
+        super().__init__(
+            func=self._trans_redis,
+            execution_mode="thread",
+            worker_limit=worker_limit,
+            unpack_task_args=unpack_task_args,
+        )
+
+        self.host = host
+        self.port = port
+        self.db = db
+        self.fetch_timeout = fetch_timeout
+        self.result_timeout = result_timeout
+
+    def init_redis(self):
+        """初始化 Redis 客户端"""
+        if not hasattr(self, "redis_client"):
+            self.redis_client = redis.Redis(
+                host=self.host, port=self.port, db=self.db, decode_responses=True
+            )
+
+    def _trans_redis(self, *task):
+        """
+        将任务写入 Redis, 并等待结果
+        """
+        self.init_redis()
+        input_key = f"{self.get_stage_tag()}:input"
+        output_key = f"{self.get_stage_tag()}:output"
+
+        # 提交任务
+        task_id = self.get_task_id(task)
+        payload = json.dumps({"id": task_id, "task": task})
+        self.redis_client.rpush(input_key, payload)
+
+        # ✅ 等待任务被 BLPOP 拿走（不在 list 中）
+        wait_start = time.time()
+        while True:
+            if self.redis_client.lpos(input_key, payload) is None:  # 已被取走
+                break
+            if time.time() - wait_start > self.fetch_timeout:
+                raise TimeoutError("Task not fetched from Redis in time")
+            time.sleep(0.1)
+
+        # ✅ 被取走后再进入结果等待阶段
+        start_time = time.time()
+        while True:
+            result = self.redis_client.hget(output_key, task_id)
+            if result:
+                self.redis_client.hdel(output_key, task_id)
+                result_obj = json.loads(result)
+                if result_obj.get("status") == "success":
+                    return result_obj.get("result")
+                elif result_obj.get("status") == "error":
+                    raise RemoteWorkerError(f"{result_obj.get('error')}")
+                else:
+                    raise ValueError(f"Unknown result status: {result_obj}")
+            if time.time() - start_time > self.result_timeout:
+                raise TimeoutError(
+                    "Redis result not returned in time after being fetched"
+                )
+            time.sleep(0.1)

celestialflow/task_progress.py

@@ -0,0 +1,57 @@
+from tqdm import tqdm
+from tqdm.asyncio import tqdm as tqdm_asy
+
+
+class NullProgress:
+    def update(self, n=1):
+        pass
+
+    def close(self):
+        pass
+
+    def refresh_total(self, total):
+        pass
+
+    def add_total(self, add_num):
+        pass
+
+
+class ProgressManager:
+    def __init__(
+        self,
+        total_tasks: int,
+        desc: str = "Processing",
+        mode: str = "normal",
+    ):
+        """
+        初始化进度条管理器
+
+        :param total_tasks: 任务总数，用于设置进度条的总长度
+        :param desc: 进度条的描述文字
+        :param mode: 任务模式，可选 "async", other
+        :param show_progress: 是否显示进度条
+        """
+        if mode == "async":
+            self.progress_bar = tqdm_asy(total=total_tasks, desc=desc)
+        else:
+            self.progress_bar = tqdm(total=total_tasks, desc=desc)
+
+    def update(self, n=1):
+        """更新进度条"""
+        self.progress_bar.update(n)
+
+    def close(self):
+        """关闭进度条"""
+        self.progress_bar.close()
+
+    def refresh_total(self, total):
+        """动态调整进度条的总任务数"""
+        self.progress_bar.total = total
+        self.progress_bar.refresh()
+
+    def add_total(self, add_num):
+        """动态增加进度条的总任务数"""
+        if not add_num:
+            return
+        total = self.progress_bar.total + add_num
+        self.refresh_total(total)

celestialflow/task_report.py

@@ -0,0 +1,162 @@
+from threading import Event, Thread
+
+import requests
+
+from .task_logging import TaskLogger
+from .task_types import TERMINATION_SIGNAL
+
+
+class TaskReporter:
+    """
+    周期性向远程服务推送任务运行状态的上报器。
+
+    - 定时从服务器拉取配置（如上报间隔、任务注入信息）
+    - 将任务图中的状态、错误、结构、拓扑等信息推送到后端接口
+    - 以后台线程方式运行，可随时 start()/stop()
+    - 主要用于可视化监控、任务远程控制与 Web UI 同步
+    """
+    def __init__(self, task_graph, logger_queue, host="127.0.0.1", port=5000):
+        from .task_graph import TaskGraph
+
+        self.task_graph: TaskGraph = task_graph
+        self.logger = TaskLogger(logger_queue)
+        self.base_url = f"http://{host}:{port}"
+        self._stop_flag = Event()
+        self._thread = None
+        self.interval = 5
+
+    def start(self):
+        if self._thread is None or not self._thread.is_alive():
+            self._stop_flag.clear()
+            self._thread = Thread(target=self._loop, daemon=True)
+            self._thread.start()
+
+    def stop(self):
+        if self._thread:
+            self.push_once()  # 最后一次
+            self._stop_flag.set()
+            self._thread.join(timeout=2)
+            self._thread = None
+            self.logger._log("DEBUG", "[Reporter] Stopped.")
+
+    def _loop(self):
+        while not self._stop_flag.is_set():
+            try:
+                self.push_once()
+            except Exception as e:
+                self.logger._log(
+                    "ERROR", f"[Reporter] Push error: {type(e).__name__}({e})."
+                )
+            self._stop_flag.wait(self.interval)
+
+    def push_once(self):
+        # 拉取逻辑
+        self._pull_interval()
+        self._pull_and_inject_tasks()
+
+        # 推送逻辑
+        self._push_errors()
+        self._push_status()
+        self._push_structure()
+        self._push_topology()
+
+    def _pull_interval(self):
+        try:
+            res = requests.get(f"{self.base_url}/api/get_interval", timeout=1)
+            if res.ok:
+                interval = res.json().get("interval", 5)
+                self.interval = max(1.0, min(interval, 60.0))
+        except Exception as e:
+            self.logger._log(
+                "WARNING", f"[Reporter] Interval fetch failed: {type(e).__name__}({e})."
+            )
+
+    def _pull_and_inject_tasks(self):
+        try:
+            res = requests.get(f"{self.base_url}/api/get_task_injection", timeout=2)
+            if res.ok:
+                tasks_list = res.json()
+                for task in tasks_list:
+                    target_node = task.get("node")
+                    task_datas = task.get("task_datas")
+
+                    if target_node not in self.task_graph.stages_status_dict:
+                        self.logger._log(
+                            "WARNING",
+                            f"[Reporter] Task injection target node {target_node} not found.",
+                        )
+                        continue
+
+                    # 这里你可以按需注入到不同的节点
+                    task_datas = [
+                        task if task != "TERMINATION_SIGNAL" else TERMINATION_SIGNAL
+                        for task in task_datas
+                    ]
+                    self.task_graph.put_stage_queue(
+                        {target_node: task_datas}, put_termination_signal=False
+                    )
+                    self.logger._log(
+                        "INFO", f"[Reporter] 注入任务到 {target_node}: {task_datas}"
+                    )
+        except Exception as e:
+            self.logger._log(
+                "WARNING",
+                f"[Reporter] Task injection fetch failed: {type(e).__name__}({e}).",
+            )
+
+    def _push_errors(self):
+        try:
+            self.task_graph.handle_fail_queue()
+            error_data = []
+            for (
+                err,
+                tag,
+            ), task_list in self.task_graph.get_error_timeline_dict().items():
+                for task, ts in task_list:
+                    error_data.append(
+                        {
+                            "error": err,
+                            "node": tag,
+                            "task_id": task if len(task) < 100 else task[:100] + "...",
+                            "timestamp": ts,
+                        }
+                    )
+            payload = {"errors": error_data}
+            requests.post(f"{self.base_url}/api/push_errors", json=payload, timeout=1)
+        except Exception as e:
+            self.logger._log(
+                "WARNING", f"[Reporter] Error push failed: {type(e).__name__}({e})."
+            )
+
+    def _push_status(self):
+        try:
+            status_data = self.task_graph.get_status_dict()
+            payload = {"status": status_data}
+            requests.post(f"{self.base_url}/api/push_status", json=payload, timeout=1)
+        except Exception as e:
+            self.logger._log(
+                "WARNING", f"[Reporter] Status push failed: {type(e).__name__}({e})."
+            )
+
+    def _push_structure(self):
+        try:
+            structure = self.task_graph.get_structure_json()
+            payload = {"items": structure}
+            requests.post(
+                f"{self.base_url}/api/push_structure", json=payload, timeout=1
+            )
+        except Exception as e:
+            self.logger._log(
+                "WARNING", f"[Reporter] Structure push failed: {type(e).__name__}({e})"
+            )
+
+    def _push_topology(self):
+        try:
+            topology = self.task_graph.get_graph_topology()
+            payload = {"topology": topology}
+            requests.post(f"{self.base_url}/api/push_topology", json=payload, timeout=1)
+        except Exception as e:
+            self.logger._log(
+                "WARNING", f"[Reporter] Topology push failed: {type(e).__name__}({e})."
+            )
+

celestialflow/task_structure.py

@@ -0,0 +1,151 @@
+from typing import List
+
+from .task_manage import TaskManager
+from .task_graph import TaskGraph
+
+
+# ========有向无环图(DAG)========
+class TaskChain(TaskGraph):
+    def __init__(self, stages: List[TaskManager], chain_mode: str = "serial"):
+        """
+        初始化 TaskChain
+        :param stages: TaskManager 列表
+        :param chain_mode: 链式模式，默认为 'serial'
+        """
+        for num, stage in enumerate(stages):
+            stage_name = f"Stage {num + 1}"
+            next_stages = [stages[num + 1]] if num < len(stages) - 1 else []
+            stage.set_graph_context(next_stages, chain_mode, stage_name)
+
+        root_stage = stages[0]
+        super().__init__([root_stage])
+
+    def start_chain(self, init_tasks_dict: dict, put_termination_signal: bool = True):
+        """
+        启动任务链
+        :param init_tasks_dict: 任务列表
+        """
+        self.start_graph(init_tasks_dict, put_termination_signal)
+
+
+class TaskCross(TaskGraph):
+    def __init__(self, layers: List[List[TaskManager]], layout_mode: str = "process"):
+        """
+        TaskCross: 多层任务交叉结构
+
+        该结构将任务按“层”组织，每层可以包含多个并行执行的 TaskManager 节点，
+        不同层之间通过依赖关系连接，形成跨层的数据流图。
+
+        :param layers: List[List[TaskManager]]
+            按层划分的任务节点列表。每个子列表代表一层，列表中的 TaskManager 将并行执行。
+            相邻层之间的所有节点将建立全连接依赖（即每个上一层节点都连接到下一层所有节点）。
+
+        :param layout_mode: str, default = 'process'
+            控制任务图的调度布局模式：
+            - 'serial'：逐层顺序执行，上一层全部完成后才启动下一层；
+            - 'process'：所有层并行启动，执行顺序由依赖关系自动调度。
+        """
+        for i in range(len(layers)):
+            curr_layer = layers[i]
+            next_layer = layers[i + 1] if i < len(layers) - 1 else []
+            for index, stage in enumerate(curr_layer[:]):
+                # 非最后一层连接为并行
+                stage.set_graph_context(
+                    next_stages=next_layer,
+                    stage_mode="process",
+                    stage_name=f"Layer{i+1}-{index+1}",
+                )
+        super().__init__(layers[0], layout_mode)
+
+    def start_cross(self, init_tasks_dict: dict, put_termination_signal: bool = True):
+        """
+        启动多层交叉结构任务图
+        """
+        self.start_graph(init_tasks_dict, put_termination_signal)
+
+
+class TaskGrid(TaskGraph):
+    def __init__(self, grid: List[List[TaskManager]], layout_mode: str = "process"):
+        rows, cols = len(grid), len(grid[0])
+        for i in range(rows):
+            for j in range(cols):
+                curr = grid[i][j]
+                nexts = []
+                if i + 1 < rows:
+                    nexts.append(grid[i + 1][j])  # down
+                if j + 1 < cols:
+                    nexts.append(grid[i][j + 1])  # right
+                curr.set_graph_context(nexts, "process", f"Grid-{i+1}-{j+1}")
+        super().__init__([grid[0][0]], layout_mode)  # 起点为左上角
+
+    def start_grid(self, init_tasks_dict: dict, put_termination_signal: bool = True):
+        """
+        启动任务网格结构
+        :param init_tasks_dict: 任务列表
+        """
+        self.start_graph(init_tasks_dict, put_termination_signal)
+
+
+# ========有环图========
+class TaskLoop(TaskGraph):
+    def __init__(self, stages: List[TaskManager]):
+        """
+        初始化 TaskLoop, 由于环的结构特性, 强制使用 'process' 节点模式
+        :param stages: TaskManager 列表
+        """
+        for num, stage in enumerate(stages):
+            stage_name = f"Stage {num + 1}"
+            next_stages = [stages[num + 1]] if num < len(stages) - 1 else [stages[0]]
+            stage.set_graph_context(next_stages, "process", stage_name)
+
+        root_stage = stages[0]
+        super().__init__([root_stage])
+
+    def start_loop(self, init_tasks_dict: dict):
+        """
+        启动任务环, 环是自锁结构, 能且仅能外部注入式停止
+        :param init_tasks_dict: 任务列表
+        """
+        self.start_graph(init_tasks_dict, False)
+
+
+class TaskWheel(TaskGraph):
+    def __init__(self, center: TaskManager, ring: List[TaskManager]):
+        # 中心连向环
+        center.set_graph_context(ring, "process", "Center")
+        # 环相连（成闭环）
+        for i, node in enumerate(ring):
+            next_stage = ring[(i + 1) % len(ring)]
+            node.set_graph_context([next_stage], "process", f"Ring-{i+1}")
+        super().__init__([center])
+
+    def start_wheel(self, init_tasks_dict: dict):
+        """
+        启动任务轮结构
+        :param init_tasks_dict: 任务列表
+        """
+        self.start_graph(init_tasks_dict, False)
+
+
+class TaskComplete(TaskGraph):
+    def __init__(self, stages: List[TaskManager]):
+        """
+        TaskComplete: 完全图结构，每个节点都连向除自己以外的所有其他节点
+        :param stages: 所有 TaskManager 节点
+        """
+        for i, stage in enumerate(stages):
+            next_stages = [s for j, s in enumerate(stages) if i != j]
+            stage.set_graph_context(
+                next_stages=next_stages,
+                stage_mode="process",
+                stage_name=f"Node {i + 1}",
+            )
+
+        super().__init__(stages)
+
+    def start_complete(self, init_tasks_dict: dict):
+        """
+        启动任务完全图
+        :param init_tasks_dict: 任务列表
+        """
+        self.start_graph(init_tasks_dict, False)