#!/usr/bin/env python3 """run_extension.py —— 扩展分析容器通用入口(宿主挂载 /opt/bioflow-entry,镜像零改动)。 执行形态(契约 §3): python -B /opt/bioflow-entry/run_extension.py --ext-dir /job/ext \ --input-dir /job/input --output-dir /job/output [--param name=value ...] 职责: 1. 读 /extension.json(command_template 唯一事实源,契约 §2 schema)。 2. render_command_template:移植 175 executor 模板渲染纯函数(白名单占位符 12 个 + {param_*};{output_dir} 必含)。 3. subprocess 执行(无 shell:shlex 切分;stdout/stderr 透传;超时进程组 SIGTERM → 3s → SIGKILL,与 176 服务侧 terminate→kill 约定同源)。 4. 输出契约校验:输出目录非空,否则退出码 3。 退出码约定(容器内零第三方依赖,仅 python 标准库): 0 成功 1.. 命令自身退出码透传(1/2/7/137…) 2 extension.json 缺失/非法、模板渲染失败、白名单外占位符、缺 {output_dir} 3 输出契约校验失败(输出目录为空) 124 执行超时(进程组已终止) """ from __future__ import annotations import argparse import json import os import re import shlex import signal import subprocess import sys from pathlib import Path from typing import Any EXIT_OK = 0 EXIT_SETUP = 2 EXIT_OUTPUT_CONTRACT = 3 EXIT_TIMEOUT = 124 DEFAULT_TIMEOUT_SEC = 36000.0 # 与 176 单阶段 10h 超时对齐 # 占位符白名单(spec_v2 PLACEHOLDERS 移植;扩展 {param_}) PLACEHOLDERS = { "code_dir", "input_dir", "input_file", "input_hmm", "input_fasta", "input_query_fasta", "input_reference_fasta", "input_pep_fasta", "input_cds_fasta", "input_id_file", "output_dir", "work_dir", } _ID_SUFFIXES = {".id", ".ids", ".geneid", ".geneids", ".txt", ".tsv", ".csv"} _FASTA_SUFFIXES = {".fa", ".fasta", ".faa", ".fsa", ".pep"} _GFF_SUFFIXES = {".gff", ".gff3"} _ACTIVE_PROC: subprocess.Popen | None = None class ExtensionRenderError(RuntimeError): """模板渲染失败(占位符缺失/参数缺值)。""" # --------------------------------------------------------------------------- # 模板渲染纯函数(executor.py 移植裁剪;不 import gfa_core,零第三方依赖) # --------------------------------------------------------------------------- def _safe_parameter_name(value: Any) -> str: text = re.sub(r"[^A-Za-z0-9_]+", "_", str(value or "").strip()) text = re.sub(r"_+", "_", text).strip("_") if not text: return "" if text[0].isdigit(): text = "p_" + text if text.startswith("param_"): text = text[6:] return text def _shell_parameter_value(value: Any) -> str: if isinstance(value, bool): return "true" if value else "false" if value is None: return "" text = str(value) if re.fullmatch(r"[A-Za-z0-9_./:+,@%-]+", text): return text return shlex.quote(text) def _first_input_by_suffix(paths: list[Path], suffixes: set[str]) -> str: for path in paths: name = path.name.lower() if path.suffix.lower() in suffixes or any(name.endswith(s) for s in suffixes): return path.as_posix() return "" def _input_files_by_suffix(paths: list[Path], suffixes: set[str]) -> list[Path]: matched = [] for path in paths: name = path.name.lower() if path.suffix.lower() in suffixes or any(name.endswith(s) for s in suffixes): matched.append(path) return matched def _first_input_by_name(paths: list[Path], tokens: set[str]) -> str: for path in paths: name = path.name.lower() if any(token in name for token in tokens): return path.as_posix() return "" def _largest_file(paths: list[Path]) -> str: if not paths: return "" return max(paths, key=lambda p: p.stat().st_size if p.exists() else -1).as_posix() def _smallest_file(paths: list[Path]) -> str: if not paths: return "" return min(paths, key=lambda p: p.stat().st_size if p.exists() else 10**18).as_posix() def _role_of(extension: dict[str, Any], placeholder: str) -> str: """input_roles 双形态兼容:契约 §2 平铺 {"input_file": "fastq"},175 旧版 {"input_file": {"role": "generic_input_file"}}。仅用于语义占位符消歧。""" roles = extension.get("input_roles") if not isinstance(roles, dict): return "" spec = roles.get(placeholder) if isinstance(spec, dict): return str(spec.get("role") or "").lower() return str(spec or "").lower() def _used_placeholders(template: str) -> set[str]: return set(re.findall(r"\{([A-Za-z_][A-Za-z0-9_]*)\}", template or "")) def _validate_template(template: str) -> str | None: """模板白名单校验(与 spec_v2 normalize 同纪律);None = 通过。""" if not template.strip(): return "extension.json 缺少 command_template" if "{output_dir}" not in template: return "command_template 必须包含 {output_dir}(输出契约硬约束)" bad = sorted( name for name in _used_placeholders(template) if name not in PLACEHOLDERS and not name.startswith("param_") ) if bad: return "command_template 含白名单外占位符: " + ", ".join(bad) return None def build_render_context( extension: dict[str, Any], *, run_dir: Path, input_dir: Path, input_files: list[Path], parameters: dict[str, Any] | None = None, output_dir: Path | None = None, code_dir: Path | None = None, ) -> dict[str, Any]: """渲染上下文:12 个白名单占位符(语义占位符按输入文件后缀/角色推导)+ param_*。""" output_dir = output_dir or (run_dir / "outputs") code_dir = code_dir or (run_dir / "code") input_file_role = _role_of(extension, "input_file") id_file = _first_input_by_suffix(input_files, _ID_SUFFIXES) annotation_file = ( _first_input_by_suffix(input_files, _GFF_SUFFIXES) if ("annotation_gff" in input_file_role or "gff" in input_file_role) else "" ) cis_element_file = ( _first_input_by_name(input_files, {"plantcare", "cis_element", "cis-elements", "cis.elements"}) if input_file_role == "generic_input_file" else "" ) input_file = ( annotation_file or cis_element_file or id_file or (input_files[0].as_posix() if input_files else input_dir.as_posix()) ) input_hmm = _first_input_by_suffix(input_files, {".hmm"}) or input_file fasta_paths = _input_files_by_suffix(input_files, _FASTA_SUFFIXES) input_fasta = fasta_paths[0].as_posix() if fasta_paths else input_file query_fasta = _largest_file(fasta_paths) or input_fasta reference_fasta = ( _smallest_file([p for p in fasta_paths if p.as_posix() != query_fasta]) if len(fasta_paths) > 1 else input_fasta ) pep_fasta = _first_input_by_name(fasta_paths, {"pep", "protein", "faa"}) or query_fasta cds_fasta = _first_input_by_name(fasta_paths, {"cds"}) or input_fasta render_context: dict[str, Any] = { "code_dir": code_dir.as_posix(), "input_dir": input_dir.as_posix(), "input_file": input_file, "input_hmm": input_hmm, "input_fasta": input_fasta, "input_query_fasta": query_fasta, "input_reference_fasta": reference_fasta, "input_pep_fasta": pep_fasta, "input_cds_fasta": cds_fasta, "input_id_file": id_file or input_file, "output_dir": output_dir.as_posix(), "work_dir": run_dir.as_posix(), } for name, value in (parameters or {}).items(): render_context["param_" + _safe_parameter_name(name)] = _shell_parameter_value(value) return render_context def render_command_template( extension: dict[str, Any], *, template: str, run_dir: Path, input_dir: Path, input_files: list[Path], parameters: dict[str, Any] | None = None, output_dir: Path | None = None, code_dir: Path | None = None, ) -> str: """模板渲染纯函数(executor.py 同签名移植)。占位符缺值 → ExtensionRenderError。""" context = build_render_context( extension, run_dir=run_dir, input_dir=input_dir, input_files=input_files, parameters=parameters, output_dir=output_dir, code_dir=code_dir, ) try: return template.format(**context) except KeyError as exc: missing = str(exc).strip("'\"") raise ExtensionRenderError(f"Missing command template value: {missing}") from exc # --------------------------------------------------------------------------- # 执行与输出契约 # --------------------------------------------------------------------------- def _output_contract_ok(output_dir: Path) -> bool: """输出契约:输出目录存在且非空(至少一个条目)。""" return output_dir.is_dir() and any(output_dir.iterdir()) def _find_code_dir(ext_dir: Path) -> Path | None: """代码目录:converted_code 优先(契约 §1 布局),raw_code 兜底。""" for name in ("converted_code", "raw_code"): candidate = ext_dir / name if candidate.is_dir(): return candidate return None def _copy_code_tree(src: Path, dst: Path) -> set[str]: """把代码树复制进可写目录;返回原始代码相对路径集合(收割时排除)。""" import shutil rels: set[str] = set() if not src.is_dir(): return rels for item in src.rglob("*"): rel = item.relative_to(src).as_posix() dest = dst / rel if item.is_dir(): dest.mkdir(parents=True, exist_ok=True) else: dest.parent.mkdir(parents=True, exist_ok=True) shutil.copy2(item, dest) rels.add(rel) return rels def _harvest_outputs(src: Path, dst: Path, exclude: set[str]) -> None: """把 src 内产物(排除原始代码路径)合并进 dst,保留相对路径。""" import shutil for item in src.rglob("*"): rel = item.relative_to(src).as_posix() if rel in exclude: continue # 排除路径是原始代码文件所在目录时也跳过其祖先目录本身 if any(rel == e or rel.startswith(e + "/") for e in exclude): continue dest = dst / rel if item.is_dir(): dest.mkdir(parents=True, exist_ok=True) else: dest.parent.mkdir(parents=True, exist_ok=True) shutil.copy2(item, dest) def _kill_group(pid: int, sig: int) -> None: try: os.killpg(pid, sig) except (ProcessLookupError, PermissionError): pass def _sigterm_handler(signum: int, frame: Any) -> None: """docker stop 的 SIGTERM:转发到子进程组后以 124 退出(日志可辨)。""" proc = _ACTIVE_PROC if proc is not None and proc.poll() is None: _kill_group(proc.pid, signal.SIGTERM) sys.exit(EXIT_TIMEOUT) def _run_command(command: list[str], cwd: Path, env: dict[str, str], timeout_sec: float) -> int: """执行命令(stdout/stderr 直接透传容器流;超时进程组 SIGTERM→3s→SIGKILL)。 可执行文件缺失(镜像缺工具)→ 可读退出码 2(不抛 traceback)。""" global _ACTIVE_PROC try: proc = subprocess.Popen(command, cwd=str(cwd), env=env, start_new_session=True) except FileNotFoundError: print(f"[run_extension] 可执行文件缺失(镜像 PATH 中不存在): {command[0]}", file=sys.stderr) return EXIT_SETUP _ACTIVE_PROC = proc try: return proc.wait(timeout=timeout_sec) except subprocess.TimeoutExpired: _kill_group(proc.pid, signal.SIGTERM) try: proc.wait(timeout=3) except subprocess.TimeoutExpired: _kill_group(proc.pid, signal.SIGKILL) proc.wait() return EXIT_TIMEOUT finally: _ACTIVE_PROC = None def _parse_params(items: list[str]) -> dict[str, Any]: parameters: dict[str, Any] = {} for item in items: name, sep, value = item.partition("=") if not sep or not name.strip(): raise ExtensionRenderError(f"invalid --param (需要 name=value): {item}") parameters[name.strip()] = value return parameters def main(argv: list[str] | None = None) -> int: parser = argparse.ArgumentParser( prog="run_extension", description="扩展分析容器通用入口(契约 EXTENSION_CONTRACT §3)", ) parser.add_argument("--ext-dir", default="/job/ext", help="扩展目录(挂载 :ro)") parser.add_argument("--input-dir", default="/job/input", help="输入目录") parser.add_argument("--output-dir", default="/job/output", help="输出目录(契约校验对象)") parser.add_argument("--work-dir", default="", help="工作目录(默认=输出目录)") parser.add_argument("--timeout", type=float, default=DEFAULT_TIMEOUT_SEC, help="执行超时(秒)") parser.add_argument("--param", action="append", default=[], help="参数 name=value(可重复)") args = parser.parse_args(argv) try: signal.signal(signal.SIGTERM, _sigterm_handler) except ValueError: # 非主线程(测试)不支持 handler pass ext_dir = Path(args.ext_dir) ext_json_path = ext_dir / "extension.json" if not ext_json_path.is_file(): print(f"[run_extension] extension.json 缺失: {ext_json_path}", file=sys.stderr) return EXIT_SETUP try: extension = json.loads(ext_json_path.read_text(encoding="utf-8")) except ValueError as exc: print(f"[run_extension] extension.json 解析失败: {exc}", file=sys.stderr) return EXIT_SETUP if not isinstance(extension, dict): print("[run_extension] extension.json 必须是 JSON 对象", file=sys.stderr) return EXIT_SETUP template = str(extension.get("command_template") or "") error = _validate_template(template) if error: print(f"[run_extension] {error}", file=sys.stderr) return EXIT_SETUP code_dir = _find_code_dir(ext_dir) if code_dir is None: print("[run_extension] 代码目录缺失(converted_code/raw_code 均不存在)", file=sys.stderr) return EXIT_SETUP input_dir = Path(args.input_dir) output_dir = Path(args.output_dir) output_dir.mkdir(parents=True, exist_ok=True) # 2026-09-15 修复:代码挂载只读,但 18 用例语义是「cd 代码目录后脚本写 # outputs/ 相对目录」。等价做法 = 把代码复制进可写 run_dir 执行, # 结束后把 run_dir 产物收割进 output_dir。显式 --work-dir 可覆盖。 work_dir = Path(args.work_dir) if args.work_dir else (output_dir / "_work") work_dir.mkdir(parents=True, exist_ok=True) code_rel_paths = _copy_code_tree(code_dir, work_dir) # 模板中 {code_dir} 指向可写副本(而非只读挂载) effective_code_dir = work_dir try: parameters = _parse_params(args.param) except ExtensionRenderError as exc: print(f"[run_extension] {exc}", file=sys.stderr) return EXIT_SETUP input_files = sorted(p for p in input_dir.rglob("*") if p.is_file()) try: command_text = render_command_template( extension, template=template, run_dir=work_dir, input_dir=input_dir, input_files=input_files, parameters=parameters, output_dir=output_dir, code_dir=effective_code_dir, ) except ExtensionRenderError as exc: print(f"[run_extension] 模板渲染失败: {exc}", file=sys.stderr) return EXIT_SETUP try: command = shlex.split(command_text) except ValueError as exc: print(f"[run_extension] 命令切分失败: {exc}", file=sys.stderr) return EXIT_SETUP if not command: print("[run_extension] 渲染后命令为空", file=sys.stderr) return EXIT_SETUP print("[run_extension] command: " + command_text) sys.stdout.flush() env = os.environ.copy() env.setdefault("PYTHONUNBUFFERED", "1") code = _run_command(command, work_dir, env, args.timeout) if code == EXIT_TIMEOUT: print(f"[run_extension] 执行超时({args.timeout}s),进程组已终止 (exit {EXIT_TIMEOUT})", file=sys.stderr) return EXIT_TIMEOUT if code != EXIT_OK: print(f"[run_extension] command exited with code {code}", file=sys.stderr) return code # 收割:work_dir 产物(排除原始代码文件)合并进 output_dir,然后删除 work 目录 try: _harvest_outputs(work_dir, output_dir, code_rel_paths) import shutil shutil.rmtree(work_dir, ignore_errors=True) except OSError as exc: print(f"[run_extension] 产物收割失败: {exc}", file=sys.stderr) return EXIT_SETUP if not _output_contract_ok(output_dir): print(f"[run_extension] 输出契约校验失败:输出目录为空 {output_dir} (exit {EXIT_OUTPUT_CONTRACT})", file=sys.stderr) return EXIT_OUTPUT_CONTRACT return EXIT_OK if __name__ == "__main__": sys.exit(main())