456 lines
20 KiB
Python
456 lines
20 KiB
Python
"""远程工作环境自动初始化：技能同步 + 依赖软件自动安装（2026-08-28）。
|
||
|
||
设计：
|
||
- 技能同步：本地技能树（公共全局 / 产线 / 项目通用 / 技能集 / 本机构组织技能）
|
||
rsync 到远程 ~/<remote_dir>/skills/。沙箱把该目录挂成 /home，技能在沙箱内
|
||
路径为 /home/skills/，agent 可直接读取。
|
||
- 依赖安装：解析同步技能 SKILL.md frontmatter 的 dependencies: 字段，
|
||
机器可解析的 pip 规格（含版本约束）装进远程 ~/<remote_dir>/.venv-skills
|
||
（沙箱内 /home/.venv-skills，技能脚本用该 venv 的 python）。
|
||
非机器可解析依赖（自然语言）不猜，跳过并汇报。
|
||
- 重量级依赖（torch/训练框架等，几 GB 且常需 GPU）默认跳过，避免远程机
|
||
被撑爆；跳过名单可通过 appbase 参数 remote_env_heavy_skip 扩展。
|
||
- pip 源：appbase 参数 remote_env_pip_index，默认阿里源。
|
||
- 全部失败不抛异常：以结构化报告返回，由调用方决定是否阻断/提示。
|
||
"""
|
||
import base64
|
||
import json
|
||
import os
|
||
import re
|
||
import shlex
|
||
import asyncio
|
||
import logging
|
||
|
||
logger = logging.getLogger("pipeline.remote_env_init")
|
||
|
||
from .work_env import (
|
||
_ssh_common_args, _ssh_target, _rel_remote_dir,
|
||
_validate_remote_dir, _remote_mkdir,
|
||
)
|
||
|
||
# 重量级 pip 包默认跳过（远程部署机通常无 GPU，torch 系单包 2GB+）。
|
||
# 名单可通过 appbase 参数 remote_env_heavy_skip 覆盖/扩展（逗号分隔）。
|
||
_DEFAULT_HEAVY = {
|
||
"torch", "torchvision", "torchaudio", "tensorflow", "jax",
|
||
"vllm", "unsloth", "deepspeed", "flash-attn", "bitsandbytes",
|
||
"llama-cpp-python", "llama_cpp_python", "transformers", "accelerate",
|
||
"peft", "trl", "diffusers", "axolotl", "lm-eval", "lm_eval",
|
||
"stable-diffusion", "comfyui", "openai-whisper", "whisper",
|
||
"segment-anything", "modal",
|
||
}
|
||
|
||
# pip 规格：名字（可带 extras）+ 可选版本约束，拒绝路径/URL（注入面）
|
||
_NAME_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$")
|
||
_SPEC_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*(\[[^\]]+\])?"
|
||
r"(\s*(==|>=|<=|~=|!=|>|<)\s*[A-Za-z0-9.*+_-]+)*$")
|
||
|
||
# Python 标准库模块：技能 frontmatter 里常把它们列进 dependencies，
|
||
# 但它们不是 pip 包，装会失败（且 venv 里自带）。
|
||
_STDLIB_SKIP = {
|
||
"asyncio", "typing", "json", "os", "sys", "re", "math", "time",
|
||
"datetime", "subprocess", "threading", "multiprocessing", "sqlite3",
|
||
"urllib", "http", "socket", "ssl", "hashlib", "hmac", "base64",
|
||
"uuid", "random", "collections", "itertools", "functools", "pathlib",
|
||
"shutil", "tempfile", "logging", "unittest", "argparse", "csv",
|
||
"struct", "ctypes", "dataclasses", "abc", "io", "string", "traceback",
|
||
"inspect", "pickle", "gzip", "zipfile", "tarfile", "email", "xml",
|
||
"html", "http.client", "socketserver", "secrets", "contextlib",
|
||
"operator", "copy", "weakref", "enum", "numbers", "decimal", "fractions",
|
||
}
|
||
|
||
|
||
# ── frontmatter 依赖解析（不引 yaml，手写轻量解析）──────────────
|
||
|
||
def _parse_skill_deps(content: str):
|
||
"""解析 SKILL.md frontmatter 的 dependencies: 列表。
|
||
|
||
支持两种写法：块式（- item 多行）与内联（[a, b]）。
|
||
只在 frontmatter 区域内解析，避免正文的散文列表混入。
|
||
返回依赖字符串列表。
|
||
"""
|
||
if not content.startswith("---"):
|
||
return []
|
||
end = content.find("\n---", 3)
|
||
if end == -1:
|
||
return []
|
||
fm = content[3:end]
|
||
lines = fm.split("\n")
|
||
deps, in_deps = [], False
|
||
for line in lines:
|
||
stripped = line.strip()
|
||
if stripped.startswith("dependencies:"):
|
||
rest = stripped.split(":", 1)[1].strip()
|
||
if rest.startswith("["):
|
||
inner = rest.strip("[]")
|
||
return [x.strip().strip("'\"") for x in inner.split(",") if x.strip()]
|
||
in_deps = True
|
||
continue
|
||
if in_deps:
|
||
if stripped.startswith("- "):
|
||
deps.append(stripped[2:].strip().strip("'\""))
|
||
elif stripped == "":
|
||
continue
|
||
else:
|
||
in_deps = False
|
||
return deps
|
||
|
||
|
||
def _pip_installable(dep: str) -> bool:
|
||
"""判断依赖是否为机器可解析的 pip 规格（名字合法 + 规格形态合法 + 非标准库）。"""
|
||
dep = (dep or "").strip()
|
||
if not dep or "/" in dep or ":" in dep:
|
||
return False
|
||
name = re.split(r"[\[<>=!~\s]", dep, 1)[0]
|
||
if name.lower() in _STDLIB_SKIP:
|
||
return False
|
||
return bool(_NAME_RE.match(name)) and bool(_SPEC_RE.match(dep))
|
||
|
||
|
||
def _specs_from_scope(scope_dir: str):
|
||
"""收集一个技能目录（含子目录）下所有技能的 pip 依赖规格。"""
|
||
out = {}
|
||
if not os.path.isdir(scope_dir):
|
||
return out
|
||
for root, _dirs, files in os.walk(scope_dir):
|
||
if "SKILL.md" not in files:
|
||
continue
|
||
try:
|
||
with open(os.path.join(root, "SKILL.md"), encoding="utf-8") as f:
|
||
content = f.read()
|
||
except Exception:
|
||
continue
|
||
for dep in _parse_skill_deps(content):
|
||
if _pip_installable(dep):
|
||
name = re.split(r"[\[<>=!~\s]", dep, 1)[0].lower()
|
||
# 带版本约束的规格优先于裸名字
|
||
if name not in out or any(c in dep for c in "=<>~"):
|
||
out[name] = dep
|
||
return out
|
||
|
||
|
||
# ── 配置（appbase params，禁硬编码）──────────────────────────
|
||
|
||
def _cfg_defaults() -> dict:
|
||
return {"pip_index": "https://mirrors.aliyun.com/pypi/simple/",
|
||
"heavy_skip": ",".join(sorted(_DEFAULT_HEAVY)),
|
||
"install_timeout": 1200}
|
||
|
||
|
||
async def _load_cfg(sor=None) -> dict:
|
||
cfg = dict(_cfg_defaults())
|
||
if sor is None:
|
||
return cfg
|
||
try:
|
||
from .workspace import get_param
|
||
idx = await get_param(sor, "remote_env_pip_index", "")
|
||
if idx:
|
||
cfg["pip_index"] = idx.strip()
|
||
heavy = await get_param(sor, "remote_env_heavy_skip", "")
|
||
if heavy:
|
||
cfg["heavy_skip"] = heavy.strip()
|
||
timeout = await get_param(sor, "remote_env_install_timeout", "")
|
||
if timeout:
|
||
cfg["install_timeout"] = int(float(timeout))
|
||
except Exception as e:
|
||
logger.warning("remote_env_init: 读取 appbase 参数失败，用默认值: %s", e)
|
||
return cfg
|
||
|
||
|
||
# ── 技能同步 ──────────────────────────────────────────────────
|
||
|
||
def _skill_scopes(skills_base: str, org_id: str):
|
||
"""要同步的技能范围（相对技能根的本地子目录）列表。
|
||
|
||
动态枚举技能根顶层目录（除 orgs/ 外全部同步，覆盖公共/产线/项目/技能集
|
||
及未来新增的范围），再加本机构组织技能（含其下项目私有技能）。
|
||
orgs/ 其他机构不同步（隔离）。
|
||
"""
|
||
scopes = []
|
||
if os.path.isdir(skills_base):
|
||
for name in sorted(os.listdir(skills_base)):
|
||
if name == "orgs" or name.startswith("."):
|
||
continue
|
||
if os.path.isdir(os.path.join(skills_base, name)):
|
||
scopes.append(name)
|
||
if org_id:
|
||
scopes.append(os.path.join("orgs", str(org_id)))
|
||
return scopes
|
||
|
||
|
||
async def sync_skills_to_remote(env: dict, org_id: str = "", skills_base: str = "") -> dict:
|
||
"""把公共/产线/项目/技能集/本机构组织技能同步到远程 ~/<remote_dir>/skills/。
|
||
|
||
幂等：--delete 覆盖。返回 {ok, synced:[...], skipped:[...], error}
|
||
"""
|
||
err = _validate_remote_dir(env.get("remote_dir", ""))
|
||
if err:
|
||
return {"ok": False, "error": err}
|
||
if not skills_base:
|
||
from pipeline_core.skill_pack import get_skills_base
|
||
skills_base = get_skills_base()
|
||
if not os.path.isdir(skills_base):
|
||
return {"ok": False, "error": f"本地技能树不存在: {skills_base}"}
|
||
|
||
rel = _rel_remote_dir(env.get("remote_dir", ""))
|
||
remote_skills = rel + "/skills"
|
||
if not await _remote_mkdir(env, remote_skills):
|
||
return {"ok": False, "error": "远程创建技能目录失败"}
|
||
|
||
target_base = _ssh_target(env) + ":~/" + remote_skills
|
||
ssh_opts = " ".join(_ssh_common_args(env))
|
||
synced, skipped = [], []
|
||
for sub in _skill_scopes(skills_base, org_id):
|
||
src = os.path.join(skills_base, sub)
|
||
if not os.path.isdir(src):
|
||
skipped.append(sub)
|
||
continue
|
||
# rsync -R + "./" 锚点：保留相对子路径（如 orgs/<id>），远程结构与本地一致
|
||
cmd = ["rsync", "-az", "--delete", "-R",
|
||
"-e", "ssh " + ssh_opts,
|
||
os.path.join(skills_base, "./" + sub + "/"),
|
||
target_base + "/"]
|
||
try:
|
||
proc = await asyncio.create_subprocess_exec(
|
||
*cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE)
|
||
_out, stderr = await asyncio.wait_for(proc.communicate(), timeout=300)
|
||
if proc.returncode != 0:
|
||
return {"ok": False, "error": f"技能同步失败({sub}): "
|
||
+ stderr.decode('utf-8', 'replace')[-500:]}
|
||
synced.append(sub)
|
||
except asyncio.TimeoutError:
|
||
return {"ok": False, "error": f"技能同步超时({sub})"}
|
||
return {"ok": True, "synced": synced, "skipped": skipped,
|
||
"remote_path": "~/" + remote_skills}
|
||
|
||
|
||
# ── 依赖收集与安装 ────────────────────────────────────────────
|
||
|
||
def collect_skill_deps(skills_base: str, org_id: str, heavy_skip: str) -> dict:
|
||
"""收集各范围技能的 pip 依赖，剔除重量级。返回 {specs, heavy_skipped, unparseable}。"""
|
||
skip = {s.strip().lower() for s in (heavy_skip or "").split(",") if s.strip()}
|
||
specs, heavy_skipped = {}, []
|
||
for sub in _skill_scopes(skills_base, org_id):
|
||
for name, spec in _specs_from_scope(os.path.join(skills_base, sub)).items():
|
||
if name in skip:
|
||
if spec not in heavy_skipped:
|
||
heavy_skipped.append(spec)
|
||
else:
|
||
specs[name] = spec
|
||
return {"specs": sorted(specs.values()), "heavy_skipped": sorted(heavy_skipped)}
|
||
|
||
|
||
async def install_remote_deps(env: dict, specs: list, pip_index: str,
|
||
timeout: int = 1200) -> dict:
|
||
"""在远程 ~/<remote_dir>/.venv-skills 建 venv 并安装依赖。幂等。
|
||
|
||
specs 经 base64 传输（含 >=[] 等特殊字符，避开 shell 展开）。
|
||
返回 {ok, installed:[...], failed:[...], error}
|
||
"""
|
||
if not specs:
|
||
return {"ok": True, "installed": [], "failed": [], "skipped_reason": "无可安装依赖"}
|
||
err = _validate_remote_dir(env.get("remote_dir", ""))
|
||
if err:
|
||
return {"ok": False, "error": err}
|
||
|
||
rel = _rel_remote_dir(env.get("remote_dir", ""))
|
||
payload = base64.b64encode(json.dumps(specs).encode()).decode()
|
||
# 不用 set -e / 管道：管道会吞掉 pip 的退出码（$? 变成 tail 的状态）。
|
||
# set -f 关 glob：pip 规格里的 * / [...] 不能被远程 shell 展开。
|
||
# 完整输出由 SSH 捕获，报告时截取尾部。
|
||
script = (
|
||
'RD="$HOME/' + rel + '"; VENV="$RD/.venv-skills"; '
|
||
'command -v python3 >/dev/null || { echo "NO_PYTHON3"; exit 3; }; '
|
||
'if [ ! -x "$VENV/bin/python" ]; then python3 -m venv "$VENV" || { echo "VENV_FAIL"; exit 4; }; fi; '
|
||
'set -f; RC=0; "$VENV/bin/pip" install --no-input --disable-pip-version-check '
|
||
+ (('-i ' + shlex.quote(pip_index) + ' ') if pip_index else '')
|
||
+ '$(echo "' + payload + '" | base64 -d | python3 -c '
|
||
'"import json,sys;print(chr(32).join(json.load(sys.stdin)))") '
|
||
'|| RC=$?; '
|
||
'echo "INSTALL_RC=$RC"'
|
||
)
|
||
ssh_cmd = ["ssh"] + _ssh_common_args(env) + [_ssh_target(env), script]
|
||
try:
|
||
proc = await asyncio.create_subprocess_exec(
|
||
*ssh_cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE)
|
||
stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=timeout + 30)
|
||
except asyncio.TimeoutError:
|
||
return {"ok": False, "error": f"依赖安装超时（{timeout}s），venv 可能装了一半，可重跑"}
|
||
except Exception as e:
|
||
return {"ok": False, "error": "SSH 执行失败: " + str(e)[:300]}
|
||
|
||
out = stdout.decode("utf-8", "replace")
|
||
if "NO_PYTHON3" in out:
|
||
return {"ok": False, "error": "远程主机缺少 python3"}
|
||
if proc.returncode not in (0, 1):
|
||
return {"ok": False, "error": "venv 创建失败: "
|
||
+ (out + stderr.decode('utf-8', 'replace'))[-600:]}
|
||
|
||
# 解析 pip 退出码（由脚本末尾 echo INSTALL_RC=$RC 输出）
|
||
rc = -1
|
||
for line in out.splitlines():
|
||
if line.startswith("INSTALL_RC="):
|
||
try:
|
||
rc = int(line.split("=", 1)[1])
|
||
except Exception:
|
||
pass
|
||
|
||
# pip 逐包结果解析：新版 pip 解析器失败时要么全装要么全不装，
|
||
# 从 ERROR 行提取包名。常见格式：
|
||
# ERROR: Could not find a version that satisfies the requirement <name> ...
|
||
# ERROR: No matching distribution found for <name>
|
||
failed = []
|
||
for line in out.splitlines():
|
||
if "ERROR:" not in line:
|
||
continue
|
||
m = (re.search(r"the requirement ([A-Za-z0-9._\[\]-]+)", line)
|
||
or re.search(r"No matching distribution found for ([A-Za-z0-9._\[\]-]+)", line))
|
||
failed.append(m.group(1) if m else line.strip()[:120])
|
||
if rc != 0 and not failed:
|
||
failed = ["部分包安装失败，详见日志末尾"]
|
||
|
||
if rc == 0:
|
||
installed = list(specs)
|
||
failed = []
|
||
else:
|
||
# 批量失败 → 逐包降级：一次 SSH 会话内逐包安装（pip 幂等，已装的秒过），
|
||
# 部分成功优于全败。
|
||
return await _install_per_package(env, specs, pip_index, rel, timeout)
|
||
return {"ok": rc == 0, "installed": installed, "failed": failed,
|
||
"venv": "~/" + rel + "/.venv-skills",
|
||
"tail": out[-800:]}
|
||
|
||
|
||
async def _install_per_package(env: dict, specs: list, pip_index: str,
|
||
rel: str, timeout: int) -> dict:
|
||
"""逐包安装（批量失败后的降级路径）。返回格式与 install_remote_deps 一致。"""
|
||
payload = base64.b64encode(json.dumps(specs).encode()).decode()
|
||
index_arg = ('-i ' + shlex.quote(pip_index) + ' ') if pip_index else ''
|
||
script = (
|
||
'RD="$HOME/' + rel + '"; VENV="$RD/.venv-skills"; '
|
||
'[ -x "$VENV/bin/pip" ] || { echo "VENV_MISSING"; exit 4; }; '
|
||
'set -f; '
|
||
'echo "' + payload + '" | base64 -d | python3 -c '
|
||
'"import json,sys;print(chr(10).join(json.load(sys.stdin)))" | '
|
||
'while IFS= read -r pkg; do '
|
||
' [ -n "$pkg" ] || continue; '
|
||
' if "$VENV/bin/pip" install --no-input --disable-pip-version-check '
|
||
+ index_arg + '"$pkg" >/dev/null 2>&1; then echo "OK:$pkg"; '
|
||
' else echo "FAILED:$pkg"; fi; '
|
||
'done'
|
||
)
|
||
ssh_cmd = ["ssh"] + _ssh_common_args(env) + [_ssh_target(env), script]
|
||
try:
|
||
proc = await asyncio.create_subprocess_exec(
|
||
*ssh_cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE)
|
||
stdout, _stderr = await asyncio.wait_for(proc.communicate(), timeout=timeout + 30)
|
||
except asyncio.TimeoutError:
|
||
return {"ok": False, "error": f"逐包安装超时（{timeout}s），可重跑"}
|
||
except Exception as e:
|
||
return {"ok": False, "error": "SSH 执行失败: " + str(e)[:300]}
|
||
|
||
out = stdout.decode("utf-8", "replace")
|
||
installed, failed = [], []
|
||
for line in out.splitlines():
|
||
if line.startswith("OK:"):
|
||
installed.append(line[3:])
|
||
elif line.startswith("FAILED:"):
|
||
failed.append(line[7:])
|
||
return {"ok": len(failed) == 0 and len(installed) > 0,
|
||
"installed": installed, "failed": failed,
|
||
"venv": "~/" + rel + "/.venv-skills",
|
||
"mode": "per-package", "tail": out[-500:]}
|
||
|
||
|
||
# ── 总编排 ────────────────────────────────────────────────────
|
||
|
||
async def init_remote_env(env: dict, org_id: str = "", sor=None,
|
||
skills_base: str = "", skip_bwrap: bool = False) -> dict:
|
||
"""远程环境一键初始化：bwrap → 目录 → 技能同步 → 依赖安装。
|
||
|
||
bwrap 缺失即中止（沙箱执行的前提）；技能/依赖失败以报告返回不中止后续步骤。
|
||
"""
|
||
from .work_env import ensure_remote_bwrap
|
||
report = {"ok": True, "steps": {}}
|
||
|
||
# 1. bwrap（前提，失败中止）
|
||
if not skip_bwrap:
|
||
bw = await ensure_remote_bwrap(env)
|
||
report["steps"]["bwrap"] = bw
|
||
if not bw.get("ok"):
|
||
report["ok"] = False
|
||
report["error"] = "远程 bwrap 部署失败: " + str(bw.get("error", ""))
|
||
return report
|
||
|
||
# 2. 远程目录
|
||
rel = _rel_remote_dir(env.get("remote_dir", ""))
|
||
if rel:
|
||
ok = await _remote_mkdir(env, rel)
|
||
report["steps"]["mkdir"] = {"ok": ok}
|
||
if not ok:
|
||
report["ok"] = False
|
||
report["error"] = "远程工作目录创建失败"
|
||
return report
|
||
|
||
# 3. 技能同步
|
||
sk = await sync_skills_to_remote(env, org_id, skills_base)
|
||
report["steps"]["skills"] = sk
|
||
if not sk.get("ok"):
|
||
report["ok"] = False
|
||
report["error"] = "技能同步失败: " + str(sk.get("error", ""))
|
||
return report
|
||
|
||
# 4. 依赖解析（快）+ 安装（慢，后台执行）
|
||
cfg = await _load_cfg(sor)
|
||
deps = collect_skill_deps(skills_base or _default_skills_base(), org_id,
|
||
cfg["heavy_skip"])
|
||
report["steps"]["deps_found"] = deps
|
||
task_key = _start_deps_install(env, deps["specs"], cfg)
|
||
report["deps_task"] = task_key
|
||
report["steps"]["deps_install"] = {"status": "running", "task": task_key}
|
||
return report
|
||
|
||
|
||
# ── 依赖安装后台任务与状态 ────────────────────────────────────
|
||
# 依赖安装可能耗时数分钟（超过 nginx 600s 代理超时），不能在 API 请求内同步等。
|
||
# 编排只启动后台任务并返回任务号，前端用 init_remote_env_status 轮询结果。
|
||
# 状态存进程内存（单进程部署；重启丢状态则重跑，幂等无副作用）。
|
||
|
||
_INSTALL_TASKS = {} # {task_key: {"status": ..., "result": ..., "started": ts}}
|
||
|
||
|
||
def _start_deps_install(env: dict, specs: list, cfg: dict) -> str:
|
||
from appPublic.uniqueID import getID
|
||
task_key = "deps_" + getID()
|
||
_INSTALL_TASKS[task_key] = {"status": "running", "result": None,
|
||
"specs": list(specs)}
|
||
|
||
async def _run():
|
||
try:
|
||
r = await install_remote_deps(env, specs, cfg["pip_index"],
|
||
cfg["install_timeout"])
|
||
_INSTALL_TASKS[task_key]["result"] = r
|
||
_INSTALL_TASKS[task_key]["status"] = "done"
|
||
except Exception as e:
|
||
_INSTALL_TASKS[task_key]["result"] = {"ok": False, "error": str(e)[:500]}
|
||
_INSTALL_TASKS[task_key]["status"] = "done"
|
||
|
||
asyncio.create_task(_run())
|
||
return task_key
|
||
|
||
|
||
def get_deps_install_status(task_key: str) -> dict:
|
||
t = _INSTALL_TASKS.get(task_key or "")
|
||
if not t:
|
||
return {"ok": False, "error": "任务不存在（可能服务重启过，请重新初始化）"}
|
||
if t["status"] == "running":
|
||
return {"ok": True, "status": "running",
|
||
"message": f"依赖安装中（{len(t.get('specs', []))} 个包），请稍候再查询"}
|
||
r = dict(t.get("result") or {})
|
||
r["status"] = "done"
|
||
return r
|
||
|
||
|
||
def _default_skills_base() -> str:
|
||
from pipeline_core.skill_pack import get_skills_base
|
||
return get_skills_base()
|