Skip to content
Merged
Show file tree
Hide file tree
Changes from 9 commits
Commits
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
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ on: # yamllint disable-line rule:truthy
- scripts/entrypoint.sh
- scripts/devinstall_software.sh
- scripts/ldpretend.sh
- scripts/devinstall_llvm.sh
- scripts/devinstall_triton.sh
- scripts/devcreate_user.sh
- scripts/devsetup.sh
Expand All @@ -20,6 +21,7 @@ on: # yamllint disable-line rule:truthy
- scripts/entrypoint.sh
- scripts/devinstall_software.sh
- scripts/ldpretend.sh
- scripts/devinstall_llvm.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
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ on: # yamllint disable-line rule:truthy
- scripts/entrypoint.sh
- scripts/devinstall_software.sh
- scripts/ldpretend.sh
- scripts/devinstall_llvm.sh
- scripts/devinstall_triton.sh
- scripts/devcreate_user.sh
- scripts/devsetup.sh
Expand All @@ -20,6 +21,7 @@ on: # yamllint disable-line rule:truthy
- scripts/entrypoint.sh
- scripts/devinstall_software.sh
- scripts/ldpretend.sh
- scripts/devinstall_llvm.sh
- scripts/devinstall_triton.sh
- scripts/devcreate_user.sh
- scripts/devsetup.sh
Expand Down
2 changes: 2 additions & 0 deletions .github/workflows/llvm-image.yml
Original file line number Diff line number Diff line change
Expand Up @@ -6,10 +6,12 @@ on: # yamllint disable-line rule:truthy
- main
paths:
- .github/workflows/llvm-image.yml
- scripts/devinstall_llvm.sh
- dockerfiles/Dockerfile.llvm
pull_request:
paths:
- .github/workflows/llvm-image.yml
- scripts/devinstall_llvm.sh
- dockerfiles/Dockerfile.llvm
schedule:
- cron: '0 0 * * 0' # Runs every Sunday at midnight UTC
Expand Down
2 changes: 2 additions & 0 deletions .github/workflows/nvidia-image.yml
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ on: # yamllint disable-line rule:truthy
- scripts/entrypoint.sh
- scripts/devinstall_software.sh
- scripts/ldpretend.sh
- scripts/devinstall_llvm.sh
- scripts/devinstall_triton.sh
- scripts/devcreate_user.sh
- scripts/devsetup.sh
Expand All @@ -20,6 +21,7 @@ on: # yamllint disable-line rule:truthy
- scripts/entrypoint.sh
- scripts/devinstall_software.sh
- scripts/ldpretend.sh
- scripts/devinstall_llvm.sh
- scripts/devinstall_triton.sh
- scripts/devcreate_user.sh
- scripts/devsetup.sh
Expand Down
7 changes: 6 additions & 1 deletion Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,9 @@ gitconfig_path ?="$(HOME)/.gitconfig"
USERNAME ?=triton
# NOTE: Requires host build system to have a valid Red Hat Subscription if true
INSTALL_NSIGHT ?=false
llvm_path ?=
user_path ?=
INSTALL_LLVM ?= skip # Options: source, skip
INSTALL_TRITON ?= source # Options: release, source, skip
INSTALL_JUPYTER ?= true
USE_CCACHE ?= 0
Expand Down Expand Up @@ -113,6 +115,9 @@ define run_container
else \
volume_arg=""; \
fi; \
if [ -n "$(llvm_path)" ]; then \
volume_arg+=" -v $(llvm_path):/workspace/llvm-project$(SELINUXFLAG)"; \
fi; \
if [ -n "$(user_path)" ]; then \
volume_arg+=" -v $(user_path):/workspace/user$(SELINUXFLAG)"; \
fi; \
Expand Down Expand Up @@ -151,7 +156,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 CUSTOM_LLVM=$(CUSTOM_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 CUSTOM_LLVM=$(CUSTOM_LLVM) -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)"; \
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
16 changes: 6 additions & 10 deletions dockerfiles/Dockerfile.llvm
Original file line number Diff line number Diff line change
Expand Up @@ -8,20 +8,18 @@ ARG TRITON_CPU_BACKEND=0
ARG LLVM_TAG
USER 0

COPY scripts/devinstall_llvm.sh devinstall_llvm

# Conditionally execute the build based on CUSTOM_LLVM
RUN if [ "$CUSTOM_LLVM" = "true" ]; then \
dnf update -y && \
dnf -y install clang rpm-build git ninja-build cmake lld && \
dnf clean all && rm -rf /var/cache/dnf && \
python3 -m pip install --upgrade pip && \
python3 -m pip install --upgrade cmake ninja sccache pybind11 && \
if [ ! -d "/llvm-project" ]; then \
git clone https://github.com/llvm/llvm-project /llvm-project; \
else \
cd /llvm-project && git fetch origin; \
fi && \
cd /llvm-project && \
python3 -m pip install -r mlir/python/requirements.txt && \
pip install --upgrade cmake ninja sccache pybind11 && \
export WORKSPACE=/usr/src && \
./devinstall_llvm source && \
cd "/usr/src/llvm-project" && \
REPO="triton"; \
PROJECTS="mlir;llvm;lld"; \
if [ "$TRITON_CPU_BACKEND" = "1" ]; then REPO="triton-cpu"; fi; \
Expand Down Expand Up @@ -63,8 +61,6 @@ RUN if [ "$CUSTOM_LLVM" = "true" ]; then \
-DLLVM_INCLUDE_BENCHMARKS=OFF \
-DMLIR_ENABLE_EXECUTION_ENGINE:bool=OFF && \
ninja install; \
else \
echo "LLVM is already up-to-date. Skipping rebuild."; \
fi \
else \
echo "Skipping LLVM build because CUSTOM_LLVM is not true"; \
Expand Down
3 changes: 2 additions & 1 deletion dockerfiles/Dockerfile.triton
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ RUN dnf -y update && \

# Stage for llvm-local-true
FROM base AS llvm-local-true
COPY --from=quay.io/triton-dev-containers/llvm:latest /install /llvm-project/install
COPY --from=quay.io/triton-dev-containers/llvm:latest /install /llvm-install

# Stage for llvm-local-false
FROM base AS llvm-local-false
Expand Down Expand Up @@ -66,6 +66,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_triton.sh /workspace/bin/devinstall_triton
COPY scripts/devcreate_user.sh /workspace/bin/devcreate_user
COPY scripts/devsetup.sh /workspace/bin/devsetup
Expand Down
3 changes: 2 additions & 1 deletion dockerfiles/Dockerfile.triton-amd
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ ENV LC_ALL=C.UTF-8 \

# LLVM Integration Stages
FROM base AS llvm-local-true
COPY --from=quay.io/triton-dev-containers/llvm:latest /install /llvm-project/install
COPY --from=quay.io/triton-dev-containers/llvm:latest /install /llvm-install

FROM base AS llvm-local-false
ENV TRITON_OFFLINE_BUILD=NO
Expand Down Expand Up @@ -84,6 +84,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_triton.sh /workspace/bin/devinstall_triton
COPY scripts/devcreate_user.sh /workspace/bin/devcreate_user
COPY scripts/devsetup.sh /workspace/bin/devsetup
Expand Down
3 changes: 2 additions & 1 deletion dockerfiles/Dockerfile.triton-cpu
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ RUN dnf update -y && \

# Stage for llvm-local-true
FROM base AS llvm-local-true
COPY --from=quay.io/triton-dev-containers/llvm:cpu-latest /install /llvm-project/install
COPY --from=quay.io/triton-dev-containers/llvm:cpu-latest /install /llvm-install

# Stage for llvm-local-false
FROM base AS llvm-local-false
Expand Down Expand Up @@ -63,6 +63,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_triton.sh /workspace/bin/devinstall_triton
COPY scripts/devcreate_user.sh /workspace/bin/devcreate_user
COPY scripts/devsetup.sh /workspace/bin/devsetup
Expand Down
25 changes: 22 additions & 3 deletions scripts/devcreate_user.sh
Original file line number Diff line number Diff line change
Expand Up @@ -110,12 +110,28 @@ create_user() {
fix_permissions() {
echo "Fixing permissions for user $USERNAME ..."
chown "$USERNAME:$GROUP_ID" -R "$HOME"
[[ -n "$WORKSPACE" && -d $WORKSPACE ]] && chown "$USERNAME:$GROUP_ID" -R "$WORKSPACE"
[[ -n "${WORKSPACE:-}" && -d $WORKSPACE ]] && chown "$USERNAME:$GROUP_ID" -R "$WORKSPACE"

mkdir -p "/run/user/$USER_ID"
chown "$USERNAME:$GROUP_ID" "/run/user/$USER_ID"
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.

get_user_home() {
grep "$USERNAME" /etc/passwd | cut -d':' -f 6
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.

create_bashrc() {
if [ ! -f "${HOME}/.bashrc" ]; then
echo "Setting up ${HOME}/.bashrc for user $USERNAME ..."
install -m 0644 -t "$HOME" \
/etc/skel/.bash_logout \
/etc/skel/.bash_profile \
/etc/skel/.bashrc

mkdir -p ${HOME}/.bashrc.d
fi
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Outdated

##
## Main
##
Expand All @@ -124,14 +140,17 @@ if [ -n "${USERNAME:-}" ] && [ "${USERNAME:-}" != "root" ]; then
echo "Creating user $USERNAME ..."
update_max_uid_gid
create_user
fix_permissions
install_sudo

if [ -n "${ROCM_VERSION:-}" ]; then
echo "Adding the user $USERNAME to the video and render groups ..."
usermod -aG video,render "$USERNAME"
fi

install_sudo
export HOME=$(get_user_home)
create_bashrc
fix_permissions
else
echo "No user specified or user is root, not creating a user ..."
create_bashrc
fi
109 changes: 109 additions & 0 deletions scripts/devinstall_llvm.sh
Original file line number Diff line number Diff line change
@@ -0,0 +1,109 @@
#! /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}}

LLVM_DIR=${WORKSPACE}/llvm-project
LLVM_REPO=https://github.com/llvm/llvm-project.git
LLVM_BUILD_PATH=$LLVM_DIR/build

pip_install() {
if command -v uv &>/dev/null; then
uv pip install "$@"
Comment thread
cmagina marked this conversation as resolved.
else
pip install "$@"
fi
}

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

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

if [ -f "${HOME}"/.bashrc.d/00-triton_llvm_gitref.sh ]; then
LLVM_GITREF=$(grep LLVM_GITREF "${HOME}/.bashrc.d/00-triton_llvm_gitref.sh" | cut -d'=' -f 2)
fi

if [ -n "${LLVM_GITREF:-}" ]; then
Comment thread
cmagina marked this conversation as resolved.
echo "Checking out LLVM_GITREF $LLVM_GITREF ..."
git checkout "$LLVM_GITREF"
fi
popd 1>/dev/null
fi
else
echo "LLVM repo already present, not cloning ..."
fi

if [ "${CUSTOM_LLVM:-false}" = "true" ] && [ -d "/llvm-install" ]; then
echo "WARNING: A CUSTOM_LLVM has been installed to /llvm-install"
echo "${HOME}/.bashrc.d/00-custom_llvm.sh sets the LLVM environment variables to use it"
echo "Override with a new file (i.e. ${HOME}/.bashrc.d/99-llvm-project.sh) or update that one to point to this directory"
else
echo "Adding LLVM_BUILD_PATH to ${HOME}/.bashrc.d/00-llvm_project.sh ..."
echo "export LLVM_BUILD_PATH=$LLVM_BUILD_PATH" >>"${HOME}/.bashrc.d/00-llvm_project.sh"
fi

}

install_build_deps() {
pushd "$LLVM_DIR" 1>/dev/null || exit 1
if [ -f mlir/python/requirements.txt ]; then
echo "Installing LLVM build dependencies ..."
pip_install -r mlir/python/requirements.txt
Comment thread
cmagina marked this conversation as resolved.
fi
Comment thread
cmagina marked this conversation as resolved.
popd 1>/dev/null
}

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

##
## Main
##

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

COMMAND=${1,,}

case $COMMAND in
source)
setup_src
install_build_deps
;;
*)
usage
exit 1
;;
esac
8 changes: 3 additions & 5 deletions scripts/devinstall_software.sh
Original file line number Diff line number Diff line change
Expand Up @@ -88,8 +88,7 @@ if ((${USE_CCACHE:-0} != 0)); then
$SUDO dnf -y install ccache

echo "Adding CCACHE environment variables to ${HOME}/.bashrc ..."
tee -a "${HOME}/.bashrc" <<EOF

tee "${HOME}/.bashrc.d/00-ccache.sh" <<EOF
# Enable CCACHE use
export CCACHE_DIR=${WORKSPACE}/.cache/ccache
export CCACHE_NOHASHDIR="true"
Expand All @@ -107,9 +106,8 @@ if [ -n "${CUDA_VERSION:-}" ]; then
"cuda-libraries-devel-$CUDA_VERSION" "cuda-nvml-devel-$CUDA_VERSION" \
"cuda-cupti-$CUDA_VERSION" cudnn cudss "libcusparse-devel-$CUDA_VERSION"

echo "Adding CUDA paths to the user environment ..."
tee -a "${HOME}/.bashrc" <<EOF

echo "Adding CUDA paths to ${HOME}/.bashrc.d/00-cuda_path.sh ..."
tee "${HOME}/.bashrc.d/00-cuda_path.sh" <<EOF
export CUDA_HOME=/usr/local/cuda
export PATH=/usr/local/cuda/bin:\$PATH
EOF
Expand Down
Loading
Loading