forked from admin/PlantAI-rag-new
feat(extension): docker 执行内核入口 + 176 本地自动镜像生成(并行链路)
- cli/run_extension.py:容器通用入口(宿主挂载 /opt/bioflow-entry,镜像零改动);
extension.json → 模板渲染(白名单 12 占位符 + param_*,{output_dir} 必含)→
subprocess 执行(无 shell,进程组 SIGTERM→3s→SIGKILL)→ 输出契约校验(空目录退出码 3);
零第三方依赖(镜像内标准库)
- extension_exec.py:服务端 docker 执行器(list-form;三数据卷+入口卷;
--cpus/--memory 白名单独立复制不 import kernel.py;terminate→kill 同约定;
返回 {exit_code,timed_out,log_tail,cmd};run_stream 逐行增量)
- extension_image.py:指纹 sha256(base+tag+排序依赖) 前 12 位 →
bioflow-ext-<fp>:dev;inspect 命中复用;Dockerfile FROM base + apt + /opt/conda/bin/pip;
失败分类(pip/apt/network/timeout/docker_missing)+ LABEL 定位中间层清理;
registry 台账 + 引用计数 + collect_garbage
- extension_streams.py:契约 §4 事件流(trial/run async 生成器,增量日志节流,
EventBus 双发,防御性失败帧,aclose 回收)
- 单测 53 用例(15 unittest + 38 pytest)全绿;现有 bioflow 226 用例回归通过
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Claude Fable 5
@@ -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. 读 <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 _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())
|
||||
@@ -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)
|
||||
@@ -0,0 +1,268 @@
|
||||
"""app/services/bioflow/extension_exec.py —— 扩展分析 docker 执行器(并行链路,零触碰 kernel.py)。
|
||||
|
||||
- 宿主挂载通用入口:run_extension.py 以 -v <bioflow_kernel/cli>:/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,<image>: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())
|
||||
@@ -0,0 +1,399 @@
|
||||
"""app/services/bioflow/extension_image.py —— 扩展镜像 176 本地自动生成(无 docker SDK)。
|
||||
|
||||
- ensure_image(base_image, requirements, tag_suffix):
|
||||
指纹 = sha256("<base>:<tag_suffix>\n" + 排序依赖行) 前 12 位;
|
||||
镜像名 = bioflow-ext-<fingerprint>:<tag_suffix>;
|
||||
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 <image_id>)。
|
||||
- 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位>:<tag_suffix>(契约 §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())
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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"] == []
|
||||
@@ -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"]
|
||||
Reference in New Issue
Block a user