- 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>
361 lines
15 KiB
Python
361 lines
15 KiB
Python
"""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())
|