Files
PlantAI-rag-new/bioflow_service/app/services/bioflow/kernel.py
T
adminandClaude Fable 5 dcb810a758 feat(bioflow): wgs/pop/gwas 26 阶段参数契约(与 rnaseq 同款机制)
- stage_image_map 加 26 个 wgs/pop/gwas 阶段 cli_contract:
  --workspace/--job-id/--prefix 公共三元组 + 上游 manifest 注入
  (prev_manifest:<stage> 解析为 workspace/final/{prefix}.{stage}.manifest.json)
  + param:<name> 用户参数 + tool:<name> 镜像内工具
- kernel _contract_args 扩展 prev_manifest/tool 两种 source
- catalog 补 wgs/pop/gwas 参数 schema(确认表单收集参考基因组/样本清单/
  表型文件/主成分等)
- 339 用例全绿;已部署 176

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-09-16 15:33:43 +08:00

361 lines
15 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""app/services/bioflow/kernel.py —— Docker 执行器(阶段→镜像映射表驱动)。
职责:
- 读 A1 交付的 stage_image_map.json(路径来自配置 BIOFLOW_IMAGE_MAP_PATH,
config 按 __file__ 推导 repo 根;零硬编码)。
- 按阶段→镜像查表拼装 docker run 命令(list 形式,无 shell;命令拼装可 mock 断言)。
- 容器内统一 `python -B -m cli.run_{stage}`(模块名来自映射表 module 字段)。
- --user 宿主 uid(BIOFLOW_HOST_UID 环境变量,默认服务进程 os.getuid())。
- -v 任务目录 + BIOFLOW_DATA 只读挂载 + -e BBA_* + --cpus/--memory + --rm。
- 超时 kill;日志全量落任务 logs/,返回尾部截断文本。
服务层纪律:纯 dataclass/dict/stdlib,无 pydantic。执行经 async subprocess
(真跑 docker 用默认实现;测试注入 fake 执行器 mock docker)。
"""
from __future__ import annotations
import asyncio
import json
import logging
import os
import shlex
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any
logger = logging.getLogger(__name__)
CONTAINER_CMD_PREFIX = ("python", "-B") # 容器内命令形态(TC-BIO-RUN-037)
class ImageMapError(Exception):
"""镜像映射表缺失/非法(服务可用但所有 docker 执行前置拦截)。"""
class UnknownStageError(Exception):
"""阶段不在映射表(catalog 校验后仍未知 → 执行前置拦截)。"""
@dataclass
class ImageMap:
"""stage_image_map.json 内存形态。"""
images: dict[str, dict[str, Any]] = field(default_factory=dict)
stages: dict[str, dict[str, Any]] = field(default_factory=dict)
family_chain: list[str] = field(default_factory=list)
def image_for(self, stage: str) -> dict[str, Any] | None:
return self.stages.get(stage)
def docker_tag(self, image: str) -> str:
return str(self.images.get(image, {}).get("docker_tag") or image)
def all_stage_names(self) -> list[str]:
return sorted(self.stages.keys())
@dataclass
class KernelConfig:
"""docker 执行配置(由 app.config 注入;测试可构造最小实例)。"""
host_uid: int = 0
docker_tag_suffix: str = "dev"
data_root: str = ""
container_data_dir: str = "/bioflow-data"
container_input_dir: str = "/job/input"
container_workspace_dir: str = "/job/workspace"
default_cpus: str = "4"
default_memory: str = "32g"
stage_timeout_sec: float = 36000.0
log_tail_lines: int = 80
bba_env: dict[str, str] = field(default_factory=dict)
cpu_allow: str = "4,8,16,32"
memory_allow: str = "16g,32g,64g"
extra_mounts: list[str] = field(default_factory=list) # 登记路径只读挂载(src==dst)
kernel_dir: str = "" # 内核树宿主路径(运行时挂载进容器)
container_kernel_dir: str = "/bioflow-kernel" # 容器内内核树路径
class Kernel:
"""Docker 执行内核:映射表加载 + 命令拼装 + 阶段执行。"""
def __init__(self, image_map: ImageMap, config: KernelConfig) -> None:
self.map = image_map
self.cfg = config
# ---- 映射表加载(类方法,供应用启动)----
@classmethod
def load(cls, image_map_path: str | Path, config: KernelConfig) -> "Kernel":
path = Path(image_map_path)
if not path.exists():
raise ImageMapError(f"stage_image_map.json 缺失: {path}")
try:
raw = json.loads(path.read_text(encoding="utf-8"))
except (ValueError, OSError) as exc:
raise ImageMapError(f"stage_image_map.json 解析失败: {exc}") from exc
image_map = ImageMap(
images=raw.get("images") or {},
stages=raw.get("stages") or {},
family_chain=raw.get("family_chain") or [],
)
return cls(image_map, config)
# ---- 查表 ----
def image_for(self, stage: str) -> dict[str, Any] | None:
return self.map.image_for(stage)
def has_stage(self, stage: str) -> bool:
return stage in self.map.stages
# ---- 命令拼装(纯函数式,可 mock 断言)----
def build_command(
self,
stage: str,
job_input_dir: str | Path,
job_workspace_dir: str | Path,
params: dict[str, Any] | None = None,
*,
cpu: str | None = None,
memory: str | None = None,
module: str | None = None,
extra_mounts: list[str] | None = None,
) -> list[str]:
"""拼装 docker run 命令(list,无 shell;TC-BIO-KERN-001 基元)。
顺序:docker,run,--rm,--user,uid,-v(×3),--cpus,--memory,-e(×N),
tag:镜像,python,-B,-m,cli.run_X,[参数…]
"""
params = params or {}
entry = self.map.stages.get(stage)
if entry is None:
raise UnknownStageError(f"阶段不在映射表: {stage}")
image = str(entry.get("image") or "")
if not image:
raise UnknownStageError(f"阶段 {stage} 无镜像映射")
mod = module or entry.get("module")
if not mod:
raise UnknownStageError(f"阶段 {stage} 无执行模块")
cmd: list[str] = ["docker", "run", "--rm", "--user", str(self.cfg.host_uid)]
cmd += ["-v", f"{job_input_dir}:{self.cfg.container_input_dir}"]
cmd += ["-v", f"{job_workspace_dir}:{self.cfg.container_workspace_dir}"]
if self.cfg.data_root:
cmd += ["-v", f"{self.cfg.data_root}:{self.cfg.container_data_dir}:ro"]
# 登记路径只读挂载(src==dst,容器内原路径直接可见;登记时已过黑名单)
for m in (extra_mounts or self.cfg.extra_mounts):
cmd += ["-v", f"{m}:{m}:ro"]
# 内核树运行时挂载(除 bioflow-gf 外其余镜像未内置内核,统一走挂载)
if self.cfg.kernel_dir:
cmd += ["-v", f"{self.cfg.kernel_dir}:{self.cfg.container_kernel_dir}:ro"]
cmd += ["-w", self.cfg.container_kernel_dir]
eff_cpu = cpu or str(params.get("cpu") or "")
eff_memory = memory or str(params.get("memory") or "")
cmd += ["--cpus", self._safe_cpu(eff_cpu or self.cfg.default_cpus)]
cmd += ["--memory", self._safe_memory(eff_memory or self.cfg.default_memory)]
for key in sorted(self.cfg.bba_env):
cmd += ["-e", f"{key}={self.cfg.bba_env[key]}"]
# 用户自定义 env 追加(params 中 env 键)
extra_env = params.get("env") if isinstance(params.get("env"), dict) else {}
for key in sorted(extra_env):
cmd += ["-e", f"{key}={extra_env[key]}"]
tag = self.map.docker_tag(image)
cmd.append(f"{tag}:{self.cfg.docker_tag_suffix}")
cmd.extend(CONTAINER_CMD_PREFIX)
cmd.extend(["-m", mod])
# 参数注入:有 cli_contract 的阶段按契约,否则沿用默认 --input-dir/--workspace
contract = entry.get("cli_contract")
if contract:
cmd += self._contract_args(contract, params, job_input_dir)
else:
cmd += ["--input-dir", self.cfg.container_input_dir]
cmd += ["--workspace", self.cfg.container_workspace_dir]
for key in sorted(params):
value = params[key]
if key in ("env", "cpu", "memory"):
continue # cpu/memory 走 --cpus/--memory;env 走 -e
if value is None:
continue
if isinstance(value, bool):
if value:
cmd.append(f"--{key.replace('_', '-')}")
elif isinstance(value, (list, tuple)):
cmd.append(f"--{key.replace('_', '-')}")
cmd.append(",".join(str(v) for v in value))
else:
cmd.append(f"--{key.replace('_', '-')}")
cmd.append(str(value))
return cmd
def _contract_args(self, contract: list[Any], params: dict[str, Any],
job_input_dir: str | Path) -> list[str]:
"""cli_contract 解析:[(flag, source), ...] → argv。
source 取值:
- input_dir/workspace:容器内 /job/input、/job/workspace
- output_final:/job/workspace/final(rnaseq 系 -o 语义:产物目录,
CLI 内部以 output.parent 为 workspace 并建 output.parent/final——
故 -o 必须指向 workspace/final,manifest 也落此处与 runner 检索对齐)
- samples_yaml:/job/input/samples.yaml(宿主侧不存在则跳过该 flag,
交由 CLI 报缺失,错误信息比服务端更可读)
- prev_manifest:<stage>:上游阶段 manifest 路径
/job/workspace/final/{prefix}.{stage}.manifest.json(缺失跳过)
- prefix/job_id:从 params 取值,缺失跳过
- param:<name>:从 params[name] 取值,缺失跳过(用户经确认表单提供)
- tool:<name>:容器内工具命令(镜像 PATH 内裸命令,如 gatk/samtools)
"""
args: list[str] = []
for flag, source in contract:
if source == "input_dir":
args += [flag, self.cfg.container_input_dir]
elif source == "workspace":
args += [flag, self.cfg.container_workspace_dir]
elif source == "output_final":
args += [flag, f"{self.cfg.container_workspace_dir}/final"]
elif str(source).startswith("prev_manifest:"):
stage_name = str(source).split(":", 1)[1]
value = str(params.get("prefix") or "").strip()
if value:
args += [flag,
f"{self.cfg.container_workspace_dir}/final/{value}.{stage_name}.manifest.json"]
elif str(source).startswith("tool:"):
name = str(source).split(":", 1)[1]
args += [flag, name]
elif source == "samples_yaml":
if (Path(job_input_dir) / "samples.yaml").is_file():
args += [flag, f"{self.cfg.container_input_dir}/samples.yaml"]
elif source == "prefix":
value = str(params.get("prefix") or "").strip()
if value:
args += [flag, value]
elif source == "job_id":
value = str(params.get("job_id") or "").strip()
if value:
args += [flag, value]
elif str(source).startswith("param:"):
key = str(source).split(":", 1)[1]
value = params.get(key)
if value is not None and str(value).strip():
args += [flag, str(value)]
return args
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("bioflow kernel: 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("bioflow kernel: memory=%s 不在白名单 %s,回落默认 %s",
memory, self.cfg.memory_allow, self.cfg.default_memory)
return self.cfg.default_memory
@staticmethod
def command_display(cmd: list[str]) -> str:
"""命令可读形态(日志用;list 执行保证无 shell 注入)。"""
return " ".join(shlex.quote(str(c)) for c in cmd)
# ---- 执行 ----
async def run_stage(
self,
stage: str,
job_input_dir: str | Path,
job_workspace_dir: str | Path,
params: dict[str, Any] | None = None,
*,
log_path: str | Path | None = None,
cpu: str | None = None,
memory: str | None = None,
extra_mounts: list[str] | None = None,
) -> dict[str, Any]:
"""执行一个阶段 docker:返回 {exit_code, timed_out, log_tail, cmd}。
- extra_mounts:本次执行额外只读挂载的宿主绝对路径(登记路径,src==dst)
- 超时:terminate → 3s → kill(TC-BIO-KERN-009)
- 日志全量写 log_path(任务 logs/),返回尾部 log_tail_lines 行
- 退出码非 0 / 超时 → 失败语义由 runner 判定(本函数只报事实)
"""
params = params or {}
cpu = cpu or str(params.get("cpu") or "")
memory = memory or str(params.get("memory") or "")
cmd = self.build_command(stage, job_input_dir, job_workspace_dir, params,
cpu=cpu or None, memory=memory or None,
extra_mounts=extra_mounts)
log_path = Path(log_path) if log_path else Path(job_workspace_dir) / "logs" / f"{stage}.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.stage_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")
log_path.write_text(text, encoding="utf-8")
return {
"exit_code": exit_code,
"timed_out": timed_out,
"log_tail": tail_lines(text, self.cfg.log_tail_lines),
"cmd": cmd,
}
def tail_lines(text: str, n: int) -> str:
"""返回文本最后 n 行(日志截断;TC-BIO-RUN-013)。"""
lines = text.splitlines()
if len(lines) <= n:
return "\n".join(lines)
return "\n".join(lines[-n:])
def _default_kernel_config() -> KernelConfig:
"""从 app.config 推导(默认 uid = 服务进程 os.getuid())。"""
from ...config import get_settings
s = get_settings()
return KernelConfig(
host_uid=s.host_uid,
docker_tag_suffix=s.docker_image_tag_suffix,
data_root=s.data_root_abs,
container_data_dir=s.container_data_dir,
container_input_dir=s.container_input_dir,
container_workspace_dir=s.container_workspace_dir,
default_cpus=s.default_cpus,
default_memory=s.default_memory,
stage_timeout_sec=s.stage_timeout_sec,
log_tail_lines=s.log_tail_lines,
bba_env=s.bba_env,
kernel_dir=s.kernel_dir_abs,
container_kernel_dir=s.container_kernel_dir,
)
def load_kernel() -> Kernel:
"""应用启动:加载映射表 + 配置推导(映射表损坏 → ImageMapError 上抛)。"""
from ...config import get_settings
return Kernel.load(get_settings().image_map_abs, _default_kernel_config())