AI预测蛋白质结构 - Protenix的快速部署和应用

应用概述

Protenix 是由字节跳动 AML AI4Science 团队开发的一款基于 PyTorch 的深度学习框架，用于蛋白质结构预测，是对 AlphaFold 3 的一个开源复现。其核心目标是降低 AlphaFold 等复杂蛋白质预测模型的技术门槛，从而更好地支持跨学科的研究需求。

使用场景

Protenix 作为一个蛋白质结构预测的开源框架，有多种应用场景，主要集中在生命科学、药物开发和基础研究领域：

1. 蛋白质结构预测：蛋白质结构预测是生命科学领域的一个重要问题，Protenix 能够提供高效、准确的蛋白质结构预测服务。
2. 药物开发：在药物研发中，Protenix 提供了一种高效工具，帮助药物设计者更快地完成药物靶点验证、小分子筛选与优化、抗体与疫苗设计等工作。
3. 基础研究：Protenix 可用于探索基础生物学问题，例如：蛋白质结构与功能的关系、蛋白质结构与生物信息学的关系、蛋白质结构与免疫学的关系等。

本应用经部署后，可以提供以下服务：

1. 蛋白质预测推理
2. 特定数据集训练或者微调模型
3. 其他研究应用

操作步骤

部署前准备工作

本次部署使用到Docker和Kubernetes，请准备好Docker环境，并确保本地有可用的Kubernestes客户端工具kubectl。
Docker环境安装请参考: 安装Docker
kubectl安装请参考：安装命令行工具（kubectl）

开通集群

资源最低要求

资源类型	数量	说明
GPU	1个起	gpu-h800
CPU	16 core起
内存	320G 起
存储	2T 起
Harbor存储	100G	用于保存镜像
对象存储	300G	用于保存模型和数据

申请开通

集群请参考：开通弹性容器集群 集群开通后，需要将集群配置信息保存到合适的位置，以便后续使用。

脚本准备

请下载模板脚本文件，并解压到本地目录。
在前面下载并解压的模板脚本文件中，对所有文件进行变量替换。

用户名密码：查看开通镜像仓库时的通知短信
镜像仓库地址：参考镜像仓库的使用
镜像仓库地址是由 镜像仓库域名/项目 组成

需要替换的变量有：

变量名	所在文件	例子
镜像仓库地址	deploy-protenix.yaml， image.sh, image.bat	registry.hd-01.alayanew.com:8443/alayanew-******-5cfd029439a8
GPU_resource_name	deploy-protenix.yaml	gpu-h800
镜像仓库域名	harbor-config.json, image.sh, image.bat	registry.hd-01.alayanew.com:8443
用户名	harbor-config.json, image.sh, image.bat	hb_abc123
密码	harbor-config.json, image.sh, image.bat	123456

secret准备

secret文件是用于保存敏感信息的，如对象存储的密钥等。本次部署中需要对harbor registry的用户名和密码进行加密，并保存到secret文件中。

• 当执行过前面的步骤后，请将harbor-config.json中的内容做base64编码。linux中可以使用base64命令进行编码。或者在线网站也可以进行编码。可以参考：https://tool.lu/encdec/
• 将编码后的内容保存到harbor-secret.yaml的.dockerconfigjson字段。 例子：

apiVersion: v1
kind: Secret
metadata:
  name: alaya-harbor-secret
type: kubernetes.io/dockerconfigjson
data:
  .dockerconfigjson: <base64编码后的内容>

镜像准备

包括以下步骤：

1. 执行刚才下载并解压的scripts目录下的Dockerfile文件，生成镜像
2. 登录镜像仓库，推送镜像

在scripts目录下, 请执行image.sh(linux)或者image.bat(windows)脚本，完成镜像的构建和推送。 当执行完毕后，请检查镜像仓库中是有镜像文件protenix-dc:v0.0.1存在。
关于如何查看镜像仓库，请参考：使用Harbor管理镜像资源

模型，数据准备

模型文件：https://af3-dev.tos-cn-beijing.volces.com/release_model/model_v1.pt
数据文件：https://af3-dev.tos-cn-beijing.volces.com/release_data.tar.gz
注意： 数据文件较大（238G），请确保本地有足够的存储空间。
注意： 由于文件较大，下载时间较长，请保证网络稳定。直接下载可能导致下载中断而需要重新下载。推荐使用下载工具下载。

如果只做推理，那么只需要下载以下两个文件即可：

• https://af3-dev.tos-cn-beijing.volces.com/release_data/components.v20240608.cif
• https://af3-dev.tos-cn-beijing.volces.com/release_data/components.v20240608.cif.rdkit_mol.pkl

准备工作包含以下步骤：

1. 下载上面的模型文件和数据文件
2. 上传到对象存储中。 对象存储的上传以及检查请参考：对象存储的使用

部署应用

包含以下步骤：

1. 执行k8s配置文件，部署应用
2. 检查k8s部署结果
3. 进入pod，并进行初始化

执行部署脚本

请在脚本目录下执行以下命令：

kubectl apply -k .

确保执行过程中没有错误。

检查部署结果

请执行以下命令：

kubectl get pods -n protenix

如果出现以下输出，说明部署成功：

NAME                               READY   STATUS    RESTARTS   AGE
deploy-protenix-7bfc98c86d-ls67t   1/1     Running   0          5h10m

其中deploy-protenix-7bfc98c86d-ls67t是pod的名字，每次部署都会产生一个新的pod名字，后面的命令都需要使用这个名字。

进入pod，并进行初始化

请执行以下命令：

kubectl exec -it deploy-protenix-7bfc98c86d-ls67t -n protenix -- /bin/bash

注意： 请使用实际的pod名字替换上面的deploy-protenix-7bfc98c86d-ls67t。

进入pod后，请执行以下命令：

/app/init.sh

从对象存储中下载相应数据：

• 下载模型文件到/af3-dev/release_model/目录
  • 模型文件：model_v1.pt
• 下载数据文件到/af3-dev/release_data/目录
  • 如果只做推理，那么只需要下载以下两个文件：components.v20240608.cif 和 components.v20240608.cif.rdkit_mol.pkl
  • 下载release_data.tar.gz并解压到/af3-dev/release_data/目录

最后请做以下检查：

• /workspace中存在目录protenix
• 模型检查：/af3-dev/release_model/model_v1.pt
• 数据检查：
  • 如果只做推理，那么只需要检查以下两个文件：
    • /af3-dev/release_data/components.v20240608.cif
    • /af3-dev/release_data/components.v20240608.cif.rdkit_mol.pkl
  • 如果做训练，那么需要在/af3-dev/release_data/目录下检查以下文件：

    ├── components.v20240608.cif [408M] # ccd source file
    ├── components.v20240608.cif.rdkit_mol.pkl [121M] # rdkit Mol object generated by ccd source file
    ├── indices [33M] # chain or interface entries
    ├── mmcif [283G]  # raw mmcif data
    ├── mmcif_bioassembly [36G] # preprocessed wwPDB structural data
    ├── mmcif_msa [450G] # msa files
    ├── posebusters_bioassembly [42M] # preprocessed posebusters structural data
    ├── posebusters_mmcif [361M] # raw mmcif data
    ├── recentPDB_bioassembly [1.5G] # preprocessed recentPDB structural data
    └── seq_to_pdb_index.json [45M] # sequence to pdb id mapping file

开始使用

推理

参考：推理

进入/workspace/protenix目录，执行以下命令

bash inference_demo.sh

用时说明：
资源配置：

• cpu: 4
• memory: 64Gi
• nvidia.com/gpu-h800: 1

用时：

• 第一次推理：约14分钟，这其中包含一些代码编译时间和数据加载时间，后续推理时间较短
• 第二次推理：约6分钟

训练和微调

参考：训练
进入/workspace/protenix目录，执行以下命令 注意： 训练需要较长时间，请保持网络畅通。推荐使用tmux等工具进行后台运行。

bash train_demo.sh

多卡训练：

torchrun --nproc_per_node=4 runner/train.py

容器中的目录说明

• /workspace/protenix：Protenix的工作目录，包含代码、执行脚本等。
• /workspace/protenix/output：Protenix的输出目录。
• /af3-dev/release_data：数据目录，包含训练数据等。
• /af3-dev/release_model：模型目录，包含训练好的模型。

关于license

请按照Protenix的版权要求使用，请参考：https://github.com/bytedance/Protenix/blob/main/LICENSE