diff --git a/infra/scripts/install.sh b/infra/scripts/install.sh index 07bf67a8..1c307fae 100644 --- a/infra/scripts/install.sh +++ b/infra/scripts/install.sh @@ -23,9 +23,7 @@ # Helpers & constants ############################################################################### SCRIPTS_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" - -SUPPORTED_FAMILIES=("2.5" "2.6" "2.7" "2.8" "2.9" "2.10") -DEFAULT_FAMILY="2.7" +source "${SCRIPTS_DIR}/pytorch_package_utils.sh" || exit 1 show_help() { cat <&2 done - INDEX_URLS+=("$candidate") + + echo "[ERROR] Could not install torch from any candidate PyTorch index." >&2 + return 1 +} + +installed_compute_matches_request() { + local installed_wheel_tag="$1" + local installed_cuda="" + + if [[ "${_CPU_ONLY}" == "1" ]]; then + [[ "${installed_wheel_tag}" == "cpu" ]] + return + fi + + if [[ -n "${_USER_CUDA}" ]]; then + if [[ "${installed_wheel_tag}" == "cpu" ]]; then + return 1 + fi + + installed_cuda="$(pytorch_cuda_tag_to_version "${installed_wheel_tag}")" || return 1 + pytorch_version_le "${installed_cuda}" "${_USER_CUDA}" + return + fi + + return 0 } ############################################################################### # Option parsing ############################################################################### _DIST=0 -_TORCH_VER="${DEFAULT_FAMILY}" +_TORCH_VER="${PYTORCH_DEFAULT_FAMILY}" +_TORCH_VER_WAS_SET=0 _USER_CUDA="" -_CPU_ONLY="" +_CPU_ONLY=0 options=$(getopt -o h --long dist,torch_ver:,cuda_ver:,cpu_only,help -- "$@") || { echo "[ERROR] Invalid command-line options" >&2; exit 1; } -eval set -- "$options" +eval set -- "${options}" while true; do case "$1" in --dist) _DIST=1 ;; - --torch_ver) _TORCH_VER="$2"; shift ;; + --torch_ver) _TORCH_VER="$2"; _TORCH_VER_WAS_SET=1; shift ;; --cuda_ver) _USER_CUDA="$2"; shift ;; --cpu_only) _CPU_ONLY=1 ;; -h|--help) show_help; exit 0 ;; @@ -83,157 +109,207 @@ while true; do shift done +if [[ "${_CPU_ONLY}" == "1" && -n "${_USER_CUDA}" ]]; then + echo "[ERROR] --cpu_only and --cuda_ver cannot be used together." >&2 + exit 1 +fi + +if [[ -n "${_USER_CUDA}" && ! "${_USER_CUDA}" =~ ^[0-9]+\.[0-9]+$ ]]; then + echo "[ERROR] Invalid --cuda_ver value '${_USER_CUDA}'. Expected MAJ.MIN." >&2 + exit 1 +fi + ############################################################################### -# Detect (and maybe keep) any existing torch installation +# Normalize the requested Torch version ############################################################################### -INSTALLED_TORCH_FULL="" -INSTALLED_TORCH_FAMILY="" -read -r INSTALLED_TORCH_FULL < <( - python3 - <<'PY' -import importlib.util, re, sys -spec = importlib.util.find_spec("torch") -if spec is None: # Torch not found → just print blanks - print() - sys.exit(0) -import torch -print(torch.__version__) -PY -) - -if [[ -n "$INSTALLED_TORCH_FULL" ]]; then - INSTALLED_TORCH_FAMILY=$(echo "$INSTALLED_TORCH_FULL" | cut -d. -f1,2) -fi +REQUEST_IS_NIGHTLY=0 +REQUEST_IS_EXACT=0 +REQUESTED_FAMILY="" +REQUESTED_EXACT_VERSION="" +REQUESTED_BUILD_TAG="" +RESOLVED_TORCH_VERSION="" -# Normalise requested spec to family / exact -REQUEST_IS_NIGHTLY="" -REQUEST_IS_EXACT="" -if [[ "$_TORCH_VER" == "nightly" ]]; then +if [[ "${_TORCH_VER}" == "nightly" ]]; then REQUEST_IS_NIGHTLY=1 -elif [[ "$_TORCH_VER" =~ ^[0-9]+\.[0-9]+$ ]]; then - : # family only -elif [[ "$_TORCH_VER" =~ ^[0-9]+\.[0-9]+\.[0-9]+ ]]; then +elif [[ "${_TORCH_VER}" =~ ^[0-9]+\.[0-9]+$ ]]; then + REQUESTED_FAMILY="${_TORCH_VER}" + if ! pytorch_is_supported_family "${REQUESTED_FAMILY}"; then + echo "[ERROR] Unsupported --torch_ver family '${_TORCH_VER}'" >&2 + exit 1 + fi + RESOLVED_TORCH_VERSION="$(pytorch_resolve_latest_stable_version "${REQUESTED_FAMILY}")" || exit 1 +elif [[ "${_TORCH_VER}" =~ ^[0-9]+\.[0-9]+\.[0-9]+(\+[A-Za-z0-9._-]+)?$ ]]; then REQUEST_IS_EXACT=1 -else - echo "[ERROR] Unsupported --torch_ver value '${_TORCH_VER}'"; exit 1 -fi - -# Respect pre-installed Torch if allowed -SKIP_TORCH_INSTALL="" -if [[ -n "$INSTALLED_TORCH_FULL" ]]; then - if [[ " ${SUPPORTED_FAMILIES[*]} " =~ " ${INSTALLED_TORCH_FAMILY} " ]]; then - if [[ -z "$REQUEST_IS_NIGHTLY" && -z "$REQUEST_IS_EXACT" ]]; then - echo "[INFO] Supported torch ${INSTALLED_TORCH_FULL} already present — keeping it" - SKIP_TORCH_INSTALL=1 - _TORCH_VER="$INSTALLED_TORCH_FAMILY" # for later requirements file pick - else - echo "[INFO] '--torch_ver' explicitly requests ${_TORCH_VER}; will override existing ${INSTALLED_TORCH_FULL}" - fi - else - echo "[WARN] Found unsupported torch ${INSTALLED_TORCH_FULL}; will install supported default" + REQUESTED_EXACT_VERSION="${_TORCH_VER}" + REQUESTED_FAMILY="$(pytorch_version_family "${_TORCH_VER}")" || exit 1 + REQUESTED_BUILD_TAG="$(pytorch_local_build_tag "${_TORCH_VER}")" + if ! pytorch_is_supported_family "${REQUESTED_FAMILY}"; then + echo "[ERROR] Unsupported --torch_ver value '${_TORCH_VER}'" >&2 + exit 1 + fi + if [[ -n "${REQUESTED_BUILD_TAG}" ]] && ! pytorch_is_supported_wheel_tag "${REQUESTED_BUILD_TAG}"; then + echo "[ERROR] Unsupported Torch build tag '+${REQUESTED_BUILD_TAG}'" >&2 + exit 1 fi + RESOLVED_TORCH_VERSION="${REQUESTED_EXACT_VERSION}" +else + echo "[ERROR] Unsupported --torch_ver value '${_TORCH_VER}'" >&2 + exit 1 fi ############################################################################### -# CUDA index-URL logic +# Detect and possibly keep an existing Torch installation ############################################################################### -get_index_url_for_cuda_version() { - local cuda_ver="$1" nightly="$2" - local maj=${cuda_ver%.*} min=${cuda_ver#*.} - echo "https://download.pytorch.org/whl${nightly:+/nightly}/cu${maj}${min}" -} - -INDEX_URL="https://download.pytorch.org/whl${REQUEST_IS_NIGHTLY:+/nightly}/cpu" -CUDA_TO_USE="" +INSTALLED_TORCH_FULL="" +INSTALLED_TORCH_BASE="" +INSTALLED_TORCH_FAMILY="" +INSTALLED_TORCH_WHEEL_TAG="" +INSTALLED_TORCH_IS_NIGHTLY="0" -if [[ -n "$_CPU_ONLY" ]]; then - echo "[INFO] Forcing CPU-only Torch installation" -else - if [[ -n "$_USER_CUDA" ]]; then - CUDA_TO_USE="$_USER_CUDA" - echo "[INFO] Using CUDA ${CUDA_TO_USE} specified with --cuda_ver" - elif command -v nvcc &>/dev/null; then - CUDA_TO_USE=$(nvcc --version | grep -oP "release \K[0-9]+\.[0-9]+") - echo "[INFO] Detected CUDA ${CUDA_TO_USE}" - elif command -v nvidia-smi &>/dev/null; then - CUDA_TO_USE=$(nvidia-smi | grep -oP "CUDA Version: \K[0-9]+\.[0-9]+") - echo "[INFO] Detected CUDA ${CUDA_TO_USE}" - fi +if TORCH_INFO="$(pytorch_get_installed_torch_info 2>/dev/null)"; then + IFS=$'\t' read -r \ + INSTALLED_TORCH_FULL \ + INSTALLED_TORCH_BASE \ + INSTALLED_TORCH_FAMILY \ + INSTALLED_TORCH_WHEEL_TAG \ + INSTALLED_TORCH_IS_NIGHTLY <<< "${TORCH_INFO}" +fi - if [[ -n "$CUDA_TO_USE" ]]; then - INDEX_URL=$(get_index_url_for_cuda_version "$CUDA_TO_USE" "$REQUEST_IS_NIGHTLY") +SKIP_TORCH_INSTALL=0 +if [[ -n "${INSTALLED_TORCH_FULL}" ]]; then + # With no explicit Torch version, preserve any supported stable family. If a + # compute-platform option requires reinstalling the wheel, reinstall the same + # family at its latest supported patch rather than falling back to the default. + if [[ "${_TORCH_VER_WAS_SET}" == "0" && \ + "${INSTALLED_TORCH_IS_NIGHTLY}" == "0" ]] && \ + pytorch_is_supported_family "${INSTALLED_TORCH_FAMILY}"; then + REQUESTED_FAMILY="${INSTALLED_TORCH_FAMILY}" + RESOLVED_TORCH_VERSION="$(pytorch_resolve_latest_stable_version \ + "${REQUESTED_FAMILY}")" || exit 1 fi -fi -INDEX_URLS=() -add_unique_index_url "$INDEX_URL" - -if [[ -z "$_CPU_ONLY" && -n "$CUDA_TO_USE" ]]; then - # PyTorch does not publish wheels for every CUDA minor version. - # Example: CUDA 13.1 maps to cu131, but cu131 may not exist. - # Try the detected/requested CUDA first, then fall back to known PyTorch - # CUDA wheel indices that are <= the detected/requested CUDA version. - PYTORCH_CUDA_FALLBACKS=("13.0" "12.8" "12.6" "12.4" "12.1" "11.8") - for cuda_fb in "${PYTORCH_CUDA_FALLBACKS[@]}"; do - if version_le "$cuda_fb" "$CUDA_TO_USE"; then - add_unique_index_url "$(get_index_url_for_cuda_version "$cuda_fb" "$REQUEST_IS_NIGHTLY")" + if installed_compute_matches_request "${INSTALLED_TORCH_WHEEL_TAG}"; then + if [[ "${_TORCH_VER_WAS_SET}" == "0" ]]; then + if [[ "${INSTALLED_TORCH_IS_NIGHTLY}" == "0" ]] && \ + pytorch_is_supported_family "${INSTALLED_TORCH_FAMILY}"; then + echo "[INFO] Supported torch ${INSTALLED_TORCH_FULL} already present — keeping it" + SKIP_TORCH_INSTALL=1 + RESOLVED_TORCH_VERSION="${INSTALLED_TORCH_BASE}" + fi + elif [[ "${REQUEST_IS_NIGHTLY}" == "1" && "${INSTALLED_TORCH_IS_NIGHTLY}" == "1" ]]; then + echo "[INFO] Requested nightly torch ${INSTALLED_TORCH_FULL} already present — keeping it" + SKIP_TORCH_INSTALL=1 + elif [[ "${REQUEST_IS_EXACT}" == "1" ]]; then + REQUESTED_BASE="$(pytorch_strip_local_version "${REQUESTED_EXACT_VERSION}")" + if [[ "${INSTALLED_TORCH_BASE}" == "${REQUESTED_BASE}" ]]; then + if [[ -z "${REQUESTED_BUILD_TAG}" || \ + "${INSTALLED_TORCH_WHEEL_TAG}" == "${REQUESTED_BUILD_TAG}" ]]; then + echo "[INFO] Requested torch ${INSTALLED_TORCH_FULL} already present — keeping it" + SKIP_TORCH_INSTALL=1 + fi + fi + elif [[ "${INSTALLED_TORCH_IS_NIGHTLY}" == "0" && \ + "${INSTALLED_TORCH_BASE}" == "${RESOLVED_TORCH_VERSION}" ]]; then + echo "[INFO] Requested torch ${INSTALLED_TORCH_FULL} already present — keeping it" + SKIP_TORCH_INSTALL=1 fi - done + fi - # Last resort: keep install working on machines without a matching CUDA wheel. - add_unique_index_url "https://download.pytorch.org/whl${REQUEST_IS_NIGHTLY:+/nightly}/cpu" + if [[ "${SKIP_TORCH_INSTALL}" == "0" ]]; then + echo "[INFO] Requested torch '${_TORCH_VER}' will replace existing ${INSTALLED_TORCH_FULL}" + fi fi ############################################################################### -# Torch installation (may be skipped) +# Resolve candidate wheel indices and install Torch ############################################################################### -install_torch() { - local spec="$1" - local index_url +if [[ "${SKIP_TORCH_INSTALL}" == "0" ]]; then + HOST_CUDA="" + PINNED_BUILD_TAG="" - for index_url in "${INDEX_URLS[@]}"; do - echo "[INFO] Installing torch (${spec}) from ${index_url}" - if python3 -m pip install ${spec} --index-url "${index_url}"; then - INDEX_URL="${index_url}" - echo "[INFO] Successfully installed torch from ${index_url}" - return 0 + if [[ "${REQUEST_IS_NIGHTLY}" == "1" ]]; then + TORCH_DEV_FILE="${SCRIPTS_DIR}/../dependency/torch_dev.txt" + PINNED_TORCH_VERSION="$(pytorch_get_pinned_requirement_version "${TORCH_DEV_FILE}" torch)" + if [[ -z "${PINNED_TORCH_VERSION}" ]]; then + echo "[ERROR] ${TORCH_DEV_FILE} must pin torch with 'torch==VERSION'." >&2 + exit 1 fi + PINNED_BUILD_TAG="$(pytorch_local_build_tag "${PINNED_TORCH_VERSION}")" + fi - echo "[WARN] Failed to install torch (${spec}) from ${index_url}; trying next candidate..." >&2 - done - - echo "[ERROR] Could not install torch (${spec}) from any candidate PyTorch index." >&2 - return 1 -} + if [[ "${_CPU_ONLY}" == "1" ]]; then + if [[ -n "${REQUESTED_BUILD_TAG}" && "${REQUESTED_BUILD_TAG}" != "cpu" ]]; then + echo "[ERROR] --cpu_only conflicts with torch build '+${REQUESTED_BUILD_TAG}'." >&2 + exit 1 + fi + if [[ -n "${PINNED_BUILD_TAG}" && "${PINNED_BUILD_TAG}" != "cpu" ]]; then + echo "[ERROR] --cpu_only conflicts with nightly torch build '+${PINNED_BUILD_TAG}'." >&2 + exit 1 + fi + REQUESTED_BUILD_TAG="cpu" + echo "[INFO] Forcing CPU-only Torch installation" + elif [[ -n "${REQUESTED_BUILD_TAG}" ]]; then + echo "[INFO] Using explicitly requested Torch wheel ${REQUESTED_BUILD_TAG}" + elif [[ -n "${PINNED_BUILD_TAG}" ]]; then + REQUESTED_BUILD_TAG="${PINNED_BUILD_TAG}" + echo "[INFO] Using Torch wheel ${REQUESTED_BUILD_TAG} pinned by torch_dev.txt" + elif [[ -n "${_USER_CUDA}" ]]; then + HOST_CUDA="${_USER_CUDA}" + echo "[INFO] Using CUDA ${HOST_CUDA} specified with --cuda_ver" + elif HOST_CUDA="$(pytorch_detect_host_cuda_version)"; then + echo "[INFO] Detected CUDA ${HOST_CUDA}" + else + HOST_CUDA="" + echo "[INFO] CUDA was not detected; using a CPU Torch wheel" + fi -if [[ -z "$SKIP_TORCH_INSTALL" ]]; then - if [[ -n "$REQUEST_IS_NIGHTLY" ]]; then - install_torch "-r ${SCRIPTS_DIR}/../dependency/torch_dev.txt" || exit 1 + if [[ "${REQUEST_IS_NIGHTLY}" == "1" ]]; then + pytorch_build_index_urls "" "${HOST_CUDA}" 1 "${REQUESTED_BUILD_TAG}" || exit 1 + install_torch -r "${TORCH_DEV_FILE}" || exit 1 else - if [[ -n "$REQUEST_IS_EXACT" ]]; then - install_torch "torch==${_TORCH_VER}" || exit 1 - else - # family only → pip’s ~= spec picks the newest patch in the family - install_torch "torch==${_TORCH_VER}.*" || exit 1 - fi + pytorch_build_index_urls \ + "${REQUESTED_FAMILY}" "${HOST_CUDA}" 0 "${REQUESTED_BUILD_TAG}" || exit 1 + echo "[INFO] Resolved torch ${_TORCH_VER} to ${RESOLVED_TORCH_VERSION}" + install_torch "torch==${RESOLVED_TORCH_VERSION}" || exit 1 fi fi +############################################################################### +# Verify the installed Torch package +############################################################################### +TORCH_INFO="$(pytorch_get_installed_torch_info)" || { + echo "[ERROR] Torch is not importable after installation." >&2 + exit 1 +} +IFS=$'\t' read -r \ + INSTALLED_TORCH_FULL \ + INSTALLED_TORCH_BASE \ + INSTALLED_TORCH_FAMILY \ + INSTALLED_TORCH_WHEEL_TAG \ + INSTALLED_TORCH_IS_NIGHTLY <<< "${TORCH_INFO}" + +echo "[INFO] Installed torch: ${INSTALLED_TORCH_FULL}" +echo "[INFO] Torch wheel build: ${INSTALLED_TORCH_WHEEL_TAG}" + ############################################################################### # Install the auxiliary Python requirements ############################################################################### REQ_FILE="${SCRIPTS_DIR}/install_requirements.txt" echo "[INFO] Installing auxiliary requirements from ${REQ_FILE##*/}" -python3 -m pip install -r "$REQ_FILE" +python3 -m pip install -r "${REQ_FILE}" || exit 1 ############################################################################### # TICO itself ############################################################################### -if [[ $_DIST -eq 1 ]]; then +if [[ "${_DIST}" -eq 1 ]]; then echo "[INFO] Installing TICO wheel from ./dist" - python3 -m pip install --force-reinstall --no-deps "${CCEX_PROJECT_PATH}"/dist/tico*.whl + python3 -m pip install --force-reinstall --no-deps \ + "${CCEX_PROJECT_PATH}"/dist/tico*.whl || exit 1 else echo "[INFO] Installing TICO in editable mode" - python3 -m pip install --editable "${CCEX_PROJECT_PATH}" + python3 -m pip install --editable "${CCEX_PROJECT_PATH}" || exit 1 fi +# TorchVision is installed by `./ccex configure test`, so a global `pip check` +# belongs at the end of that command rather than between the two setup stages. echo "[SUCCESS] ./ccex install completed" diff --git a/infra/scripts/pytorch_package_utils.sh b/infra/scripts/pytorch_package_utils.sh new file mode 100644 index 00000000..6643cb83 --- /dev/null +++ b/infra/scripts/pytorch_package_utils.sh @@ -0,0 +1,315 @@ +#!/bin/bash + +# Copyright (c) 2025 Samsung Electronics Co., Ltd. All Rights Reserved +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Shared PyTorch package metadata and helper functions for ccex scripts. + +if [[ -n "${TICO_PYTORCH_PACKAGE_UTILS_SOURCED:-}" ]]; then + return 0 +fi +TICO_PYTORCH_PACKAGE_UTILS_SOURCED=1 + +PYTORCH_DEFAULT_FAMILY="2.7" +PYTORCH_SUPPORTED_FAMILIES=("2.5" "2.6" "2.7" "2.8" "2.9" "2.10") + +# Keep family requests deterministic. A family such as 2.7 resolves to the +# latest stable patch explicitly supported by TICO instead of asking pip to +# backtrack over every 2.7.x release. +declare -A PYTORCH_LATEST_STABLE_VERSION=( + ["2.5"]="2.5.1" + ["2.6"]="2.6.0" + ["2.7"]="2.7.1" + ["2.8"]="2.8.0" + ["2.9"]="2.9.1" + ["2.10"]="2.10.0" +) + +# Official CUDA wheel variants for each supported stable Torch family. Values +# are ordered from newest to oldest so that the best compatible wheel is tried +# first. The CPU index is appended separately as the final fallback. +declare -A PYTORCH_STABLE_CUDA_WHEELS=( + ["2.5"]="12.4 12.1 11.8" + ["2.6"]="12.6 12.4 11.8" + ["2.7"]="12.8 12.6 11.8" + ["2.8"]="12.9 12.8 12.6" + ["2.9"]="13.0 12.8 12.6" + ["2.10"]="13.0 12.8 12.6" +) + +# Used only when a nightly requirement does not contain an explicit local +# build tag such as +cpu or +cu130. +PYTORCH_NIGHTLY_CUDA_FALLBACKS=("13.0" "12.8" "12.6") + +pytorch_is_supported_family() { + local family="$1" + local supported + + for supported in "${PYTORCH_SUPPORTED_FAMILIES[@]}"; do + [[ "${supported}" == "${family}" ]] && return 0 + done + return 1 +} + +pytorch_version_le() { + local lhs="$1" + local rhs="$2" + + [[ "$(printf '%s\n%s\n' "${lhs}" "${rhs}" | sort -V | head -n1)" == "${lhs}" ]] +} + +pytorch_strip_local_version() { + local version="$1" + echo "${version%%+*}" +} + +pytorch_version_family() { + local version + version="$(pytorch_strip_local_version "$1")" + + if [[ "${version}" =~ ^([0-9]+\.[0-9]+) ]]; then + echo "${BASH_REMATCH[1]}" + return 0 + fi + + return 1 +} + +pytorch_is_nightly_version() { + local version + version="$(pytorch_strip_local_version "$1")" + [[ "${version}" == *".dev"* ]] +} + +pytorch_local_build_tag() { + local version="$1" + + if [[ "${version}" == *+* ]]; then + echo "${version#*+}" + fi +} + +pytorch_resolve_latest_stable_version() { + local family="$1" + + if ! pytorch_is_supported_family "${family}"; then + return 1 + fi + + echo "${PYTORCH_LATEST_STABLE_VERSION[${family}]}" +} + +pytorch_resolve_torchvision_version() { + local torch_version + local major + local minor + local patch + local family + local vision_minor + + torch_version="$(pytorch_strip_local_version "$1")" + if [[ ! "${torch_version}" =~ ^([0-9]+)\.([0-9]+)\.([0-9]+)$ ]]; then + return 1 + fi + + major="${BASH_REMATCH[1]}" + minor="${BASH_REMATCH[2]}" + patch="${BASH_REMATCH[3]}" + family="${major}.${minor}" + + if [[ "${major}" != "2" ]] || ! pytorch_is_supported_family "${family}"; then + return 1 + fi + + # Torch 2.N.P is paired with TorchVision 0.(N+15).P for the supported + # release families (2.5 through 2.10). + vision_minor=$((10#${minor} + 15)) + echo "0.${vision_minor}.${patch}" +} + +pytorch_cuda_version_to_tag() { + local cuda_version="$1" + local major + local minor + + if [[ ! "${cuda_version}" =~ ^([0-9]+)\.([0-9]+)$ ]]; then + return 1 + fi + + major="${BASH_REMATCH[1]}" + minor="${BASH_REMATCH[2]}" + echo "cu${major}${minor}" +} + +pytorch_cuda_tag_to_version() { + local tag="$1" + local digits + local length + + if [[ ! "${tag}" =~ ^cu([0-9]+)$ ]]; then + return 1 + fi + + digits="${BASH_REMATCH[1]}" + length=${#digits} + if (( length < 3 )); then + return 1 + fi + + # Supported CUDA releases use a two-digit major version (11, 12, 13, ...). + # Keep the remaining digits as the minor version so a future cu1310 tag maps + # to 13.10 instead of 131.0. + echo "${digits:0:2}.${digits:2}" +} + +pytorch_is_supported_wheel_tag() { + local tag="$1" + [[ "${tag}" == "cpu" || "${tag}" =~ ^cu[0-9]+$ ]] +} + +pytorch_index_url() { + local wheel_tag="$1" + local is_nightly="${2:-0}" + local channel="" + + if [[ "${is_nightly}" == "1" ]]; then + channel="/nightly" + fi + + echo "https://download.pytorch.org/whl${channel}/${wheel_tag}" +} + +pytorch_detect_host_cuda_version() { + local cuda_version="" + + if command -v nvcc &>/dev/null; then + cuda_version="$(nvcc --version | sed -nE 's/.*release ([0-9]+\.[0-9]+).*/\1/p' | head -n1)" + fi + + if [[ -z "${cuda_version}" ]] && command -v nvidia-smi &>/dev/null; then + cuda_version="$(nvidia-smi | sed -nE 's/.*CUDA Version: ([0-9]+\.[0-9]+).*/\1/p' | head -n1)" + fi + + [[ -n "${cuda_version}" ]] || return 1 + echo "${cuda_version}" +} + +pytorch_add_unique_index_url() { + local candidate="$1" + local existing + + for existing in "${PYTORCH_INDEX_URLS[@]}"; do + [[ "${existing}" == "${candidate}" ]] && return 0 + done + + PYTORCH_INDEX_URLS+=("${candidate}") +} + +pytorch_build_index_urls() { + local family="$1" + local host_cuda="$2" + local is_nightly="${3:-0}" + local requested_wheel_tag="${4:-}" + local cuda_version + local cuda_tag + local -a cuda_candidates=() + + PYTORCH_INDEX_URLS=() + + if [[ -n "${requested_wheel_tag}" ]]; then + if ! pytorch_is_supported_wheel_tag "${requested_wheel_tag}"; then + echo "[ERROR] Unsupported PyTorch wheel tag '${requested_wheel_tag}'" >&2 + return 1 + fi + + pytorch_add_unique_index_url "$(pytorch_index_url "${requested_wheel_tag}" "${is_nightly}")" + return 0 + fi + + if [[ -n "${host_cuda}" ]]; then + if [[ "${is_nightly}" == "1" ]]; then + # Try the detected CUDA index first because nightly availability changes + # independently of stable release families. + cuda_tag="$(pytorch_cuda_version_to_tag "${host_cuda}")" || return 1 + pytorch_add_unique_index_url "$(pytorch_index_url "${cuda_tag}" 1)" + cuda_candidates=("${PYTORCH_NIGHTLY_CUDA_FALLBACKS[@]}") + else + if ! pytorch_is_supported_family "${family}"; then + echo "[ERROR] Unsupported Torch family '${family}'" >&2 + return 1 + fi + read -r -a cuda_candidates <<< "${PYTORCH_STABLE_CUDA_WHEELS[${family}]}" + fi + + for cuda_version in "${cuda_candidates[@]}"; do + if pytorch_version_le "${cuda_version}" "${host_cuda}"; then + cuda_tag="$(pytorch_cuda_version_to_tag "${cuda_version}")" || return 1 + pytorch_add_unique_index_url "$(pytorch_index_url "${cuda_tag}" "${is_nightly}")" + fi + done + fi + + # A CPU wheel keeps installation usable on machines without a compatible + # published CUDA wheel. An explicit +cuXXX request does not reach this path. + pytorch_add_unique_index_url "$(pytorch_index_url "cpu" "${is_nightly}")" +} + +pytorch_get_pinned_requirement_version() { + local requirement_file="$1" + local package_name="$2" + + [[ -f "${requirement_file}" ]] || return 1 + + sed -nE \ + "s/^[[:space:]]*${package_name}==([^[:space:]#]+).*$/\\1/p" \ + "${requirement_file}" | head -n1 +} + +pytorch_get_installed_torch_info() { + python3 - <<'PY' +import importlib.util +import re +import sys + +if importlib.util.find_spec("torch") is None: + sys.exit(1) + +try: + import torch +except Exception as exc: + print(f"Failed to import torch: {exc}", file=sys.stderr) + sys.exit(2) + +full_version = str(torch.__version__) +base_version = full_version.split("+", 1)[0] +family_match = re.match(r"^(\d+\.\d+)", base_version) +if family_match is None: + print(f"Unrecognized torch version: {full_version}", file=sys.stderr) + sys.exit(3) + +family = family_match.group(1) +is_nightly = "1" if ".dev" in base_version else "0" +cuda_version = getattr(torch.version, "cuda", None) +hip_version = getattr(torch.version, "hip", None) + +if cuda_version: + wheel_tag = "cu" + str(cuda_version).replace(".", "") +elif hip_version: + wheel_tag = "rocm" + str(hip_version) +else: + wheel_tag = "cpu" + +print("\t".join((full_version, base_version, family, wheel_tag, is_nightly))) +PY +} diff --git a/infra/scripts/test_configure.sh b/infra/scripts/test_configure.sh index 04de5038..0feda55e 100755 --- a/infra/scripts/test_configure.sh +++ b/infra/scripts/test_configure.sh @@ -24,126 +24,263 @@ ############################################################################### SCRIPTS_DIR="${CCEX_PROJECT_PATH}/infra/scripts" TEST_DIR="${CCEX_PROJECT_PATH}/test" - -SUPPORTED_FAMILIES=("2.5" "2.6" "2.7" "2.8" "2.9" "2.10") -DEFAULT_FAMILY="2.7" +source "${SCRIPTS_DIR}/pytorch_package_utils.sh" || exit 1 show_help() { -cat <&2; exit 1; } -eval set -- "$options" +eval set -- "${options}" while true; do case "$1" in - --torch_ver) _TORCH_VER="$2"; shift 2;; - --cuda_ver) _USER_CUDA="$2"; shift 2;; - --cpu_only) _CPU_ONLY=1; shift;; - -h|--help) show_help; exit 0;; - --) shift; break;; - *) echo "[ERROR] Unknown option $1"; exit 1;; + --torch_ver) _TORCH_VER="$2"; _TORCH_VER_WAS_SET=1; shift ;; + --cuda_ver) _USER_CUDA="$2"; shift ;; + --cpu_only) _CPU_ONLY=1 ;; + -h|--help) show_help; exit 0 ;; + --) shift; break ;; + *) echo "[ERROR] Unknown option $1"; exit 1 ;; esac + shift done +if [[ "${_CPU_ONLY}" == "1" && -n "${_USER_CUDA}" ]]; then + echo "[ERROR] --cpu_only and --cuda_ver cannot be used together." >&2 + exit 1 +fi + +if [[ -n "${_USER_CUDA}" && ! "${_USER_CUDA}" =~ ^[0-9]+\.[0-9]+$ ]]; then + echo "[ERROR] Invalid --cuda_ver value '${_USER_CUDA}'. Expected MAJ.MIN." >&2 + exit 1 +fi + ############################################################################### -# Torch version analysis +# Inspect the installed Torch package ############################################################################### -REQUEST_IS_NIGHTLY="" -REQUEST_IS_EXACT="" -if [[ "${_TORCH_VER}" == "nightly" ]]; then - REQUEST_IS_NIGHTLY=1 +TORCH_INFO="$(pytorch_get_installed_torch_info)" || { + echo "[ERROR] Torch is not installed or cannot be imported." >&2 + echo "[ERROR] Run './ccex install' before './ccex configure test'." >&2 + exit 1 +} +IFS=$'\t' read -r \ + INSTALLED_TORCH_FULL \ + INSTALLED_TORCH_BASE \ + INSTALLED_TORCH_FAMILY \ + INSTALLED_TORCH_WHEEL_TAG \ + INSTALLED_TORCH_IS_NIGHTLY <<< "${TORCH_INFO}" + +echo "[INFO] Installed torch: ${INSTALLED_TORCH_FULL}" +echo "[INFO] Torch wheel build: ${INSTALLED_TORCH_WHEEL_TAG}" + +if [[ "${INSTALLED_TORCH_WHEEL_TAG}" != "cpu" && \ + ! "${INSTALLED_TORCH_WHEEL_TAG}" =~ ^cu[0-9]+$ ]]; then + echo "[ERROR] Unsupported installed Torch wheel '${INSTALLED_TORCH_WHEEL_TAG}'." >&2 + echo "[ERROR] ccex currently supports CPU and CUDA PyTorch wheels." >&2 + exit 1 fi -FAMILY="$(echo "${_TORCH_VER}" | grep -oE '^[0-9]+\.[0-9]+' || echo "${DEFAULT_FAMILY}")" -if [[ -z "${REQUEST_IS_NIGHTLY}" && " ${SUPPORTED_FAMILIES[*]} " != *" ${FAMILY} "* ]]; then - echo "[ERROR] Unsupported --torch_ver ${_TORCH_VER}"; exit 1 + +if [[ "${INSTALLED_TORCH_IS_NIGHTLY}" == "0" ]] && \ + ! pytorch_is_supported_family "${INSTALLED_TORCH_FAMILY}"; then + echo "[ERROR] Installed torch ${INSTALLED_TORCH_FULL} is not supported." >&2 + exit 1 fi ############################################################################### -# Index-url (same rules as install.sh) +# Validate optional version and compute-platform requests ############################################################################### -get_index_url() { - local cuda="$1" nightly="$2" - local maj=${cuda%.*} min=${cuda#*.} - echo "https://download.pytorch.org/whl${nightly:+/nightly}/cu${maj}${min}" -} +if [[ "${_TORCH_VER_WAS_SET}" == "1" ]]; then + if [[ "${_TORCH_VER}" == "nightly" ]]; then + if [[ "${INSTALLED_TORCH_IS_NIGHTLY}" != "1" ]]; then + echo "[ERROR] --torch_ver nightly was requested, but ${INSTALLED_TORCH_FULL} is installed." >&2 + exit 1 + fi + elif [[ "${_TORCH_VER}" =~ ^[0-9]+\.[0-9]+$ ]]; then + if ! pytorch_is_supported_family "${_TORCH_VER}"; then + echo "[ERROR] Unsupported --torch_ver family '${_TORCH_VER}'" >&2 + exit 1 + fi + if [[ "${INSTALLED_TORCH_IS_NIGHTLY}" == "1" || \ + "${INSTALLED_TORCH_FAMILY}" != "${_TORCH_VER}" ]]; then + echo "[ERROR] Requested torch family ${_TORCH_VER}, but ${INSTALLED_TORCH_FULL} is installed." >&2 + echo "[ERROR] Re-run './ccex install --torch_ver ${_TORCH_VER}' first." >&2 + exit 1 + fi + elif [[ "${_TORCH_VER}" =~ ^[0-9]+\.[0-9]+\.[0-9]+(\+[A-Za-z0-9._-]+)?$ ]]; then + REQUESTED_BASE="$(pytorch_strip_local_version "${_TORCH_VER}")" + REQUESTED_BUILD_TAG="$(pytorch_local_build_tag "${_TORCH_VER}")" + REQUESTED_FAMILY="$(pytorch_version_family "${_TORCH_VER}")" || exit 1 -INDEX_URL="https://download.pytorch.org/whl${REQUEST_IS_NIGHTLY:+/nightly}/cpu" -if [[ -z "${_CPU_ONLY}" ]]; then - if [[ -n "${_USER_CUDA}" ]]; then - CUDA="${_USER_CUDA}" - echo "[INFO] Using user-specified CUDA ${CUDA}" - else - if command -v nvcc &>/dev/null; then CUDA="$(nvcc --version | grep -oP 'release \K[0-9]+\.[0-9]+')" ; fi - if [[ -z "${CUDA:-}" ]] && command -v nvidia-smi &>/dev/null; then - CUDA="$(nvidia-smi | grep -oP 'CUDA Version: \K[0-9]+\.[0-9]+')" + if ! pytorch_is_supported_family "${REQUESTED_FAMILY}"; then + echo "[ERROR] Unsupported --torch_ver value '${_TORCH_VER}'" >&2 + exit 1 fi + if [[ "${INSTALLED_TORCH_BASE}" != "${REQUESTED_BASE}" ]]; then + echo "[ERROR] Requested torch ${_TORCH_VER}, but ${INSTALLED_TORCH_FULL} is installed." >&2 + exit 1 + fi + if [[ -n "${REQUESTED_BUILD_TAG}" && \ + "${INSTALLED_TORCH_WHEEL_TAG}" != "${REQUESTED_BUILD_TAG}" ]]; then + echo "[ERROR] Requested Torch wheel ${REQUESTED_BUILD_TAG}, but ${INSTALLED_TORCH_WHEEL_TAG} is installed." >&2 + exit 1 + fi + else + echo "[ERROR] Unsupported --torch_ver value '${_TORCH_VER}'" >&2 + exit 1 fi - if [[ -n "${CUDA:-}" ]]; then - echo "[INFO] Detected CUDA ${CUDA}" - INDEX_URL="$(get_index_url "${CUDA}" "${REQUEST_IS_NIGHTLY}")" +fi + +if [[ "${_CPU_ONLY}" == "1" && "${INSTALLED_TORCH_WHEEL_TAG}" != "cpu" ]]; then + echo "[ERROR] --cpu_only was requested, but ${INSTALLED_TORCH_WHEEL_TAG} Torch is installed." >&2 + echo "[ERROR] Re-run './ccex install --cpu_only' first." >&2 + exit 1 +fi + +if [[ -n "${_USER_CUDA}" ]]; then + if [[ "${INSTALLED_TORCH_WHEEL_TAG}" == "cpu" ]]; then + echo "[INFO] Installed Torch is CPU-only; --cuda_ver does not change TorchVision selection." + else + INSTALLED_WHEEL_CUDA="$(pytorch_cuda_tag_to_version "${INSTALLED_TORCH_WHEEL_TAG}")" || exit 1 + if ! pytorch_version_le "${INSTALLED_WHEEL_CUDA}" "${_USER_CUDA}"; then + echo "[ERROR] CUDA ${_USER_CUDA} cannot run the installed ${INSTALLED_TORCH_WHEEL_TAG} Torch wheel." >&2 + exit 1 + fi + echo "[INFO] CUDA ${_USER_CUDA} is compatible with ${INSTALLED_TORCH_WHEEL_TAG}." fi -else - echo "[INFO] Forcing CPU-only test deps" fi ############################################################################### -# torchvision version mapping per torch family +# Install the TorchVision build matching the installed Torch package ############################################################################### -declare -A TORCHVISION_FAMILY=( - ["2.5"]="0.20.*" - ["2.6"]="0.21.*" - ["2.7"]="0.22.*" - ["2.8"]="0.23.*" - ["2.9"]="0.24.*" - ["2.10"]="0.25.*" - ["2.11"]="0.26.*" - ["2.12"]="0.27.*" - ["2.13"]="0.28.*" - ["2.14"]="0.29.*" -) +INDEX_URL="$(pytorch_index_url \ + "${INSTALLED_TORCH_WHEEL_TAG}" "${INSTALLED_TORCH_IS_NIGHTLY}")" +EXPECTED_VISION_VERSION="" -############################################################################### -# 1) Install torchvision -############################################################################### -if [[ -n "${REQUEST_IS_NIGHTLY}" ]]; then - echo "[INFO] Installing torchvision (nightly) from ${INDEX_URL}" - python3 -m pip install -r ${SCRIPTS_DIR}/../dependency/torchvision_dev.txt \ - --index-url "${INDEX_URL}" +if [[ "${INSTALLED_TORCH_IS_NIGHTLY}" == "1" ]]; then + TORCHVISION_DEV_FILE="${SCRIPTS_DIR}/../dependency/torchvision_dev.txt" + EXPECTED_VISION_FULL="$(pytorch_get_pinned_requirement_version \ + "${TORCHVISION_DEV_FILE}" torchvision)" + if [[ -z "${EXPECTED_VISION_FULL}" ]]; then + echo "[ERROR] ${TORCHVISION_DEV_FILE} must pin torchvision with 'torchvision==VERSION'." >&2 + exit 1 + fi + + PINNED_VISION_BUILD_TAG="$(pytorch_local_build_tag "${EXPECTED_VISION_FULL}")" + if [[ -n "${PINNED_VISION_BUILD_TAG}" && \ + "${PINNED_VISION_BUILD_TAG}" != "${INSTALLED_TORCH_WHEEL_TAG}" ]]; then + echo "[ERROR] torchvision_dev.txt pins +${PINNED_VISION_BUILD_TAG}, but Torch uses ${INSTALLED_TORCH_WHEEL_TAG}." >&2 + exit 1 + fi + + EXPECTED_VISION_VERSION="$(pytorch_strip_local_version "${EXPECTED_VISION_FULL}")" + echo "[INFO] Installing torchvision ${EXPECTED_VISION_FULL} from ${INDEX_URL}" + python3 -m pip install -r "${TORCHVISION_DEV_FILE}" \ + --index-url "${INDEX_URL}" || exit 1 else - VISION_VER="${TORCHVISION_FAMILY[${FAMILY}]}" - echo "[INFO] Installing torchvision==${VISION_VER} from ${INDEX_URL}" - python3 -m pip install "torchvision==${VISION_VER}" --index-url "${INDEX_URL}" -fi + EXPECTED_VISION_VERSION="$(pytorch_resolve_torchvision_version \ + "${INSTALLED_TORCH_BASE}")" || { + echo "[ERROR] No TorchVision mapping for torch ${INSTALLED_TORCH_BASE}." >&2 + exit 1 + } + echo "[INFO] Installing torchvision==${EXPECTED_VISION_VERSION} from ${INDEX_URL}" + python3 -m pip install "torchvision==${EXPECTED_VISION_VERSION}" \ + --index-url "${INDEX_URL}" || exit 1 +fi ############################################################################### -# 2) Install additional test-only requirements +# Install additional test-only requirements ############################################################################### EXTRA_REQ_FILE="${TEST_DIR}/requirements.txt" EXTRA_REQ_PRE_FILE="${TEST_DIR}/requirements_pre.txt" -python3 -m pip install -r ${EXTRA_REQ_FILE} -python3 -m pip install -r ${EXTRA_REQ_PRE_FILE} --pre +python3 -m pip install -r "${EXTRA_REQ_FILE}" || exit 1 +python3 -m pip install -r "${EXTRA_REQ_PRE_FILE}" --pre || exit 1 + +############################################################################### +# Validate the final Torch and TorchVision pair +############################################################################### +EXPECTED_TORCH_BASE="${INSTALLED_TORCH_BASE}" \ +EXPECTED_VISION_BASE="${EXPECTED_VISION_VERSION}" \ +EXPECTED_WHEEL_TAG="${INSTALLED_TORCH_WHEEL_TAG}" \ +python3 - <<'PY' +import os +import sys + +try: + import torch + import torchvision +except Exception as exc: + print(f"[ERROR] Failed to import torch/torchvision: {exc}", file=sys.stderr) + sys.exit(1) + +torch_full = str(torch.__version__) +vision_full = str(torchvision.__version__) +torch_base = torch_full.split("+", 1)[0] +vision_base = vision_full.split("+", 1)[0] +expected_torch = os.environ["EXPECTED_TORCH_BASE"] +expected_vision = os.environ["EXPECTED_VISION_BASE"] +expected_tag = os.environ["EXPECTED_WHEEL_TAG"] + +if torch_base != expected_torch: + print( + f"[ERROR] Torch changed during configure: expected {expected_torch}, got {torch.__version__}", + file=sys.stderr, + ) + sys.exit(1) + +if vision_base != expected_vision: + print( + f"[ERROR] Unexpected TorchVision: expected {expected_vision}, got {vision_full}", + file=sys.stderr, + ) + sys.exit(1) + +vision_tag = vision_full.split("+", 1)[1] if "+" in vision_full else "" +# x86_64 PyTorch wheel indices encode +cpu/+cuXXX in TorchVision's version, +# while some other architectures publish a tagless but index-specific wheel. +if vision_tag and vision_tag != expected_tag: + print( + f"[ERROR] TorchVision wheel mismatch: expected {expected_tag}, got {vision_full}", + file=sys.stderr, + ) + sys.exit(1) + +actual_tag = "cpu" if torch.version.cuda is None else "cu" + torch.version.cuda.replace(".", "") +if actual_tag != expected_tag: + print( + f"[ERROR] Torch wheel changed during configure: expected {expected_tag}, got {actual_tag}", + file=sys.stderr, + ) + sys.exit(1) + +print(f"[INFO] Verified torch {torch_full} with torchvision {vision_full}") +PY + +python3 -m pip check || { + echo "[ERROR] Installed Python packages have incompatible dependencies." >&2 + exit 1 +} echo "[SUCCESS] ./ccex configure test completed"