Skip to content

Commit 60db4e0

Browse files
FindHaoFindHao
committed
Refactor Triton installation script to improve cache handling and logging; update GitHub Actions workflow to cache Triton source and build more effectively
1 parent 5e1c78d commit 60db4e0

File tree

2 files changed

+43
-46
lines changed

2 files changed

+43
-46
lines changed

‎.ci/install-triton.sh

Lines changed: 39 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -26,28 +26,25 @@ conda activate "$CONDA_ENV"
2626
# Create cache directory
2727
mkdir -p "$TRITON_CACHE_DIR"
2828

29-
# Check if Triton is already installed and working
30-
if python -c "import triton; print(f'Triton version: {triton.__version__}')" 2>/dev/null; then
31-
echo "Triton is already installed, checking commit compatibility..."
32-
33-
# Check if the cached commit matches the target commit
34-
if [ -f "$TRITON_CACHE_DIR/commit" ]; then
35-
CACHED_COMMIT=$(cat "$TRITON_CACHE_DIR/commit")
36-
if [ "$CACHED_COMMIT" = "$TRITON_COMMIT" ] && [ "$TRITON_COMMIT" != "main" ]; then
37-
echo "Triton is already installed with correct commit ($CACHED_COMMIT), skipping installation"
38-
exit 0
39-
elif [ "$TRITON_COMMIT" = "main" ]; then
40-
echo "Target is 'main' branch (API fallback), will reinstall to get latest"
41-
echo "Cached commit: $CACHED_COMMIT"
42-
else
43-
echo "Triton installed but commit mismatch: cached=$CACHED_COMMIT, target=$TRITON_COMMIT"
44-
echo "Will reinstall Triton..."
45-
fi
29+
# Check if we have cached source with correct commit
30+
if [ -f "$TRITON_CACHE_DIR/commit" ] && [ -d "$TRITON_SOURCE_DIR" ]; then
31+
CACHED_COMMIT=$(cat "$TRITON_CACHE_DIR/commit")
32+
if [ "$CACHED_COMMIT" = "$TRITON_COMMIT" ] && [ "$TRITON_COMMIT" != "main" ]; then
33+
echo "Found cached Triton source with correct commit ($CACHED_COMMIT)"
34+
echo "Will use cached source and re-install to new conda environment"
35+
USE_CACHED_SOURCE=true
36+
elif [ "$TRITON_COMMIT" = "main" ]; then
37+
echo "Target is 'main' branch (API fallback), will rebuild from scratch"
38+
echo "Cached commit: $CACHED_COMMIT"
39+
USE_CACHED_SOURCE=false
4640
else
47-
echo "Triton installed but no commit info found, will reinstall..."
41+
echo "Cached source commit mismatch: cached=$CACHED_COMMIT, target=$TRITON_COMMIT"
42+
echo "Will rebuild from scratch"
43+
USE_CACHED_SOURCE=false
4844
fi
4945
else
50-
echo "Triton not installed or not working, proceeding with installation..."
46+
echo "No cached source found or no commit info, will build from scratch"
47+
USE_CACHED_SOURCE=false
5148
fi
5249

5350
# Update libstdc++ to match system version
@@ -71,25 +68,28 @@ if [ -n "$TRITON_PKG_DIR" ] && [ -d "$TRITON_PKG_DIR" ]; then
7168
rm -rf "$TRITON_PKG_DIR"
7269
fi
7370

74-
# Clone or update Triton repository
75-
echo "Setting up Triton repository..."
76-
if [ -d "$TRITON_SOURCE_DIR" ]; then
71+
# Setup Triton repository based on cache status
72+
if [ "$USE_CACHED_SOURCE" = "true" ]; then
7773
echo "Using cached Triton source..."
7874
cd "$TRITON_SOURCE_DIR"
79-
# Reset to clean state and fetch latest
80-
git reset --hard HEAD
81-
git clean -fd
82-
git fetch origin
75+
ACTUAL_COMMIT=$(git rev-parse HEAD)
76+
echo "Using cached Triton commit: $ACTUAL_COMMIT"
8377
else
78+
echo "Setting up Triton repository from scratch..."
79+
if [ -d "$TRITON_SOURCE_DIR" ]; then
80+
echo "Removing existing source directory..."
81+
rm -rf "$TRITON_SOURCE_DIR"
82+
fi
83+
8484
echo "Cloning Triton repository..."
8585
git clone https://github.com/triton-lang/triton.git "$TRITON_SOURCE_DIR"
8686
cd "$TRITON_SOURCE_DIR"
87-
fi
8887

89-
# Checkout specific commit for reproducibility
90-
git checkout "$TRITON_COMMIT"
91-
ACTUAL_COMMIT=$(git rev-parse HEAD)
92-
echo "Using Triton commit: $ACTUAL_COMMIT"
88+
# Checkout specific commit for reproducibility
89+
git checkout "$TRITON_COMMIT"
90+
ACTUAL_COMMIT=$(git rev-parse HEAD)
91+
echo "Using Triton commit: $ACTUAL_COMMIT"
92+
fi
9393

9494
# Install build dependencies
9595
echo "Installing build dependencies..."
@@ -107,7 +107,13 @@ echo "Using CC: $CC"
107107
echo "Using CXX: $CXX"
108108

109109
# Install Triton in editable mode with clang
110-
echo "Installing Triton in editable mode with clang..."
110+
if [ "$USE_CACHED_SOURCE" = "true" ]; then
111+
echo "Installing cached Triton to new conda environment..."
112+
echo "This should be fast since build artifacts are cached"
113+
else
114+
echo "Compiling and installing Triton from scratch..."
115+
echo "This will take 30-50 minutes for compilation"
116+
fi
111117
pip install -e .
112118

113119
# Verify Triton installation
@@ -124,6 +130,6 @@ python -c "import triton; print(f'Triton version: {triton.__version__}')" || {
124130
python -c "import triton; print(f'Triton path: {triton.__file__}')"
125131

126132
# Save commit info for cache validation
127-
echo "$ACTUAL_COMMIT" > "$TRITON_CACHE_DIR/commit"
133+
echo "$ACTUAL_COMMIT" >"$TRITON_CACHE_DIR/commit"
128134

129135
echo "Triton installation completed successfully!"

‎.github/workflows/test.yml

Lines changed: 4 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -42,14 +42,6 @@ jobs:
4242
restore-keys: |
4343
${{ runner.os }}-pip-3.11-
4444
45-
- name: Cache Conda
46-
uses: actions/cache@v3
47-
with:
48-
path: /opt/miniconda3
49-
key: ${{ runner.os }}-conda-${{ hashFiles('.ci/setup.sh') }}
50-
restore-keys: |
51-
${{ runner.os }}-conda-
52-
5345
- name: Get Triton latest commit
5446
id: triton-commit
5547
run: |
@@ -67,17 +59,16 @@ jobs:
6759
echo "Using fallback cache key with timestamp: main-fallback-$TIMESTAMP"
6860
fi
6961
70-
- name: Cache Triton build
62+
- name: Cache Triton source and build
7163
uses: actions/cache@v3
7264
with:
7365
path: |
7466
/tmp/triton
7567
/tmp/triton-cache
76-
/opt/miniconda3/envs/tritonparse/lib/python3.11/site-packages/triton*
77-
key: ${{ runner.os }}-triton-${{ hashFiles('.ci/install-triton.sh') }}-${{ steps.triton-commit.outputs.cache-key }}
68+
key: ${{ runner.os }}-triton-source-${{ hashFiles('.ci/install-triton.sh') }}-${{ steps.triton-commit.outputs.cache-key }}
7869
restore-keys: |
79-
${{ runner.os }}-triton-${{ hashFiles('.ci/install-triton.sh') }}-
80-
${{ runner.os }}-triton-
70+
${{ runner.os }}-triton-source-${{ hashFiles('.ci/install-triton.sh') }}-
71+
${{ runner.os }}-triton-source-
8172
8273
- name: Setup environment
8374
env:

0 commit comments

Comments
 (0)