Cleanup based on other comments

- Moved to using local variables inside install_whl cleaning up old globals from when it wasn't a function
- Fixed the devinstall_triton torch dep install to respect the user's INSTALL_TORCH

Signed-off-by: Craig Magina <cmagina@redhat.com>

Author: cmagina

scripts/devinstall_torch.sh

@@ -24,9 +24,6 @@ WORKSPACE=${WORKSPACE:-${HOME}}
 TORCH_DIR=${WORKSPACE}/torch
 TORCH_REPO=https://github.com/pytorch/pytorch.git
 
-declare -a PIP_INSTALL_ARGS
-PIP_TORCH_INDEX_URL_BASE=https://download.pytorch.org/whl
-
 SUDO=''
 if ((EUID != 0)) && command -v sudo &>/dev/null; then
 	SUDO="sudo"
@@ -102,7 +99,22 @@ install_deps() {
 }
 
 install_whl() {
-	echo "Installing Torch ${PIP_TORCH_INDEX_URL_BUILD:-release} from PyPI ..."
+	local pip_build="$1"
+	local pip_torch_index_url_base
+	local -a pip_install_args
+	local compute_platform
+
+
+	pip_torch_index_url_base="https://download.pytorch.org/whl"
+
+	case "$pip_build" in
+		release) ;;
+		nightly | test)
+			pip_torch_index_url_base="${pip_torch_index_url_base}/${pip_build}"
+			;;
+	esac
+
+	echo "Installing Torch $pip_build from PyPI ..."
 
 	if [ -n "${PIP_TORCH_VERSION:-}" ]; then
 		echo "Using the specified version $PIP_TORCH_VERSION of torch"
@@ -127,34 +139,33 @@ install_whl() {
 
 	if [ -n "${PIP_TORCH_INDEX_URL:-}" ]; then
 		echo "Using the specified index, $PIP_TORCH_INDEX_URL"
-		PIP_INSTALL_ARGS+=("--index-url" "$PIP_TORCH_INDEX_URL")
+		pip_install_args+=("--index-url" "$PIP_TORCH_INDEX_URL")
 	elif command -v uv &>/dev/null && [ -n "${UV_TORCH_BACKEND:-}" ]; then
 		echo "Using the specified uv backend, $UV_TORCH_BACKEND"
-		PIP_INSTALL_ARGS+=("--torch-backend" "$UV_TORCH_BACKEND")
+		pip_install_args+=("--torch-backend" "$UV_TORCH_BACKEND")
 	elif ! command -v uv &>/dev/null && [ -n "${UV_TORCH_BACKEND:-}" ]; then
 		echo "Error: UV_TORCH_BACKEND is set to $UV_TORCH_BACKEND but uv is not available."
 		exit 1
 	else
 		# Set compute platform for torch wheel installation
 		if [ -n "${ROCM_VERSION:-}" ]; then
 			echo "Using the ROCm version $ROCM_VERSION backend"
-			COMPUTE_PLATFORM="rocm$(get_rocm_version)"
+			compute_platform="rocm$(get_rocm_version)"
 		elif ((${TRITON_CPU_BACKEND:-0} == 1)); then
 			echo "Using the CPU backend"
-			COMPUTE_PLATFORM="cpu"
+			compute_platform="cpu"
 		elif [ -n "${CUDA_VERSION:-}" ]; then
 			echo "Using the CUDA version $CUDA_VERSION backend"
-			COMPUTE_PLATFORM="cu${CUDA_VERSION/[.-]/}"
+			compute_platform="cu${CUDA_VERSION/[.-]/}"
 		fi
 
-		if [ -n "${COMPUTE_PLATFORM:-}" ]; then
-			[[ -n "${PIP_TORCH_INDEX_URL_BUILD:-}" ]] && PIP_TORCH_INDEX_URL_BUILD="/${PIP_TORCH_INDEX_URL_BUILD}"
-			PIP_TORCH_INDEX_URL="${PIP_TORCH_INDEX_URL_BASE}${PIP_TORCH_INDEX_URL_BUILD:-}/${COMPUTE_PLATFORM}"
-			PIP_INSTALL_ARGS+=("--index-url" "$PIP_TORCH_INDEX_URL")
+		if [ -n "${compute_platform:-}" ]; then
+			PIP_TORCH_INDEX_URL="${pip_torch_index_url_base}/${compute_platform}"
+			pip_install_args+=("--index-url" "$PIP_TORCH_INDEX_URL")
 		fi
 	fi
 
-	pip_install -U --force-reinstall "${PIP_INSTALL_ARGS[@]}" "${TORCH_PACKAGES[@]}"
+	pip_install -U --force-reinstall "${pip_install_args[@]}" "${TORCH_PACKAGES[@]}"
 
 	# Fix up LD_LIBRARY_PATH for CUDA
 	ldpretend
@@ -186,14 +197,9 @@ source)
 	install_build_deps
 	install_deps
 	;;
-release)
-	install_deps
-	install_whl
-	;;
-nightly | test)
-	PIP_TORCH_INDEX_URL_BUILD=$COMMAND
+nightly | release | test)
 	install_deps
-	install_whl
+	install_whl "$COMMAND"
 	;;
 *)
 	usage

scripts/devinstall_triton.sh

@@ -122,8 +122,15 @@ install_deps() {
 	pip_install cmake ctypeslib2 matplotlib ninja \
 		numpy pandas pybind11 pytest pyyaml scipy tabulate wheel
 
-	echo "Installing Torch as a Triton dependency ..."
-	devinstall_torch release
+	if [ "${INSTALL_TORCH:-}" != "source" ]; then
+		if [ -n "${INSTALL_TORCH:-}" ]; then
+			echo "Installing Torch $INSTALL_TORCH as a dependency ..."
+			devinstall_torch "${INSTALL_TORCH}"
+		else
+			echo "Installing Torch as a dependency ..."
+			devinstall_torch release
+		fi
+	fi
 }
 
 install_whl() {