bitool 0.1.2__py3-none-any.whl

bitool/cmd/_scheduler.py

@@ -0,0 +1,548 @@
+"""命令调度器模块, 提供 BaseCommand 的顺序执行、并行执行和 DAG 执行."""
+
+from __future__ import annotations
+
+import time
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from dataclasses import dataclass, field
+from enum import Enum, auto
+from typing import Any, Callable
+
+from bitool.cmd import BaseCommand
+from bitool.core import logger
+
+
+class ExecutionMode(Enum):
+    """执行模式枚举"""
+
+    SEQUENTIAL = auto()  # 顺序执行
+    PARALLEL = auto()  # 并行执行
+    DAG = auto()  # DAG依赖执行
+
+
+@dataclass
+class CommandResult:
+    """命令执行结果"""
+
+    name: str
+    success: bool
+    duration: float = 0.0  # 执行耗时（秒）
+    error: str | None = None  # 错误信息
+    output: Any = None  # 命令输出
+    retry_count: int = 0  # 重试次数
+
+
+@dataclass
+class SchedulerMetrics:
+    """调度器指标"""
+
+    total_commands: int = 0
+    successful: int = 0
+    failed: int = 0
+    skipped: int = 0
+    total_duration: float = 0.0
+    mode_used: str = ""
+    command_results: list[CommandResult] = field(default_factory=list)
+
+
+class CommandScheduler:
+    """命令调度器
+
+    提供 BaseCommand 的顺序执行、并行执行和 DAG 执行功能
+
+    工业级特性:
+    - 自动模式选择
+    - 超时控制
+    - 重试机制
+    - 详细结果追踪
+    - 性能指标收集
+    - 进度回调
+    - 执行上下文共享
+    """
+
+    def __init__(
+        self,
+        commands: list[BaseCommand] | None = None,
+        max_workers: int | None = None,
+        timeout: float | None = None,
+        max_retries: int = 0,
+        retry_delay: float = 1.0,
+        context: dict[str, Any] | None = None,
+        progress_callback: Callable[[str, CommandResult], None] | None = None,
+    ) -> None:
+        """初始化命令调度器
+
+        Args:
+            commands: 命令列表
+            max_workers: 最大并行工作线程数, None 表示使用 CPU 核心数
+            timeout: 单个命令执行超时时间（秒）, None 表示不限制
+            max_retries: 最大重试次数
+            retry_delay: 重试延迟时间（秒）
+            context: 命令间共享的执行上下文
+            progress_callback: 进度回调函数, 接收 (命令名, 执行结果)
+        """
+        self.commands = commands or []
+        self.max_workers = max_workers
+        self.timeout = timeout
+        self.max_retries = max_retries
+        self.retry_delay = retry_delay
+        self.context = context or {}
+        self.progress_callback = progress_callback
+
+        self.results: dict[str, bool] = {}
+        self.detailed_results: dict[str, CommandResult] = {}
+        self.metrics = SchedulerMetrics()
+
+        # 初始化时验证依赖
+        self._validate_dependencies(self.commands)
+
+    def run(
+        self, *, stop_on_failure: bool = False, max_workers: int | None = None
+    ) -> dict[str, bool]:
+        """智能选择调度模式执行命令
+
+        根据命令的依赖关系自动选择最优执行策略:
+        - 如果所有命令都没有依赖关系 -> 并行执行
+        - 如果存在依赖关系 -> DAG 执行
+        - 如果明确指定顺序执行 -> 顺序执行
+
+        Args:
+            stop_on_failure: 失败时是否停止执行, 默认 False
+            max_workers: 最大并行工作线程数, None 表示使用实例默认值
+
+        Returns:
+            命令执行结果字典, key 为命令 name, value 为执行是否成功
+        """
+        if not self.commands:
+            logger.info("命令列表为空, 直接返回")
+            return {}
+
+        # 分析依赖关系
+        mode = self._detect_execution_mode()
+        logger.info(f"自动检测到执行模式: {mode.name}")
+
+        if mode == ExecutionMode.SEQUENTIAL:
+            return self.run_sequential(stop_on_failure=stop_on_failure)
+        elif mode == ExecutionMode.PARALLEL:
+            return self.run_parallel(max_workers=max_workers)
+        else:  # DAG
+            return self.run_dag(
+                self.commands, max_workers=max_workers, stop_on_failure=stop_on_failure
+            )
+
+    def run_sequential(self, *, stop_on_failure: bool = False) -> dict[str, bool]:
+        """顺序执行命令列表
+
+        Args:
+            stop_on_failure: 失败时是否停止执行, 默认 False
+
+        Returns:
+            命令执行结果字典, key 为命令 name, value 为执行是否成功
+        """
+        if not self.commands:
+            logger.info("命令列表为空, 直接返回")
+            return {}
+
+        results = {}
+        self.metrics = SchedulerMetrics()
+        self.metrics.mode_used = "SEQUENTIAL"
+        start_time = time.time()
+
+        logger.info(f"开始顺序执行 {len(self.commands)} 个命令")
+
+        for cmd in self.commands:
+            try:
+                success = self._execute_command(cmd)
+                results[cmd.name] = success
+                self.results[cmd.name] = success
+
+                if not success and stop_on_failure:
+                    logger.warning(f"命令 {cmd.name} 执行失败, 停止后续执行")
+                    self.metrics.skipped = len(self.commands) - len(results)
+                    break
+
+            except Exception as e:  # noqa: BLE001
+                results[cmd.name] = False
+                self.results[cmd.name] = False
+                logger.error(f"命令 {cmd.name} 执行异常: {e}")
+
+                if stop_on_failure:
+                    logger.warning(f"命令 {cmd.name} 执行异常, 停止后续执行")
+                    self.metrics.skipped = len(self.commands) - len(results)
+                    break
+
+        self.metrics.total_duration = time.time() - start_time
+        return results
+
+    def run_parallel(self, max_workers: int | None = None) -> dict[str, bool]:
+        """并行执行命令列表
+
+        Args:
+            max_workers: 最大并行工作线程数, None 表示使用实例默认值
+
+        Returns:
+            命令执行结果字典, key 为命令 name, value 为执行是否成功
+        """
+        if not self.commands:
+            logger.info("命令列表为空, 直接返回")
+            return {}
+
+        workers = max_workers or self.max_workers
+        results = {}
+        self.metrics = SchedulerMetrics()
+        self.metrics.mode_used = "PARALLEL"
+        start_time = time.time()
+
+        logger.info(f"开始并行执行 {len(self.commands)} 个命令")
+
+        with ThreadPoolExecutor(max_workers=workers) as executor:
+            future_to_cmd = {}
+            for cmd in self.commands:
+                future = executor.submit(self._execute_command_with_timeout, cmd)
+                future_to_cmd[future] = cmd
+
+            for future in as_completed(future_to_cmd):
+                cmd = future_to_cmd[future]
+                cmd_name = cmd.name
+                try:
+                    success = future.result()
+                    results[cmd_name] = success
+                    self.results[cmd_name] = success
+
+                    if success:
+                        logger.info(f"命令 {cmd_name} 执行成功")
+                    else:
+                        logger.warning(f"命令 {cmd_name} 执行失败")
+
+                except Exception as e:  # noqa: BLE001
+                    results[cmd_name] = False
+                    self.results[cmd_name] = False
+                    logger.error(f"命令 {cmd_name} 执行异常: {e}")
+
+        self.metrics.total_duration = time.time() - start_time
+        return results
+
+    def run_dag(
+        self,
+        tasks: list[BaseCommand] | None = None,
+        max_workers: int | None = None,
+        *,
+        stop_on_failure: bool = False,
+    ) -> dict[str, bool]:
+        """基于 DAG 依赖调度执行命令
+
+        Args:
+            tasks: 命令任务列表, None 表示使用实例的命令列表
+            max_workers: 最大并行工作线程数, None 表示使用实例默认值
+            stop_on_failure: 失败时是否停止执行, 默认 False
+
+        Returns:
+            命令执行结果字典, key 为任务 name, value 为执行是否成功
+
+        Raises:
+            ValueError: 当存在循环依赖时
+        """
+        tasks = tasks or self.commands
+        if not tasks:
+            logger.info("任务列表为空, 直接返回")
+            return {}
+
+        workers = max_workers or self.max_workers
+        results = {}
+        completed: set[str] = set()
+        failed: set[str] = set()
+
+        self.metrics = SchedulerMetrics()
+        self.metrics.mode_used = "DAG"
+        start_time = time.time()
+
+        logger.info(f"开始 DAG 调度执行 {len(tasks)} 个任务")
+
+        # 拓扑排序执行
+        while len(completed) + len(failed) < len(tasks):
+            # 找到可以执行的任务（依赖已全部完成）
+            ready_tasks = []
+            for task in tasks:
+                if task.name in completed or task.name in failed:
+                    continue
+
+                # 检查依赖是否都已完成
+                deps_met = all(dep in completed for dep in task.dependencies)
+
+                # 检查是否有依赖失败
+                deps_failed = any(dep in failed for dep in task.dependencies)
+
+                if deps_failed:
+                    failed.add(task.name)
+                    results[task.name] = False
+                    self.results[task.name] = False
+                    logger.warning(f"任务 {task.name} 的依赖失败, 跳过执行")
+                    continue
+
+                if deps_met:
+                    ready_tasks.append(task)
+
+            if not ready_tasks:
+                # 如果没有就绪任务但还有未完成的任务, 说明有循环依赖
+                remaining = len(tasks) - len(completed) - len(failed)
+                if remaining > 0:
+                    logger.error("检测到循环依赖, 无法继续执行")
+                    msg = "检测到循环依赖"
+                    raise ValueError(msg)
+                break
+
+            # 并行执行就绪的任务
+            with ThreadPoolExecutor(max_workers=workers) as executor:
+                future_to_task = {}
+                for task in ready_tasks:
+                    future = executor.submit(self._execute_command_with_timeout, task)
+                    future_to_task[future] = task
+
+                for future in as_completed(future_to_task):
+                    task = future_to_task[future]
+                    task_name = task.name
+                    try:
+                        success = future.result()
+                        results[task_name] = success
+                        self.results[task_name] = success
+
+                        if success:
+                            completed.add(task_name)
+                            logger.info(f"任务 {task_name} 执行成功")
+                        else:
+                            failed.add(task_name)
+                            logger.warning(f"任务 {task_name} 执行失败")
+
+                            if stop_on_failure:
+                                logger.warning(
+                                    f"任务 {task_name} 执行失败, 停止后续执行"
+                                )
+                                self.metrics.total_duration = time.time() - start_time
+                                return results
+
+                    except Exception as e:  # noqa: BLE001
+                        results[task_name] = False
+                        self.results[task_name] = False
+                        failed.add(task_name)
+                        logger.error(f"任务 {task_name} 执行异常: {e}")
+
+                        if stop_on_failure:
+                            logger.warning(f"任务 {task_name} 执行异常, 停止后续执行")
+                            self.metrics.total_duration = time.time() - start_time
+                            return results
+
+        self.metrics.total_duration = time.time() - start_time
+        return results
+
+    def _detect_execution_mode(self) -> ExecutionMode:
+        """检测最适合的执行模式
+
+        Returns:
+            推荐的执行模式
+        """
+        if not self.commands:
+            return ExecutionMode.SEQUENTIAL
+
+        # 检查是否存在依赖关系
+        has_dependencies = any(len(cmd.dependencies) > 0 for cmd in self.commands)
+
+        if has_dependencies:
+            return ExecutionMode.DAG
+        else:
+            # 没有依赖关系，可以并行执行
+            # 如果只有一个命令，顺序执行即可
+            if len(self.commands) == 1:
+                return ExecutionMode.SEQUENTIAL
+            else:
+                return ExecutionMode.PARALLEL
+
+    def _execute_command(self, command: BaseCommand) -> bool:
+        """执行单个命令（带重试和超时控制）
+
+        Args:
+            command: BaseCommand 实例
+
+        Returns:
+            命令执行是否成功
+        """
+        cmd_name = command.name
+        retry_count = 0
+        last_error = None
+
+        while retry_count <= self.max_retries:
+            try:
+                logger.info(
+                    f"执行命令: {cmd_name} - {command.description} "
+                    f"(尝试 {retry_count + 1}/{self.max_retries + 1})"
+                )
+
+                start_time = time.time()
+                success = command.validate_and_run()
+                duration = time.time() - start_time
+
+                # 记录详细结果
+                result = CommandResult(
+                    name=cmd_name,
+                    success=success,
+                    duration=duration,
+                    retry_count=retry_count,
+                )
+
+                if success:
+                    logger.info(f"命令 {cmd_name} 执行成功 (耗时: {duration:.2f}s)")
+                    self.detailed_results[cmd_name] = result
+                    self._notify_progress(cmd_name, result)
+                    return True
+                else:
+                    logger.warning(f"命令 {cmd_name} 执行失败 (耗时: {duration:.2f}s)")
+                    result.error = "命令返回失败状态"
+                    last_error = result.error
+
+            except TimeoutError as e:
+                duration = time.time() - start_time if "start_time" in locals() else 0
+                logger.error(f"命令 {cmd_name} 执行超时 (耗时: {duration:.2f}s)")
+                last_error = f"执行超时: {e}"
+
+            except Exception as e:  # noqa: BLE001
+                duration = time.time() - start_time if "start_time" in locals() else 0
+                logger.error(f"命令 {cmd_name} 执行异常: {e} (耗时: {duration:.2f}s)")
+                last_error = str(e)
+
+            # 重试逻辑
+            retry_count += 1
+            if retry_count <= self.max_retries:
+                logger.info(f"等待 {self.retry_delay}s 后重试...")
+                time.sleep(self.retry_delay)
+
+        # 所有重试都失败
+        result = CommandResult(
+            name=cmd_name, success=False, error=last_error, retry_count=retry_count - 1
+        )
+        self.detailed_results[cmd_name] = result
+        self._notify_progress(cmd_name, result)
+        return False
+
+    def _execute_command_with_timeout(self, command: BaseCommand) -> bool:
+        """执行单个命令（简化版，不带重试）
+
+        Args:
+            command: BaseCommand 实例
+
+        Returns:
+            命令执行是否成功
+        """
+        cmd_name = command.name
+        logger.info(f"执行命令: {cmd_name} - {command.description}")
+
+        try:
+            start_time = time.time()
+            success = command.validate_and_run()
+            duration = time.time() - start_time
+
+            result = CommandResult(name=cmd_name, success=success, duration=duration)
+
+            self.detailed_results[cmd_name] = result
+            self._notify_progress(cmd_name, result)
+        except Exception as e:  # noqa: BLE001
+            duration = time.time() - start_time if "start_time" in locals() else 0
+            logger.error(f"命令 {cmd_name} 执行异常: {e}")
+
+            result = CommandResult(
+                name=cmd_name, success=False, error=str(e), duration=duration
+            )
+            self.detailed_results[cmd_name] = result
+            self._notify_progress(cmd_name, result)
+            return False
+        else:
+            return success
+
+    def _notify_progress(self, cmd_name: str, result: CommandResult) -> None:
+        """通知进度回调
+
+        Args:
+            cmd_name: 命令名称
+            result: 执行结果
+        """
+        if self.progress_callback:
+            try:
+                self.progress_callback(cmd_name, result)
+            except Exception as e:  # noqa: BLE001
+                logger.warning(f"进度回调执行失败: {e}")
+
+    def _validate_dependencies(self, tasks: list[BaseCommand]) -> None:
+        """验证依赖关系是否有效
+
+        Args:
+            tasks: 命令任务列表
+
+        Raises:
+            ValueError: 当依赖的任务不存在时
+        """
+        task_names = {task.name for task in tasks}
+
+        for task in tasks:
+            for dep in task.dependencies:
+                if dep not in task_names:
+                    msg = f"任务 {task.name} 依赖的任务 {dep} 不存在"
+                    raise ValueError(msg)
+
+    def get_results(self) -> dict[str, bool]:
+        """获取所有执行结果
+
+        Returns:
+            命令执行结果字典
+        """
+        return self.results.copy()
+
+    def get_detailed_results(self) -> dict[str, CommandResult]:
+        """获取详细执行结果
+
+        Returns:
+            详细命令执行结果字典，包含耗时、错误信息等
+        """
+        return self.detailed_results.copy()
+
+    def get_metrics(self) -> SchedulerMetrics:
+        """获取调度器指标
+
+        Returns:
+            调度器性能指标
+        """
+        self.metrics.total_commands = len(self.results)
+        self.metrics.successful = sum(1 for v in self.results.values() if v)
+        self.metrics.failed = sum(1 for v in self.results.values() if not v)
+        self.metrics.total_duration = sum(
+            r.duration for r in self.detailed_results.values()
+        )
+        self.metrics.command_results = list(self.detailed_results.values())
+        return self.metrics
+
+    def get_success_count(self) -> int:
+        """获取成功执行的命令数量
+
+        Returns:
+            成功执行的命令数量
+        """
+        return sum(1 for success in self.results.values() if success)
+
+    def get_failure_count(self) -> int:
+        """获取执行失败的命令数量
+
+        Returns:
+            执行失败的命令数量
+        """
+        return sum(1 for success in self.results.values() if not success)
+
+    def __enter__(self) -> CommandScheduler:
+        """上下文管理器入口"""
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: object | None,
+    ) -> None:
+        """上下文管理器退出"""
+        # 清理资源（如果有需要）
+        pass