Skip to content

About

ai_protein_design,AI Protein Design Pipeline for BCBM,面向乳腺癌脑转移(BCBM)的完整计算蛋白设计项目

Resources

Stars

2 stars

Watchers

0 watching

Forks

Latest commit

 

History

2 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

AI Protein Design Pipeline for BCBM — Resilient Edition 2.1

1

面向乳腺癌脑转移(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

1

支持的硬件

RTX 5090 / Blackwell

安装器会识别 RTX 50 系或 compute capability 12.x,并优先尝试:

  1. PyTorch stable cu130(驱动报告 CUDA 13.x 时);
  2. PyTorch stable cu128;
  3. PyTorch nightly cu128,仅作为恢复方案。

Blackwell 不会错误回退到 cu121 或 cu126。如果驱动不支持 CUDA 12.8,安装器会记录清晰错误,然后保留 CPU 兼容环境。系统驱动不会被脚本擅自替换。

RTX 4090 / Ada,24 GB

优先选择驱动能够支持的最高 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 风险。

48 GB 或以上

使用 gpu_ultra 配置:更高 ProteinMPNN batch、完整 Chai 参数和更多验证候选。

没有 GPU

完整自动流程仍可运行:

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/;
  • 已完成组件不会被无条件删除。

强制安装全部 GPU 组件

bash install.sh \
  --with-proteinmpnn \
  --with-rfdiffusion \
  --with-chai \
  --include-optional-rfdiffusion-models

仅 CPU

bash 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 路径;
  • 中断后重新运行,从检查点继续。

4090 / 5090 全面运行

bash run_pipeline.sh \
  --step all \
  --mode scaffold \
  --backend rfdiffusion \
  --max-retries 6

Binder:

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 6

无 GPU 全面运行

bash 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/

下载保护机制

RFdiffusion 模型

使用 HTTP Range:

python scripts/download_rfdiffusion_models.py --attempts 12

下载中的文件:

Base_ckpt.pt.part
Base_ckpt.pt.part.json

只有下载完整并通过最低大小检查后,才原子改名为:

Base_ckpt.pt

RCSB PDB

每个 PDB 同样使用续传下载,并逐个写入分析 CSV。第 200 个 PDB 失败时,前 199 个结果不会丢失。

ESMFold API

POST 请求不能按字节续传,但会执行:

  • 连接重试;
  • 读取重试;
  • 429 / 5xx 重试;
  • Retry-After;
  • 指数退避;
  • 单序列结果缓存。

已生成的 ESMFold PDB 不会重复请求。

GPU OOM 恢复

检测到以下错误时:

CUDA out of memory
OutOfMemoryError
CUBLAS_STATUS_ALLOC_FAILED

运行器会:

  1. 结束失败进程,释放其显存;
  2. 启动全新进程;
  3. 设置 PIPELINE_LOW_MEMORY=1;
  4. 将 ProteinMPNN batch 降至 1;
  5. 将 CUDA allocator 的 max_split_size_mb 降至 64;
  6. 减少 Chai recycle 和 diffusion steps;
  7. 从已有主链/序列检查点继续。

输出真实性

程序严格区分以下后端:

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。
  • 输出是计算候选,不代表药物安全性、有效性或临床结论。

上游资料

About

ai_protein_design,AI Protein Design Pipeline for BCBM,面向乳腺癌脑转移(BCBM)的完整计算蛋白设计项目

Resources

Stars

2 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages