RenJW/RNAediting

1import os2import subprocess3import shutil4from pathlib import Path5import threading6import time7 8class TaskExecutor:9    """执行NGS分析任务 - 专为 Hugging Face Spaces 优化"""10    11    def __init__(self):12        pass13    14    def execute_tasks(self, project_path, samples, ref_mapping):15        """16        执行所有样本的分析任务 (直接执行,不使用 sbatch)17        18        Args:19            project_path: 项目路径20            samples: 样本列表21            ref_mapping: 样本到参考序列的映射22            23        Returns:24            str: 执行结果消息25        """26        project_path = Path(project_path)27        28        try:29            return self._execute_directly(project_path, samples, ref_mapping)30        except Exception as e:31            return f"❌ 任务执行失败: {str(e)}"32    33    def _execute_directly(self, project_path, samples, ref_mapping):34        """35        直接执行任务(不使用 sbatch)36        适用于 Hugging Face Spaces 环境37        """38        try:39            # 检查必要的工具40            missing_tools = self._check_required_tools()41            if missing_tools:42                return f"❌ 缺少必要的工具: {', '.join(missing_tools)}\n请确保已安装: fastqc, fastp, bowtie2, samtools"43            44            # 创建任务执行脚本45            from script_generator import ScriptGenerator46            generator = ScriptGenerator()47            direct_script = generator.generate_direct_pipeline_script(str(project_path))48            49            if not direct_script:50                return "❌ 生成执行脚本失败: generate_direct_pipeline_script 返回空值"51            52            direct_script_path = project_path / "pipeline_direct.sh"53            54            try:55                with open(direct_script_path, 'w') as f:56                    f.write(direct_script)57                os.chmod(direct_script_path, 0o755)58            except Exception as e:59                return f"❌ 写入执行脚本失败: {str(e)}\n   路径: {direct_script_path}"60            61            # 验证脚本文件确实存在62            if not direct_script_path.exists():63                return f"❌ 脚本文件创建失败: {direct_script_path}"64            65            print(f"✅ 创建执行脚本: {direct_script_path}")66            67            # 创建日志目录68            log_dir = project_path / ".job.out"69            log_dir.mkdir(exist_ok=True)70            71            # 为每个样本执行任务72            results = []73            threads = []  # 跟踪所有线程74            75            for sample in samples:76                ref_name = ref_mapping.get(sample, "ref")77                78                results.append(f"🔄 开始处理样本: {sample} (参考序列: {ref_name})")79                80                # 检查参考序列文件是否存在81                ref_file = project_path / "ref" / f"{ref_name}.fa"82                if not ref_file.exists():83                    results.append(f"  ❌ 参考序列文件不存在: {ref_file}")84                    continue85                86                # 构建参考序列索引(如果不存在)87                if not self._check_bowtie2_index(ref_file):88                    results.append(f"  📚 构建 bowtie2 索引...")89                    index_result = self._build_bowtie2_index(ref_file)90                    if not index_result:91                        results.append(f"  ❌ 构建索引失败")92                        continue93                    results.append(f"  ✅ 索引构建完成")94                95                # 执行分析流程(后台运行)96                log_file = log_dir / f"{sample}.log"97                results.append(f"  ▶️ 启动分析任务(日志: {log_file.name})")98                99                # 在后台启动任务,并保存线程引用100                thread = self._run_sample_analysis_background(101                    direct_script_path,102                    sample,103                    ref_name,104                    project_path105                )106                threads.append(thread)107            108            # 等待所有任务完成109            results.append(f"\n⏳ 等待 {len(threads)} 个分析任务完成...")110            for i, thread in enumerate(threads, 1):111                thread.join()112                results.append(f"  ✓ 任务 {i}/{len(threads)} 完成")113            114            # 检查每个样本的执行结果115            results.append(f"\n📋 检查分析结果:")116            for sample in samples:117                log_file = log_dir / f"{sample}.log"118                reditools_dir = project_path / "05-reditools2" / sample119                120                if log_file.exists():121                    # 读取日志最后几行122                    with open(log_file, 'r') as f:123                        log_lines = f.readlines()124                        last_lines = ''.join(log_lines[-5:])  # 最后5行125                    126                    # 检查输出目录127                    if reditools_dir.exists() and any(reditools_dir.iterdir()):128                        results.append(f"  ✅ {sample}: 分析成功,已生成结果")129                    else:130                        results.append(f"  ❌ {sample}: 输出目录为空或不存在")131                        results.append(f"     日志摘要:\n{last_lines}")132                else:133                    results.append(f"  ⚠️ {sample}: 未找到日志文件")134 135            136            message = "\n".join(results)137            message += f"\n\n✅ 已完成 {len(samples)} 个样本的分析"138            message += f"\n📁 日志目录: {log_dir}"139            140            return message141            142        except Exception as e:143            return f"❌ 直接执行失败: {str(e)}"144    145    def _check_required_tools(self):146        """检查必要的工具是否可用"""147        required_tools = ["fastqc", "fastp", "bowtie2", "samtools"]148        missing = []149        150        for tool in required_tools:151            try:152                result = subprocess.run(153                    [tool, "--version"],154                    capture_output=True,155                    timeout=5156                )157                if result.returncode != 0:158                    # 某些工具用 -v 而不是 --version159                    result = subprocess.run(160                        [tool, "-v"],161                        capture_output=True,162                        timeout=5163                    )164            except:165                missing.append(tool)166        167        return missing168    169    def _check_bowtie2_index(self, ref_file):170        """检查 bowtie2 索引是否存在"""171        # bowtie2 的索引文件不需要 .fa 后缀172        ref_path = Path(ref_file)173        # 索引文件直接使用参考序列文件名(不带.fa)174        base_name = str(ref_path.parent / ref_path.stem)175        index_files = [176            Path(f"{base_name}.1.bt2"),177            Path(f"{base_name}.2.bt2"),178            Path(f"{base_name}.3.bt2"),179            Path(f"{base_name}.4.bt2"),180        ]181        return all(f.exists() for f in index_files)182    183    def _build_bowtie2_index(self, ref_file):184        """构建 bowtie2 索引"""185        try:186            ref_path = Path(ref_file)187            # 输出前缀不包含.fa后缀188            output_prefix = str(ref_path.parent / ref_path.stem)189            result = subprocess.run(190                ["bowtie2-build", str(ref_file), output_prefix],191                capture_output=True,192                text=True,193                timeout=600  # 10分钟超时194            )195            return result.returncode == 0196        except:197            return False198    199    def _run_sample_analysis_background(self, script_path, sample, ref_name, project_path):200        """在后台运行样本分析,返回线程对象"""201        def run():202            try:203                # 验证脚本文件存在204                if not Path(script_path).exists():205                    print(f"❌ 样本 {sample}: 脚本文件不存在: {script_path}")206                    return207                208                # 创建日志文件路径209                log_dir = project_path / ".job.out"210                log_file = log_dir / f"{sample}.log"211                212                print(f"🔄 样本 {sample}: 开始执行 {Path(script_path).name}")213                214                # 将输出写入日志文件而不是捕获215                # 使用绝对路径避免相对路径问题216                abs_script_path = Path(script_path).resolve()217                218                with open(log_file, 'w') as f:219                    result = subprocess.run(220                        ["bash", str(abs_script_path), sample, ref_name],221                        cwd=str(project_path),222                        stdout=f,223                        stderr=subprocess.STDOUT,224                        text=True225                    )226                    227                # 检查返回码228                if result.returncode != 0:229                    print(f"❌ 样本 {sample} 分析失败 (退出码: {result.returncode})")230                    print(f"   查看日志: {log_file}")231                else:232                    print(f"✅ 样本 {sample} 分析完成")233                    234            except Exception as e:235                print(f"❌ 样本 {sample} 分析异常: {str(e)}")236        237        # 在新线程中启动238        thread = threading.Thread(target=run, daemon=False)  # 改为非daemon,确保等待完成239        thread.start()240        return thread  # 返回线程对象241    242    def _run_sample_analysis_sync(self, script_path, sample, ref_name, project_path, log_file):243        """同步运行样本分析(用于测试)"""244        try:245            with open(log_file, 'w') as f:246                result = subprocess.run(247                    ["bash", str(script_path), sample, ref_name],248                    cwd=str(project_path),249                    stdout=f,250                    stderr=subprocess.STDOUT,251                    text=True,252                    timeout=3600  # 1小时超时253                )254            return result.returncode == 0255        except subprocess.TimeoutExpired:256            return False257        except Exception:258            return False259