OneScience/install.sh at main · onescience-ai/OneScience · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
#!/bin/bash

# 开启严格模式：遇到错误、未定义变量或管道错误时立即退出
set -euo pipefail
pip install uv
# ==========================================
# 1. 参数校验与环境准备
# ==========================================
DOMAIN="${1:-all}" # 获取用户输入的参数，默认安装全部领域

# 处理帮助命令
if [[ "$DOMAIN" == "-h" || "$DOMAIN" == "--help" ]]; then
    echo "Usage: bash $0 [earth|cfd|bio|matchem|all]"
    echo "If no domain is specified, all domains will be installed by default."
    exit 0
fi

# 校验输入的领域是否支持
case "$DOMAIN" in
    earth|cfd|bio|matchem|all)
        # 校验通过，继续执行
        ;;
    *)
        echo "Error: Unsupported domain '$DOMAIN'"
        echo "Usage: bash $0 [earth|cfd|bio|matchem|all]"
        exit 1
        ;;
esac

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
CONSTRAINTS_FILE="$SCRIPT_DIR/constraints.txt"

# ==========================================
# 2. 依赖配置区 (后续维护仅需修改此处)
# ==========================================
# 约定：
# - CORE_WHEELS: 所有领域都依赖的基础轮子
# - EARTH/CFD/BIO/MATCHEM_WHEELS: 特定领域的定制轮子

CORE_WHEELS=(
    # 基础依赖占位：如果所有领域都需要某个轮子，添加在这里
)

EARTH_WHEELS=(
    "https://download.sourcefind.cn:65024/directlink/4/pytorch/DAS1.7/torch-2.5.1+das.opt1.dtk25042-cp311-cp311-manylinux_2_28_x86_64.whl"
    "https://download.sourcefind.cn:65024/directlink/4/vision/DAS1.7/torchvision-0.20.1+das.opt1.dtk25042-cp311-cp311-manylinux_2_28_x86_64.whl"
    "https://download.sourcefind.cn:65024/file/4/apex/DAS1.6/apex-1.4.0+das.opt1.dtk25041-cp311-cp311-manylinux_2_28_x86_64.whl"
    "https://download.sourcefind.cn:65024/directlink/4/onnxruntime/DAS1.7/onnxruntime-1.19.2+das.opt1.dtk25042-cp311-cp311-manylinux_2_28_x86_64.whl"
    "https://download.sourcefind.cn:65024/directlink/4/dgl/DAS1.7/dgl-2.2.1+das.opt1.dtk25042-cp311-cp311-manylinux_2_28_x86_64.whl"
    "https://download.sourcefind.cn:65024/file/4/flash_attn/DAS1.6/flash_attn-2.6.1+das.opt1.dtk25041-cp311-cp311-manylinux_2_28_x86_64.whl"
    "https://download.sourcefind.cn:65024/file/4/triton/DAS1.7/triton-3.1.0+das.opt1.dtk25042-cp311-cp311-manylinux_2_28_x86_64.whl"
)

CFD_WHEELS=(
    "https://download.sourcefind.cn:65024/directlink/4/pytorch/DAS1.7/torch-2.5.1+das.opt1.dtk25042-cp311-cp311-manylinux_2_28_x86_64.whl"
    "https://download.sourcefind.cn:65024/directlink/4/vision/DAS1.7/torchvision-0.20.1+das.opt1.dtk25042-cp311-cp311-manylinux_2_28_x86_64.whl"
    "https://download.sourcefind.cn:65024/directlink/4/fastpt/DAS1.7/fastpt-2.1.1+das.dtk25042-cp311-cp311-linux_x86_64.whl"
    # "https://download.sourcefind.cn:65024/directlink/4/fastpt/DAS1.7/fastpt-2.2.0+das.dtk25042-cp311-cp311-linux_x86_64.whl"
    "https://download.sourcefind.cn:65024/directlink/4/dgl/DAS1.7/dgl-2.2.1+das.opt1.dtk25042-cp311-cp311-manylinux_2_28_x86_64.whl"
    "https://download.sourcefind.cn:65024/file/4/tensorflow/DAS1.7/tensorflow-2.18.0+das.opt1.dtk25042-cp311-cp311-manylinux_2_28_x86_64.whl"
    "https://download.sourcefind.cn:65024/directlink/9/onesicence/dtk-25.04.2/torch_scatter-2.1.0+das.opt1.dtk25043-cp311-cp311-linux_x86_64.whl"
    "https://download.sourcefind.cn:65024/directlink/9/onesicence/dtk-25.04.2/torch_cluster-1.6.3+das.opt1.dtk25043-cp311-cp311-linux_x86_64.whl"
)

BIO_WHEELS=(
    "https://download.sourcefind.cn:65024/directlink/4/pytorch/DAS1.7/torch-2.5.1+das.opt1.dtk25042-cp311-cp311-manylinux_2_28_x86_64.whl"
    "https://download.sourcefind.cn:65024/directlink/4/vision/DAS1.7/torchvision-0.20.1+das.opt1.dtk25042-cp311-cp311-manylinux_2_28_x86_64.whl"
    "https://download.sourcefind.cn:65024/file/4/jax/DAS1.7/jax-0.4.34+das.opt1.dtk25042-py3-none-any.whl"
    "https://download.sourcefind.cn:65024/file/4/jax/DAS1.7/jax_rocm60_pjrt-0.4.34+das.opt1.dtk25042-py3-none-manylinux2014_x86_64.whl"
    "https://download.sourcefind.cn:65024/file/4/jax/DAS1.7/jax_rocm60_plugin-0.4.34+das.opt1.dtk25042-cp311-cp311-manylinux2014_x86_64.whl"
    "https://download.sourcefind.cn:65024/file/4/jax/DAS1.7/jaxlib-0.4.34+das.opt1.dtk25042-cp311-cp311-manylinux2014_x86_64.whl"
    "https://download.sourcefind.cn:65024/directlink/4/transformer_engine/DAS1.7/transformer_engine-2.5.0+das.opt1.dtk25042-cp311-cp311-manylinux_2_28_x86_64.whl"
    "https://download.sourcefind.cn:65024/directlink/4/flash_attn/DAS1.7/flash_attn-2.6.1+das.opt1.dtk25042-cp311-cp311-manylinux_2_28_x86_64.whl"
    "https://download.sourcefind.cn:65024/directlink/9/onesicence/dtk-25.04.2/bionemo_noodles-0.1.2-cp311-cp311-manylinux_2_28_x86_64.whl"
    "https://download.sourcefind.cn:65024/file/9/onesicence/dtk-26.04/nemo_toolkit-2.5.0rc0-py3-none-any.whl"
    "https://download.sourcefind.cn:65024/directlink/9/onesicence/dtk-25.04.2/openmm-8.3.1-cp311-cp311-linux_x86_64.whl"
    "https://download.sourcefind.cn:65024/file/4/triton/DAS1.7/triton-3.1.0+das.opt1.dtk25042-cp311-cp311-manylinux_2_28_x86_64.whl"
    "https://download.sourcefind.cn:65024/file/4/jax_triton/DAS1.7/jax_triton-0.2.0+das.opt1.dtk25042-py3-none-any.whl"
    "https://download.sourcefind.cn:65024/file/4/tensorflow/DAS1.7/tensorflow-2.18.0+das.opt1.dtk25042-cp311-cp311-manylinux_2_28_x86_64.whl"
    "https://download.sourcefind.cn:65024/directlink/4/dgl/DAS1.7/dgl-2.2.1+das.opt1.dtk25042-cp311-cp311-manylinux_2_28_x86_64.whl"
    # "https://download.sourcefind.cn:65024/file/4/dgl/DAS1.7/dgl-2.2.1+das.opt1.dtk25042-cp311-cp311-manylinux_2_28_x86_64.whl"
    "https://download.sourcefind.cn:65024/file/9/onesicence/dtk-25.04.2/megatron_core-0.15.0-cp311-cp311-linux_x86_64.whl"
    "https://download.sourcefind.cn:65024/file/9/onesicence/dtk-25.04.2/causal_conv1d-1.5.2-cp311-cp311-linux_x86_64.whl"
    "https://download.sourcefind.cn:65024/directlink/4/fastpt/DAS1.7/fastpt-2.1.1+das.dtk25042-cp311-cp311-linux_x86_64.whl"
    "https://download.sourcefind.cn:65024/directlink/9/onesicence/dtk-25.04.2/torch_scatter-2.1.0+das.opt1.dtk25043-cp311-cp311-linux_x86_64.whl"
    "https://download.sourcefind.cn:65024/directlink/9/onesicence/dtk-25.04.2/torch_cluster-1.6.3+das.opt1.dtk25043-cp311-cp311-linux_x86_64.whl"
    "https://download.sourcefind.cn:65024/file/9/onesicence/dtk-25.04.2/torch_sparse-0.6.18-cp311-cp311-linux_x86_64.whl"
    "https://download.sourcefind.cn:65024/file/9/onesicence/dtk-25.04.2/torch_spline_conv-1.2.2-cp311-cp311-linux_x86_64.whl"
    "https://download.sourcefind.cn:65024/file/9/onesicence/dtk-26.04/numba_mwu-0.1.1+dcucompat.1-py3-none-any.whl"
    "https://download.sourcefind.cn:65024/file/9/onesicence/dtk-26.04/pdex-0.2.5+dcucompat.1-py3-none-any.whl"
    "https://download.sourcefind.cn:65024/file/9/onesicence/dtk-26.04/cell_eval-0.8.1+dcucompat.1-py3-none-any.whl"
)

MATCHEM_WHEELS=(
    "https://download.sourcefind.cn:65024/directlink/4/pytorch/DAS1.7/torch-2.5.1+das.opt1.dtk25042-cp311-cp311-manylinux_2_28_x86_64.whl"
    "https://download.sourcefind.cn:65024/directlink/4/vision/DAS1.7/torchvision-0.20.1+das.opt1.dtk25042-cp311-cp311-manylinux_2_28_x86_64.whl"
    "https://download.sourcefind.cn:65024/directlink/4/dgl/DAS1.7/dgl-2.2.1+das.opt1.dtk25042-cp311-cp311-manylinux_2_28_x86_64.whl"
    "https://download.sourcefind.cn:65024/directlink/4/fastpt/DAS1.7/fastpt-2.1.1+das.dtk25042-cp311-cp311-linux_x86_64.whl"
    "https://download.sourcefind.cn:65024/file/4/triton/DAS1.7/triton-3.1.0+das.opt1.dtk25042-cp311-cp311-manylinux_2_28_x86_64.whl"
    "https://download.sourcefind.cn:65024/file/4/tensorflow/DAS1.7/tensorflow-2.18.0+das.opt1.dtk25042-cp311-cp311-manylinux_2_28_x86_64.whl" # 跑DP模型会有JIT编译错误，需切换到DAS1.8版本
    "https://download.sourcefind.cn:65024/directlink/9/onesicence/dtk-25.04.2/torch_scatter-2.1.2-cp311-cp311-linux_x86_64.whl"
    "https://download.sourcefind.cn:65024/directlink/9/onesicence/dtk-25.04.2/torch_cluster-1.6.3-cp311-cp311-linux_x86_64.whl"
    "https://download.sourcefind.cn:65024/directlink/9/onesicence/dtk-25.04.2/torch_sparse-0.6.18-cp311-cp311-linux_x86_64.whl"
    #"https://download.sourcefind.cn:65024/file/4/tensorflow/DAS1.8/tensorflow-2.18.0+das.opt1.dtk2604-cp311-cp311-manylinux_2_28_x86_64.whl"
#torch-geometric torch-scatter torch-sparse torch-cluster torch-spline-conv
    )

# ==========================================
# 3. 核心安装逻辑
# ==========================================

# 对轮子链接按首次出现顺序去重，兼容不支持关联数组的 Bash 版本
dedupe_args() {
    local deduped=()
    local item
    local existing
    local found

    for item in "$@"; do
        found=0
        for existing in "${deduped[@]}"; do
            if [[ "$existing" == "$item" ]]; then
                found=1
                break
            fi
        done

        if [[ "$found" -eq 0 ]]; then
            deduped+=("$item")
        fi
    done

    printf '%s\n' "${deduped[@]}"
}

# 统一的轮子安装函数
install_wheels() {
    local label="$1"
    shift # 移除第一个参数 (label)，剩下的全都是轮子链接
    local deduped_wheels=()

    # 如果没有传入具体的轮子，则跳过
    if [[ "$#" -eq 0 ]]; then
        echo ">>> No $label DAS wheels to install. Skipping."
        return 0
    fi

    while IFS= read -r wheel; do
        [[ -n "$wheel" ]] && deduped_wheels+=("$wheel")
    done < <(dedupe_args "$@")

    echo ">>> Installing ${label} DAS wheels..."
    # 优化：一次性传递所有轮子给 pip，避免 for 循环，提升安装速度并让 pip 更好地处理依赖关系
    uv pip install "${deduped_wheels[@]}" --system
}

# 1. 安装核心基础依赖
install_wheels "core" "${CORE_WHEELS[@]}"

# 2. 根据用户传入的 domain 安装专属依赖
case "$DOMAIN" in
    earth)
        install_wheels "earth" "${EARTH_WHEELS[@]}"
        ;;
    cfd)
        install_wheels "cfd" "${CFD_WHEELS[@]}"
        ;;
    bio)
        install_wheels "bio" "${BIO_WHEELS[@]}"
        ;;
    matchem)
        install_wheels "matchem" "${MATCHEM_WHEELS[@]}"
        echo ">>> Installing cuequivariance base package"
        uv pip install cuequivariance==0.8.0
        ;;
    all)
        install_wheels "all" \
            "${EARTH_WHEELS[@]}" \
            "${CFD_WHEELS[@]}" \
            "${BIO_WHEELS[@]}" \
            "${MATCHEM_WHEELS[@]}"
        ;;
esac

EXTRAS_SPEC="$DOMAIN"
if [[ "$DOMAIN" == "all" ]]; then
    EXTRAS_SPEC="earth,cfd,bio,matchem"
fi

if [[ "$DOMAIN" != "bio" && "$DOMAIN" != "all" ]]; then
    export ONESCIENCE_SKIP_AF3_BUILD=1
fi
if [[ "$DOMAIN" == "bio" || "$DOMAIN" == "all" ]]; then
    export ALPHAFOLD3_TMP_ROOT="${ALPHAFOLD3_TMP_ROOT:-$SCRIPT_DIR/af3_build}"
    export ALPHAFOLD3_CIFPP_COMPONENTS="${ALPHAFOLD3_CIFPP_COMPONENTS:-$ALPHAFOLD3_TMP_ROOT/components.cif}"
    export ALPHAFOLD3_CIFPP_DATA_DIR="${ALPHAFOLD3_CIFPP_DATA_DIR:-$ALPHAFOLD3_TMP_ROOT}"
    mkdir -p "$ALPHAFOLD3_TMP_ROOT"
    # wheel 构建阶段跳过 AF3 编译，后置步骤单独编译一次
    export ONESCIENCE_SKIP_AF3_BUILD=1
fi

echo ">>> Installing OneScience [$DOMAIN] with constraints"
uv pip install -c "$CONSTRAINTS_FILE" ".[${EXTRAS_SPEC}]" --system

# ==========================================
# 4. 后置环境初始化 (环境完全就绪后执行)
# ==========================================
# 如果选用离线依赖包，可预先设置：
#
# export ALPHAFOLD3_DEP_DIR=$ONESCIENCE_MODELS_DIR/AlphaFold3/_dep
# export ALPHAFOLD3_CIFPP_DATA_DIR=$ONESCIENCE_MODELS_DIR/AlphaFold3/_dep/libcifpp/rsrc
# export ALPHAFOLD3_CIFPP_COMPONENTS=$ONESCIENCE_MODELS_DIR/AlphaFold3/_dep/libcifpp/rsrc/components.cif
#
# 注意：
# - ALPHAFOLD3_CIFPP_DATA_DIR 必须是目录，不是 components.cif 文件路径
# - 若未设置，则会回退到本地 af3_build，并在缺文件时尝试自动下载
# 如果没有设置自动下载
if [[ "$DOMAIN" == "bio" || "$DOMAIN" == "all" ]]; then
    echo ">>> Building AlphaFold3 C++ extensions and data files"
    unset ONESCIENCE_SKIP_AF3_BUILD  # 取消跳过标志，允许编译
    python setup.py build_py

    # Copy built artifacts directly into site-packages (no reinstall needed)
    AF3_SRC="$SCRIPT_DIR/src/onescience/flax_models/alphafold3"
    AF3_DEST="$(python -c 'import site; print(site.getsitepackages()[0])')/onescience/flax_models/alphafold3"
    echo ">>> Copying built artifacts to $AF3_DEST"
    find "$AF3_SRC" -name "cpp*.so" -exec cp {} "$AF3_DEST/" \;
    cp -r "$AF3_SRC/_tools" "$AF3_DEST/" 2>/dev/null || true

    # cpp.so (libcifpp) looks for components.cif in site-packages/share/libcifpp/
    # at import time. Copy it there before building data files.
    if [[ -n "${ALPHAFOLD3_CIFPP_COMPONENTS:-}" && -f "$ALPHAFOLD3_CIFPP_COMPONENTS" ]]; then
        CIFPP_DEST="$(python -c 'import site; print(site.getsitepackages()[0])')/share/libcifpp"
        mkdir -p "$CIFPP_DEST"
        cp $ALPHAFOLD3_CIFPP_COMPONENTS "$CIFPP_DEST/components.cif"
        echo ">>> Copied components.cif to $CIFPP_DEST"
    fi
fi

if [[ "$DOMAIN" == "cfd" || "$DOMAIN" == "matchem" || "$DOMAIN" == "all" ]]; then
    FASTPT_SITE="$(pip show fastpt | sed -n 's/^Location: //p')"
    TORCH_SITE="$(pip show torch | sed -n 's/^Location: //p')"

    FASTPT_TORCH_LIB="$FASTPT_SITE/fastpt/torch/lib"
    TORCH_LIB="$TORCH_SITE/torch/lib"

    [[ -d "$FASTPT_TORCH_LIB" ]]
    [[ -d "$TORCH_SITE/torch" ]]

    mkdir -p "$TORCH_LIB"
    cp -a "$FASTPT_TORCH_LIB"/. "$TORCH_LIB"/
fi

# export LD_LIBRARY_PATH="$CONDA_PREFIX/lib:$LD_LIBRARY_PATH"
echo ">>> Installation completed successfully!"