diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 4f256d1..e5b99c8 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -11,105 +11,9 @@ permissions: contents: read jobs: - test-channel-args: - name: Test channel args (${{ matrix.name }}) - runs-on: ubuntu-latest - strategy: - fail-fast: false - matrix: - include: - - name: default - country: "" - rc: false - expected_channels: "--channel=conda-forge --channel=njzjz" - - name: default RC - country: "" - rc: true - expected_channels: "--channel=conda-forge --channel=conda-forge/label/deepmd-kit_rc --channel=njzjz" - - name: China - country: CN - rc: false - expected_channels: "--channel=https://mirrors.ustc.edu.cn/anaconda/cloud/conda-forge/ --channel=njzjz" - - name: China RC - country: CN - rc: true - expected_channels: "--channel=https://mirrors.ustc.edu.cn/anaconda/cloud/conda-forge/ --channel=conda-forge/label/deepmd-kit_rc --channel=njzjz" - env: - DP1S_COUNTRY: ${{ matrix.country }} - DP1S_NO_PATH_UPDATE: 1 - steps: - - name: Checkout code - uses: actions/checkout@v4 - - - name: Install mocks - run: | - set -euo pipefail - mock_bin="${RUNNER_TEMP}/mock-bin" - mkdir -p "${mock_bin}" - cat > "${mock_bin}/curl" <<'EOF' - #!/bin/sh - cat <<'INSTALLER' - #!/bin/sh - set -eu - mkdir -p "${PIXI_HOME}/bin" - cat > "${PIXI_HOME}/bin/pixi" <<'PIXI' - #!/bin/sh - set -eu - printf '%s\n' "$@" > "${DP1S_PIXI_ARGS_FILE}" - mkdir -p "${PIXI_HOME}/bin" - cat > "${PIXI_HOME}/bin/dp" <<'DP' - #!/bin/sh - echo "DeepMD-kit mock" - DP - chmod +x "${PIXI_HOME}/bin/dp" - cat > "${PIXI_HOME}/bin/lmp" <<'LMP' - #!/bin/sh - echo "pair deepmd" - LMP - chmod +x "${PIXI_HOME}/bin/lmp" - cat > "${PIXI_HOME}/bin/mpirun" <<'MPIRUN' - #!/bin/sh - echo "mpirun mock" - MPIRUN - chmod +x "${PIXI_HOME}/bin/mpirun" - PIXI - chmod +x "${PIXI_HOME}/bin/pixi" - INSTALLER - EOF - chmod +x "${mock_bin}/curl" - echo "${mock_bin}" >> "${GITHUB_PATH}" - - - name: Run script with mocked installer - run: | - set -euo pipefail - export DP1S_HOME="${RUNNER_TEMP}/dp1s" - export DP1S_PIXI_ARGS_FILE="${RUNNER_TEMP}/pixi-args.txt" - if [[ "${{ matrix.rc }}" == "true" ]]; then - export DP1S_DEEPMD_RC=1 - fi - bash dp1s.sh - - actual_channels=$(grep '^--channel=' "${DP1S_PIXI_ARGS_FILE}" | tr '\n' ' ' | sed 's/ $//') - expected_channels='${{ matrix.expected_channels }}' - if [[ "${actual_channels}" != "${expected_channels}" ]]; then - echo "Expected channels: ${expected_channels}" - echo "Actual channels: ${actual_channels}" - exit 1 - fi - test: name: Test dp1s.sh (${{ matrix.name }}) runs-on: ubuntu-latest - strategy: - fail-fast: false - matrix: - include: - - name: default - country: "" - - name: China - country: CN - env: - DP1S_COUNTRY: ${{ matrix.country }} steps: - name: Checkout code uses: actions/checkout@v4 diff --git a/README.md b/README.md index fca9025..2d98f26 100644 --- a/README.md +++ b/README.md @@ -26,4 +26,3 @@ The installation script has several options that can be manipulated through envi - `DP1S_NO_PATH_UPDATE`: If set the `$PATH` will not be updated to add DeePMD-kit to it. - `DP1S_DEEPMD_RC`: If set, add the `conda-forge/label/deepmd-kit_rc` channel to install DeePMD-kit release candidates. - `DEEPMD_VERSION`: The version of DeePMD-kit getting installed, can be used to up- or down-grade. - diff --git a/dp1s.sh b/dp1s.sh index 648e184..d734a16 100755 --- a/dp1s.sh +++ b/dp1s.sh @@ -25,33 +25,219 @@ logging "This script will automatically download and install DeePMD-kit (${DEEPM DP1S_HOME=${DP1S_HOME:-~/.dp1s} export PIXI_HOME=$DP1S_HOME DP1S_BIN_PATH=$DP1S_HOME/bin +DP1S_CACHE_DIR=${DP1S_CACHE_DIR:-"$DP1S_HOME/cache"} +DP1S_MIRROR_CACHE=${DP1S_MIRROR_CACHE:-"$DP1S_CACHE_DIR/mirror-choice"} +DP1S_MIRROR_CACHE_TTL=${DP1S_MIRROR_CACHE_TTL:-86400} -# 1. check the location of the machine +deepmd_rc_channel="conda-forge/label/deepmd-kit_rc" + +conda_mirror_names() { + if [[ -n "${DP1S_CONDA_MIRRORS:-}" ]]; then + echo "${DP1S_CONDA_MIRRORS}" | tr ',' ' ' + else + echo "ustc tuna bfsu nju conda-forge" + fi +} + +pixi_mirror_names() { + if [[ -n "${DP1S_PIXI_MIRRORS:-}" ]]; then + echo "${DP1S_PIXI_MIRRORS}" | tr ',' ' ' + else + echo "github ghfast" + fi +} + +conda_mirror_base() { + case "$1" in + conda-forge|official) echo "conda-forge" ;; + ustc) echo "https://mirrors.ustc.edu.cn/anaconda/cloud/conda-forge/" ;; + tuna) echo "https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge/" ;; + bfsu) echo "https://mirrors.bfsu.edu.cn/anaconda/cloud/conda-forge/" ;; + nju) echo "https://mirror.nju.edu.cn/anaconda/cloud/conda-forge/" ;; + http://*|https://*) echo "$1" ;; + *) return 1 ;; + esac +} + +conda_mirror_probe() { + local base="$1" + if [[ "${base}" == "conda-forge" ]]; then + echo "https://conda.anaconda.org/conda-forge/noarch/current_repodata.json" + else + echo "${base%/}/noarch/current_repodata.json" + fi +} + +pixi_mirror_base() { + case "$1" in + github|official) echo "https://github.com/prefix-dev/pixi" ;; + ghfast) echo "https://ghfast.top/https://github.com/prefix-dev/pixi" ;; + http://*|https://*) echo "$1" ;; + *) return 1 ;; + esac +} + +pixi_mirror_probe() { + local base="$1" + case "$base" in + *) echo "${base%/}/releases/latest/download/pixi-x86_64-unknown-linux-musl.tar.gz" ;; + esac +} + +probe_url() { + local url="$1" + curl -fsSL --connect-timeout "${DP1S_MIRROR_CONNECT_TIMEOUT:-60}" --max-time "${DP1S_MIRROR_TIMEOUT:-60}" \ + --range 0-8191 -o /dev/null -w '%{time_total}' "$url" +} + +pick_fastest_mirror() { + local kind="$1" + local tmpdir result_file name base probe_time probe_url_path mirror_list pending pids pid + tmpdir=$(mktemp -d) + + case "$kind" in + conda) mirror_list=$(conda_mirror_names) ;; + pixi) mirror_list=$(pixi_mirror_names) ;; + *) return 1 ;; + esac + + pids="" + pending=0 + for name in $mirror_list; do + case "$kind" in + conda) + base=$(conda_mirror_base "$name" 2>/dev/null || true) + [[ -n "$base" ]] || continue + probe_url_path=$(conda_mirror_probe "$base") + ;; + pixi) + base=$(pixi_mirror_base "$name" 2>/dev/null || true) + [[ -n "$base" ]] || continue + probe_url_path=$(pixi_mirror_probe "$base") + ;; + *) return 1 ;; + esac + result_file="$tmpdir/${kind}-${pending}-${name//[^A-Za-z0-9_.-]/_}" + ( + if probe_time=$(probe_url "$probe_url_path" 2>/dev/null); then + printf 'ok\t%s\t%s\t%s\t%s\n' "$probe_time" "$name" "$base" "$probe_url_path" > "$result_file" + else + printf 'failed\t%s\t%s\t%s\n' "$name" "$base" "$probe_url_path" > "$result_file" + fi + ) & + pids="$pids $!" + pending=$((pending + 1)) + done + + if [[ "$pending" -eq 0 ]]; then + rm -rf "$tmpdir" + return 1 + fi + + local finished success_line completed=0 status probe_seconds result_name result_base result_probe_url + while (( completed < pending )); do + for result_file in "$tmpdir"/${kind}-*; do + [[ -f "$result_file" ]] || continue + finished="$result_file.done" + [[ -e "$finished" ]] && continue + : > "$finished" + completed=$((completed + 1)) + + IFS=$'\t' read -r status probe_seconds result_name result_base result_probe_url < "$result_file" || true + if [[ "$status" == "ok" ]]; then + success_line=$(cat "$result_file") + for pid in $pids; do + kill "$pid" 2>/dev/null || true + done + wait 2>/dev/null || true + IFS=$'\t' read -r status probe_seconds result_name result_base result_probe_url <<< "$success_line" + printf '%s\t%s\t%s\n' "$result_name" "$result_base" "$probe_seconds" + rm -rf "$tmpdir" + return 0 + fi + done + sleep 0.05 + done + + wait 2>/dev/null || true + rm -rf "$tmpdir" + return 1 +} + +read_mirror_cache() { + [[ -f "$DP1S_MIRROR_CACHE" ]] || return 1 + local now ts cache_conda_name cache_conda_channel cache_pixi_name cache_pixi_repo + now=$(date +%s) + IFS=$'\t' read -r ts cache_conda_name cache_conda_channel cache_pixi_name cache_pixi_repo < "$DP1S_MIRROR_CACHE" || return 1 + [[ "$ts" =~ ^[0-9]+$ ]] || return 1 + (( now - ts < DP1S_MIRROR_CACHE_TTL )) || return 1 + conda_channel="$cache_conda_channel" + conda_channel_name="$cache_conda_name" + if [[ -z "${PIXI_REPOURL:-}" ]]; then + export PIXI_REPOURL="$cache_pixi_repo" + pixi_repo_name="$cache_pixi_name" + fi +} + +write_mirror_cache() { + mkdir -p "$DP1S_CACHE_DIR" + printf '%s\t%s\t%s\t%s\t%s\n' \ + "$(date +%s)" "${conda_channel_name:-manual}" "${conda_channel}" \ + "${pixi_repo_name:-manual}" "${PIXI_REPOURL:-}" > "$DP1S_MIRROR_CACHE" +} + +# 1. select mirrors ((progress++)) || : -if [[ -v DP1S_COUNTRY ]]; then - country=${DP1S_COUNTRY} -else - country=$(curl -fsSL --connect-timeout 5 --max-time 10 https://ipinfo.io/country || :) +conda_channel_name="" +pixi_repo_name="" +if [[ -n "${DP1S_CONDA_CHANNEL:-}" ]]; then + conda_channel="$DP1S_CONDA_CHANNEL" + conda_channel_name="manual" + logging "Use conda channel from DP1S_CONDA_CHANNEL: ${conda_channel}" +elif [[ "${DP1S_CHANNEL_AUTO:-1}" != "0" ]] && read_mirror_cache; then + logging "Use cached conda channel: ${conda_channel_name} (${conda_channel})" + if [[ -n "${pixi_repo_name:-}" ]]; then + logging "Use cached fastest pixi repo: ${pixi_repo_name} (${PIXI_REPOURL})" + fi +elif [[ "${DP1S_CHANNEL_AUTO:-1}" != "0" ]]; then + logging "Benchmark conda channel mirrors" + if fastest_conda=$(pick_fastest_mirror conda); then + IFS=$'\t' read -r conda_channel_name conda_channel conda_probe_time <<< "$fastest_conda" + logging "Use fastest conda channel: ${conda_channel_name} (${conda_channel}, ${conda_probe_time}s)" + fi fi -deepmd_rc_channel="conda-forge/label/deepmd-kit_rc" +if [[ -z "${conda_channel:-}" ]]; then + conda_channel="conda-forge" + conda_channel_name="conda-forge" + logging "Use fallback conda channel: ${conda_channel_name} (${conda_channel})" +fi -case "${country}" in - CN) - logging "Location: ${country}" - conda_channel="https://mirrors.ustc.edu.cn/anaconda/cloud/conda-forge/" - export PIXI_REPOURL=https://ghfast.top/https://github.com/prefix-dev/pixi - ;; - "") - logging "Location detection failed; falling back to conda-forge" - conda_channel="conda-forge" - ;; - *) - logging "Location: ${country}" - conda_channel="conda-forge" - ;; -esac +if [[ -n "${DP1S_PIXI_REPOURL:-}" ]]; then + export PIXI_REPOURL="$DP1S_PIXI_REPOURL" + pixi_repo_name="manual" + logging "Use pixi repo from DP1S_PIXI_REPOURL: ${PIXI_REPOURL}" +elif [[ -n "${PIXI_REPOURL:-}" ]]; then + pixi_repo_name="manual" +elif [[ "${DP1S_CHANNEL_AUTO:-1}" != "0" ]]; then + logging "Benchmark pixi installer mirrors" + if fastest_pixi=$(pick_fastest_mirror pixi); then + IFS=$'\t' read -r pixi_repo_name PIXI_REPOURL pixi_probe_time <<< "$fastest_pixi" + export PIXI_REPOURL + logging "Use fastest pixi repo: ${pixi_repo_name} (${PIXI_REPOURL}, ${pixi_probe_time}s)" + fi +fi + +if [[ -z "${PIXI_REPOURL:-}" ]]; then + export PIXI_REPOURL=https://github.com/prefix-dev/pixi + pixi_repo_name="github" + logging "Use fallback pixi repo: ${PIXI_REPOURL}" +fi + +if [[ "${DP1S_CHANNEL_AUTO:-1}" != "0" && "${conda_channel_name}" != "manual" && "${pixi_repo_name:-}" != "manual" ]]; then + write_mirror_cache || true +fi channel_args=(--channel="$conda_channel") if [[ -v DP1S_DEEPMD_RC ]]; then