388 lines
13 KiB
Python
388 lines
13 KiB
Python
|
import asyncio
|
||
|
|
import json
|
||
|
|
import subprocess
|
||
|
|
from dataclasses import dataclass, field
|
||
|
|
from datetime import datetime
|
||
|
|
from enum import Enum
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
from gfa_core.mas.tools.registry import ToolMetadata, ToolRegistry, create_default_registry
|
||
|
|
|
||
|
|
|
||
|
|
class ExecutionStatus(Enum):
|
||
|
|
PENDING = "pending"
|
||
|
|
RUNNING = "running"
|
||
|
|
COMPLETED = "completed"
|
||
|
|
FAILED = "failed"
|
||
|
|
TIMEOUT = "timeout"
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class ExecutionResult:
|
||
|
|
tool_id: str
|
||
|
|
status: ExecutionStatus
|
||
|
|
return_code: int = 0
|
||
|
|
stdout: str = ""
|
||
|
|
stderr: str = ""
|
||
|
|
outputs: dict[str, str] = field(default_factory=dict)
|
||
|
|
started_at: datetime | None = None
|
||
|
|
completed_at: datetime | None = None
|
||
|
|
error: str | None = None
|
||
|
|
|
||
|
|
|
||
|
|
class ToolExecutor:
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
registry: ToolRegistry | None = None,
|
||
|
|
work_dir: str = "/tmp/gene-family-tools",
|
||
|
|
use_docker: bool = False,
|
||
|
|
simulation_mode: bool = False,
|
||
|
|
):
|
||
|
|
self.registry = registry or create_default_registry()
|
||
|
|
self.work_dir = Path(work_dir)
|
||
|
|
self.work_dir.mkdir(parents=True, exist_ok=True)
|
||
|
|
self.use_docker = use_docker
|
||
|
|
self.simulation_mode = simulation_mode
|
||
|
|
|
||
|
|
def get_tool(self, tool_id: str) -> ToolMetadata | None:
|
||
|
|
return self.registry.get(tool_id)
|
||
|
|
|
||
|
|
def list_tools(self, category: str | None = None) -> list[ToolMetadata]:
|
||
|
|
if category:
|
||
|
|
return self.registry.list_by_category(category)
|
||
|
|
return self.registry.list_all()
|
||
|
|
|
||
|
|
def execute(
|
||
|
|
self,
|
||
|
|
tool_id: str,
|
||
|
|
inputs: dict[str, str],
|
||
|
|
parameters: dict[str, Any] | None = None,
|
||
|
|
work_subdir: str | None = None,
|
||
|
|
) -> ExecutionResult:
|
||
|
|
tool = self.registry.get(tool_id)
|
||
|
|
if not tool:
|
||
|
|
return ExecutionResult(
|
||
|
|
tool_id=tool_id,
|
||
|
|
status=ExecutionStatus.FAILED,
|
||
|
|
error=f"Tool not found: {tool_id}",
|
||
|
|
)
|
||
|
|
|
||
|
|
parameters = parameters or {}
|
||
|
|
work_path = self.work_dir / (work_subdir or tool_id)
|
||
|
|
work_path.mkdir(parents=True, exist_ok=True)
|
||
|
|
|
||
|
|
started_at = datetime.now()
|
||
|
|
|
||
|
|
if self.simulation_mode:
|
||
|
|
return self._simulate_execution(tool, inputs, parameters, work_path, started_at)
|
||
|
|
|
||
|
|
try:
|
||
|
|
if self.use_docker and tool.docker_image:
|
||
|
|
cmd = self._build_docker_command(tool, inputs, parameters, work_path)
|
||
|
|
else:
|
||
|
|
cmd = self._build_command(tool, inputs, parameters, work_path)
|
||
|
|
|
||
|
|
result = subprocess.run(
|
||
|
|
cmd,
|
||
|
|
cwd=str(work_path),
|
||
|
|
capture_output=True,
|
||
|
|
text=True,
|
||
|
|
timeout=tool.timeout,
|
||
|
|
)
|
||
|
|
|
||
|
|
if tool.id == "mafft" and result.returncode == 0:
|
||
|
|
mafft_output = work_path / "aligned.fasta"
|
||
|
|
with open(mafft_output, "w") as f:
|
||
|
|
f.write(result.stdout)
|
||
|
|
|
||
|
|
outputs = self._collect_outputs(tool, work_path)
|
||
|
|
|
||
|
|
status = ExecutionStatus.COMPLETED if result.returncode == 0 else ExecutionStatus.FAILED
|
||
|
|
|
||
|
|
return ExecutionResult(
|
||
|
|
tool_id=tool_id,
|
||
|
|
status=status,
|
||
|
|
return_code=result.returncode,
|
||
|
|
stdout=result.stdout,
|
||
|
|
stderr=result.stderr,
|
||
|
|
outputs=outputs,
|
||
|
|
started_at=started_at,
|
||
|
|
completed_at=datetime.now(),
|
||
|
|
)
|
||
|
|
|
||
|
|
except subprocess.TimeoutExpired:
|
||
|
|
return ExecutionResult(
|
||
|
|
tool_id=tool_id,
|
||
|
|
status=ExecutionStatus.TIMEOUT,
|
||
|
|
error=f"Tool timed out after {tool.timeout}s",
|
||
|
|
started_at=started_at,
|
||
|
|
completed_at=datetime.now(),
|
||
|
|
)
|
||
|
|
except FileNotFoundError:
|
||
|
|
return ExecutionResult(
|
||
|
|
tool_id=tool_id,
|
||
|
|
status=ExecutionStatus.FAILED,
|
||
|
|
error=f"Executable not found: {tool.executable}",
|
||
|
|
started_at=started_at,
|
||
|
|
completed_at=datetime.now(),
|
||
|
|
)
|
||
|
|
except Exception as e:
|
||
|
|
return ExecutionResult(
|
||
|
|
tool_id=tool_id,
|
||
|
|
status=ExecutionStatus.FAILED,
|
||
|
|
error=str(e),
|
||
|
|
started_at=started_at,
|
||
|
|
completed_at=datetime.now(),
|
||
|
|
)
|
||
|
|
|
||
|
|
async def aexecute(
|
||
|
|
self,
|
||
|
|
tool_id: str,
|
||
|
|
inputs: dict[str, str],
|
||
|
|
parameters: dict[str, Any] | None = None,
|
||
|
|
work_subdir: str | None = None,
|
||
|
|
) -> ExecutionResult:
|
||
|
|
return await asyncio.to_thread(
|
||
|
|
self.execute, tool_id, inputs, parameters, work_subdir
|
||
|
|
)
|
||
|
|
|
||
|
|
def _build_command(
|
||
|
|
self,
|
||
|
|
tool: ToolMetadata,
|
||
|
|
inputs: dict[str, str],
|
||
|
|
parameters: dict[str, Any],
|
||
|
|
work_path: Path,
|
||
|
|
) -> list[str]:
|
||
|
|
cmd = [tool.executable]
|
||
|
|
|
||
|
|
if tool.id == "hmmsearch":
|
||
|
|
output_tbl = work_path / "output.tbl"
|
||
|
|
dom_tbl = work_path / "dom.tbl"
|
||
|
|
cmd.extend([
|
||
|
|
"--tblout", str(output_tbl),
|
||
|
|
"--domtblout", str(dom_tbl),
|
||
|
|
"-E", str(parameters.get("e_value", 1e-5)),
|
||
|
|
"--cpu", str(parameters.get("cpu", 4)),
|
||
|
|
inputs.get("query", ""),
|
||
|
|
inputs.get("target", ""),
|
||
|
|
])
|
||
|
|
|
||
|
|
elif tool.id == "blastp":
|
||
|
|
output_file = work_path / "blast_results.txt"
|
||
|
|
cmd.extend([
|
||
|
|
"-query", inputs.get("query", ""),
|
||
|
|
"-db", inputs.get("db", ""),
|
||
|
|
"-evalue", str(parameters.get("e_value", 1e-5)),
|
||
|
|
"-max_target_seqs", str(parameters.get("max_target_seqs", 500)),
|
||
|
|
"-outfmt", str(parameters.get("outfmt", 6)),
|
||
|
|
"-out", str(output_file),
|
||
|
|
])
|
||
|
|
|
||
|
|
elif tool.id == "mafft":
|
||
|
|
strategy = parameters.get("strategy", "auto")
|
||
|
|
if strategy != "auto":
|
||
|
|
cmd.append(f"--{strategy}")
|
||
|
|
cmd.extend([
|
||
|
|
"--maxiterate", str(parameters.get("maxiterate", 1000)),
|
||
|
|
"--reorder",
|
||
|
|
"--thread", str(parameters.get("threads", 4)),
|
||
|
|
inputs.get("sequences", ""),
|
||
|
|
])
|
||
|
|
|
||
|
|
elif tool.id == "iqtree":
|
||
|
|
cmd.extend([
|
||
|
|
"-s", inputs.get("alignment", ""),
|
||
|
|
"-m", parameters.get("model", "MFP"),
|
||
|
|
"-bb", str(parameters.get("bootstrap", 1000)),
|
||
|
|
"-nt", str(parameters.get("threads", 4)),
|
||
|
|
"-pre", str(work_path / "iqtree"),
|
||
|
|
])
|
||
|
|
|
||
|
|
elif tool.id == "meme":
|
||
|
|
output_dir = work_path / "meme_out"
|
||
|
|
cmd.extend([
|
||
|
|
inputs.get("sequences", ""),
|
||
|
|
"-nmotifs", str(parameters.get("nmotifs", 10)),
|
||
|
|
"-minw", str(parameters.get("min_width", 6)),
|
||
|
|
"-maxw", str(parameters.get("max_width", 200)),
|
||
|
|
"-oc", str(output_dir),
|
||
|
|
"-mod", "anr",
|
||
|
|
"-protein",
|
||
|
|
])
|
||
|
|
|
||
|
|
elif tool.id == "mcscanx":
|
||
|
|
cmd.extend([
|
||
|
|
inputs.get("gff", "").replace(".gff", ""),
|
||
|
|
])
|
||
|
|
|
||
|
|
return cmd
|
||
|
|
|
||
|
|
def _build_docker_command(
|
||
|
|
self,
|
||
|
|
tool: ToolMetadata,
|
||
|
|
inputs: dict[str, str],
|
||
|
|
parameters: dict[str, Any],
|
||
|
|
work_path: Path,
|
||
|
|
) -> list[str]:
|
||
|
|
cmd = [
|
||
|
|
"docker", "run", "--rm",
|
||
|
|
"-v", f"{work_path}:/data",
|
||
|
|
"-w", "/data",
|
||
|
|
tool.docker_image,
|
||
|
|
tool.executable,
|
||
|
|
]
|
||
|
|
|
||
|
|
base_cmd = self._build_command(tool, inputs, parameters, work_path)
|
||
|
|
cmd.extend(base_cmd[1:])
|
||
|
|
|
||
|
|
return cmd
|
||
|
|
|
||
|
|
def _simulate_execution(
|
||
|
|
self,
|
||
|
|
tool: ToolMetadata,
|
||
|
|
inputs: dict[str, str],
|
||
|
|
parameters: dict[str, Any],
|
||
|
|
work_path: Path,
|
||
|
|
started_at: datetime,
|
||
|
|
) -> ExecutionResult:
|
||
|
|
import random
|
||
|
|
import time
|
||
|
|
|
||
|
|
time.sleep(0.5)
|
||
|
|
|
||
|
|
outputs = {}
|
||
|
|
stdout_lines = [f"[SIMULATION] Executing {tool.name}"]
|
||
|
|
stdout_lines.append(f"[SIMULATION] Input files: {list(inputs.values())}")
|
||
|
|
stdout_lines.append(f"[SIMULATION] Parameters: {parameters}")
|
||
|
|
|
||
|
|
if tool.id == "hmmsearch":
|
||
|
|
output_tbl = work_path / "output.tbl"
|
||
|
|
dom_tbl = work_path / "dom.tbl"
|
||
|
|
|
||
|
|
target_file = inputs.get("target", "")
|
||
|
|
num_hits = random.randint(3, 8)
|
||
|
|
|
||
|
|
with open(output_tbl, "w") as f:
|
||
|
|
f.write(f"# hmmsearch output (simulated)\n")
|
||
|
|
f.write(f"# target: {target_file}\n")
|
||
|
|
for i in range(num_hits):
|
||
|
|
evalue = 10 ** (-random.randint(5, 50))
|
||
|
|
score = random.uniform(50, 300)
|
||
|
|
f.write(f"gene{i+1}\t-\tSPL\t-\t{evalue:.2e}\t{score:.1f}\n")
|
||
|
|
|
||
|
|
with open(dom_tbl, "w") as f:
|
||
|
|
f.write(f"# domain hits (simulated)\n")
|
||
|
|
|
||
|
|
outputs["tblout"] = str(output_tbl)
|
||
|
|
outputs["domtblout"] = str(dom_tbl)
|
||
|
|
stdout_lines.append(f"[SIMULATION] Found {num_hits} hits with E-value < {parameters.get('e_value', 1e-5)}")
|
||
|
|
|
||
|
|
elif tool.id == "blastp":
|
||
|
|
output_file = work_path / "blast_results.txt"
|
||
|
|
|
||
|
|
num_hits = random.randint(5, 15)
|
||
|
|
with open(output_file, "w") as f:
|
||
|
|
for i in range(num_hits):
|
||
|
|
evalue = 10 ** (-random.randint(3, 30))
|
||
|
|
identity = random.uniform(30, 95)
|
||
|
|
f.write(f"query\tsubject{i+1}\t{identity:.1f}\t100\t85\t0\t1\t100\t{evalue:.2e}\t200\n")
|
||
|
|
|
||
|
|
outputs["results"] = str(output_file)
|
||
|
|
stdout_lines.append(f"[SIMULATION] Found {num_hits} similar sequences")
|
||
|
|
|
||
|
|
elif tool.id == "mafft":
|
||
|
|
output_file = work_path / "aligned.fasta"
|
||
|
|
|
||
|
|
seq_file = inputs.get("sequences", "")
|
||
|
|
aligned_seqs = ">gene1_aligned\nM-A-V-Q-K-L-I-R\n>gene2_aligned\nM-A-G-R-K-L-C-F\n>gene3_aligned\nM-S-D-R-Q-Q-L-V"
|
||
|
|
|
||
|
|
with open(output_file, "w") as f:
|
||
|
|
f.write(aligned_seqs)
|
||
|
|
|
||
|
|
outputs["alignment"] = str(output_file)
|
||
|
|
stdout_lines.append(f"[SIMULATION] Aligned sequences from {seq_file}")
|
||
|
|
|
||
|
|
elif tool.id == "iqtree":
|
||
|
|
tree_file = work_path / "iqtree.treefile"
|
||
|
|
log_file = work_path / "iqtree.log"
|
||
|
|
|
||
|
|
tree_content = "(gene1:0.1,gene2:0.2,(gene3:0.15,gene4:0.18):0.05);"
|
||
|
|
with open(tree_file, "w") as f:
|
||
|
|
f.write(tree_content)
|
||
|
|
|
||
|
|
with open(log_file, "w") as f:
|
||
|
|
f.write(f"IQ-TREE (simulated)\n")
|
||
|
|
f.write(f"Best-fit model: {parameters.get('model', 'MFP')}\n")
|
||
|
|
f.write(f"Bootstrap: {parameters.get('bootstrap', 1000)}\n")
|
||
|
|
|
||
|
|
outputs["tree"] = str(tree_file)
|
||
|
|
outputs["log"] = str(log_file)
|
||
|
|
stdout_lines.append(f"[SIMULATION] Built phylogenetic tree")
|
||
|
|
|
||
|
|
elif tool.id == "meme":
|
||
|
|
meme_dir = work_path / "meme_out"
|
||
|
|
meme_dir.mkdir(exist_ok=True)
|
||
|
|
|
||
|
|
xml_file = meme_dir / "meme.xml"
|
||
|
|
with open(xml_file, "w") as f:
|
||
|
|
f.write('<?xml version="1.0"?>\n<meme>\n')
|
||
|
|
for i in range(parameters.get("nmotifs", 10)):
|
||
|
|
f.write(f'<motif id="motif_{i+1}" width="{random.randint(6, 20)}" evalue="{10**-random.randint(5,20)}"/>\n')
|
||
|
|
f.write('</meme>\n')
|
||
|
|
|
||
|
|
outputs["motifs"] = str(xml_file)
|
||
|
|
stdout_lines.append(f"[SIMULATION] Found {parameters.get('nmotifs', 10)} motifs")
|
||
|
|
|
||
|
|
elif tool.id == "mcscanx":
|
||
|
|
collinearity = work_path / "output.collinearity"
|
||
|
|
tandem = work_path / "output.tandem"
|
||
|
|
|
||
|
|
with open(collinearity, "w") as f:
|
||
|
|
f.write("### Collinearity blocks (simulated)\n")
|
||
|
|
f.write("chr1\tgene1\tgene2\tchr2\tgene3\tgene4\n")
|
||
|
|
|
||
|
|
with open(tandem, "w") as f:
|
||
|
|
f.write("# Tandem duplications (simulated)\n")
|
||
|
|
|
||
|
|
outputs["collinearity"] = str(collinearity)
|
||
|
|
outputs["tandem"] = str(tandem)
|
||
|
|
stdout_lines.append(f"[SIMULATION] Analyzed synteny")
|
||
|
|
|
||
|
|
return ExecutionResult(
|
||
|
|
tool_id=tool.id,
|
||
|
|
status=ExecutionStatus.COMPLETED,
|
||
|
|
return_code=0,
|
||
|
|
stdout="\n".join(stdout_lines),
|
||
|
|
stderr="",
|
||
|
|
outputs=outputs,
|
||
|
|
started_at=started_at,
|
||
|
|
completed_at=datetime.now(),
|
||
|
|
)
|
||
|
|
|
||
|
|
def _collect_outputs(self, tool: ToolMetadata, work_path: Path) -> dict[str, str]:
|
||
|
|
outputs = {}
|
||
|
|
|
||
|
|
output_mappings = {
|
||
|
|
"tblout": "output.tbl",
|
||
|
|
"domtblout": "dom.tbl",
|
||
|
|
"tree": "iqtree.treefile",
|
||
|
|
"log": "iqtree.log",
|
||
|
|
"alignment": "aligned.fasta",
|
||
|
|
"results": "blast_results.txt",
|
||
|
|
"motifs": "meme_out/meme.xml",
|
||
|
|
"collinearity": "output.collinearity",
|
||
|
|
"tandem": "output.tandem",
|
||
|
|
}
|
||
|
|
|
||
|
|
for output_def in tool.outputs:
|
||
|
|
if output_def.name in output_mappings:
|
||
|
|
path = work_path / output_mappings[output_def.name]
|
||
|
|
if path.exists():
|
||
|
|
outputs[output_def.name] = str(path)
|
||
|
|
|
||
|
|
for f in work_path.glob("*.treefile"):
|
||
|
|
outputs["tree"] = str(f)
|
||
|
|
|
||
|
|
return outputs
|