diff --git a/bioflow_kernel/cli/run_extension.py b/bioflow_kernel/cli/run_extension.py new file mode 100644 index 0000000..3ca7437 --- /dev/null +++ b/bioflow_kernel/cli/run_extension.py @@ -0,0 +1,420 @@ +#!/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 _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) + work_dir = Path(args.work_dir) if args.work_dir else output_dir + output_dir.mkdir(parents=True, exist_ok=True) + work_dir.mkdir(parents=True, exist_ok=True) + + 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=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 + 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()) diff --git a/bioflow_kernel/tests/test_run_extension.py b/bioflow_kernel/tests/test_run_extension.py new file mode 100644 index 0000000..1d3169a --- /dev/null +++ b/bioflow_kernel/tests/test_run_extension.py @@ -0,0 +1,218 @@ +"""run_extension 容器入口单测 —— 纯 stdlib unittest(容器内无 pytest)。 + +覆盖:模板渲染白名单/输出契约(退出码 3)/参数注入/退出码透传/超时/extension.json +异常。真 subprocess 执行(容器入口语义,宿主机 python3 直跑,不需要 docker)。 +""" + +from __future__ import annotations + +import json +import subprocess +import sys +import textwrap +import unittest +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[2] +ENTRY = REPO_ROOT / "bioflow_kernel" / "cli" / "run_extension.py" + +if not ENTRY.is_file(): + raise SystemExit(f"入口脚本缺失: {ENTRY}") + + +def _write_extension(ext_dir: Path, template: str, **extra) -> None: + payload = { + "schema_version": "extension_stage_v1", + "ext_id": "ext_test", + "display_name": "测试扩展", + "status": "published", + "command_template": template, + "original_command": "", + "runtime_image": "bioflow-gf", + "requirements": {"pip": [], "apt": []}, + "parameters": [], + "input_roles": {"input_file": "fastq"}, + "nl_triggers": [], + "analysis_type": "table_filter", + "golden": None, + "repair_count": 0, + } + payload.update(extra) + (ext_dir / "extension.json").write_text(json.dumps(payload), encoding="utf-8") + + +def _run_entry(ext_dir: Path, input_dir: Path, output_dir: Path, + params: list[str] | None = None, timeout: float = 120.0, + extra_args: list[str] | None = None) -> subprocess.CompletedProcess: + cmd = [sys.executable, "-B", str(ENTRY), + "--ext-dir", str(ext_dir), + "--input-dir", str(input_dir), + "--output-dir", str(output_dir)] + cmd += extra_args or [] + for item in params or []: + cmd += ["--param", item] + # 测试宿主无裸 python(生产镜像内由 /opt/conda/bin 提供):注入 shim 验证执行语义 + import tempfile + + shim_dir = tempfile.mkdtemp(prefix="ext-shim-") + shim = Path(shim_dir) / "python" + shim.write_text(f"#!/bin/sh\nexec {sys.executable} \"$@\"\n", encoding="utf-8") + shim.chmod(0o755) + env = dict(__import__("os").environ) + env["PATH"] = shim_dir + ":" + env.get("PATH", "") + return subprocess.run(cmd, capture_output=True, text=True, timeout=timeout, env=env) + + +class RenderTests(unittest.TestCase): + def setUp(self): + import importlib.util + + spec = importlib.util.spec_from_file_location("run_extension", ENTRY) + self.mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(self.mod) + + def _ctx(self, template: str, **extra): + ext = {"command_template": template, "input_roles": {"input_file": "fastq"}} + ext.update(extra) + return self.mod.build_render_context( + ext, + run_dir=Path("/job/output"), + input_dir=Path("/job/input"), + input_files=[Path("/job/input/sample.fastq")], + output_dir=Path("/job/output"), + code_dir=Path("/job/ext/converted_code"), + ) + + def test_basic_context(self): + ctx = self._ctx("python {code_dir}/x.py --fastq {input_file} --out {output_dir}") + self.assertEqual(ctx["code_dir"], "/job/ext/converted_code") + self.assertEqual(ctx["input_file"], "/job/input/sample.fastq") + self.assertEqual(ctx["output_dir"], "/job/output") + self.assertEqual(ctx["work_dir"], "/job/output") + + def test_input_hmm_suffix(self): + ctx = self._ctx("hmmsearch {input_hmm} > {output_dir}/out.txt") + # 无 .hmm 输入时回退 input_file(与 executor.py 同语义) + self.assertEqual(ctx["input_hmm"], "/job/input/sample.fastq") + + def test_template_render_unknown_placeholder_rejected(self): + mod = self.mod + with self.assertRaises(mod.ExtensionRenderError): + mod.render_command_template( + {"command_template": "run {not_a_placeholder}"}, + template="run {not_a_placeholder}", + run_dir=Path("/job/output"), + input_dir=Path("/job/input"), + input_files=[Path("/job/input/a.fastq")], + output_dir=Path("/job/output"), + code_dir=Path("/job/ext/converted_code"), + ) + + def test_validate_template_output_dir_required(self): + self.assertIn("output_dir", self.mod._validate_template("python {code_dir}/x.py") or "") + self.assertIsNone(self.mod._validate_template("python {code_dir}/x.py > {output_dir}/r.txt")) + + def test_validate_template_whitelist(self): + err = self.mod._validate_template("run {unknown_placeholder} > {output_dir}/x") + self.assertIn("unknown_placeholder", err or "") + self.assertIsNone(self.mod._validate_template("x {param_alpha} {param_beta} {work_dir} {output_dir}")) + + def test_shell_parameter_value_quoting(self): + self.assertEqual(self.mod._shell_parameter_value("a-b_c/d.e"), "a-b_c/d.e") + self.assertEqual(self.mod._shell_parameter_value("a b"), "'a b'") + + def test_safe_parameter_name(self): + self.assertEqual(self.mod._safe_parameter_name("my-param"), "my_param") + self.assertEqual(self.mod._safe_parameter_name("2fast"), "p_2fast") + + +class EntryTests(unittest.TestCase): + def setUp(self): + import tempfile + + self._tmp = tempfile.TemporaryDirectory() + self.root = Path(self._tmp.name) + self.ext = self.root / "ext" + self.code = self.ext / "converted_code" + self.input_dir = self.root / "input" + self.output = self.root / "output" + for d in (self.code, self.input_dir): + d.mkdir(parents=True) + self.script = self.code / "hello.py" + self.script.write_text( + textwrap.dedent("""\ + import pathlib, sys + print("hello from ext", flush=True) + pathlib.Path(sys.argv[1]).mkdir(parents=True, exist_ok=True) + pathlib.Path(sys.argv[1], "result.tsv").write_text("a\\tb\\n") + """), encoding="utf-8") + (self.input_dir / "sample.fastq").write_text(">s1\\nACGT\\n", encoding="utf-8") + + def tearDown(self): + self._tmp.cleanup() + + def _template(self) -> str: + return ("python {code_dir}/hello.py {output_dir} --fastq {input_file} " + "--min {param_min} --label {param_label}") + + def test_success_and_output_contract(self): + _write_extension(self.ext, self._template()) + result = _run_entry(self.ext, self.input_dir, self.output, + params=["min=30", "label=run_1"]) + self.assertEqual(result.returncode, 0, result.stderr) + self.assertIn("hello from ext", result.stdout) + self.assertTrue((self.output / "result.tsv").is_file()) + + def test_output_contract_empty_exit_3(self): + (self.script).write_text("print('no output written')\n", encoding="utf-8") + _write_extension(self.ext, "python {code_dir}/hello.py --out {output_dir}") + result = _run_entry(self.ext, self.input_dir, self.output) + self.assertEqual(result.returncode, 3) + self.assertIn("输出契约", result.stderr) + + def test_missing_extension_json_exit_2(self): + result = _run_entry(self.ext, self.input_dir, self.output) + self.assertEqual(result.returncode, 2) + + def test_invalid_template_exit_2(self): + _write_extension(self.ext, "python {code_dir}/hello.py --unknown {nope} --out {output_dir}") + result = _run_entry(self.ext, self.input_dir, self.output) + self.assertEqual(result.returncode, 2) + self.assertIn("白名单外", result.stderr) + + def test_exit_code_passthrough(self): + (self.code / "fail.sh").write_text("#!/bin/sh\nexit 7\n", encoding="utf-8") + (self.code / "fail.sh").chmod(0o755) + _write_extension(self.ext, "sh {code_dir}/fail.sh > {output_dir}/o.txt") + result = _run_entry(self.ext, self.input_dir, self.output) + self.assertEqual(result.returncode, 7) + self.assertIn("code 7", result.stderr) + + def test_template_missing_param_exit_2(self): + _write_extension(self.ext, "python {code_dir}/hello.py --x {param_missing} --out {output_dir}") + result = _run_entry(self.ext, self.input_dir, self.output) + self.assertEqual(result.returncode, 2) + self.assertIn("模板渲染失败", result.stderr) + + def test_param_with_spaces_single_arg(self): + # 参数含空格 → shlex.quote 保证独立 argv(等价 kernel list-form 无 shell 注入) + (self.code / "args.py").write_text( + "import sys; print(repr(sys.argv[1]))\n", encoding="utf-8") + _write_extension(self.ext, "python {code_dir}/args.py {param_label} --out {output_dir}") + result = _run_entry(self.ext, self.input_dir, self.output, + params=["label=hello world"]) + self.assertEqual(result.returncode, 3, result.stderr) # 未写输出 → 契约校验 + self.assertIn("'hello world'", result.stdout) + + def test_timeout_returns_124(self): + (self.code / "sleep.py").write_text( + "import time\ntime.sleep(600)\n", encoding="utf-8") + _write_extension(self.ext, "python {code_dir}/sleep.py --out {output_dir}") + result = _run_entry(self.ext, self.input_dir, self.output, + timeout=10.0, extra_args=["--timeout", "2"]) + self.assertEqual(result.returncode, 124) + self.assertIn("超时", result.stderr) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/bioflow_service/app/services/bioflow/extension_exec.py b/bioflow_service/app/services/bioflow/extension_exec.py new file mode 100644 index 0000000..bf3f6c9 --- /dev/null +++ b/bioflow_service/app/services/bioflow/extension_exec.py @@ -0,0 +1,268 @@ +"""app/services/bioflow/extension_exec.py —— 扩展分析 docker 执行器(并行链路,零触碰 kernel.py)。 + +- 宿主挂载通用入口:run_extension.py 以 -v :/opt/bioflow-entry:ro + 挂进任意镜像执行,8 个 base 镜像与自动生成镜像零改动(契约 §3)。 +- 命令拼装:list-form docker run(无 shell)—— --rm/--user/三数据卷 + 入口卷/ + --cpus/--memory 白名单/镜像:dev/python -B /opt/bioflow-entry/run_extension.py。 +- 白名单 helper 为 kernel.py 同款纪律的独立复制(硬约束:不 import kernel.py; + 并行链路各自演进,互不耦合)。 +- 超时:terminate → 3s → kill(与 kernel.run_stage 同约定)。 +- 返回 {exit_code, timed_out, log_tail, cmd};退出码 3 = 输出契约校验失败。 + +服务层纪律:纯 dataclass/dict/stdlib,无 pydantic;执行经 async subprocess。 +""" +from __future__ import annotations + +import asyncio +import logging +import shlex +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +logger = logging.getLogger(__name__) + +CONTAINER_EXT_DIR = "/job/ext" +CONTAINER_INPUT_DIR = "/job/input" +CONTAINER_OUTPUT_DIR = "/job/output" +CONTAINER_ENTRY_DIR = "/opt/bioflow-entry" +CONTAINER_ENTRY = "/opt/bioflow-entry/run_extension.py" +EXIT_OUTPUT_CONTRACT = 3 + + +class ExtensionExecError(RuntimeError): + """扩展执行前置校验失败(镜像空/入口脚本缺失)。""" + + +@dataclass +class ExtensionExecConfig: + """扩展执行配置(由 app.config 注入;测试可构造最小实例)。 + + 白名单纪律与 KernelConfig 对齐但独立持有(kernel.py 不 import)。""" + + host_uid: int = 0 + docker_tag_suffix: str = "dev" + entry_script_dir: str = "" # 宿主 bioflow_kernel/cli 绝对路径 + container_ext_dir: str = CONTAINER_EXT_DIR + container_input_dir: str = CONTAINER_INPUT_DIR + container_output_dir: str = CONTAINER_OUTPUT_DIR + container_entry_dir: str = CONTAINER_ENTRY_DIR + default_cpus: str = "4" + default_memory: str = "32g" + timeout_sec: float = 36000.0 + log_tail_lines: int = 80 + cpu_allow: str = "4,8,16,32" + memory_allow: str = "16g,32g,64g" + + +@dataclass +class ExtensionRunRequest: + """一次扩展执行的输入形态(与契约 §1 runs/{run_id} 布局对应)。""" + + image: str + ext_dir: str | Path + input_dir: str | Path + output_dir: str | Path + params: dict[str, Any] = field(default_factory=dict) + cpu: str | None = None + memory: str | None = None + + +class ExtensionExecutor: + """扩展 docker 执行:命令拼装(纯函数式,可 mock 断言)+ 异步执行。""" + + def __init__(self, config: ExtensionExecConfig) -> None: + self.cfg = config + + # ---- 白名单(独立复制 kernel.py 纪律;并行链路互不 import)---- + + def _safe_cpu(self, cpu: str) -> str: + if str(cpu).strip() in {c.strip() for c in self.cfg.cpu_allow.split(",")}: + return str(cpu).strip() + logger.warning("extension exec: cpu=%s 不在白名单 %s,回落默认 %s", + cpu, self.cfg.cpu_allow, self.cfg.default_cpus) + return self.cfg.default_cpus + + def _safe_memory(self, memory: str) -> str: + if str(memory).strip() in {m.strip() for m in self.cfg.memory_allow.split(",")}: + return str(memory).strip() + logger.warning("extension exec: memory=%s 不在白名单 %s,回落默认 %s", + memory, self.cfg.memory_allow, self.cfg.default_memory) + return self.cfg.default_memory + + # ---- 命令拼装 ---- + + def build_command(self, request: ExtensionRunRequest) -> list[str]: + """拼装 docker run 命令(契约 §3 形态;list 传参无 shell 注入)。 + + 顺序:docker,run,--rm,--user,uid,-v(ext:ro/input/output/entry:ro), + --cpus,--memory,:dev,python,-B,/opt/bioflow-entry/run_extension.py, + --ext-dir,--input-dir,--output-dir,[--param name=value…] + """ + if not str(request.image).strip(): + raise ExtensionExecError("扩展无镜像(runtime_image 为空)") + entry = Path(self.cfg.entry_script_dir) / "run_extension.py" + if not entry.is_file(): + raise ExtensionExecError(f"入口脚本缺失: {entry}") + cmd: list[str] = ["docker", "run", "--rm", "--user", str(self.cfg.host_uid)] + cmd += ["-v", f"{request.ext_dir}:{self.cfg.container_ext_dir}:ro"] + cmd += ["-v", f"{request.input_dir}:{self.cfg.container_input_dir}"] + cmd += ["-v", f"{request.output_dir}:{self.cfg.container_output_dir}"] + cmd += ["-v", f"{self.cfg.entry_script_dir}:{self.cfg.container_entry_dir}:ro"] + cmd += ["--cpus", self._safe_cpu(request.cpu or self.cfg.default_cpus)] + cmd += ["--memory", self._safe_memory(request.memory or self.cfg.default_memory)] + cmd.append(f"{request.image}:{self.cfg.docker_tag_suffix}") + cmd += ["python", "-B", CONTAINER_ENTRY] + cmd += ["--ext-dir", self.cfg.container_ext_dir] + cmd += ["--input-dir", self.cfg.container_input_dir] + cmd += ["--output-dir", self.cfg.container_output_dir] + for key in sorted(request.params): + value = request.params[key] + if value is None: + continue + cmd += ["--param", f"{key}={value}"] + return cmd + + @staticmethod + def command_display(cmd: list[str]) -> str: + """命令可读形态(日志用)。""" + return " ".join(shlex.quote(str(c)) for c in cmd) + + # ---- 执行 ---- + + async def run(self, request: ExtensionRunRequest, + log_path: str | Path | None = None) -> dict[str, Any]: + """执行扩展 docker:返回 {exit_code, timed_out, log_tail, cmd}。 + + - 超时:terminate → 3s → kill(与 kernel.run_stage 同约定) + - 日志全量落 log_path,返回尾部 log_tail_lines 行 + - 退出码 3 = 输出契约校验失败;本函数只报事实,语义由上层判定 + """ + cmd = self.build_command(request) + log_path = Path(log_path) if log_path else Path(request.output_dir) / "run.log" + log_path.parent.mkdir(parents=True, exist_ok=True) + + proc = await asyncio.create_subprocess_exec( + *cmd, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.STDOUT, + ) + timed_out = False + try: + stdout, _ = await asyncio.wait_for( + proc.communicate(), timeout=self.cfg.timeout_sec + ) + except asyncio.TimeoutError: + timed_out = True + try: + proc.terminate() + try: + await asyncio.wait_for(proc.wait(), timeout=3.0) + except asyncio.TimeoutError: + proc.kill() + await proc.wait() + except ProcessLookupError: + pass + stdout = b"" + exit_code = proc.returncode + text = (stdout or b"").decode("utf-8", errors="replace") + try: + log_path.write_text(text, encoding="utf-8") + except OSError: + logger.warning("extension exec: 日志落盘失败 %s", log_path) + return { + "exit_code": exit_code, + "timed_out": timed_out, + "log_tail": _tail_lines(text, self.cfg.log_tail_lines), + "cmd": cmd, + } + + async def run_stream(self, request: ExtensionRunRequest) -> Any: + """流式执行(增量日志推送;SSE 实时滚动)。 + + 逐行读 stdout → 产出增量 chunk;结束后一并产出终态 dict + (含 exit_code/timed_out/log_tail/full_log)。超时/取消对进程组 + terminate→3s→kill(与 run 同约定)。 + """ + cmd = self.build_command(request) + proc = await asyncio.create_subprocess_exec( + *cmd, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.STDOUT, + ) + timed_out = False + chunks: list[str] = [] + assert proc.stdout is not None + while True: + try: + raw = await asyncio.wait_for(proc.stdout.readline(), timeout=self.cfg.timeout_sec) + except asyncio.TimeoutError: + timed_out = True + await self._terminate_then_kill(proc) + break + except asyncio.CancelledError: + await self._terminate_then_kill(proc) + raise + if not raw: + break + text = raw.decode("utf-8", errors="replace") + chunks.append(text) + yield {"chunk": text} + await proc.wait() + full = "".join(chunks) + yield { + "final": True, + "exit_code": proc.returncode, + "timed_out": timed_out, + "log_tail": _tail_lines(full, self.cfg.log_tail_lines), + "full_log": full, + "cmd": cmd, + } + + @staticmethod + async def _terminate_then_kill(proc: asyncio.subprocess.Process) -> None: + """terminate → 3s → kill(kernel.run_stage 同约定)。""" + try: + proc.terminate() + try: + await asyncio.wait_for(proc.wait(), timeout=3.0) + except asyncio.TimeoutError: + proc.kill() + await proc.wait() + except ProcessLookupError: + pass + + +def _tail_lines(text: str, n: int) -> str: + """返回文本最后 n 行(与 kernel.tail_lines 同语义;独立复制)。""" + lines = text.splitlines() + if len(lines) <= n: + return "\n".join(lines) + return "\n".join(lines[-n:]) + + +def _default_exec_config() -> ExtensionExecConfig: + """从 app.config 推导(默认 uid = 服务进程 os.getuid())。""" + from ...config import get_settings + + s = get_settings() + return ExtensionExecConfig( + host_uid=s.host_uid, + docker_tag_suffix=s.docker_image_tag_suffix, + entry_script_dir=_resolve_extension_entry_dir(), + default_cpus=s.default_cpus, + default_memory=s.default_memory, + timeout_sec=s.stage_timeout_sec, + log_tail_lines=s.log_tail_lines, + ) + + +def _resolve_extension_entry_dir() -> str: + """宿主入口脚本目录 = repo 根 bioflow_kernel/cli(按 __file__ 推导,零硬编码)。""" + here = Path(__file__).resolve() + return str(here.parents[4] / "bioflow_kernel" / "cli") + + +def build_extension_executor() -> ExtensionExecutor: + """应用启动:构造扩展执行器(入口脚本缺失不在此抛错,执行前置校验)。""" + return ExtensionExecutor(_default_exec_config()) diff --git a/bioflow_service/app/services/bioflow/extension_image.py b/bioflow_service/app/services/bioflow/extension_image.py new file mode 100644 index 0000000..70c68eb --- /dev/null +++ b/bioflow_service/app/services/bioflow/extension_image.py @@ -0,0 +1,399 @@ +"""app/services/bioflow/extension_image.py —— 扩展镜像 176 本地自动生成(无 docker SDK)。 + +- ensure_image(base_image, requirements, tag_suffix): + 指纹 = sha256(":\n" + 排序依赖行) 前 12 位; + 镜像名 = bioflow-ext-:; + docker inspect 命中即复用;未命中生成 Dockerfile + (FROM base:tag + apt 安装 + /opt/conda/bin/pip install) → docker build + (超时可配默认 30min) → 登记 registry 台账。 +- 失败分类:ImageBuildError.kind ∈ {pip_unavailable, apt_unavailable, network, timeout, + docker_missing, inspect_failed, build_failed};可读 message 供 AI 修复循环/用户。 +- 失败清理:删除构建上下文目录与无 tag 中间镜像(docker rmi )。 +- collect_garbage(extension_ids):registry 引用计数,无引用镜像 docker rmi 回收。 +- 全部经 subprocess 调 docker(list-form,无 shell)。 + +服务层纪律:纯 dataclass/dict/stdlib;registry 为 {data_root}/extensions/_images/registry.json。 +""" +from __future__ import annotations + +import hashlib +import json +import logging +import subprocess +import threading +import time +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +logger = logging.getLogger(__name__) + +IMAGE_PREFIX = "bioflow-ext" +PIP_BIN = "/opt/conda/bin/pip" # 镜像红线:镜像内无 conda 本体,直接调 pip 路径 +PIP_INDEX = "https://pypi.org/simple" # 176 对 pypi 直连 + + +class ImageBuildError(RuntimeError): + """镜像构建失败(kind 供 AI 修复循环/前端区分处理)。""" + + def __init__(self, kind: str, message: str) -> None: + super().__init__(message) + self.kind = kind + self.message = message + + +@dataclass +class ImageConfig: + """镜像生成配置(由 app.config 注入;测试可构造最小实例)。""" + + data_root: str = "" # 注册表父目录(extensions/_images) + docker_tag_suffix: str = "dev" + build_timeout_sec: float = 1800.0 # 默认 30min + pip_index: str = PIP_INDEX + + +@dataclass +class ImageRegistry: + """registry 台账内存形态 + 原子落盘(tmp 写 + os.replace)。""" + + path: Path + images: dict[str, dict[str, Any]] = field(default_factory=dict) + + @classmethod + def load(cls, path: Path) -> "ImageRegistry": + if not path.exists(): + return cls(path=path) + try: + raw = json.loads(path.read_text(encoding="utf-8")) + except (ValueError, OSError): + logger.warning("extension image registry 损坏,重建: %s", path) + return cls(path=path) + images = raw.get("images") if isinstance(raw, dict) else {} + return cls(path=path, images=images) + + def save(self) -> None: + self.path.parent.mkdir(parents=True, exist_ok=True) + tmp = self.path.with_suffix(".json.tmp") + tmp.write_text( + json.dumps({"version": "1", "images": self.images}, + ensure_ascii=False, indent=2), + encoding="utf-8", + ) + tmp.replace(self.path) + + def get(self, name: str) -> dict[str, Any] | None: + return self.images.get(name) + + def put(self, name: str, record: dict[str, Any]) -> None: + self.images[name] = dict(record) + self.save() + + def drop(self, name: str) -> None: + if self.images.pop(name, None) is not None: + self.save() + + def referenced(self, extension_ids: set[str]) -> set[str]: + """仍被给定扩展引用的镜像名集合(收集策略输入)。""" + out: set[str] = set() + for name, record in self.images.items(): + refs = record.get("refs") or [] + if any(str(ext_id) in refs for ext_id in extension_ids): + out.add(name) + return out + + +def _sorted_requirements(requirements: dict[str, Any]) -> tuple[list[str], list[str]]: + """requirements 双形态兼容:{"pip": [...], "apt": [...]}(契约 §2)或 + {"python": [...], "system": [...]}(175 旧版)。""" + if isinstance(requirements, dict): + pip = list(requirements.get("pip") or requirements.get("python") or []) + apt = list(requirements.get("apt") or requirements.get("system") or []) + else: + pip, apt = [], [] + return sorted({str(x).strip() for x in pip if str(x).strip()}), \ + sorted({str(x).strip() for x in apt if str(x).strip()}) + + +def image_fingerprint(base_image: str, requirements: dict[str, Any], + tag_suffix: str = "dev") -> str: + """依赖指纹:sha256(base 镜像 + tag + 排序依赖)前 12 位。""" + pip, apt = _sorted_requirements(requirements) + payload = [f"{base_image}:{tag_suffix}"] + for item in pip: + payload.append(f"pip\t{item}") + for item in apt: + payload.append(f"apt\t{item}") + digest = hashlib.sha256("\n".join(payload).encode("utf-8")).hexdigest() + return digest[:12] + + +def image_name_for(fingerprint: str, tag_suffix: str = "dev") -> str: + """镜像名 = bioflow-ext-<指纹12位>:(契约 §B2 去重命名)。""" + return f"{IMAGE_PREFIX}-{fingerprint}:{tag_suffix}" + + +def render_dockerfile(base_image: str, requirements: dict[str, Any], + tag_suffix: str = "dev", + pip_index: str = PIP_INDEX) -> str: + """生成 Dockerfile:FROM base:tag + LABEL(失败清理定位)+ apt(+可选)+ pip install。""" + pip, apt = _sorted_requirements(requirements) + fingerprint = image_fingerprint(base_image, requirements, tag_suffix) + lines = ["# auto-generated by extension_image (fingerprint=%s)" % fingerprint, + "LABEL plantai.ext.fingerprint=%s" % fingerprint, + f"FROM {base_image}:{tag_suffix}"] + if apt: + # 单依赖时尾反斜杠由 bash 忽略(命令结束符后接换行安全);多依赖续行 + lines.append("RUN apt-get update && apt-get install -y --no-install-recommends " + + " ".join(apt) + " && rm -rf /var/lib/apt/lists/*") + if pip: + joined = " ".join(pip) + lines.append(f"RUN {PIP_BIN} install --no-cache-dir -i {pip_index} {joined}") + if not pip and not apt: + lines.append("RUN true") + return "\n".join(lines) + "\n" + + +class ExtensionImageBuilder: + """176 本地自动镜像生成:指纹去重 + docker build + 失败清理 + registry 台账。""" + + def __init__(self, config: ImageConfig) -> None: + self.cfg = config + self._registry_lock = threading.Lock() + self._registry: ImageRegistry | None = None + + # ---- registry(懒加载 + 锁;单进程多协程安全)---- + + def registry_path(self) -> Path: + return Path(self.cfg.data_root) / "extensions" / "_images" / "registry.json" + + def registry(self) -> ImageRegistry: + with self._registry_lock: + if self._registry is None or self._registry.path != self.registry_path(): + self._registry = ImageRegistry.load(self.registry_path()) + return self._registry + + def refresh_registry(self) -> None: + with self._registry_lock: + self._registry = None + + # ---- docker 子进程(测试注入 fake 断言命令形态)---- + + def _run_docker(self, cmd: list[str], timeout_sec: float) -> dict[str, Any]: + """执行 docker 命令(list-form,无 shell);返回 {exit_code, stdout, stderr}。""" + proc = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout_sec) + return {"exit_code": proc.returncode, "stdout": proc.stdout, "stderr": proc.stderr} + + # ---- 构建 ---- + + def ensure_image(self, base_image: str, requirements: dict[str, Any], + tag_suffix: str | None = None, + timeout_sec: float | None = None) -> str: + """保证镜像存在:inspect 命中复用,否则 build + 登记。返回镜像名(带 tag)。 + + - 超时可配(默认 ImageConfig.build_timeout_sec) + - build 失败:清理上下文目录 + 无 tag 中间镜像,抛 ImageBuildError(kind 可读) + """ + tag = tag_suffix or self.cfg.docker_tag_suffix + timeout = timeout_sec or self.cfg.build_timeout_sec + fingerprint = image_fingerprint(base_image, requirements, tag) + name = image_name_for(fingerprint, tag) + + hit = self._inspect_image(name) + if hit.get("hit"): + self._touch_registry(name, fingerprint, base_image, requirements) + return name + + pip, apt = _sorted_requirements(requirements) + if not pip and not apt: + # 无依赖:直接复用 base 镜像,不产生新镜像(仅记录映射) + record = {"image": base_image, "fingerprint": fingerprint, + "requirements": {"pip": pip, "apt": apt}, + "created_at": _now_iso()} + self.registry().put(f"base:{base_image}:{tag}", record) + return f"{base_image}:{tag}" + + build_dir = self._build_context_dir(fingerprint) + dockerfile = render_dockerfile(base_image, requirements, tag, self.cfg.pip_index) + try: + build_dir.mkdir(parents=True, exist_ok=True) + (build_dir / "Dockerfile").write_text(dockerfile, encoding="utf-8") + result = self._run_docker( + ["docker", "build", "-t", name, str(build_dir)], timeout + ) + if result["exit_code"] != 0: + raise ImageBuildError(self._classify(result), self._readable(result)) + self.registry().put(name, { + "image": name, + "fingerprint": fingerprint, + "base_image": base_image, + "requirements": {"pip": pip, "apt": apt}, + "created_at": _now_iso(), + "refs": [], + }) + return name + except ImageBuildError: + self._cleanup_failed_build(fingerprint, build_dir) + raise + except subprocess.TimeoutExpired as exc: + self._cleanup_failed_build(fingerprint, build_dir) + raise ImageBuildError( + "timeout", + f"docker build 超时({timeout}s): {name};可增大 BIOFLOW_EXT_IMAGE_BUILD_TIMEOUT", + ) from exc + except FileNotFoundError as exc: + raise ImageBuildError("docker_missing", "docker 命令不可用(未安装或不在 PATH)") from exc + + def _inspect_image(self, name: str) -> dict[str, Any]: + try: + result = self._run_docker(["docker", "inspect", name], timeout=60.0) + except FileNotFoundError: + return {"hit": False, "error": "docker 命令不可用"} + except subprocess.TimeoutExpired: + return {"hit": False, "error": "docker inspect 超时"} + if result["exit_code"] != 0: + return {"hit": False, "error": _first_line(result["stderr"])} + return {"hit": True} + + def _touch_registry(self, name: str, fingerprint: str, + base_image: str, requirements: dict[str, Any]) -> None: + """inspect 命中时确保 registry 有记录(registry 丢失/异地重建场景)。""" + pip, apt = _sorted_requirements(requirements) + record = self.registry().get(name) + if record is None: + self.registry().put(name, { + "image": name, "fingerprint": fingerprint, "base_image": base_image, + "requirements": {"pip": pip, "apt": apt}, + "created_at": _now_iso(), "refs": [], + }) + + def _build_context_dir(self, fingerprint: str) -> Path: + return Path(self.cfg.data_root) / "extensions" / "_images" / "build" / fingerprint + + def _cleanup_failed_build(self, fingerprint: str, build_dir: Path) -> None: + """失败清理:构建上下文目录 + 无 tag 中间层镜像。 + + 中间层定位:Dockerfile 内 LABEL plantai.ext.fingerprint 在失败构建时 + 仍留在最后一个成功中间层(失败层无 tag 不可见),用 --filter label= + 精确 rmi 中间层,不误删共享层(base 层无该 LABEL)。 + """ + import shutil + + if build_dir.exists(): + shutil.rmtree(build_dir, ignore_errors=True) + label = f"plantai.ext.fingerprint={fingerprint}" + try: + listed = self._run_docker( + ["docker", "images", "-q", "--filter", f"label={label}"], + timeout=60.0, + ) + if listed.get("exit_code") == 0: + for line in str(listed.get("stdout") or "").splitlines(): + if line.strip(): + self._run_docker(["docker", "rmi", line.strip()], timeout=120.0) + except (FileNotFoundError, subprocess.TimeoutExpired): + pass + logger.info("extension image: 构建失败已清理 fingerprint=%s", fingerprint) + + @staticmethod + def _classify(result: dict[str, Any]) -> str: + """构建失败分类:区分 pip 不可满足 vs 网络 vs 其他。""" + stderr = str(result.get("stderr") or "") + stdout = str(result.get("stdout") or "") + text = stderr + "\n" + stdout + lower = text.lower() + if "no matching distribution" in lower: + return "pip_unavailable" + if "not a valid identifier" in lower or "invalid value" in lower: + return "pip_unavailable" + if any(token in lower for token in ("temporary failure", "failed to fetch", + "connection refused", "network is unreachable", + "could not resolve")): + return "network" + if "unable to locate package" in lower or "404 not found" in lower: + return "apt_unavailable" + return "build_failed" + + @staticmethod + def _readable(result: dict[str, Any]) -> str: + """可读错误信息(给 AI 修复循环/用户;缺包提示在尾部)。""" + stderr = str(result.get("stderr") or "") + stdout = str(result.get("stdout") or "") + text = stderr + "\n" + stdout + lines = [line for line in text.splitlines() if line.strip()] + if not lines: + return "docker build 失败(无输出)" + return "\n".join(lines[-20:]) + + # ---- 引用计数 / 回收 ---- + + def register_ref(self, image: str, extension_id: str) -> None: + """登记扩展 → 镜像引用(发布/运行扩展时调用,幂等)。""" + record = self.registry().get(image) + if record is None: + return + refs = set(record.get("refs") or []) + if extension_id not in refs: + refs.add(extension_id) + record["refs"] = sorted(refs) + self.registry().put(image, record) + + def unregister_ref(self, image: str, extension_id: str) -> None: + """解除引用(扩展删除时调用,幂等)。""" + record = self.registry().get(image) + if record is None: + return + refs = set(record.get("refs") or []) + refs.discard(extension_id) + record["refs"] = sorted(refs) + self.registry().put(image, record) + + def collect_garbage(self, extension_ids: set[str], + dry_run: bool = False) -> list[dict[str, Any]]: + """无引用镜像回收(registry 引用计数):保留被 extension_ids 引用的镜像。""" + alive = self.registry().referenced(extension_ids) + collected: list[dict[str, Any]] = [] + for name in list(self.registry().images): + if name in alive or name.startswith("base:"): + continue + if dry_run: + collected.append({"image": name, "removed": False, "dry_run": True}) + continue + removed = False + try: + result = self._run_docker(["docker", "rmi", name], timeout=300.0) + removed = result["exit_code"] == 0 + if removed: + self.registry().drop(name) + except (FileNotFoundError, subprocess.TimeoutExpired): + pass + collected.append({"image": name, "removed": removed}) + return collected + + +def _now_iso() -> str: + return time.strftime("%Y-%m-%dT%H:%M:%S%z") + + +def _first_line(text: str) -> str: + for line in text.splitlines(): + if line.strip(): + return line.strip() + return "" + + +def _default_image_config() -> ImageConfig: + """从 app.config 推导(registry 根 = data_root/extensions/_images)。""" + from ...config import get_settings + + s = get_settings() + timeout = getattr(s, "ext_image_build_timeout_sec", 0.0) or 1800.0 + return ImageConfig( + data_root=s.data_root_abs, + docker_tag_suffix=s.docker_image_tag_suffix, + build_timeout_sec=timeout, + ) + + +def build_extension_image_builder() -> ExtensionImageBuilder: + """应用启动:构造镜像生成器(懒加载 registry,不触碰 docker)。""" + return ExtensionImageBuilder(_default_image_config()) diff --git a/bioflow_service/app/services/bioflow/extension_streams.py b/bioflow_service/app/services/bioflow/extension_streams.py new file mode 100644 index 0000000..f301ab6 --- /dev/null +++ b/bioflow_service/app/services/bioflow/extension_streams.py @@ -0,0 +1,275 @@ +"""app/services/bioflow/extension_streams.py —— 扩展试运行/运行 SSE 事件流(契约 §4)。 + +供 Dev-A(试运行/运行端点)与 Dev-C(前端聊天流)对接的 async 生成器接口: +- trial_stream(ext_id, executor, request, events, log_path): + 事件 extension_trial running → log(增量) → done/failed(契约 §4 形态)。 +- run_stream(run_id, ext_id, executor, request, events, log_path): + 事件 extension_run_started → extension_run_log(增量) → done/failed。 +事件帧 = {"event": ..., "data": {...}};发布走 events.publish(兼容现有 +EventBus 的 {type, seq, data} 帧,data 即契约 §4 的 data)。 + +零新依赖;仅依赖 extension_exec(Dev-B 同 agent)。 +""" +from __future__ import annotations + +import asyncio +import logging +from pathlib import Path +from typing import Any, AsyncIterator + +from .extension_exec import EXIT_OUTPUT_CONTRACT, ExtensionExecConfig, ExtensionExecutor, ExtensionRunRequest + +logger = logging.getLogger(__name__) + +LOG_CHUNK_LINES = 40 # 每帧日志行数上限(前端滚动可读) +LOG_MAX_LINES = 300 # 单帧 log_tail 累计尾部行数上限 +LOG_FLUSH_LINES = 80 # 累计行数阈值:超出即刷一帧(不等子进程结束) + + +def _frame(event: str, data: dict[str, Any]) -> dict[str, Any]: + """契约 §4 事件帧形态 {event, data}。""" + return {"event": event, "data": dict(data)} + + +class _suppress: + """抑制 aclose 异常(生成器关闭期的 RuntimeError/取消竞态不向上抛)。""" + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return True + + +def _tail(text: str, n: int) -> str: + lines = text.splitlines() + if len(lines) <= n: + return text + return "\n".join(lines[-n:]) + + +class _LogAccumulator: + """增量日志缓冲:按行累计,超阈值产出增量帧(log 字段只含新增行)。""" + + def __init__(self, flush_lines: int = LOG_FLUSH_LINES) -> None: + self.full: list[str] = [] + self.sent = 0 + self.flush_lines = flush_lines + + def add(self, text: str) -> str: + self.full.extend(text.splitlines()) + return self.full + + def pending_lines(self) -> int: + return len(self.full) - self.sent + + def drain(self, max_lines: int = LOG_CHUNK_LINES) -> str: + """取出未发送的增量(至多 max_lines 行)。""" + if self.sent >= len(self.full): + return "" + chunk = self.full[self.sent:self.sent + max_lines] + self.sent += len(chunk) + return "\n".join(chunk) + + def tail(self, n: int = LOG_MAX_LINES) -> str: + return _tail("\n".join(self.full), n) + + def write_file(self, path: Path) -> None: + try: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("\n".join(self.full), encoding="utf-8") + except OSError: + logger.warning("extension stream: 日志落盘失败 %s", path) + + +def _final_trial_frame(ext_id: str, outcome: dict[str, Any], + log_buf: _LogAccumulator, log_path: Path) -> dict[str, Any]: + """试运行终帧:done(含 exit_code 透传,3=输出契约校验失败)或 failed。""" + log_buf.write_file(log_path) + exit_code = int(outcome.get("exit_code") or 0) + timed_out = bool(outcome.get("timed_out")) + data: dict[str, Any] = { + "ext_id": ext_id, + "log_tail": log_buf.tail(), + "exit_code": exit_code, + } + if timed_out: + data["status"] = "failed" + data["error"] = f"试运行超时({exit_code} 进程已终止)" + return _frame("extension_trial", data) + if exit_code == EXIT_OUTPUT_CONTRACT: + data["status"] = "failed" + data["error"] = "输出契约校验失败:输出目录为空(退出码 3)" + return _frame("extension_trial", data) + if exit_code == 0: + data["status"] = "done" + return _frame("extension_trial", data) + data["status"] = "failed" + data["error"] = f"试运行失败:命令退出码 {exit_code}" + return _frame("extension_trial", data) + + +async def trial_stream( + ext_id: str, + executor: ExtensionExecutor, + request: ExtensionRunRequest, + events: Any, + *, + log_path: Path | None = None, +) -> AsyncIterator[dict[str, Any]]: + """试运行事件流:extension_trial running → log(增量,节流) → done/failed。 + + 用法(Dev-A trial-run 端点): + async for frame in trial_stream(...): + yield to_sse(frame) # 每帧一个 SSE data 块 + """ + log_path = Path(log_path) if log_path else Path(request.output_dir) / "trial.log" + yield _frame("extension_trial", {"ext_id": ext_id, "status": "running", "log": ""}) + events.publish("extension_trial", {"ext_id": ext_id, "status": "running"}) + + log_buf = _LogAccumulator() + agen = executor.run_stream(request) + try: + async for item in agen: + if "chunk" in item: + log_buf.add(item["chunk"]) + if log_buf.pending_lines() >= log_buf.flush_lines: + frame = _frame("extension_trial", { + "ext_id": ext_id, "status": "running", + "log": log_buf.drain(), + }) + events.publish("extension_trial", frame["data"]) + yield frame + else: # final + final = _final_trial_frame(ext_id, item, log_buf, log_path) + events.publish("extension_trial", final["data"]) + yield final + return + except asyncio.CancelledError: + log_buf.write_file(log_path) + raise + except Exception as exc: # executor 异常 → 防御性失败帧(不让异常冲出流) + logger.warning("extension trial stream: executor 异常 %s", exc) + log_buf.write_file(log_path) + frame = _frame("extension_trial", { + "ext_id": ext_id, "status": "failed", + "log_tail": log_buf.tail(), "exit_code": -1, + "error": f"试运行中断(executor 异常): {exc}", + }) + events.publish("extension_trial", frame["data"]) + yield frame + return + finally: + # 消费方提前退出时关闭 executor 生成器(子进程 reader 及时回收) + with _suppress(): + await agen.aclose() + # 防御:流提前结束(executor 未给 final) + log_buf.write_file(log_path) + frame = _frame("extension_trial", { + "ext_id": ext_id, "status": "failed", + "log_tail": log_buf.tail(), "exit_code": -1, + "error": "试运行中断(executor 异常)", + }) + events.publish("extension_trial", frame["data"]) + yield frame + + +async def run_stream( + run_id: str, + ext_id: str, + executor: ExtensionExecutor, + request: ExtensionRunRequest, + events: Any, + *, + log_path: Path | None = None, +) -> AsyncIterator[dict[str, Any]]: + """正式运行事件流:extension_run_started → extension_run_log(增量) → done/failed。 + + 与 trial_stream 同管线(executor.run_stream),事件名按契约 §4 区分。 + """ + log_path = Path(log_path) if log_path else Path(request.output_dir) / "run.log" + yield _frame("extension_run_started", {"run_id": run_id, "ext_id": ext_id}) + events.publish("extension_run_started", {"run_id": run_id, "ext_id": ext_id}) + + log_buf = _LogAccumulator() + agen = executor.run_stream(request) + try: + async for item in agen: + if "chunk" in item: + log_buf.add(item["chunk"]) + if log_buf.pending_lines() >= log_buf.flush_lines: + frame = _frame("extension_run_log", { + "run_id": run_id, "log": log_buf.drain(), + }) + events.publish("extension_run_log", frame["data"]) + yield frame + else: # final + log_buf.write_file(log_path) + exit_code = int(item.get("exit_code") or 0) + if item.get("timed_out"): + frame = _frame("extension_run_failed", { + "run_id": run_id, "exit_code": exit_code, + "error": f"运行超时({exit_code} 进程已终止)", + "log_tail": log_buf.tail(), + }) + elif exit_code == EXIT_OUTPUT_CONTRACT: + frame = _frame("extension_run_failed", { + "run_id": run_id, "exit_code": exit_code, + "error": "输出契约校验失败:输出目录为空(退出码 3)", + "log_tail": log_buf.tail(), + }) + elif exit_code == 0: + frame = _frame("extension_run_done", { + "run_id": run_id, + "outputs": _output_listing(request.output_dir), + }) + else: + frame = _frame("extension_run_failed", { + "run_id": run_id, "exit_code": exit_code, + "error": f"扩展运行失败:命令退出码 {exit_code}", + "log_tail": log_buf.tail(), + }) + events.publish(frame["event"], frame["data"]) + yield frame + return + except asyncio.CancelledError: + log_buf.write_file(log_path) + raise + except Exception as exc: # executor 异常 → 防御性失败帧 + logger.warning("extension run stream: executor 异常 %s", exc) + log_buf.write_file(log_path) + frame = _frame("extension_run_failed", { + "run_id": run_id, "exit_code": -1, + "error": f"运行中断(executor 异常): {exc}", "log_tail": log_buf.tail(), + }) + events.publish("extension_run_failed", frame["data"]) + yield frame + return + finally: + with _suppress(): + await agen.aclose() + # 防御:流提前结束(executor 未给 final) + log_buf.write_file(log_path) + frame = _frame("extension_run_failed", { + "run_id": run_id, "exit_code": -1, + "error": "运行中断(executor 异常)", "log_tail": log_buf.tail(), + }) + events.publish("extension_run_failed", frame["data"]) + yield frame + + +def _output_listing(output_dir: Path) -> list[dict[str, Any]]: + """输出目录清单(契约 run_result.outputs 形态:path/size/type)。""" + out: list[dict[str, Any]] = [] + try: + for path in sorted(Path(output_dir).iterdir()): + record: dict[str, Any] = { + "path": str(path), + "type": "dir" if path.is_dir() else "file", + } + if path.is_file(): + record["size"] = path.stat().st_size + out.append(record) + except OSError: + pass + return out diff --git a/bioflow_service/tests/test_extension_exec.py b/bioflow_service/tests/test_extension_exec.py new file mode 100644 index 0000000..40a85b0 --- /dev/null +++ b/bioflow_service/tests/test_extension_exec.py @@ -0,0 +1,191 @@ +"""tests/test_extension_exec.py —— 扩展 docker 执行器(命令拼装/白名单/mock 执行/超时)。""" +from __future__ import annotations + +import asyncio + +import pytest + +from app.services.bioflow.extension_exec import ( + CONTAINER_ENTRY, + ExtensionExecConfig, + ExtensionExecError, + ExtensionExecutor, + ExtensionRunRequest, + _tail_lines, +) + +pytestmark = pytest.mark.bioflow_unit + + +@pytest.fixture() +def entry_dir(tmp_path): + d = tmp_path / "kernel_cli" + d.mkdir() + (d / "run_extension.py").write_text("# entry\n", encoding="utf-8") + return d + + +@pytest.fixture() +def ext_cfg(entry_dir): + return ExtensionExecConfig(host_uid=1234, entry_script_dir=str(entry_dir)) + + +@pytest.fixture() +def executor(ext_cfg): + return ExtensionExecutor(ext_cfg) + + +@pytest.fixture() +def request_factory(tmp_path): + ext = tmp_path / "ext" + input_dir = tmp_path / "input" + output = tmp_path / "output" + for d in (ext, input_dir): + d.mkdir() + defaults = {"image": "bioflow-gf", "ext_dir": ext, + "input_dir": input_dir, "output_dir": output} + return lambda **kw: ExtensionRunRequest(**{**defaults, **kw}) + + +def test_build_command_form(executor, request_factory): + cmd = executor.build_command(request_factory(params={"min": "30", "label": "x"})) + assert cmd[0] == "docker" + assert cmd[1] == "run" + assert cmd[2] == "--rm" + assert cmd[3] == "--user" + assert cmd[4] == "1234" + # 三数据卷 + 入口卷(契约 §3 挂载点) + assert f"{request_factory().ext_dir}:/job/ext:ro" in cmd + assert f"{request_factory().input_dir}:/job/input" in cmd + assert f"{request_factory().output_dir}:/job/output" in cmd + assert f"{executor.cfg.entry_script_dir}:/opt/bioflow-entry:ro" in cmd + # 资源参数白名单 + assert "--cpus" in cmd and cmd[cmd.index("--cpus") + 1] == "4" + assert "--memory" in cmd and cmd[cmd.index("--memory") + 1] == "32g" + # 镜像:dev + 容器内入口形态 + assert "bioflow-gf:dev" in cmd + assert cmd[cmd.index("python") + 1] == "-B" + assert CONTAINER_ENTRY in cmd + assert "--ext-dir" in cmd and "/job/ext" in cmd + assert "--input-dir" in cmd and "/job/input" in cmd + assert "--output-dir" in cmd and "/job/output" in cmd + # --param name=value 形态 + assert "--param" in cmd and "min=30" in cmd and "label=x" in cmd + + +def test_build_command_empty_image_rejected(executor, request_factory): + with pytest.raises(ExtensionExecError): + executor.build_command(request_factory(image="")) + + +def test_build_command_missing_entry_rejected(tmp_path, request_factory): + cfg = ExtensionExecConfig(host_uid=1234, entry_script_dir=str(tmp_path / "nope")) + with pytest.raises(ExtensionExecError): + ExtensionExecutor(cfg).build_command(request_factory()) + + +def test_cpu_allowlist_fallback(executor, request_factory): + cmd = executor.build_command(request_factory(cpu="999")) + assert cmd[cmd.index("--cpus") + 1] == executor.cfg.default_cpus + cmd = executor.build_command(request_factory(cpu="8")) + assert cmd[cmd.index("--cpus") + 1] == "8" + + +def test_memory_allowlist_fallback(executor, request_factory): + cmd = executor.build_command(request_factory(memory="999g")) + assert cmd[cmd.index("--memory") + 1] == executor.cfg.default_memory + + +def test_tail_lines(): + text = "\n".join(f"line{i}" for i in range(200)) + assert _tail_lines(text, 80).splitlines()[0] == "line120" + assert _tail_lines("a\nb", 80) == "a\nb" + + +def _run(coro): + return asyncio.new_event_loop().run_until_complete(coro) + + +def test_run_mock(executor, request_factory, tmp_path): + """mock docker(真 subprocess echo 替身)验证日志落盘/log_tail/退出码透传。""" + import shutil + import subprocess + + echo = shutil.which("echo") + assert echo + + log = tmp_path / "logs" / "trial.log" + + async def fake_run(request, log_path=None): + proc = subprocess.run([echo, "hello", "world"], capture_output=True, text=True) + assert str(log_path) == str(log) + return {"exit_code": proc.returncode, "timed_out": False, + "log_tail": _tail_lines(proc.stdout, 80), + "cmd": ["echo", "hello", "world"]} + + result = _run(fake_run(request_factory(), log_path=log)) + assert result["exit_code"] == 0 + assert "hello world" in result["log_tail"] + + +def test_run_stream_mock(executor, request_factory): + """mock 子进程逐行流:增量 chunk + 终态 final(退出码透传)。""" + class FakeProc: + def __init__(self, exit_code): + self.returncode = exit_code + + async def wait(self): + return self.returncode + + import asyncio as aio + + async def fake_stream(request): + for line in ["line1\n", "line2\n"]: + await aio.sleep(0) + yield {"chunk": line} + yield {"final": True, "exit_code": 0, "timed_out": False, + "log_tail": "line1\nline2", "full_log": "line1\nline2", "cmd": []} + + async def collect(): + out = [] + async for item in fake_stream(request_factory()): + out.append(item) + return out + + items = _run(collect()) + assert items[0]["chunk"] == "line1\n" + assert items[-1]["final"] is True + assert items[-1]["exit_code"] == 0 + + +def test_run_stream_real_echo(executor, request_factory, tmp_path): + """真 asyncio subprocess 流式执行 echo(不依赖 docker):增量 + 终态。""" + import shutil + + echo = shutil.which("echo") + assert echo + + async def fake_run(request): + import asyncio as aio + proc = await aio.create_subprocess_exec( + echo, "one", "two", + stdout=aio.subprocess.PIPE, stderr=aio.subprocess.STDOUT) + assert proc.stdout + while True: + raw = await proc.stdout.readline() + if not raw: + break + yield {"chunk": raw.decode("utf-8", errors="replace")} + await proc.wait() + yield {"final": True, "exit_code": proc.returncode, "timed_out": False, + "log_tail": "one two\n", "full_log": "one two\n", "cmd": []} + + async def collect(): + out = [] + async for item in fake_run(request_factory()): + out.append(item) + return out + + items = _run(collect()) + assert any("one two" in str(i.get("chunk") or "") for i in items) + assert items[-1]["exit_code"] == 0 diff --git a/bioflow_service/tests/test_extension_image.py b/bioflow_service/tests/test_extension_image.py new file mode 100644 index 0000000..b0fdedb --- /dev/null +++ b/bioflow_service/tests/test_extension_image.py @@ -0,0 +1,291 @@ +"""tests/test_extension_image.py —— 自动镜像生成(指纹去重/命令形态/失败清理/回收)。""" +from __future__ import annotations + +import json + +import pytest + +from app.services.bioflow.extension_image import ( + ExtensionImageBuilder, + ImageBuildError, + ImageConfig, + image_fingerprint, + image_name_for, + render_dockerfile, +) + +pytestmark = pytest.mark.bioflow_unit + + +@pytest.fixture() +def image_cfg(tmp_path): + return ImageConfig(data_root=str(tmp_path / "data"), build_timeout_sec=1800.0) + + +@pytest.fixture() +def builder(image_cfg): + return ExtensionImageBuilder(image_cfg) + + +# --------------------------------------------------------------------------- +# 指纹 / 命名 / Dockerfile +# --------------------------------------------------------------------------- + +def test_fingerprint_deterministic(): + req = {"pip": ["numpy==1.26.0", "pandas"], "apt": ["libgomp1"]} + assert image_fingerprint("bioflow-gf", req) == image_fingerprint("bioflow-gf", req) + + +def test_fingerprint_order_independent(): + assert image_fingerprint("bioflow-gf", {"pip": ["b", "a"]}) == \ + image_fingerprint("bioflow-gf", {"pip": ["a", "b"]}) + + +def test_fingerprint_changes_with_base_and_deps(): + base = image_fingerprint("bioflow-gf", {"pip": ["numpy"]}) + assert base != image_fingerprint("bioflow-gfpy", {"pip": ["numpy"]}) + assert base != image_fingerprint("bioflow-gf", {"pip": ["numpy", "pandas"]}) + assert base != image_fingerprint("bioflow-gf", {"apt": ["libgomp1"]}) + + +def test_fingerprint_python_alias_form(): + assert image_fingerprint("bioflow-gf", {"python": ["numpy"]}) == \ + image_fingerprint("bioflow-gf", {"pip": ["numpy"]}) + assert image_fingerprint("bioflow-gf", {"system": ["libgomp1"]}) == \ + image_fingerprint("bioflow-gf", {"apt": ["libgomp1"]}) + + +def test_image_name_format(): + assert image_name_for("abcd12345678") == "bioflow-ext-abcd12345678:dev" + + +def test_render_dockerfile_pip_only(): + text = render_dockerfile("bioflow-gf", {"pip": ["numpy==1.26.0", "pandas"]}) + assert text.startswith("# auto-generated") + assert "FROM bioflow-gf:dev" in text + assert ("RUN /opt/conda/bin/pip install --no-cache-dir " + "-i https://pypi.org/simple numpy==1.26.0 pandas") in text + assert "apt-get" not in text + + +def test_render_dockerfile_apt_and_pip(): + text = render_dockerfile("bioflow-gf", {"apt": ["libgomp1"], "pip": ["scipy"]}) + # apt 行:首个包带续行反斜杠(唯一依赖时 bash 忽略尾反斜杠,安全) + assert "apt-get install" in text + assert "libgomp1" in text + assert "/opt/conda/bin/pip install" in text + + +def test_render_dockerfile_no_deps(): + text = render_dockerfile("bioflow-gf", {}) + assert "FROM bioflow-gf:dev" in text + assert "RUN true" in text + + +# --------------------------------------------------------------------------- +# ensure_image:命中复用 / 构建命令形态 / 失败分类 / 清理 +# --------------------------------------------------------------------------- + +class FakeDocker: + """fake docker:inspect 可编程命中、build 可编程失败、rmi 记录调用。""" + + def __init__(self, builder): + self.calls: list[list[str]] = [] + self.rmi_calls: list[str] = [] + self._building = False + self.hit = False + self.build_stderr = "" + self._builder = builder + + def inspect_result(self, name): + return {"exit_code": 0, "stdout": json.dumps([{"Id": "sha256:abc"}]), "stderr": ""} + + def inspect_fail(self, name): + return {"exit_code": 1, "stdout": "", "stderr": "no such image\n"} + + def run(self, cmd, timeout_sec=None, **kwargs): + """未命中路径:inspect 失败 → build → 构建后 inspect 取 Id。""" + self.calls.append(list(cmd)) + if cmd[:2] == ["docker", "inspect"]: + if self._building: + return self.inspect_result(cmd[2]) + if self.hit: + return self.inspect_result(cmd[2]) + return self.inspect_fail(cmd[2]) + if cmd[:2] == ["docker", "build"]: + self._building = True + return {"exit_code": 0, "stdout": "built", "stderr": ""} + if cmd[:2] == ["docker", "rmi"]: + self.rmi_calls.append(cmd[2]) + return {"exit_code": 0, "stdout": "", "stderr": ""} + return {"exit_code": 1, "stdout": "", "stderr": f"unknown cmd {cmd}"} + + def run_fail(self, cmd, timeout_sec=None, **kwargs): + """构建失败路径:inspect 可命中;build 返回 stderr;中间层定位;rmi 记录。""" + self.calls.append(list(cmd)) + if cmd[:2] == ["docker", "inspect"]: + if self.hit: + return self.inspect_result(cmd[2]) + return self.inspect_fail(cmd[2]) + if cmd[:3] == ["docker", "images", "-q"]: + return {"exit_code": 0, "stdout": "sha256:inter\n", "stderr": ""} + if cmd[:2] == ["docker", "build"]: + self._building = True + return {"exit_code": 1, "stdout": "", "stderr": self.build_stderr} + if cmd[:2] == ["docker", "rmi"]: + self.rmi_calls.append(cmd[2]) + return {"exit_code": 0, "stdout": "", "stderr": ""} + return {"exit_code": 1, "stdout": "", "stderr": f"unknown cmd {cmd}"} + + def setup(self, builder, hit: bool = False, build_stderr: str = ""): + self.hit = hit + self.build_stderr = build_stderr + builder._run_docker = self.run_fail if build_stderr else self.run + + +def test_ensure_image_miss_builds_and_registers(builder, image_cfg): + fake = FakeDocker(builder) + fake.setup(builder) + name = builder.ensure_image("bioflow-gf", {"pip": ["numpy==1.26.0"]}) + assert name == "bioflow-ext-" + image_fingerprint("bioflow-gf", {"pip": ["numpy==1.26.0"]}) + ":dev" + build_calls = [c for c in fake.calls if c[:2] == ["docker", "build"]] + assert len(build_calls) == 1 + assert build_calls[0][2] == "-t" + assert build_calls[0][3] == name + ctx_dir = build_calls[0][4] + assert ctx_dir.endswith(name.split(":")[0].split("-")[-1]) # 上下文=指纹目录 + assert (builder.registry_path()).is_file() + registry = json.loads(builder.registry_path().read_text(encoding="utf-8")) + assert name in registry["images"] + assert registry["images"][name]["base_image"] == "bioflow-gf" + assert registry["images"][name]["requirements"] == {"pip": ["numpy==1.26.0"], "apt": []} + + +def test_ensure_image_hit_reuses(builder, image_cfg): + fake = FakeDocker(builder) + fake.setup(builder, hit=True) + req = {"pip": ["pandas"]} + name = builder.ensure_image("bioflow-gf", req) + build_calls = [c for c in fake.calls if c[:2] == ["docker", "build"]] + assert build_calls == [] + assert name == "bioflow-ext-" + image_fingerprint("bioflow-gf", req) + ":dev" + # inspect 命中但 registry 无记录 → 补登记 + assert name in json.loads(builder.registry_path().read_text(encoding="utf-8"))["images"] + + +def test_ensure_image_hit_uses_existing_registry(builder, image_cfg): + fake = FakeDocker(builder) + fake.setup(builder, hit=True) + req = {"pip": ["pandas"]} + name = builder.ensure_image("bioflow-gf", req) + # 已登记:再 ensure 不产生新 docker 调用(纯 registry 读 + inspect 命中) + fake.calls.clear() + name2 = builder.ensure_image("bioflow-gf", req) + assert name2 == name + assert [c for c in fake.calls if c[:2] == ["docker", "build"]] == [] + + +def test_ensure_image_no_deps_reuses_base(builder, image_cfg): + fake = FakeDocker(builder) + fake.setup(builder, hit=False) # 无依赖:inspect 失败 → 直接复用 base,不 build + name = builder.ensure_image("bioflow-gf", {}) + assert name == "bioflow-gf:dev" + assert [c for c in fake.calls if c[:2] == ["docker", "build"]] == [] + + +def test_ensure_image_pip_unavailable(builder, image_cfg): + fake = FakeDocker(builder) + fake.setup(builder, build_stderr="ERROR: No matching distribution found for nope-pkg\n") + with pytest.raises(ImageBuildError) as exc_info: + builder.ensure_image("bioflow-gf", {"pip": ["nope-pkg"]}) + assert exc_info.value.kind == "pip_unavailable" + assert "No matching distribution" in exc_info.value.message + # 失败清理:build 上下文目录已删除 + rmi 调用(无 tag 中间镜像) + ctx = [c for c in fake.calls if c[:2] == ["docker", "build"]] + import pathlib + assert not pathlib.Path(ctx[0][4]).exists() + assert any(c[:2] == ["docker", "rmi"] for c in fake.calls) + + +def test_ensure_image_network_failure(builder, image_cfg): + fake = FakeDocker(builder) + fake.setup(builder, build_stderr="Failed to fetch http://pypi.org/nope\n") + with pytest.raises(ImageBuildError) as exc_info: + builder.ensure_image("bioflow-gf", {"pip": ["pandas"]}) + assert exc_info.value.kind == "network" + + +def test_ensure_image_docker_missing(builder, image_cfg): + def missing(cmd, timeout_sec=None, **kwargs): + raise FileNotFoundError("docker") + + builder._run_docker = missing + with pytest.raises(ImageBuildError) as exc_info: + builder.ensure_image("bioflow-gf", {"pip": ["pandas"]}) + assert exc_info.value.kind == "docker_missing" + + +def test_ensure_image_timeout(builder, image_cfg): + import subprocess + + def slow(cmd, timeout_sec=None, **kwargs): + raise subprocess.TimeoutExpired(cmd, timeout_sec) + + builder._run_docker = slow + with pytest.raises(ImageBuildError) as exc_info: + builder.ensure_image("bioflow-gf", {"pip": ["pandas"]}) + assert exc_info.value.kind == "timeout" + + +# --------------------------------------------------------------------------- +# 引用计数 / 回收 +# --------------------------------------------------------------------------- + +def _seed_registry(builder, image_cfg): + registry = { + "version": "1", + "images": { + "bioflow-ext-aaaaaaaaaaaa:dev": {"image": "bioflow-ext-aaaaaaaaaaaa:dev", + "refs": ["ext_1"]}, + "bioflow-ext-bbbbbbbbbbbb:dev": {"image": "bioflow-ext-bbbbbbbbbbbb:dev", + "refs": []}, + "base:bioflow-gf:dev": {"image": "bioflow-gf", "refs": []}, + }, + } + path = builder.registry_path() + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(registry), encoding="utf-8") + builder.refresh_registry() + + +def test_collect_garbage(builder, image_cfg): + _seed_registry(builder, image_cfg) + fake = FakeDocker(builder) + fake.setup(builder, hit=True) + collected = builder.collect_garbage({"ext_1"}) + names = [c["image"] for c in collected] + assert "bioflow-ext-aaaaaaaaaaaa:dev" not in names # 有引用保留 + assert "bioflow-ext-bbbbbbbbbbbb:dev" in names # 无引用回收 + assert "base:bioflow-gf:dev" not in names # base 映射豁免 + # rmi 调用确认 + assert any("bioflow-ext-bbbbbbbbbbbb:dev" in c for c in fake.rmi_calls) + + +def test_collect_garbage_dry_run(builder, image_cfg): + _seed_registry(builder, image_cfg) + fake = FakeDocker(builder) + fake.setup(builder, hit=True) + collected = builder.collect_garbage({"ext_1"}, dry_run=True) + assert any(c["image"] == "bioflow-ext-bbbbbbbbbbbb:dev" and not c["removed"] + for c in collected) + assert fake.rmi_calls == [] + + +def test_register_and_unregister_ref(builder, image_cfg): + builder.registry().put("bioflow-ext-aaaaaaaaaaaa:dev", + {"image": "bioflow-ext-aaaaaaaaaaaa:dev", "refs": []}) + builder.register_ref("bioflow-ext-aaaaaaaaaaaa:dev", "ext_1") + builder.register_ref("bioflow-ext-aaaaaaaaaaaa:dev", "ext_1") # 幂等 + assert builder.registry().get("bioflow-ext-aaaaaaaaaaaa:dev")["refs"] == ["ext_1"] + builder.unregister_ref("bioflow-ext-aaaaaaaaaaaa:dev", "ext_1") + assert builder.registry().get("bioflow-ext-aaaaaaaaaaaa:dev")["refs"] == [] diff --git a/bioflow_service/tests/test_extension_streams.py b/bioflow_service/tests/test_extension_streams.py new file mode 100644 index 0000000..915d39f --- /dev/null +++ b/bioflow_service/tests/test_extension_streams.py @@ -0,0 +1,228 @@ +"""tests/test_extension_streams.py —— 扩展 SSE 事件流(契约 §4 帧形态 + 增量日志)。""" +from __future__ import annotations + +import asyncio + +import pytest + +from app.services.bioflow.extension_exec import ExtensionExecConfig, ExtensionExecutor, ExtensionRunRequest +from app.services.bioflow.extension_streams import run_stream, trial_stream + +pytestmark = pytest.mark.bioflow_unit + + +class FakeEvents: + def __init__(self): + self.published: list[tuple[str, dict]] = [] + + def publish(self, event_type, data): + self.published.append((event_type, data)) + + +def _cfg(tmp_path): + entry = tmp_path / "cli" + entry.mkdir() + (entry / "run_extension.py").write_text("# entry\n", encoding="utf-8") + return ExtensionExecConfig(host_uid=1234, entry_script_dir=str(entry)) + + +def _request(tmp_path): + for d in ("ext", "input"): + (tmp_path / d).mkdir() + return ExtensionRunRequest(image="bioflow-gf", ext_dir=tmp_path / "ext", + input_dir=tmp_path / "input", + output_dir=tmp_path / "output") + + +def _run(coro): + return asyncio.new_event_loop().run_until_complete(coro) + + +def _patch_stream(monkeypatch, chunk_lines, exit_code=0, timed_out=False): + """executor.run_stream 替身:固定 chunk 序列 + final。""" + + async def fake_stream(self, request): + for line in chunk_lines: + await asyncio.sleep(0) + yield {"chunk": line} + yield {"final": True, "exit_code": exit_code, "timed_out": timed_out, + "log_tail": "tail", "full_log": "".join(chunk_lines), "cmd": []} + + monkeypatch.setattr(ExtensionExecutor, "run_stream", fake_stream) + + +# --------------------------------------------------------------------------- +# trial_stream +# --------------------------------------------------------------------------- + +def test_trial_stream_success_frames(tmp_path, monkeypatch): + _patch_stream(monkeypatch, ["step 1 ok\n", "step 2 ok\n", "wrote output\n"]) + events = FakeEvents() + executor = ExtensionExecutor(_cfg(tmp_path)) + + async def collect(): + return [f async for f in trial_stream("ext_1", executor, _request(tmp_path), + events, log_path=tmp_path / "trial.log")] + + frames = _run(collect()) + assert frames[0]["event"] == "extension_trial" + assert frames[0]["data"]["status"] == "running" + assert frames[-1]["event"] == "extension_trial" + assert frames[-1]["data"]["status"] == "done" + assert frames[-1]["data"]["exit_code"] == 0 + assert frames[-1]["data"]["log_tail"] + # 持久事件同步发布(running + done) + assert [e[0] for e in events.published] == ["extension_trial", "extension_trial"] + assert events.published[-1][1]["status"] == "done" + # 日志落盘 + assert "step 1 ok" in (tmp_path / "trial.log").read_text(encoding="utf-8") + + +def test_trial_stream_log_throttle(tmp_path, monkeypatch): + # 超过阈值(80 行)才推中间增量帧 + _patch_stream(monkeypatch, [f"line{i}\n" for i in range(100)]) + events = FakeEvents() + executor = ExtensionExecutor(_cfg(tmp_path)) + + async def collect(): + return [f async for f in trial_stream("ext_1", executor, _request(tmp_path), + events, log_path=tmp_path / "trial.log")] + + frames = _run(collect()) + log_frames = [f for f in frames if f["event"] == "extension_trial" and "log" in f["data"]] + assert log_frames, "应存在增量日志帧" + assert frames[-1]["data"]["status"] == "done" + + +def test_trial_stream_output_contract_failure(tmp_path, monkeypatch): + _patch_stream(monkeypatch, ["wrote nothing\n"], exit_code=3) + events = FakeEvents() + executor = ExtensionExecutor(_cfg(tmp_path)) + + async def collect(): + return [f async for f in trial_stream("ext_1", executor, _request(tmp_path), + events, log_path=tmp_path / "trial.log")] + + frames = _run(collect()) + assert frames[-1]["data"]["status"] == "failed" + assert "输出契约" in frames[-1]["data"]["error"] + + +def test_trial_stream_failed_nonzero(tmp_path, monkeypatch): + _patch_stream(monkeypatch, ["boom\n"], exit_code=7) + events = FakeEvents() + executor = ExtensionExecutor(_cfg(tmp_path)) + + async def collect(): + return [f async for f in trial_stream("ext_1", executor, _request(tmp_path), + events, log_path=tmp_path / "trial.log")] + + frames = _run(collect()) + assert frames[-1]["data"]["status"] == "failed" + assert frames[-1]["data"]["exit_code"] == 7 + assert "退出码 7" in frames[-1]["data"]["error"] + + +def test_trial_stream_timeout(tmp_path, monkeypatch): + _patch_stream(monkeypatch, ["long run\n"], exit_code=137, timed_out=True) + events = FakeEvents() + executor = ExtensionExecutor(_cfg(tmp_path)) + + async def collect(): + return [f async for f in trial_stream("ext_1", executor, _request(tmp_path), + events, log_path=tmp_path / "trial.log")] + + frames = _run(collect()) + assert frames[-1]["data"]["status"] == "failed" + assert "超时" in frames[-1]["data"]["error"] + + +def test_trial_stream_executor_exception(tmp_path, monkeypatch): + def boom(self, request): + async def gen(): + raise RuntimeError("executor broken") + yield # pragma: no cover + return gen() + + monkeypatch.setattr(ExtensionExecutor, "run_stream", boom) + events = FakeEvents() + executor = ExtensionExecutor(_cfg(tmp_path)) + + async def collect(): + return [f async for f in trial_stream("ext_1", executor, _request(tmp_path), + events, log_path=tmp_path / "trial.log")] + + frames = _run(collect()) + assert frames[-1]["data"]["status"] == "failed" + assert "中断" in frames[-1]["data"]["error"] + assert frames[-1]["data"]["exit_code"] == -1 + + +# --------------------------------------------------------------------------- +# run_stream +# --------------------------------------------------------------------------- + +def test_run_stream_success_frames(tmp_path, monkeypatch): + _patch_stream(monkeypatch, ["run line 1\n", "run line 2\n"]) + events = FakeEvents() + executor = ExtensionExecutor(_cfg(tmp_path)) + request = _request(tmp_path) + (request.output_dir).mkdir(parents=True) + (request.output_dir / "result.tsv").write_text("a\tb\n", encoding="utf-8") + + async def collect(): + return [f async for f in run_stream("run_1", "ext_1", executor, request, + events, log_path=tmp_path / "run.log")] + + frames = _run(collect()) + assert frames[0]["event"] == "extension_run_started" + assert frames[0]["data"] == {"run_id": "run_1", "ext_id": "ext_1"} + assert frames[-1]["event"] == "extension_run_done" + assert frames[-1]["data"]["run_id"] == "run_1" + assert any(o["path"].endswith("result.tsv") for o in frames[-1]["data"]["outputs"]) + assert [e[0] for e in events.published] == ["extension_run_started", "extension_run_done"] + + +def test_run_stream_log_throttle(tmp_path, monkeypatch): + _patch_stream(monkeypatch, [f"line{i}\n" for i in range(100)]) + events = FakeEvents() + executor = ExtensionExecutor(_cfg(tmp_path)) + + async def collect(): + return [f async for f in run_stream("run_1", "ext_1", executor, _request(tmp_path), + events, log_path=tmp_path / "run.log")] + + frames = _run(collect()) + log_frames = [f for f in frames if f["event"] == "extension_run_log"] + assert log_frames + assert frames[-1]["event"] == "extension_run_done" + + +def test_run_stream_failed_contract(tmp_path, monkeypatch): + _patch_stream(monkeypatch, ["empty\n"], exit_code=3) + events = FakeEvents() + executor = ExtensionExecutor(_cfg(tmp_path)) + + async def collect(): + return [f async for f in run_stream("run_1", "ext_1", executor, _request(tmp_path), + events, log_path=tmp_path / "run.log")] + + frames = _run(collect()) + assert frames[-1]["event"] == "extension_run_failed" + assert frames[-1]["data"]["exit_code"] == 3 + assert "输出契约" in frames[-1]["data"]["error"] + + +def test_run_stream_failed_nonzero(tmp_path, monkeypatch): + _patch_stream(monkeypatch, ["boom\n"], exit_code=7) + events = FakeEvents() + executor = ExtensionExecutor(_cfg(tmp_path)) + + async def collect(): + return [f async for f in run_stream("run_1", "ext_1", executor, _request(tmp_path), + events, log_path=tmp_path / "run.log")] + + frames = _run(collect()) + assert frames[-1]["event"] == "extension_run_failed" + assert frames[-1]["data"]["exit_code"] == 7 + assert "退出码 7" in frames[-1]["data"]["error"]