Prhub

#1899 Snapshot the external commands of every shell launch script

原始 PR 作者 fzyzcjy 合并时间 2026-08-09 18:33 文件变更 71 提交数 6 评论 3 代码增减 +16330 / -5

执行摘要

为全部 shell 启动脚本建立外部命令快照测试

该 PR 属于 #1837 追踪的启动脚本重构与增强系列,目标是给每个 shell 启动脚本的外部命令做快照,防止后续对启动命令格式、参数顺序或 ray job submit 调用链的改动在不知不觉中破坏已有脚本。issue 评论中 yueming-yuan 也提出了对 CI 复杂度与 .sh 脚本长期维护的疑虑,促使团队在自动化防护与维护成本之间做权衡。

值得精读,尤其是 sh_harness.py 的 shim 与冻结环境设计,以及 TestDiscovery 的发现-快照双射校验,这些模式可以复用到其他 CLI/脚本类仓库。对于维护者,建议关注快照更新机制的滥用风险,并参考 PR 2279 的做法合理规划快照测试的 CI 执行频率。

讨论亮点

issue 评论中 yueming-yuan 提出两个疑虑:一是担心保留此快照测试会让未来 CI 更复杂,询问合入后是否移除;二是建议把所有 .sh 脚本转换成 .py。作者未在 PR 内直接回应,快照测试随 PR 合并;后续 PR 2279 将快照测试改为手动执行、PR 1910 开始把模型配置 shell 脚本改写为 Python,两条疑虑分别得到回应。

实现拆解

实现分为 5 个步骤:

  1. 扩展测试底座 tests/fast/launch_scripts/sh_harness.py:新增 iter_launch_scripts()ray job submit 关键字发现脚本;run_launch_script() 新增 args 透传;扩充 shim 命令白名单(awkpsscpssh 等),并让 python/python3 的 shim 拦截 cluster_resources 查询与 import 探测,其余参数透传给真实 Python,避免资源查询依赖真实集群。同时冻结日期、进程列表等输出,保证快照可复现。

  2. 新增主测试文件 tests/fast/launch_scripts/test_sh_launch_scripts.py:用 LaunchScriptCase 描述需要显式参数或环境变量的脚本(如 p2p 参数、BASE_FOLDEROUTPUT_DIR 等);recorded fixture 参数化执行每个脚本;TestEveryLaunchScript 校验调用序列与快照一致,且恰好有一次 ray job submitTestDiscovery 校验发现集合与快照集合双向一致,防止脚本悄悄脱离覆盖。

  3. 批量生成快照:将每条外部命令的 argv 与返回码写入 tests/snapshots/launch_scripts/sh/ 下与脚本同路径的 .txt 文件,作为后续比较基准;快照可通过 MILES_UPDATE_LAUNCH_SCRIPT_SNAPSHOTS 环境变量重新生成。

  4. 合入系列前置修复:分支合并中带入了 PYTHONBUFFERED 拼写修正、缺失行连接导致的路径不可解析、硬编码 checkout 位置改为自动推导等变更,这些修复是快照能够稳定落盘的前提。

  5. 配套测试:除主测试外,harness 自身行为(fork 顺序、shim 输出、冻结环境拒绝解冻)也有针对性测试覆盖,确保底座本身可信。

文件 模块 状态 重要度
tests/fast/launch_scripts/test_sh_launch_scripts.py 启动脚本 added 7.31
tests/fast/launch_scripts/sh_harness.py 测试底座 modified 6.28
tests/snapshots/launch_scripts/sh/examples/on_policy_distillation/run-qwen3-8B-opd.sh.txt 快照数据 added 4.24

关键符号

iter_launch_scripts assert_matches_snapshot run_launch_script format_invocations recorded test_invocations_match_snapshot test_submits_exactly_one_ray_job test_every_discovered_script_has_a_snapshot_and_vice_versa

关键源码片段

tests/fast/launch_scripts/test_sh_launch_scripts.py test-coverage

新增的 shell 启动脚本快照主测试,是本次变更的核心。遍历发现的全部启动脚本,校验外部命令调用序列与快照的一致性,并确保恰好一次 `ray job submit`。

from dataclasses import dataclass, fieldimport pytest
from tests.fast.launch_scripts.sh_harness import (
    REPO_ROOT,
    assert_matches_snapshot,
    format_invocations,
    iter_launch_scripts,
    run_launch_script,
)_SNAPSHOT_DIR = REPO_ROOT / "tests" / "snapshots" / "launch_scripts" / "sh"# 快照基准路径与脚本相对仓库根路径保持一致
​
​
@dataclass(frozen=True)
class LaunchScriptCase:
    # 部分脚本必须有显式参数或环境变量才肯继续执行,
    # 测试需要为它们提供最小可信输入。
    args: tuple[str, ...] = ()
    env: dict[str, str] = field(default_factory=dict)
​
​
_CHECKPOINT_DIR = "/frozen/checkpoints"
_HEAD_NODE_IP = "10.0.0.1"# 需要显式输入才能跑完的脚本清单,按相对仓库根路径索引;
# `{workdir}` 会被替换成沙箱临时目录,保证快照可复现。
_SCRIPTS_REFUSING_TO_RUN_WITHOUT_EXPLICIT_INPUTS: dict[str, LaunchScriptCase] = {
    "examples/lora/run-qwen2.5-3B-megatron-lora-disaggregated-multi-node.sh": LaunchScriptCase(args=("p2p", "0")),
    "examples/infra_features/p2p_weight_transfer/run-glm5-disagg-profile.sh": LaunchScriptCase(
        args=("GLM-5", "p2p", "0", _HEAD_NODE_IP), env={"MILES_LOG_DIR": "{workdir}"}
    ),
    "scripts/run-qwen3-235B-A22B.sh": LaunchScriptCase(env={"BASE_FOLDER": _CHECKPOINT_DIR}),
    "scripts/run-qwen3.6-27B.sh": LaunchScriptCase(env={"OUTPUT_DIR": _CHECKPOINT_DIR}),
}# 发现所有会执行 `ray job submit` 的 shell 启动脚本
_SCRIPTS = [script.relative_to(REPO_ROOT).as_posix() for script in iter_launch_scripts()]
​
​
@pytest.fixture(params=_SCRIPTS, scope="module")
def recorded(request, tmp_path_factory):
    rel = request.param
    case = _SCRIPTS_REFUSING_TO_RUN_WITHOUT_EXPLICIT_INPUTS.get(rel, LaunchScriptCase())
    tmp_path = tmp_path_factory.mktemp("launch_script")
    workdir = tmp_path / "workdir"
    run = run_launch_script(
        REPO_ROOT / rel,
        sandbox=tmp_path,
        args=case.args,
        extra_env={key: value.format(workdir=workdir) for key, value in case.env.items()},
    )
    return rel, run
​
​
class TestEveryLaunchScript:
    def test_invocations_match_snapshot(self, recorded):
        # 把脚本真实发出的外部命令序列和快照逐条对比,
        # 任何改动(哪怕只是参数顺序)都会在这里暴露。
        rel, run = recorded
        snapshot = _SNAPSHOT_DIR / f"{rel}.txt"
        actual = f"# returncode: {run.returncode}\n\n{format_invocations(run.invocations)}"
        assert_matches_snapshot(snapshot, actual, rel)
​
    def test_submits_exactly_one_ray_job(self, recorded):
        # 启动脚本必须恰好提交一次 ray job,否则说明调用链已断。
        _, run = recorded
        assert run.returncode == 0
        assert len(run.ray_job_submit_argv()) > 10
​
​
class TestDiscovery:
    def test_every_discovered_script_has_a_snapshot_and_vice_versa(self):
        # 发现集合与快照集合必须双向一致,
        # 避免脚本悄悄脱离覆盖范围也没有任何信号。
        discovered = {f"{rel}.txt" for rel in _SCRIPTS}
        recorded = {path.relative_to(_SNAPSHOT_DIR).as_posix() for path in _SNAPSHOT_DIR.rglob("*.txt")}
        assert discovered == recorded
        assert len(discovered) > 60
tests/fast/launch_scripts/sh_harness.py test-coverage

沙箱测试底座,定义了脚本发现、命令 shim、冻结环境、快照比较 / 再生成机制,所有快照测试都依赖它。

import os
import sys
from pathlib import Path_ARG_SEPARATOR = "\x1f"
_RECORD_SEPARATOR = "\x1e"_SYSTEM_PATH = "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin"# 冻结环境:确保快照不依赖开发机或 CI 宿主机的注入变量
_FROZEN_ENV = {
    "HOME": "/root",
    "LANG": "C",
    "LC_ALL": "C",
    "TERM": "dumb",
    "MASTER_ADDR": "127.0.0.1",
    "NODE_RANK": "0",
    "WANDB_KEY": "frozen-wandb-key",
    "WANDB_API_KEY": "frozen-wandb-api-key",
}# 需要被替换成假命令的可执行文件白名单;
# 新增真实命令前先确认是否会打破沙箱隔离。
_SHIMMED_COMMANDS = (
    "apt", "apt-get", "awk", "curl", "date", "docker", "git", "hf", "ip", "mkdir",
    "nc", "nvidia-smi", "pip", "pip3", "pkill", "ps", "python", "python3", "ray",
    "rm", "rsync", "scp", "sleep", "ssh", "torchrun", "wget",
)# 让“等待 N 个 GPU 加入”的轮询立即退出,而不是死等
_GPU_COUNT_LARGER_THAN_ANY_WAIT_LOOP_EXPECTS = "1000000"# 冻结 `ps` 输出,避免录制到宿主机的真实进程列表
_FROZEN_RAY_DASHBOARD_PROCESS = "root 1 0.0 0.0 ray dashboard --node-ip-address=10.0.0.1 --dashboard-port=8265"_SHIM_STDOUT = {
    "date": "20260101_000000",
    "ps": _FROZEN_RAY_DASHBOARD_PROCESS,
}# python shim:拦截集群资源查询与 import 探测,其余透传真实 Python
_PYTHON_SHIM_BODY = """case "${1:-}" in
-c)
    case "$2" in
    *cluster_resources*) printf '%s\\n' 'REPLACE_GPU_COUNT' ;;
    *import*) ;;
    *) "$MILES_SH_HARNESS_REAL_PYTHON" "$@" ;;
    esac
    ;;
esac
""".replace("REPLACE_GPU_COUNT", _GPU_COUNT_LARGER_THAN_ANY_WAIT_LOOP_EXPECTS)_SHIM_BODY = {"python": _PYTHON_SHIM_BODY, "python3": _PYTHON_SHIM_BODY}
​
​
def iter_launch_scripts() -> list[Path]:
    # 用“包含 `ray job submit`”作为发现条件,避免把普通工具脚本纳入
    roots = [REPO_ROOT / "scripts", REPO_ROOT / "examples"]
    return sorted(
        path for root in roots for path in root.rglob("*.sh") if "ray job submit" in path.read_text(errors="replace")
    )
​
​
SNAPSHOT_UPDATE_ENV_VAR = "MILES_UPDATE_LAUNCH_SCRIPT_SNAPSHOTS"
​
​
def assert_matches_snapshot(snapshot: Path, actual: str, subject: str) -> None:
    # 设置更新环境变量时直接写回快照,用于批量再生成;
    # 平时则严格相等对比,任何漂移都视为失败。
    if os.environ.get(SNAPSHOT_UPDATE_ENV_VAR):
        snapshot.parent.mkdir(parents=True, exist_ok=True)
        snapshot.write_text(actual)
        return
    assert snapshot.exists(), f"missing snapshot for {subject}; regenerate with {SNAPSHOT_UPDATE_ENV_VAR}=1"
    assert actual == snapshot.read_text()

评论区精华

快照测试在未来 CI 中的去留 设计

yueming-yuan 在 issue 评论中担心此类快照测试会让未来 CI 更复杂,建议合入后移除。

结论:作者未在 PR 内回应,快照测试保留;后续 PR 2279 将其改为手动执行,回应了该担忧。 · 待处理

将 .sh 脚本转换为 .py 设计

yueming-yuan 建议趁系列合并期把所有 .sh 脚本迁移为 .py,减少双套维护成本。

结论:本 PR 未做转换;后续系列(PR 1910 等)将模型配置 shell 脚本重写为 Python,方向部分吻合。 · 待处理

风险与影响

  1. 快照严格相等比对对顺序敏感:任何命令行顺序调整都会触发失败,这是设计目标,但也意味着无关联的重排需要人工确认。
  2. shim 命令白名单有限:tests/fast/launch_scripts/sh_harness.py 中未列出的外部命令会在沙箱中真实执行,可能出现网络访问、挂起或污染快照;当前通过 120 秒超时与进程组清理兜底。
  3. 快照更新模式(MILES_UPDATE_LAUNCH_SCRIPT_SNAPSHOTS)会把当前输出直接写回快照,误用时可能掩盖真实回归,需配合 review 人工确认。
  4. CI 时长影响:60+ 脚本逐一沙箱执行,超时上限 120 秒/脚本,后续 PR 2279 将该测试移出手工 CI 以缓解压力。

对用户无直接行为影响;对团队是在 scripts/examples/ 启动脚本上新增硬性回归门槛,任何改动启动命令的 PR 都需要同步更新快照。测试底座被后续多个 PR(如 1901、1906、1908)复用,成为启动脚本系列重构的公共基础设施,长期降低启动命令回归的成本。

快照严格比对易误报 沙箱 shim 覆盖有限 快照更新模式可能掩盖回归 CI 时长增长

关联 Issue

#1837 Tracking issue for refactoring and enhancements

完整报告

参与讨论