Skip to content

feat: add PyTorch setup script for modular installation #117

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
cmagina merged 6 commits into from
Apr 9, 2026
Merged

feat: add PyTorch setup script for modular installation #117

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
0964c2e
feat: add PyTorch setup script for modular installation
cmagina Feb 7, 2026
c9aa1c0
Fix for lack of sudo bailing out even for pip install only runs
cmagina Apr 8, 2026
a879dc2
Re-order git submodule update
cmagina Apr 8, 2026
6672bc2
Added default pip index url for nightly and test with no compute plat…
cmagina Apr 8, 2026
5602f7a
Handle INSTALL_TORCH=skip as triton dep install
cmagina Apr 8, 2026
b88845f
Added new torch environment variables to SAVE_VARS
cmagina Apr 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions .github/workflows/amd-image.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ on: # yamllint disable-line rule:truthy
- scripts/devinstall_software.sh
- scripts/ldpretend.sh
- scripts/devinstall_llvm.sh
- scripts/devinstall_torch.sh
- scripts/devinstall_triton.sh
- scripts/devcreate_user.sh
- scripts/devsetup.sh
Expand All @@ -22,6 +23,7 @@ on: # yamllint disable-line rule:truthy
- scripts/devinstall_software.sh
- scripts/ldpretend.sh
- scripts/devinstall_llvm.sh
- scripts/devinstall_torch.sh
- scripts/devinstall_triton.sh
- scripts/devcreate_user.sh
- scripts/devsetup.sh
Expand Down
2 changes: 2 additions & 0 deletions .github/workflows/cpu-image.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ on: # yamllint disable-line rule:truthy
- scripts/devinstall_software.sh
- scripts/ldpretend.sh
- scripts/devinstall_llvm.sh
- scripts/devinstall_torch.sh
- scripts/devinstall_triton.sh
- scripts/devcreate_user.sh
- scripts/devsetup.sh
Expand All @@ -22,6 +23,7 @@ on: # yamllint disable-line rule:truthy
- scripts/devinstall_software.sh
- scripts/ldpretend.sh
- scripts/devinstall_llvm.sh
- scripts/devinstall_torch.sh
- scripts/devinstall_triton.sh
- scripts/devcreate_user.sh
- scripts/devsetup.sh
Expand Down
2 changes: 2 additions & 0 deletions .github/workflows/nvidia-image.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ on: # yamllint disable-line rule:truthy
- scripts/devinstall_software.sh
- scripts/ldpretend.sh
- scripts/devinstall_llvm.sh
- scripts/devinstall_torch.sh
- scripts/devinstall_triton.sh
- scripts/devcreate_user.sh
- scripts/devsetup.sh
Expand All @@ -22,6 +23,7 @@ on: # yamllint disable-line rule:truthy
- scripts/devinstall_software.sh
- scripts/ldpretend.sh
- scripts/devinstall_llvm.sh
- scripts/devinstall_torch.sh
- scripts/devinstall_triton.sh
- scripts/devcreate_user.sh
- scripts/devsetup.sh
Expand Down
7 changes: 6 additions & 1 deletion Makefile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,12 +36,14 @@ TRITON_CPU_BACKEND ?=0
TRITON_TAG ?= latest
triton_path ?=$(source_dir)
llvm_path ?=
torch_path ?=
user_path ?=
gitconfig_path ?="$(HOME)/.gitconfig"
USERNAME ?=triton
# NOTE: Requires host build system to have a valid Red Hat Subscription if true
INSTALL_NSIGHT ?=false
INSTALL_LLVM ?= skip # Options: source, skip
INSTALL_TORCH ?= skip # Options: nightly, release, source, skip, test
INSTALL_TRITON ?= source # Options: release, source, skip
INSTALL_JUPYTER ?= true
USE_CCACHE ?= 0
Expand Down Expand Up @@ -97,6 +99,9 @@ define run_container
if [ -n "$(llvm_path)" ]; then \
volume_arg+=" -v $(llvm_path):/workspace/llvm-project$(SELINUXFLAG)"; \
fi; \
if [ -n "$(torch_path)" ]; then \
volume_arg+=" -v $(torch_path):/workspace/torch$(SELINUXFLAG)"; \
fi; \
if [ -n "$(user_path)" ]; then \
volume_arg+=" -v $(user_path):/workspace/user$(SELINUXFLAG)"; \
fi; \
Expand Down Expand Up @@ -135,7 +140,7 @@ define run_container
else \
port_arg=""; \
fi; \
env_vars="-e USERNAME=$(USER) -e USER_UID=`id -u $(USER)` -e USER_GID=`id -g $(USER)` -e TORCH_VERSION=$(torch_version) -e INSTALL_LLVM=$(INSTALL_LLVM) -e INSTALL_TOOLS=$(DEMO_TOOLS) -e INSTALL_JUPYTER=$(INSTALL_JUPYTER) -e NOTEBOOK_PORT=$(NOTEBOOK_PORT) -e INSTALL_TRITON=$(INSTALL_TRITON) -e USE_CCACHE=$(USE_CCACHE) -e MAX_JOBS=$(MAX_JOBS)"; \
env_vars="-e USERNAME=$(USER) -e USER_UID=`id -u $(USER)` -e USER_GID=`id -g $(USER)` -e TORCH_VERSION=$(torch_version) -e INSTALL_LLVM=$(INSTALL_LLVM) -e INSTALL_TOOLS=$(DEMO_TOOLS) -e INSTALL_JUPYTER=$(INSTALL_JUPYTER) -e NOTEBOOK_PORT=$(NOTEBOOK_PORT) -e INSTALL_TORCH=$(INSTALL_TORCH) -e INSTALL_TRITON=$(INSTALL_TRITON) -e USE_CCACHE=$(USE_CCACHE) -e MAX_JOBS=$(MAX_JOBS)"; \
if [ "$(STRIPPED_CMD)" = "docker" ]; then \
$(CTR_CMD) run $$env_vars $$gpu_args $$profiling_args $$port_arg \
-ti $$volume_arg $$gitconfig_arg $(IMAGE_REPO)/$(strip $(1)):$(TRITON_TAG) bash; \
Expand Down
1 change: 1 addition & 0 deletions dockerfiles/Dockerfile.triton
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,7 @@ ENV BASH_ENV=/workspace/bin/activate \
COPY --from=quay.io/triton-dev-containers/gosu /usr/local/bin/gosu /usr/local/bin/gosu

COPY scripts/devinstall_llvm.sh /workspace/bin/devinstall_llvm
COPY scripts/devinstall_torch.sh /workspace/bin/devinstall_torch
COPY scripts/devinstall_triton.sh /workspace/bin/devinstall_triton
COPY scripts/devcreate_user.sh /workspace/bin/devcreate_user
COPY scripts/devsetup.sh /workspace/bin/devsetup
Expand Down
1 change: 1 addition & 0 deletions dockerfiles/Dockerfile.triton-amd
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,7 @@ WORKDIR /workspace
COPY --from=quay.io/triton-dev-containers/gosu /usr/local/bin/gosu /usr/local/bin/gosu

COPY scripts/devinstall_llvm.sh /workspace/bin/devinstall_llvm
COPY scripts/devinstall_torch.sh /workspace/bin/devinstall_torch
COPY scripts/devinstall_triton.sh /workspace/bin/devinstall_triton
COPY scripts/devcreate_user.sh /workspace/bin/devcreate_user
COPY scripts/devsetup.sh /workspace/bin/devsetup
Expand Down
1 change: 1 addition & 0 deletions dockerfiles/Dockerfile.triton-cpu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@ WORKDIR /workspace
COPY --from=quay.io/triton-dev-containers/gosu /usr/local/bin/gosu /usr/local/bin/gosu

COPY scripts/devinstall_llvm.sh /workspace/bin/devinstall_llvm
COPY scripts/devinstall_torch.sh /workspace/bin/devinstall_torch
COPY scripts/devinstall_triton.sh /workspace/bin/devinstall_triton
COPY scripts/devcreate_user.sh /workspace/bin/devcreate_user
COPY scripts/devsetup.sh /workspace/bin/devsetup
Expand Down
219 changes: 219 additions & 0 deletions scripts/devinstall_torch.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,219 @@
#! /bin/bash -e

trap "echo -e '\nScript interrupted. Exiting gracefully.'; exit 1" SIGINT

# Copyright (C) 2024-2025 Red Hat, Inc.
#
# 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.
#
# SPDX-License-Identifier: Apache-2.0
set -euo pipefail

WORKSPACE=${WORKSPACE:-${HOME}}

TORCH_DIR=${WORKSPACE}/torch
TORCH_REPO=https://github.com/pytorch/pytorch.git

SUDO=''
if ((EUID != 0)) && command -v sudo &>/dev/null; then
SUDO="sudo"
fi
Comment on lines +27 to +30
Copy link
Copy Markdown

@coderabbitai coderabbitai bot Apr 8, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Privilege check is too broad and blocks wheel-only installs.

This exits early for non-root users without sudo even when running release/nightly/test, which don’t need system package installation.

Suggested fix 
-SUDO=''
-if ((EUID != 0)) && command -v sudo &>/dev/null; then
-	SUDO="sudo"
-elif ((EUID != 0)); then
-	echo "ERROR: $(basename "$0") requires root privileges or sudo." >&2
-	exit 1
-fi
+SUDO=''
+if ((EUID != 0)) && command -v sudo &>/dev/null; then
+	SUDO="sudo"
+fi
 install_build_deps() {
 	echo "Installing Torch build dependencies ..."
+	if ((EUID != 0)) && [ -z "${SUDO}" ]; then
+		echo "ERROR: $(basename "$0") requires root privileges or sudo for build dependencies." >&2
+		exit 1
+	fi
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@scripts/devinstall_torch.sh` around lines 27 - 33, The current privilege
check unconditionally exits for non-root users without sudo by guarding SUDO,
EUID and command -v sudo; change it so non-root without sudo is allowed when
doing wheel-only installs (e.g., install modes 'release', 'nightly', 'test') by
first detecting the requested install mode/argument and only enforcing the sudo
requirement for modes that need system package installation. Concretely, inspect
the install mode variable/argument used earlier in the script and wrap the
existing SUDO/EUID/command -v sudo logic so it runs only for non-wheel modes;
keep the SUDO variable, EUID check and command -v sudo usage but skip the exit
branch for wheel-only modes.


pip_install() {
if command -v uv &>/dev/null; then
uv pip install "$@"
else
pip install "$@"
fi
}

# Remove the dashes or periods from the CUDA version, e.g. 128 from 12-8
get_cuda_version() {
echo "${CUDA_VERSION//[.-]/}"
}

# Extract the major.minor version from ROCM_VERSION, e.g. 6.4 from 6.4.4
get_rocm_version() {
[[ "$ROCM_VERSION" =~ ^([0-9]+\.[0-9]+) ]] && echo "${BASH_REMATCH[1]}" ||
echo "$ROCM_VERSION"
}

setup_src() {
echo "Downloading Torch source code and setting up the environment for building from source..."

if [ ! -d "$TORCH_DIR" ]; then
echo "Cloning the Torch repo $TORCH_REPO to $TORCH_DIR ..."
git clone "$TORCH_REPO" "$TORCH_DIR"
if [ ! -d "$TORCH_DIR" ]; then
echo "$TORCH_DIR not found. ERROR Cloning repository..."
exit 1
else
pushd "$TORCH_DIR" 1>/dev/null || exit 1

if [ -n "${TORCH_GITREF:-}" ]; then
git checkout "$TORCH_GITREF"
fi

git submodule sync
git submodule update --init --recursive

echo "Install pre-commit hooks into your local Torch git repo (one-time)"
pip_install pre-commit
pre-commit install
popd 1>/dev/null
fi
else
echo "Torch repo already present, not cloning ..."
fi
}

install_build_deps() {
echo "Installing Torch build dependencies ..."

pushd "$TORCH_DIR" 1>/dev/null || exit 1

if [ -f requirements.txt ]; then
pip_install --group dev
pip_install mkl-static mkl-include
fi
Comment on lines +85 to +88
Copy link
Copy Markdown

@coderabbitai coderabbitai bot Apr 8, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

tmp="$(mktemp -d)"
trap 'rm -rf "$tmp"' EXIT

git clone --depth 1 https://github.com/pytorch/pytorch.git "$tmp/pytorch"

echo "== Check PyTorch dependency metadata =="
if [ -f "$tmp/pytorch/pyproject.toml" ]; then
  rg -n '^\[dependency-groups\]|^dev\s*=' "$tmp/pytorch/pyproject.toml" || true
fi
rg -n '^' "$tmp/pytorch/requirements.txt" | head -n 30 || true

echo "== Check local pip/uv support for --group =="
pip --help | rg -n -- '--group' || true
if command -v uv >/dev/null 2>&1; then
  uv pip install --help | rg -n -- '--group' || true
fi

Repository: redhat-et/triton-dev-containers

Length of output: 1826

🏁 Script executed:

cat -n scripts/devinstall_torch.sh

Repository: redhat-et/triton-dev-containers

Length of output: 7064

🏁 Script executed:

ls -la | head -20

Repository: redhat-et/triton-dev-containers

Length of output: 1257

🏁 Script executed:

find . -name "requirements.txt" -type f 2>/dev/null | head -10

Repository: redhat-et/triton-dev-containers

Length of output: 57

🏁 Script executed:

grep -n "pip_install" scripts/devinstall_torch.sh | head -20

Repository: redhat-et/triton-dev-containers

Length of output: 522

🏁 Script executed:

grep -r "\-\-group" scripts/

Repository: redhat-et/triton-dev-containers

Length of output: 128

🏁 Script executed:

grep -r "pip_install" . --include="*.sh" 2>/dev/null

Repository: redhat-et/triton-dev-containers

Length of output: 1858

--group dev flag incompatible with fallback to standard pip.

Line 88 uses pip_install --group dev, which relies on uv-specific syntax. Since pip_install falls back to standard pip when uv is unavailable (line 39), this command will fail with an unrecognized flag error if uv is not installed. The file check at line 87 suggests intent to install from requirements.txt, but the command should use -r requirements.txt for portability across both pip and uv.

Suggested fix 
 	if [ -f requirements.txt ]; then
-		pip_install --group dev
+		pip_install -r requirements.txt
 		pip_install mkl-static mkl-include
 	fi
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
if [ -f requirements.txt ]; then
pip_install --group dev
pip_install mkl-static mkl-include
fi
if [ -f requirements.txt ]; then
pip_install -r requirements.txt
pip_install mkl-static mkl-include
fi
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@scripts/devinstall_torch.sh` around lines 87 - 90, The pip_install invocation
uses an uv-specific flag (--group dev) which will fail when pip_install falls
back to standard pip; update the block that checks for requirements.txt to call
pip_install -r requirements.txt (so it works with both uv and pip) and keep the
subsequent pip_install mkl-static mkl-include lines unchanged; locate the
pip_install call in the block that references requirements.txt and replace the
uv-specific --group dev usage with the portable -r requirements.txt form.

Copy link
Copy Markdown
Collaborator Author

@cmagina cmagina Apr 8, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Incorrect, pip install --group dev is supported by both.

Copy link
Copy Markdown

@coderabbitai coderabbitai bot Apr 8, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@cmagina, you're correct — my claim that --group is uv-specific was wrong. pip install --group support for PEP 735 Dependency Groups was added in pip 25.1, so it is indeed supported by both pip (25.1+) and uv.

That said, there's still a subtle issue worth noting: --group installs from pyproject.toml dependency groups, accepting arguments of the form group or path:group — it reads from pyproject.toml, not from requirements.txt. The surrounding if [ -f requirements.txt ] guard checks for a different file format entirely, so the two may not be logically consistent (the guard passes if requirements.txt exists, but the install reads from pyproject.toml). This is separate from the uv/pip question and may be worth a quick look to ensure the intent is correct.

I apologize for the incorrect original comment.

📝 Learnings were identified and not saved (knowledge base disabled). Enable


if ((EUID == 0)) || [ -n "${SUDO:-}" ]; then
$SUDO dnf -y install numactl-devel
else
echo "ERROR: Can't install some build deps without root or sudo permissions." >&2
exit 1
fi

if [ -n "${ROCM_VERSION:-}" ]; then
python tools/amd_build/build_amd.py
fi

popd 1>/dev/null
}

install_deps() {
echo "Installing Torch dependencies ..."
pip_install numpy
}

install_whl() {
local pip_build="$1"

local compute_platform
local pip_torch_index_url_base
local -a pip_install_args

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"
PIP_TORCH_VERSION="==$PIP_TORCH_VERSION"
fi

if [ -n "${PIP_TORCHVISION_VERSION:-}" ]; then
echo "Installing the specified version $PIP_TORCHVISION_VERSION of torchvision"
PIP_TORCHVISION_VERSION="==$PIP_TORCHVISION_VERSION"
fi

if [ -n "${PIP_TORCHAUDIO_VERSION:-}" ]; then
echo "Installing the specified version $PIP_TORCHAUDIO_VERSION of torchaudio"
PIP_TORCHAUDIO_VERSION="==$PIP_TORCHAUDIO_VERSION"
fi

declare -a TORCH_PACKAGES=(
"torch${PIP_TORCH_VERSION:-}"
"torchvision${PIP_TORCHVISION_VERSION:-}"
"torchaudio${PIP_TORCHAUDIO_VERSION:-}"
)

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")
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")
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)"
elif ((${TRITON_CPU_BACKEND:-0} == 1)); then
echo "Using the CPU backend"
compute_platform="cpu"
elif [ -n "${CUDA_VERSION:-}" ]; then
echo "Using the CUDA version $CUDA_VERSION backend"
compute_platform="cu$(get_cuda_version)"
fi

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")
else
PIP_TORCH_INDEX_URL="${pip_torch_index_url_base}"
pip_install_args+=("--index-url" "$PIP_TORCH_INDEX_URL")
fi
Comment on lines +170 to +176
Copy link
Copy Markdown

@coderabbitai coderabbitai bot Apr 8, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

nightly/test may silently miss the intended index.

If no CUDA/ROCm/CPU backend is inferred, no index URL is set, so nightly/test can resolve from default PyPI instead of the nightly/test channel.

Suggested fix 
 		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")
+		elif [ "$pip_build" != "release" ]; then
+			# Ensure nightly/test always use the intended channel
+			pip_install_args+=("--index-url" "${pip_torch_index_url_base}/cpu")
 		fi
 	fi
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
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
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")
elif [ "$pip_build" != "release" ]; then
# Ensure nightly/test always use the intended channel
pip_install_args+=("--index-url" "${pip_torch_index_url_base}/cpu")
fi
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@scripts/devinstall_torch.sh` around lines 167 - 170, The current block only
sets PIP_TORCH_INDEX_URL when compute_platform is non-empty, which lets
"nightly"/"test" builds fall back to PyPI; update the logic around
compute_platform/PIP_TORCH_INDEX_URL so that when the requested channel is a
non-stable one (detect via TORCH_VERSION containing "nightly" or "test" or a
TORCH_CHANNEL env var) you still set PIP_TORCH_INDEX_URL (use
pip_torch_index_url_base with compute_platform if present, otherwise use
pip_torch_index_url_base alone) and push it into pip_install_args; change the
code around the compute_platform check that currently sets PIP_TORCH_INDEX_URL
and pip_install_args to handle the empty-compute_platform case for nightly/test
to avoid resolving from default PyPI.

fi

pip_install -U --force-reinstall "${pip_install_args[@]}" "${TORCH_PACKAGES[@]}"

# Fix up LD_LIBRARY_PATH for CUDA
ldpretend
}

usage() {
cat >&2 <<EOF
Usage: $(basename "$0") [COMMAND]
source Download Torch's source (if needed) and install the build deps
release Install Torch
nightly Install the Torch nightly wheel
test Install the Torch test wheel
EOF
}

##
## Main
##
if [ $# -ne 1 ]; then
usage
exit 1
fi

COMMAND=${1,,}

case $COMMAND in
source)
setup_src
install_build_deps
install_deps
;;
nightly | release | test)
install_deps
install_whl "$COMMAND"
;;
*)
usage
exit 1
;;
esac
37 changes: 12 additions & 25 deletions scripts/devinstall_triton.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,9 +19,6 @@ trap "echo -e '\nScript interrupted. Exiting gracefully.'; exit 1" SIGINT
# SPDX-License-Identifier: Apache-2.0
set -euo pipefail

declare -a PIP_INSTALL_ARGS
PIP_TORCH_INDEX_URL_BASE=https://download.pytorch.org/whl

WORKSPACE=${WORKSPACE:-${HOME}}

TRITON_DIR=${WORKSPACE}/triton
Expand Down Expand Up @@ -108,30 +105,20 @@ install_deps() {
pip_install cmake ctypeslib2 matplotlib ninja \
numpy pandas pybind11 pytest pyyaml scipy tabulate wheel

if [ -n "${TORCH_VERSION:-}" ]; then
echo "Installing the specified version $TORCH_VERSION of torch"
PIP_TORCH_VERSION="==$TORCH_VERSION"
fi

if [ -n "${ROCM_VERSION:-}" ]; then
echo "Installing torch for ROCm version $ROCM_VERSION"
pip_install "torch${PIP_TORCH_VERSION:-}" \
--index-url "${PIP_TORCH_INDEX_URL_BASE}/rocm$(get_rocm_version)"
elif ((${TRITON_CPU_BACKEND:-0} == 1)); then
echo "Installing torch for CPU"
pip_install "torch${PIP_TORCH_VERSION:-}" \
--index-url "${PIP_TORCH_INDEX_URL_BASE}/cpu"
elif [ -n "${CUDA_VERSION:-}" ]; then
echo "Installing torch for CUDA version $CUDA_VERSION"
pip_install "torch${PIP_TORCH_VERSION:-}" \
--index-url "${PIP_TORCH_INDEX_URL_BASE}/cu$(get_cuda_version)"
else
echo "Installing torch ..."
pip_install "torch${PIP_TORCH_VERSION:-}"
if [ "${INSTALL_TORCH:-}" != "source" ]; then
if [ -n "${INSTALL_TORCH:-}" ] && [ "${INSTALL_TORCH}" != "skip" ]; 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() {
local -a pip_install_args

echo "Installing Triton from PyPI ..."

if command -v uv &>/dev/null; then
Expand All @@ -151,7 +138,7 @@ install_whl() {
UV_TORCH_BACKEND=auto
fi

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
Expand All @@ -162,7 +149,7 @@ install_whl() {
PIP_TRITON_VERSION="==$PIP_TRITON_VERSION"
fi

pip_install -U --force-reinstall "${PIP_INSTALL_ARGS[@]}" "triton${PIP_TRITON_VERSION:-}"
pip_install -U --force-reinstall "${pip_install_args[@]}" "triton${PIP_TRITON_VERSION:-}"

# Fix up LD_LIBRARY_PATH for CUDA
ldpretend
Expand Down
Loading
Loading