471 lines
17 KiB
Python
471 lines
17 KiB
Python
#!/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. 读 <ext-dir>/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_<name>})
|
|
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())
|