面向乳腺癌脑转移(BCBM)的完整计算蛋白设计项目。此版本在原六步流程上增加:
- RTX 4090、RTX 5090、24 GB 及以上 NVIDIA GPU 自动识别与高性能配置。
- RTX 50 系 Blackwell 使用 CUDA 12.8 或更高的 PyTorch 构建,并执行真实 CUDA 矩阵运算验证。
- 无 NVIDIA GPU 时自动进入 CPU 兼容路径,完整完成“模板结构改造 → ProteinMPNN → ESMFold API → 排序 → 报告”。
- HTTP Range 断点续传、
.part文件保留、原子重命名、指数退避和随机抖动。 - Git 增量下载与重试;Git 不可用时自动回退 GitHub 源码压缩包。
- pip 持久缓存、超时控制和多次重试。
- 步骤级、PDB 级、主链级和序列级检查点。
- GPU OOM 自动清理到新进程并使用低显存配置重试。
- 不再生成随机伪序列;不可完成的科学步骤会明确降级或报错。
ai_protein_design_complete/
├── install.sh # 一键安装入口
├── installer.py # GPU/CPU 自适应安装器
├── run_pipeline.sh # 一键运行入口
├── pipeline_runner.py # 检查点、重试、OOM 恢复
├── doctor.py # 环境与 CUDA 实算诊断
├── configs/config.yaml
├── pipeline/
│ ├── runtime.py # 硬件识别和运行参数
│ ├── resilience.py # 续传、重试、Git 和命令恢复
│ └── checkpoints.py # 项目检查点
├── scripts/
│ ├── detect_hardware.py
│ ├── download_rfdiffusion_models.py # RFdiffusion 权重续传
│ └── write_install_state.py
├── 01_rcsb_filter.py
├── 02_download_and_analyze.py
├── 03_rfdiffusion_design.py
├── 04_proteinmpnn_sequence.py
├── 05_validate_and_rank.py
└── 06_pipeline_report.py
安装器会识别 RTX 50 系或 compute capability 12.x,并优先尝试:
- PyTorch stable
cu130(驱动报告 CUDA 13.x 时); - PyTorch stable
cu128; - PyTorch nightly
cu128,仅作为恢复方案。
Blackwell 不会错误回退到 cu121 或 cu126。如果驱动不支持 CUDA 12.8,安装器会记录清晰错误,然后保留 CPU 兼容环境。系统驱动不会被脚本擅自替换。
优先选择驱动能够支持的最高 PyTorch CUDA wheel,通常依次尝试 cu130、cu128、cu126、cu121。安装后执行 GPU 矩阵乘法,只有实算通过才把 CUDA 标记为可用。
24 GB 配置默认:
- 自动安装 ProteinMPNN;
- 自动尝试 RFdiffusion;
- Linux 下自动尝试 Chai-1;
- ProteinMPNN batch size 默认 4;
- RFdiffusion 每个主链独立任务,保障断点恢复;
- Chai-1 默认 1 个样本、1 次 recycle、100 diffusion steps,降低 24 GB OOM 风险。
使用 gpu_ultra 配置:更高 ProteinMPNN batch、完整 Chai 参数和更多验证候选。
完整自动流程仍可运行:
RCSB 筛选
→ 下载目标 PDB
→ CPU 模板结构改造后端
→ ProteinMPNN CPU 序列设计
→ ESMFold API 结构预测
→ RMSD / 稳定性 / BBB 评分
→ 报告
CPU 后端不会假称完成了 de novo RFdiffusion。它在输出中标记为:
cpu_template_fallback
该路径属于“结构保持的模板重设计”,适合无 GPU 环境进行完整软件验证和初步序列改造。Binder、无条件新骨架和小分子 motif 的真实主链生成仍需要 RFdiffusion GPU 后端。
unzip ai_protein_design_resilient_v2_1.zip
cd ai_protein_design_complete
bash install.sh安装过程可随时中断。重新执行相同命令时:
- 已完成 pip wheel 保留在
.cache/pip/; - 已完成 Git objects 保留在
external/<repo>/.git/; - RFdiffusion 模型保留为
*.part并从已有字节继续; - Chai 权重保留在
.cache/chai/; - 已完成组件不会被无条件删除。
bash install.sh \
--with-proteinmpnn \
--with-rfdiffusion \
--with-chai \
--include-optional-rfdiffusion-modelsbash install.sh \
--cpu-only \
--with-proteinmpnn \
--without-rfdiffusion \
--without-chai安装过程中网络断开或某个组件失败后,直接再次执行:
bash install.sh --repair未完成下载会继续,不会主动删除 .part、pip cache 或 Git objects。
source .venv/bin/activate
python doctor.py
bash run_pipeline.sh --smoke重点检查:
{
"torch_probe": {
"cuda_available": true,
"device": "NVIDIA GeForce RTX 4090",
"capability": "8.9",
"cuda_math_ok": true
}
}5090 应显示 capability 12.x,并且 cuda_math_ok 必须为 true。
最推荐命令:
bash run_pipeline.sh --step all --mode auto自动逻辑:
- 有可用 RFdiffusion GPU:使用
scaffold; - 没有可用 RFdiffusion GPU:使用
redesign+ CPU template backend; - Chai-1 未安装时自动跳过,不阻塞 ESMFold 路径;
- 中断后重新运行,从检查点继续。
bash run_pipeline.sh \
--step all \
--mode scaffold \
--backend rfdiffusion \
--max-retries 6Binder:
bash run_pipeline.sh --step 1
bash run_pipeline.sh --step 2
bash run_pipeline.sh --step 3 --mode binder --target EGFR --backend rfdiffusion
bash run_pipeline.sh --step 4 --mode binder --target EGFR
bash run_pipeline.sh --step 5
bash run_pipeline.sh --step 6bash run_pipeline.sh \
--step all \
--mode auto \
--backend auto \
--skip-chai也可以明确指定:
bash run_pipeline.sh \
--step all \
--mode redesign \
--backend cpu-template \
--skip-chai默认开启恢复:
bash run_pipeline.sh --step all --mode auto --resume强制从某一步重新执行:
bash run_pipeline.sh --step all --mode auto --restart-step 5关闭项目检查点:
bash run_pipeline.sh --step 3 --mode scaffold --no-resume增加重试次数:
bash run_pipeline.sh --step all --mode auto --max-retries 8 --retry-delay 5检查点目录:
.pipeline/checkpoints/
日志目录:
.pipeline/logs/
下载缓存:
.cache/
使用 HTTP Range:
python scripts/download_rfdiffusion_models.py --attempts 12下载中的文件:
Base_ckpt.pt.part
Base_ckpt.pt.part.json
只有下载完整并通过最低大小检查后,才原子改名为:
Base_ckpt.pt
每个 PDB 同样使用续传下载,并逐个写入分析 CSV。第 200 个 PDB 失败时,前 199 个结果不会丢失。
POST 请求不能按字节续传,但会执行:
- 连接重试;
- 读取重试;
- 429 / 5xx 重试;
- Retry-After;
- 指数退避;
- 单序列结果缓存。
已生成的 ESMFold PDB 不会重复请求。
检测到以下错误时:
CUDA out of memory
OutOfMemoryError
CUBLAS_STATUS_ALLOC_FAILED
运行器会:
- 结束失败进程,释放其显存;
- 启动全新进程;
- 设置
PIPELINE_LOW_MEMORY=1; - 将 ProteinMPNN batch 降至 1;
- 将 CUDA allocator 的
max_split_size_mb降至 64; - 减少 Chai recycle 和 diffusion steps;
- 从已有主链/序列检查点继续。
程序严格区分以下后端:
rfdiffusion
cpu_template_fallback
不会把 CPU 模板复制包装成 RFdiffusion 新骨架,也不会在 ProteinMPNN 缺失时生成随机氨基酸序列。
- “保证在所有机器上绝不出错”在第三方 CUDA、驱动、网络和上游仓库不断变化的条件下并不现实。本项目采用的是可验证、可恢复、可降级的工程机制。
- RTX 5090 需要支持 CUDA 12.8 或更高的驱动和 PyTorch build。
- RFdiffusion 上游依赖历史 SE(3)-Transformer/DGL 环境,安装器会自动尝试并运行 import/help probe;若该组件仍与当前系统不兼容,会保留 CPU template fallback,而不是伪造成功。
- Chai-1 官方要求 Linux、CUDA 和 bfloat16 GPU;24 GB 消费卡适合较小复合物,较大复合物可能仍需要 48–80 GB。
- 输出是计算候选,不代表药物安全性、有效性或临床结论。
- PyTorch installation: https://pytorch.org/get-started/locally/
- NVIDIA Blackwell compatibility: https://docs.nvidia.com/cuda/blackwell-compatibility-guide/
- RFdiffusion: https://github.com/RosettaCommons/RFdiffusion
- ProteinMPNN: https://github.com/dauparas/ProteinMPNN
- Chai-1: https://github.com/chaidiscovery/chai-lab

