complexa-slurm

Complexa SLURM 技能

使用 slurm_utils/launch_protein_binder_search.sh、 slurm_utils/launch_laproteina_design_pipeline.sh、 slurm_utils/launch_monomer_eval_from_pdb_dir.sh 和 slurm_utils/launch_laproteina_train.sh，向远程 SLURM 集群提交 Proteina-Complexa 作业。先探查本地和集群环境，收集所需标志，始终先使用 --dry-run 预览，并在提交前向用户展示解析后的 sbatch 命令，然后提交、捕获作业 ID，并生成 slurm_manifest.json 以供重放。

不涉及 complexa CLI。 SLURM 提交只能通过 bash 脚本完成——没有 complexa slurm 子命令。启动脚本会加载 .env、通过 rsync 同步仓库、生成 sbatch 脚本，并通过 SSH 调用 sbatch。不要尝试用 ssh ... sbatch <<EOF 替代它们——那样会失去 rsync 门控、sweeper 展开和 --dry-run 安全保障。

此技能支持的功能

• 在集群 GPU 上启动四阶段 binder 搜索流水线（生成 → 筛选 → 评估 → 分析），支持单目标或通过 --targets-file 指定多个目标。 通过向启动脚本传入匹配的 configs/search_*_pipeline.yaml，驱动 Complexa 的全部三种设计流水线——蛋白 binder（默认）、配体 binder、AME / 酶支架设计。
• 启动 LaProteina 单体设计（launch_laproteina_design_pipeline.sh， 默认配置 design_monomer_pipeline.yaml），用于无条件蛋白质生成。
• 根据 YAML 配置调整多节点分布式训练规模（nnodes_、ngpus_per_node_、ncpus_per_task_train_、run_name）。
• 通过 --runtime 为每次提交选择 Docker（Pyxis/Enroot） 或 UV venv 运行环境。
• 在集群上编排 Hydra sweeps（--sweeper）和每个配置的覆盖项（--override）。
• 默认安全——每次提交都必须先通过 dry-run 预览，并由用户批准。
• 通过 rsync 自动下载结果（可使用 --skip-download 跳过）。

步骤 1：集群预检

本地预检（_shared/scripts/preflight.sh）会检查本地主机上的 GPU / 检查点 / 工具。SLURM 提交还需要集群可连接、有 sbatch 可执行文件，并且分区处于运行状态。请运行两项预检。

bash .claude/skills/_shared/scripts/preflight.sh           # local
bash .claude/skills/complexa-slurm/scripts/cluster_preflight.sh   # cluster

cluster_preflight.sh 会写入 cluster_preflight.json 并检查：

1. .env 第 5 节（CLUSTER_USER、CLUSTER_HOST、CLUSTER_ACCOUNT、 CLUSTER_PARTITION、CLUSTER_ROOT_REMOTE、CLUSTER_DATA_PATH、 CLUSTER_RUNTIME）已填写。运行环境对应的变量： CLUSTER_CONTAINER_IMAGE（docker）或 CLUSTER_UV_VENV + CLUSTER_CACHE_DIR （uv）。
2. ssh -o BatchMode=yes -o ConnectTimeout=5 $CLUSTER_USER@$CLUSTER_HOST true 执行成功（必须已配置免密码 SSH）。
3. 登录节点的 $PATH 中有 sbatch。
4. sinfo -p $CLUSTER_PARTITION -h 至少返回一个节点。
5. 集群上存在 CLUSTER_CKPT_PATH 和 CLUSTER_SHARED_MODELS_PATH（如果已设置）。
6. 本地仓库的 git SHA（用于清单）。

如果任何必需检查失败，请勿继续步骤 2——先修复 .env、SSH 或权限问题。各变量的含义请参阅 reference/cluster_env.md。

步骤 2：选择工作负载

意图	启动脚本
Complexa 设计流水线（蛋白 binder、配体 binder、AME）	slurm_utils/launch_protein_binder_search.sh（传入匹配的 configs/search_*_pipeline.yaml）
LaProteina 无条件单体设计	slurm_utils/launch_laproteina_design_pipeline.sh
从 PDB 目录进行单体评估	slurm_utils/launch_monomer_eval_from_pdb_dir.sh
分布式训练（微调、RL、基础训练）	slurm_utils/launch_laproteina_train.sh

完整的标志列表见 reference/slurm_workloads.md。

步骤 3：通过 AskUserQuestion 收集参数

只询问上下文中缺少的信息。建议顺序：

1. 工作负载 — binder 搜索 / LaProteina 单体设计 / 训练。
2. 目标 / 任务（binder）— 例如 22_DerF21、02_PDL1、39_7V11_LIGAND， 或 M0096_1chm；也可以提供批处理模式所用的 --targets-file 路径。
3. 配置 YAML 路径（训练）— 默认为 configs/training_local_latents.yaml；如果用户要用其他配置，请询问。
4. 运行时覆盖 — 使用 .env 中的默认值，还是强制使用 docker / uv？
5. Sweep — 是否传入 --sweeper FILE？（仅适用于 binder）
6. 目标文件 — 是否传入 --targets-file FILE 和 --num-runs N？（仅适用于 binder）
7. 是否已在集群上？ — 如果用户当前位于登录节点，则传入 --on-cluster（跳过 SSH/rsync）。仅适用于 binder 脚本。
8. 是否将结果保留在远程端？ — 传入 --skip-download，将输出留在集群上。默认情况下会通过 rsync 将输出同步回来。

步骤 4：务必先进行试运行

构建命令并使用 --dry-run 运行。向用户展示完整输出 — 尤其是生成的 sbatch 脚本正文、Hydra 覆盖项、rsync 目标以及资源配置块（--nodes、--gres=gpu、--time、 --partition）。提交前必须明确征得用户同意。

./slurm_utils/launch_protein_binder_search.sh --dry-run 22_DerF21
./slurm_utils/launch_laproteina_design_pipeline.sh --dry-run
./slurm_utils/launch_laproteina_train.sh --dry-run --config configs/training_local_latents.yaml

试运行日志会显示 "DRY RUN MODE — No changes will be made" 和 "Would create slurm_batch_script_*.sh with content: …"。原样展示该代码块。

此关卡不可省略的原因： 每次实际提交都会占用账户配额、安排可能长达许多节点小时的任务、可能导致其他用户的任务被驱逐，并会在 $CLUSTER_ROOT_REMOTE 上创建运行目录，且脚本拒绝覆盖该目录。提交前先检查，是发现分区错误、运行名称错误或 Hydra 覆盖项中的拼写错误，成本最低的方式。

步骤 5：提交

使用相同的命令重新运行，但不带 --dry-run。启动脚本会读取 .env，在 .rsync_lock/ 下同步代码，将 slurm_batch_script_*.sh 写入集群上的运行目录，并通过 SSH 调用 sbatch。

./slurm_utils/launch_protein_binder_search.sh 22_DerF21
./slurm_utils/launch_laproteina_design_pipeline.sh
./slurm_utils/launch_laproteina_train.sh --config configs/training_local_latents.yaml

查看标准输出中是否出现 Submitted batch job <ID>（辅助程序会记录 Submitted <stage> [i/N] with Job ID: <ID>）。训练任务会重复提交 --dependency=singleton，次数为 --num-jobs（默认值为 5）— 请记录所有 ID。Binder 搜索和 LaProteina 设计流水线会按顺序提交各阶段，并在阶段之间调用 wait_for_job，因此调用方会一直阻塞，直到获取下一阶段的 ID。

步骤 6：监控

在提交时获取一次状态快照并展示给用户。不要轮询 — 重新调用此技能的成本很低。

ssh "$CLUSTER_USER@$CLUSTER_HOST" "squeue -j <id1>,<id2>,... -u $CLUSTER_USER"

集群上的日志位置（位于运行目录 $CLUSTER_ROOT_REMOTE/<run_name>/ 下）：

脚本	阶段	日志路径
binder 搜索	生成	slurm_run_outputs/inf/slurm_<jobid>_<array>.{out,err}
binder 搜索	过滤	slurm_run_outputs/filter/slurm_<jobid>_<array>.{out,err}
binder 搜索	评估	slurm_run_outputs/eval/slurm_<jobid>_<array>.{out,err}
binder 搜索	分析	slurm_run_outputs/agg/slurm_<jobid>_<array>.{out,err}
LaProteina 设计流水线	gen/eval/agg	slurm_run_outputs/{gen,eval,agg}/...
训练	全部	slurm_run_outputs/slurm_<jobid>.{out,err}

后续跟进（任务仍在运行？是否失败？）时，请让用户重新调用此技能，或自行运行 ssh $CLUSTER_HOST，然后执行 squeue -u $CLUSTER_USER / sacct -j <id>。

第 7 步：下载结果

用户报告作业已完成时：

• 默认（不带标志）：launch_*.sh 会在流水线结束时自动使用 rsync 将结果同步回来。Binder 搜索会将结果下载到 <run_name>-<YYYY_MM_DD_HH>/{inference,evaluation_results}，并通过 script_utils/aggregate_successful_samples/aggregate_successful_samples.py 汇总。 LaProteina 设计流水线会将结果下载到 results_downloaded/<run_name>-<YYYY_MM_DD_HH>/。
• 传入了 --skip-download：结果保留在集群上的 $CLUSTER_ROOT_REMOTE/<run_name>/。请手动拉取： rsync -az $CLUSTER_USER@$CLUSTER_HOST:$CLUSTER_ROOT_REMOTE/<run_name>/ ./local_results/。
• 传入了 --on-cluster：流水线在登录节点上运行；结果保留在集群上的当前工作目录中。

训练不会自动下载——检查点位于 $CLUSTER_ROOT_REMOTE/<run_name>/（或配置了 $CLUSTER_CKPT_PATH 时所在的路径）。

第 8 步：生成 slurm_manifest.json

在 ./complexa_slurm/ 中放置一个清单文件，让用户可以用一个文件获取重放所需的全部内容。

{
  "timestamp": "2026-05-15T18:32:01Z",
  "workload": "binder_search",
  "launch_cmd": "./slurm_utils/launch_protein_binder_search.sh --runtime docker 22_DerF21",
  "cluster_host": "login.mycluster.example.com",
  "cluster_partition": "gpu,compute",
  "cluster_account": "my_slurm_account",
  "runtime": "docker",
  "run_name": "search_binder_pipeline-search-22_DerF21",
  "remote_run_dir": "/lustre/.../runs/search_binder_pipeline-search-22_DerF21",
  "sbatch_scripts": [
    "slurm_batch_script_gen.sh",
    "slurm_batch_script_filter.sh",
    "slurm_batch_script_eval.sh",
    "slurm_batch_script_analyze.sh"
  ],
  "slurm_job_ids": {"generation": "1234567", "filter": "...", "evaluation": "...", "analyze": "..."},
  "git_sha": "abc1234"
}

启动脚本已经会将 .githash 写入远程运行目录；请将该 SHA 同步到本地清单文件中。

硬件

工作负载	每个作业的节点数	每个节点的 GPU 数	每个任务的 CPU 数	墙钟时间（默认）	备注
Binder 搜索（任意阶段）	1	1（gpus-per-node 1）	来自流水线 YAML 中的 ncpus_	04:00:00	作业数组，每个配置对应一个作业
LaProteina 设计流水线（任意阶段）	1	1	来自流水线 YAML 中的 ncpus_	02:00:00	作业数组，每个配置对应一个作业
训练	来自 nnodes_（例如 12）	来自 ngpus_per_node_（例如 8）	来自 ncpus_per_task_train_	04:00:00	单作业重排队链（典型配置：12 × 8 = 96 个 H100，用于 Binder 微调）

默认值来自 slurm_utils/slurm_helper.sh 中的 generate_array_slurm_header / generate_train_slurm_header。各脚本的完整说明见 reference/slurm_workloads.md。

故障排查

症状	原因	解决方法
ssh: connect: Permission denied (publickey)	未配置免密 SSH / CLUSTER_SSH_KEY 错误	运行 ssh-copy-id，或在 .env 中设置 CLUSTER_SSH_KEY=/abs/path/key
Directory ... already exists on remote	集群上已有使用相同 run_name 的先前运行	运行 ssh $CLUSTER_HOST 'rm -rf $RUN_DIR'，或修改 YAML 中的 run_name
Container sqsh file not found on cluster	CLUSTER_CONTAINER_IMAGE 路径错误	运行 ssh $CLUSTER_HOST 'ls -lh <image>' 验证，或改用仓库 URL
Missing required fields in <pipeline_config>	YAML 中没有 gen_njobs / eval_njobs / run_name	根据规范流水线 YAML 重新推导；不要编辑下游副本
INVALID_QOS / account ... not allowed in partition	CLUSTER_ACCOUNT 无权访问 CLUSTER_PARTITION	选择该账户有权限使用的分区；运行 ssh $CLUSTER_HOST 'sshare -A $CLUSTER_ACCOUNT'
所有数组作业都无限期处于 PENDING (Resources) 状态	分区超额订阅或预留配置错误	尝试其他分区；运行 sinfo -p $CLUSTER_PARTITION 检查

更长的故障排查列表（抢占、容器拉取、缺少检查点、rsync 故障、目标文件格式、--on-cluster 与默认模式的区别）见 reference/troubleshooting.md。

---

参考资料：

• reference/cluster_env.md — .env 第 5 节中的变量、默认值，以及缺少各变量时会出现的问题。
• reference/slurm_workloads.md — 各启动脚本的完整标志矩阵和 sbatch 结构，并附有可直接复制粘贴的示例。
• reference/troubleshooting.md — 常见故障模式及其诊断和解决方法。