93 lines
2.7 KiB
Bash
93 lines
2.7 KiB
Bash
#!/usr/bin/env bash
|
|
set -euo pipefail
|
|
|
|
ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
|
|
export UV_PROJECT_ENVIRONMENT="${UV_PROJECT_ENVIRONMENT:-$ROOT_DIR/.venv}"
|
|
|
|
METAX_PYPI_INDEX="${METAX_PYPI_INDEX:-https://repos.metax-tech.com/r/maca-pypi/simple}"
|
|
METAX_PYPI_HOST="${METAX_PYPI_HOST:-repos.metax-tech.com}"
|
|
METAX_INSTALL_GPU_DEPS="${METAX_INSTALL_GPU_DEPS:-false}"
|
|
|
|
METAX_TORCH_SPEC="${METAX_TORCH_SPEC:-torch}"
|
|
METAX_TORCHAUDIO_SPEC="${METAX_TORCHAUDIO_SPEC:-}"
|
|
METAX_TORCHVISION_SPEC="${METAX_TORCHVISION_SPEC:-}"
|
|
METAX_VLLM_SPEC="${METAX_VLLM_SPEC:-vllm}"
|
|
|
|
usage() {
|
|
cat <<EOF
|
|
Usage:
|
|
./scripts/sync_metax_env.sh [--list-versions PACKAGE] [uv sync args...]
|
|
|
|
Environment overrides:
|
|
METAX_TORCH_SPEC=torch==x.y.z
|
|
METAX_TORCHAUDIO_SPEC=torchaudio==x.y.z # optional
|
|
METAX_TORCHVISION_SPEC=torchvision==x.y.z # optional
|
|
METAX_VLLM_SPEC=vllm==x.y.z
|
|
METAX_INSTALL_GPU_DEPS=false
|
|
METAX_PYPI_INDEX=${METAX_PYPI_INDEX}
|
|
EOF
|
|
}
|
|
|
|
cd "$ROOT_DIR"
|
|
|
|
LIST_PACKAGE=""
|
|
if [[ "${1:-}" == "-h" || "${1:-}" == "--help" ]]; then
|
|
usage
|
|
exit 0
|
|
fi
|
|
if [[ "${1:-}" == "--list-versions" ]]; then
|
|
LIST_PACKAGE="${2:-torch}"
|
|
shift
|
|
[[ $# -gt 0 ]] && shift
|
|
fi
|
|
|
|
if [[ -n "$LIST_PACKAGE" ]]; then
|
|
if command -v pip3 >/dev/null 2>&1; then
|
|
pip3 index versions "$LIST_PACKAGE" \
|
|
-i "$METAX_PYPI_INDEX" \
|
|
--trusted-host "$METAX_PYPI_HOST"
|
|
elif command -v pip >/dev/null 2>&1; then
|
|
pip index versions "$LIST_PACKAGE" \
|
|
-i "$METAX_PYPI_INDEX" \
|
|
--trusted-host "$METAX_PYPI_HOST"
|
|
else
|
|
python3 -m pip index versions "$LIST_PACKAGE" \
|
|
-i "$METAX_PYPI_INDEX" \
|
|
--trusted-host "$METAX_PYPI_HOST"
|
|
fi
|
|
exit 0
|
|
fi
|
|
|
|
uv venv "$UV_PROJECT_ENVIRONMENT"
|
|
uv sync --project "$ROOT_DIR/environments/metax" --no-dev --no-install-project "$@"
|
|
|
|
gpu_install_args=(pip install --reinstall)
|
|
if [[ "$METAX_INSTALL_GPU_DEPS" != "true" ]]; then
|
|
gpu_install_args+=(--no-deps)
|
|
fi
|
|
gpu_install_args+=(
|
|
--index-url "$METAX_PYPI_INDEX" \
|
|
--trusted-host "$METAX_PYPI_HOST" \
|
|
)
|
|
for package_spec in "$METAX_TORCH_SPEC" "$METAX_TORCHAUDIO_SPEC" "$METAX_TORCHVISION_SPEC" "$METAX_VLLM_SPEC"; do
|
|
[[ -n "$package_spec" ]] && gpu_install_args+=("$package_spec")
|
|
done
|
|
|
|
uv "${gpu_install_args[@]}"
|
|
|
|
if [[ "$METAX_INSTALL_GPU_DEPS" != "true" ]]; then
|
|
echo "[sync_metax_env] Installed MetaX GPU stack with --no-deps to avoid PyPI/NVIDIA fallback."
|
|
echo "[sync_metax_env] Set METAX_INSTALL_GPU_DEPS=true only if the MetaX index hosts all required GPU-stack dependencies."
|
|
fi
|
|
|
|
"$UV_PROJECT_ENVIRONMENT/bin/python" - <<'PY'
|
|
import importlib
|
|
import torch
|
|
|
|
print("torch", torch.__version__)
|
|
print("torch.cuda.is_available", torch.cuda.is_available())
|
|
for name in ("vllm",):
|
|
module = importlib.import_module(name)
|
|
print(name, getattr(module, "__version__", "unknown"))
|
|
PY
|