40 lines
1.9 KiB
Bash
40 lines
1.9 KiB
Bash
|
#!/usr/bin/env bash
|
|||
|
|
# 在「已激活的 fishserver」中，将 torch/torchvision 换成指定 CUDA 索引的构建。
|
|||
|
|
# 若从未单独装过 fishserver，请直接跑 bootstrap_fishserver.sh（已内含 GPU 轮子）。
|
|||
|
|
set -euo pipefail
|
|||
|
|
|
|||
|
|
TORCH_INDEX="${TORCH_INDEX:-https://download.pytorch.org/whl/cu124}"
|
|||
|
|
export PIP_DEFAULT_TIMEOUT="${PIP_DEFAULT_TIMEOUT:-120}"
|
|||
|
|
|
|||
|
|
# 必须用「目标环境」里的 python：仅靠 `command -v python` 在激活失败时仍会指向 base。
|
|||
|
|
if [[ -n "${FISHSERVER_PY:-}" ]]; then
|
|||
|
|
PY="$FISHSERVER_PY"
|
|||
|
|
elif [[ -n "${CONDA_PREFIX:-}" && -x "${CONDA_PREFIX}/bin/python" ]]; then
|
|||
|
|
PY="${CONDA_PREFIX}/bin/python"
|
|||
|
|
else
|
|||
|
|
PY="$(command -v python)"
|
|||
|
|
fi
|
|||
|
|
|
|||
|
|
echo "Using PyTorch wheel index: $TORCH_INDEX"
|
|||
|
|
echo "conda env: ${CONDA_DEFAULT_ENV:-"(unset)"} CONDA_PREFIX: ${CONDA_PREFIX:-"(unset)"}"
|
|||
|
|
echo "Using python: $PY ($("$PY" --version 2>&1))"
|
|||
|
|
|
|||
|
|
if [[ -z "${FISHSERVER_PY:-}" && "${CONDA_DEFAULT_ENV:-}" == "base" ]]; then
|
|||
|
|
echo "" >&2
|
|||
|
|
echo "WARNING: 当前 conda 仍是 base，刚才那类输出会把 torch 装进 base（常见为 py3.12），不是 fishserver（py3.11）。" >&2
|
|||
|
|
echo " 请先: conda activate fishserver 再跑本脚本，或:" >&2
|
|||
|
|
echo " FISHSERVER_PY=\$HOME/miniforge3/envs/fishserver/bin/python bash packaging/patch_cuda_torch.sh" >&2
|
|||
|
|
echo "" >&2
|
|||
|
|
fi
|
|||
|
|
|
|||
|
|
# 已装过 PyPI/cu130 的 torch 时，不加 uninstall / force 会一律 "already satisfied"，版本不会变。
|
|||
|
|
# 全局 pip.conf 里的 extra-index-url（如 pypi.ngc.nvidia.com）也会参与解析，用空配置避开。
|
|||
|
|
echo "Removing any existing torch / torchvision …"
|
|||
|
|
"$PY" -m pip uninstall -y torch torchvision 2>/dev/null || true
|
|||
|
|
|
|||
|
|
echo "Installing torch / torchvision from $TORCH_INDEX only …"
|
|||
|
|
PIP_CONFIG_FILE=/dev/null PIP_EXTRA_INDEX_URL="" \
|
|||
|
|
"$PY" -m pip install --no-cache-dir torch torchvision --index-url "$TORCH_INDEX"
|
|||
|
|
|
|||
|
|
"$PY" -c "import torch; print('torch', torch.__version__, 'cuda?', torch.cuda.is_available())"
|