410 lines
18 KiB
Python
410 lines
18 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 launch_remote_install(env: dict, specs: list, pip_index: str, rel: str) -> bool:
|
||
"""在远程以 nohup 分离进程启动逐包安装（脱离 SSH 会话，不怕断连/重启）。
|
||
|
||
进度写 ~/<rel>/.venv-skills-install.progress（OK:/FAILED: 行），
|
||
完成写 ~/<rel>/.venv-skills-install.done（汇总）。
|
||
本调用秒级返回，是否成功启动由返回值给出。
|
||
"""
|
||
payload = base64.b64encode(json.dumps(specs).encode()).decode()
|
||
index_arg = ('-i ' + shlex.quote(pip_index) + ' ') if pip_index else ''
|
||
# 分离脚本：建 venv → 逐包安装 → 写进度/完成标记。
|
||
# 用 set -f 防 pip 规格里 * / [...] 被远程 shell 展开。
|
||
inner = (
|
||
'RD="$HOME/' + rel + '"; VENV="$RD/.venv-skills"; '
|
||
'PROG="$RD/.venv-skills-install.progress"; DONE="$RD/.venv-skills-install.done"; '
|
||
'if [ ! -x "$VENV/bin/python" ]; then '
|
||
'python3 -m venv "$VENV" || { echo "VENV_FAIL" > "$DONE"; exit 4; }; fi; '
|
||
'set -f; : > "$PROG"; '
|
||
'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" >> "$PROG"; '
|
||
'else echo "FAILED:$pkg" >> "$PROG"; fi; done; '
|
||
'OKC=$(grep -c "^OK:" "$PROG" 2>/dev/null || echo 0); '
|
||
'FC=$(grep -c "^FAILED:" "$PROG" 2>/dev/null || echo 0); '
|
||
'echo "OK=$OKC FAILED=$FC" > "$DONE"'
|
||
)
|
||
script = (
|
||
'RD="$HOME/' + rel + '"; '
|
||
'command -v python3 >/dev/null || { echo "NO_PYTHON3"; exit 3; }; '
|
||
'mkdir -p "$RD"; '
|
||
'nohup bash -c ' + shlex.quote(inner) +
|
||
' </dev/null >"$RD/.venv-skills-install.log" 2>&1 & echo LAUNCHED'
|
||
)
|
||
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=60)
|
||
except Exception as e:
|
||
logger.warning("launch_remote_install 启动失败: %s", e)
|
||
return False
|
||
out = stdout.decode("utf-8", "replace")
|
||
if "NO_PYTHON3" in out:
|
||
return False
|
||
return "LAUNCHED" in out
|
||
|
||
|
||
async def poll_remote_install(env: dict, rel: str) -> dict:
|
||
"""轮询远程安装进度。短 SSH 调用。
|
||
|
||
返回 {done:bool, ok:bool, installed:[], failed:[], venv, ...}
|
||
"""
|
||
script = (
|
||
'RD="$HOME/' + rel + '"; DONE="$RD/.venv-skills-install.done"; '
|
||
'PROG="$RD/.venv-skills-install.progress"; '
|
||
'if [ -f "$DONE" ]; then echo "STATE=done"; cat "$DONE"; echo "===PROG==="; '
|
||
'cat "$PROG" 2>/dev/null; '
|
||
'else echo "STATE=running"; cat "$PROG" 2>/dev/null; fi'
|
||
)
|
||
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=30)
|
||
except Exception as e:
|
||
return {"done": False, "error": "轮询失败: " + str(e)[:200]}
|
||
|
||
out = stdout.decode("utf-8", "replace")
|
||
done = "STATE=done" in out
|
||
body = out.split("STATE=", 1)[1] if "STATE=" in out else out
|
||
if "===PROG===" in body:
|
||
summary, prog = body.split("===PROG===", 1)
|
||
else:
|
||
summary, prog = body, ""
|
||
|
||
installed, failed = [], []
|
||
for line in prog.splitlines():
|
||
line = line.strip()
|
||
if line.startswith("OK:"):
|
||
installed.append(line[3:])
|
||
elif line.startswith("FAILED:"):
|
||
failed.append(line[7:])
|
||
|
||
result = {
|
||
"done": done,
|
||
"ok": done and len(failed) == 0 and len(installed) > 0,
|
||
"installed": installed,
|
||
"failed": failed,
|
||
"venv": "~/" + rel + "/.venv-skills",
|
||
"mode": "per-package-detached",
|
||
}
|
||
if done:
|
||
result["summary"] = summary.strip()
|
||
return result
|
||
|
||
|
||
# ── 总编排 ────────────────────────────────────────────────────
|
||
|
||
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. 依赖解析（快）+ 启动远程分离式安装（nohup，脱离 SSH 会话，
|
||
# 进度落远程文件，轮询无状态——服务重启/断连都不影响）
|
||
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
|
||
launched = await launch_remote_install(env, deps["specs"], cfg["pip_index"], rel)
|
||
report["steps"]["deps_launch"] = {"ok": launched, "count": len(deps["specs"])}
|
||
if not launched:
|
||
report["warning"] = "依赖安装进程启动失败（远程可能缺 python3），请重新运行初始化"
|
||
return report
|
||
|
||
|
||
# ── 依赖安装状态查询（无状态，轮询远程进度文件）────────────────
|
||
# 安装进程由远程 nohup 托管，进度写 ~/<remote_dir>/.venv-skills-install.*。
|
||
# 查询只需 SSH 读文件，不依赖本进程内存，服务重启后仍可查。
|
||
|
||
async def get_deps_install_status(sor, owner_type: str, owner_id: str) -> dict:
|
||
"""查询远程依赖安装进度。从 sd_work_envs 取远程配置后轮询进度文件。"""
|
||
from .work_env import get_work_env
|
||
env = await get_work_env(
|
||
sor, owner_id if owner_type == "user" else "",
|
||
owner_id if owner_type == "org" else "")
|
||
if env.get("mode") != "remote" or not env.get("remote_dir"):
|
||
return {"ok": False, "error": "当前非远程模式，无安装任务"}
|
||
rel = _rel_remote_dir(env.get("remote_dir", ""))
|
||
return await poll_remote_install(env, rel)
|
||
|
||
|
||
def _default_skills_base() -> str:
|
||
from pipeline_core.skill_pack import get_skills_base
|
||
return get_skills_base()
|