refactor: modularize setup scripts for better maintainability (#115)

* refactor: modularize setup scripts for better maintainability

Replace monolithic entrypoint.sh and user.sh scripts with a modular
architecture that separates concerns:

- Add scripts/setup.sh as the main orchestrator
- Add scripts/setup_user.sh for user creation and configuration
- Add scripts/install_software.sh for runtime package installation
- Add scripts/entrypoint.sh as a simplified container entry point
- Add scripts/ldpretend.sh as a helper utility for cuda python libraries
- Add scripts/setup_triton.sh for Triton-specific setup tasks

Remove old monolithic scripts:
- dockerfiles/entrypoint.sh (replaced by modular scripts)
- dockerfiles/user.sh (replaced by setup_user.sh)

Update from pip to uv for virtual environment management.

This modularization provides a foundation for framework-specific
setup scripts and makes the codebase more maintainable and testable.

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

Author: cmagina

.github/workflows/amd-image.yml

@@ -7,14 +7,22 @@ on: # yamllint disable-line rule:truthy
     paths:
       - .github/workflows/amd-image.yml
       - dockerfiles/Dockerfile.triton-amd
-      - entrypoint.sh
-      - users.sh
+      - scripts/entrypoint.sh
+      - scripts/devinstall_software.sh
+      - scripts/ldpretend.sh
+      - scripts/devinstall_triton.sh
+      - scripts/devcreate_user.sh
+      - scripts/devsetup.sh
   pull_request:
     paths:
       - .github/workflows/amd-image.yml
       - dockerfiles/Dockerfile.triton-amd
-      - entrypoint.sh
-      - users.sh
+      - scripts/entrypoint.sh
+      - scripts/devinstall_software.sh
+      - scripts/ldpretend.sh
+      - scripts/devinstall_triton.sh
+      - scripts/devcreate_user.sh
+      - scripts/devsetup.sh
   schedule:
     - cron: '0 0 * * 0' # Runs every Sunday at midnight UTC
   workflow_dispatch:

.github/workflows/cpu-image.yml

@@ -7,14 +7,22 @@ on: # yamllint disable-line rule:truthy
     paths:
       - .github/workflows/cpu-image.yml
       - dockerfiles/Dockerfile.triton-cpu
-      - entrypoint.sh
-      - users.sh
+      - scripts/entrypoint.sh
+      - scripts/devinstall_software.sh
+      - scripts/ldpretend.sh
+      - scripts/devinstall_triton.sh
+      - scripts/devcreate_user.sh
+      - scripts/devsetup.sh
   pull_request:
     paths:
       - .github/workflows/cpu-image.yml
       - dockerfiles/Dockerfile.triton-cpu
-      - entrypoint.sh
-      - users.sh
+      - scripts/entrypoint.sh
+      - scripts/devinstall_software.sh
+      - scripts/ldpretend.sh
+      - scripts/devinstall_triton.sh
+      - scripts/devcreate_user.sh
+      - scripts/devsetup.sh
   schedule:
     - cron: '0 0 * * 0' # Runs every Sunday at midnight UTC
   workflow_dispatch:

.github/workflows/nvidia-image.yml

@@ -7,14 +7,22 @@ on: # yamllint disable-line rule:truthy
     paths:
       - .github/workflows/nvidia-image.yml
       - dockerfiles/Dockerfile.triton
-      - entrypoint.sh
-      - users.sh
+      - scripts/entrypoint.sh
+      - scripts/devinstall_software.sh
+      - scripts/ldpretend.sh
+      - scripts/devinstall_triton.sh
+      - scripts/devcreate_user.sh
+      - scripts/devsetup.sh
   pull_request:
     paths:
       - .github/workflows/nvidia-image.yml
       - dockerfiles/Dockerfile.triton
-      - entrypoint.sh
-      - users.sh
+      - scripts/entrypoint.sh
+      - scripts/devinstall_software.sh
+      - scripts/ldpretend.sh
+      - scripts/devinstall_triton.sh
+      - scripts/devcreate_user.sh
+      - scripts/devsetup.sh
   schedule:
     - cron: '0 0 * * 0' # Runs every Sunday at midnight UTC
   workflow_dispatch:
@@ -29,7 +37,6 @@ jobs:
       build-args: |
         USERNAME=triton
         CUSTOM_LLVM=false
-        INSTALL_CUDNN=true
       platforms: linux/amd64
     secrets:
       qt_username: ${{ secrets.qt_username }}

Makefile

@@ -40,10 +40,15 @@ TRITON_TAG ?= latest
 triton_path ?=$(source_dir)
 gitconfig_path ?="$(HOME)/.gitconfig"
 USERNAME ?=triton
-create_user ?=true
 # NOTE: Requires host build system to have a valid Red Hat Subscription if true
 INSTALL_NSIGHT ?=false
 user_path ?=
+INSTALL_TRITON ?= source # Options: release, source, skip
+INSTALL_JUPYTER ?= true
+USE_CCACHE ?= 0
+CUDA_VERSION ?= 12-8
+ROCM_VERSION ?= 6.2
+MAX_JOBS ?= $(shell nproc --all)
 
 # Modify image tag if CUSTOM_LLVM is enabled
 ifeq ($(CUSTOM_LLVM),true)
@@ -80,19 +85,20 @@ gosu-image: image-builder-check ## Build the Triton gosu image
 .PHONY: triton-image
 triton-image: image-builder-check gosu-image llvm-image ## Build the Triton devcontainer image
 	$(CTR_CMD) build -t $(IMAGE_REPO)/$(NVIDIA_IMAGE_NAME):$(TRITON_TAG) \
-		--build-arg CUSTOM_LLVM=$(CUSTOM_LLVM) -f dockerfiles/Dockerfile.triton .
+		--build-arg CUSTOM_LLVM=$(CUSTOM_LLVM) --build-arg BUILD_CUDA_VERSION=$(CUDA_VERSION) \
+		-f dockerfiles/Dockerfile.triton .
 
 .PHONY: triton-cpu-image
 triton-cpu-image: image-builder-check gosu-image ## Build the Triton CPU image
 	$(MAKE) llvm-image CUSTOM_LLVM=$(CUSTOM_LLVM) TRITON_CPU_BACKEND=1 LLVM_IMAGE_LABEL=cpu-latest
 	$(CTR_CMD) build -t $(IMAGE_REPO)/$(CPU_IMAGE_NAME):$(TRITON_TAG) \
-		--build-arg CUSTOM_LLVM=$(CUSTOM_LLVM) --build-arg TRITON_CPU_BACKEND=1 \
-		-f dockerfiles/Dockerfile.triton-cpu .
+		--build-arg CUSTOM_LLVM=$(CUSTOM_LLVM) -f dockerfiles/Dockerfile.triton-cpu .
 
 .PHONY: triton-amd-image
 triton-amd-image: image-builder-check gosu-image llvm-image ## Build the Triton AMD devcontainer image
 	$(CTR_CMD) build -t $(IMAGE_REPO)/$(AMD_IMAGE_NAME):$(TRITON_TAG) \
-		--build-arg CUSTOM_LLVM=$(CUSTOM_LLVM) -f dockerfiles/Dockerfile.triton-amd .
+		--build-arg CUSTOM_LLVM=$(CUSTOM_LLVM) --build-arg BUILD_ROCM_VERSION=$(ROCM_VERSION) \
+		-f dockerfiles/Dockerfile.triton-amd .
 
 ##@ Container Run
 # If you are on an OS that has the user in /etc/passwd then we can pass
@@ -110,7 +116,7 @@ define run_container
 	if [ -n "$(user_path)" ]; then \
 		volume_arg+=" -v $(user_path):/workspace/user$(SELINUXFLAG)"; \
 	fi; \
-	if [ "$(OS)" != "Darwin" ] && ! getent passwd $(USER) > /dev/null && [ "$(create_user)" = "false" ]; then \
+	if [ "$(OS)" != "Darwin" ] && ! getent passwd $(USER) > /dev/null; then \
 		volume_arg+=" -v /etc/passwd:/etc/passwd:ro -v /etc/group:/etc/group:ro"; \
 	fi; \
 	if [ -f "$(gitconfig_path)" ]; then \
@@ -120,7 +126,6 @@ define run_container
 	fi; \
 	if [ "$(strip $(1))" = "$(AMD_IMAGE_NAME)" ]; then \
 		gpu_args="--device=/dev/kfd --device=/dev/dri --security-opt seccomp=unconfined --group-add=video --cap-add=SYS_PTRACE --ipc=host --env HIP_VISIBLE_DEVICES=$(HIP_DEVICES)"; \
-		profiling_args=""; \
 	elif [ "$(strip $(1))" = "$(NVIDIA_IMAGE_NAME)" ]; then \
 		if command -v nvidia-ctk >/dev/null 2>&1 && nvidia-ctk cdi list | grep -q "nvidia.com/gpu=all"; then \
 			gpu_args="--device nvidia.com/gpu=all"; \
@@ -146,16 +151,12 @@ define run_container
 	else \
 		port_arg=""; \
 	fi; \
-	env_vars="-e USERNAME=$(USER) -e TORCH_VERSION=$(torch_version) -e CUSTOM_LLVM=$(CUSTOM_LLVM) -e DEMO_TOOLS=$(DEMO_TOOLS) -e NOTEBOOK_PORT=$(NOTEBOOK_PORT)"; \
-	if [ "$(create_user)" = "true" ]; then \
-		$(CTR_CMD) run -e CREATE_USER=$(create_user) $$env_vars $$port_arg \
-		-e USER_UID=`id -u $(USER)` -e USER_GID=`id -g $(USER)` $$gpu_args $$profiling_args $$keep_ns_arg \
-		-ti $$volume_arg $$gitconfig_arg $(IMAGE_REPO)/$(strip $(1)):$(TRITON_TAG) bash; \
-	elif [ "$(STRIPPED_CMD)" = "docker" ]; then \
-		$(CTR_CMD) run --user $(shell id -u):$(shell id -g) $$env_vars $$gpu_args $$profiling_args $$port_arg \
+	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)"; \
+	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; \
 	elif [ "$(STRIPPED_CMD)" = "podman" ]; then \
-		$(CTR_CMD) run --user $(USER) $$env_vars $$keep_ns_arg $$gpu_args $$profiling_args $$port_arg \
+		$(CTR_CMD) run $$env_vars $$keep_ns_arg $$gpu_args $$profiling_args $$port_arg \
 		-ti $$volume_arg $$gitconfig_arg $(IMAGE_REPO)/$(strip $(1)):$(TRITON_TAG) bash; \
 	fi
 endef

dockerfiles/Dockerfile.triton

@@ -1,4 +1,4 @@
-# Copyright (C) 2024-2025 Red Hat, Inc.
+# Copyright (C) 2024-2026 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.
@@ -13,7 +13,9 @@
 # limitations under the License.
 #
 # SPDX-License-Identifier: Apache-2.0
+
 ARG CUSTOM_LLVM=false
+ARG BUILD_CUDA_VERSION=12-8
 
 FROM registry.access.redhat.com/ubi9/python-312 AS base
 ARG CUSTOM_LLVM
@@ -26,7 +28,7 @@ USER 0
 RUN dnf -y update && \
     dnf -y install clang cmake lld ninja-build openblas openblas-devel \
     llvm llvm-libs libomp libomp-devel sudo && \
-    dnf clean all
+    dnf clean all && rm -rf /var/cache/dnf
 
 # Stage for llvm-local-true
 FROM base AS llvm-local-true
@@ -38,18 +40,18 @@ ENV TRITON_OFFLINE_BUILD=NO
 
 # Use intermediate stage selection
 FROM llvm-local-${CUSTOM_LLVM} AS final
+ARG BUILD_CUDA_VERSION
 
 # Create the /workspace directory and set permissions
-RUN mkdir -p /workspace && \
-    python -m venv /workspace && \
+WORKDIR /workspace
+RUN python -m pip install uv && \
+    uv venv --python 3.12 --seed /workspace && \
     echo "unset BASH_ENV PROMPT_COMMAND ENV" >> /workspace/bin/activate && \
     chmod -R 777 /workspace
 
-# Create a symlink to the installed version of CUDA
-RUN ln -sf /usr/local/cuda-${CUDA_VERSION/-/.} /usr/local/cuda
-
 ENV BASH_ENV=/workspace/bin/activate \
     ENV=/workspace/bin/activate \
+    CUDA_VERSION=${BUILD_CUDA_VERSION} \
     PROMPT_COMMAND=". /workspace/bin/activate" \
     PYTHON_VERSION=3.12 \
     PATH=/workspace/bin:$PATH \
@@ -58,15 +60,22 @@ ENV BASH_ENV=/workspace/bin/activate \
     PYTHONPATH=/workspace/lib/python$PYTHON_VERSION/site-packages \
     XDG_CACHE_HOME=/workspace \
     TRITON_CACHE_DIR=/workspace/.triton/cache \
-    TRITON_HOME=/workspace/
+    TRITON_HOME=/workspace/ \
+    WORKSPACE=/workspace \
+    UV_HTTP_TIMEOUT=60
 
-WORKDIR /workspace
-RUN echo 'export LD_LIBRARY_PATH=/usr/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
-RUN echo "export MAX_JOBS=$(nproc --all)" >> "${HOME}"/.bashrc
 COPY --from=quay.io/triton-dev-containers/gosu /usr/local/bin/gosu /usr/local/bin/gosu
-COPY dockerfiles/user.sh user.sh
-COPY dockerfiles/entrypoint.sh /entrypoint.sh
+
+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
+COPY scripts/devinstall_software.sh /workspace/bin/devinstall_software
+COPY scripts/ldpretend.sh /workspace/bin/ldpretend
+
+COPY scripts/entrypoint.sh /entrypoint.sh
+
 COPY hack/triton-gpu-check.py triton-gpu-check.py
 COPY examples/flash_attention_demo.ipynb flash_attention_demo.ipynb
+
 ENTRYPOINT ["/entrypoint.sh"]
 CMD ["tail", "-f", "/dev/null"]

dockerfiles/Dockerfile.triton-amd

@@ -1,4 +1,4 @@
-# Copyright (C) 2024-2025 Red Hat, Inc.
+# Copyright (C) 2024-2026 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.
@@ -15,17 +15,19 @@
 # SPDX-License-Identifier: Apache-2.0
 
 ARG CUSTOM_LLVM=false
-ARG ROCM_VERSION=6.2
+ARG BUILD_ROCM_VERSION=6.2
 
 # Base Stage with ROCm Installation
 FROM registry.access.redhat.com/ubi9/python-312 AS base
-ARG ROCM_VERSION=6.2
+ARG BUILD_ROCM_VERSION
+
 USER 0
 
 # Install the ROCm rpms
-RUN echo "[ROCm]" > /etc/yum.repos.d/rocm.repo && \
+RUN RHEL_VERSION=$(rpm --eval '%{rhel}') && \
+    echo "[ROCm]" > /etc/yum.repos.d/rocm.repo && \
     echo "name=ROCm" >> /etc/yum.repos.d/rocm.repo && \
-    echo "baseurl=https://repo.radeon.com/rocm/rhel9/$ROCM_VERSION/main" >> /etc/yum.repos.d/rocm.repo && \
+    echo "baseurl=https://repo.radeon.com/rocm/rhel${RHEL_VERSION}/${BUILD_ROCM_VERSION}/main" >> /etc/yum.repos.d/rocm.repo && \
     echo "enabled=1" >> /etc/yum.repos.d/rocm.repo && \
     echo "gpgcheck=0" >> /etc/yum.repos.d/rocm.repo && \
     dnf install -y llvm clang clang-libs lld && \
@@ -40,10 +42,10 @@ RUN echo "[ROCm]" > /etc/yum.repos.d/rocm.repo && \
 
 # Set environment variables for ROCm
 ENV LC_ALL=C.UTF-8 \
-LANG=C.UTF-8 \
-ROCM_PATH=/opt/rocm \
-LD_LIBRARY_PATH=/usr/lib64:/usr/lib:/opt/rocm/lib:/opt/rocm/llvm/lib \
-PATH=/opt/rocm/bin:/opt/rocm/llvm/bin:$PATH
+    LANG=C.UTF-8 \
+    ROCM_PATH=/opt/rocm \
+    LD_LIBRARY_PATH=/usr/lib64:/usr/lib:/opt/rocm/lib:/opt/rocm/llvm/lib \
+    PATH=/opt/rocm/bin:/opt/rocm/llvm/bin:$PATH
 
 # LLVM Integration Stages
 FROM base AS llvm-local-true
@@ -54,36 +56,45 @@ ENV TRITON_OFFLINE_BUILD=NO
 
 # Final Stage
 FROM llvm-local-${CUSTOM_LLVM} AS final
-ARG ROCM_VERSION=6.2
+ARG BUILD_ROCM_VERSION
 
 # Create the /workspace directory and set permissions
 RUN mkdir -p /workspace && \
-    python -m venv /workspace && \
+    python -m pip install uv && \
+    uv venv --python 3.12 --seed /workspace && \
     echo "unset BASH_ENV PROMPT_COMMAND ENV" >> /workspace/bin/activate && \
     chmod -R 777 /workspace
 
 ENV BASH_ENV=/workspace/bin/activate \
     ENV=/workspace/bin/activate \
     PROMPT_COMMAND=". /workspace/bin/activate" \
-    PYTHON_VERSION=3.11 \
+    PYTHON_VERSION=3.12 \
     PATH=/workspace/bin:$PATH \
     PYTHONUNBUFFERED=1 \
-    AMD=true \
-    ROCM_VERSION=$ROCM_VERSION \
+    ROCM_VERSION=$BUILD_ROCM_VERSION \
     PIP_PREFIX=/workspace \
-    PYTHONPATH=/workspace/lib/python$PYTHON_VERSION/site-packages \
+    PYTHONPATH=/workspace/lib/python3.12/site-packages \
     XDG_CACHE_HOME=/workspace \
     TRITON_CACHE_DIR=/workspace/.triton/cache \
-    TRITON_HOME=/workspace/
+    TRITON_HOME=/workspace/ \
+    WORKSPACE=/workspace \
+    UV_HTTP_TIMEOUT=60
 
 WORKDIR /workspace
-RUN echo 'export LD_LIBRARY_PATH=/usr/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
-RUN echo "export MAX_JOBS=$(nproc --all)" >> "${HOME}"/.bashrc
+
 COPY --from=quay.io/triton-dev-containers/gosu /usr/local/bin/gosu /usr/local/bin/gosu
-COPY dockerfiles/user.sh user.sh
-COPY dockerfiles/entrypoint.sh /entrypoint.sh
+
+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
+COPY scripts/devinstall_software.sh /workspace/bin/devinstall_software
+COPY scripts/ldpretend.sh /workspace/bin/ldpretend
+
+COPY scripts/entrypoint.sh /entrypoint.sh
+
 COPY hack/triton-gpu-check.py triton-gpu-check.py
-COPY examples/triton-vector-add-rocm.py triton-vector-add-rocm.py
 COPY examples/flash_attention_demo.ipynb flash_attention_demo.ipynb
+COPY examples/triton-vector-add-rocm.py triton-vector-add-rocm.py
+
 ENTRYPOINT ["/entrypoint.sh"]
 CMD ["tail", "-f", "/dev/null"]

dockerfiles/Dockerfile.triton-cpu

@@ -1,4 +1,4 @@
-# Copyright (C) 2024-2025 Red Hat, Inc.
+# Copyright (C) 2024-2026 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.
@@ -39,9 +39,11 @@ FROM llvm-local-${CUSTOM_LLVM} AS final
 
 # Create the /workspace directory and set permissions
 RUN mkdir -p /workspace && \
-    python -m venv /workspace && \
+    python -m pip install uv && \
+    uv venv --python 3.12 --seed /workspace && \
     echo "unset BASH_ENV PROMPT_COMMAND ENV" >> /workspace/bin/activate && \
     chmod -R 777 /workspace
+
 ENV BASH_ENV=/workspace/bin/activate \
     ENV=/workspace/bin/activate \
     PROMPT_COMMAND=". /workspace/bin/activate" \
@@ -53,14 +55,22 @@ ENV BASH_ENV=/workspace/bin/activate \
     PYTHONPATH=/workspace/lib/python$PYTHON_VERSION/site-packages \
     XDG_CACHE_HOME=/workspace \
     TRITON_CACHE_DIR=/workspace/.triton/cache \
-    TRITON_HOME=/workspace/
+    TRITON_HOME=/workspace/ \
+    WORKSPACE=/workspace \
+    UV_HTTP_TIMEOUT=60
 
 WORKDIR /workspace
-RUN echo 'export LD_LIBRARY_PATH=/usr/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
-RUN echo "export MAX_JOBS=$(nproc --all)" >> "${HOME}"/.bashrc
+
 COPY --from=quay.io/triton-dev-containers/gosu /usr/local/bin/gosu /usr/local/bin/gosu
-COPY dockerfiles/user.sh user.sh
-COPY dockerfiles/entrypoint.sh /entrypoint.sh
+
+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
+COPY scripts/devinstall_software.sh /workspace/bin/devinstall_software
+COPY scripts/ldpretend.sh /workspace/bin/ldpretend
+
+COPY scripts/entrypoint.sh /entrypoint.sh
 COPY hack/triton-gpu-check.py triton-gpu-check.py
+
 ENTRYPOINT ["/entrypoint.sh"]
 CMD ["tail", "-f", "/dev/null"]