[BE] Add Parallel CI for Pip-Installed Triton and Update Docs

.ci/install-triton-pip.sh

@@ -0,0 +1,51 @@
+#!/bin/bash
+
+# Install Triton from pip
+# This script installs the latest version of Triton from PyPI.
+
+set -e
+
+echo "🚀 Installing Triton from pip..."
+START_TIME=$(date +%s)
+
+# Function to show elapsed time
+show_elapsed() {
+    CURRENT_TIME=$(date +%s)
+    ELAPSED=$((CURRENT_TIME - START_TIME))
+    echo "⏱️ Elapsed time: ${ELAPSED}s"
+}
+
+# Pre-flight checks
+echo "🔍 Running pre-flight checks..."
+
+# Ensure we're in the conda environment
+if [ -z "$CONDA_ENV" ]; then
+    echo "ERROR: CONDA_ENV is not set"
+    exit 1
+fi
+
+# Activate conda environment
+source /opt/miniconda3/etc/profile.d/conda.sh
+conda activate "$CONDA_ENV"
+
+# Uninstall existing triton to avoid conflicts
+echo "Uninstalling existing Triton versions..."
+pip uninstall -y pytorch-triton triton || true
+
+# Install Triton from pip
+echo "Installing the latest Triton from PyPI..."
+pip install triton
+
+show_elapsed
+
+# Verify Triton installation
+echo "Verifying Triton installation..."
+if python -c "import triton; print(f'Triton version: {triton.__version__}')" 2>/dev/null; then
+    python -c "import triton; print(f'Triton path: {triton.__file__}')"
+    echo "✅ Triton installation verified successfully"
+    show_elapsed
+    echo "🎉 Triton installation completed successfully!"
+else
+    echo "❌ ERROR: Failed to import Triton"
+    exit 1
+fi

.github/workflows/test.yml

@@ -56,7 +56,7 @@ jobs:
               run: |
                   make lint-check || (echo "❌ Linting failed. Please run 'make format' to fix formatting issues, then commit the changes." && echo "📖 For detailed formatting guide, see: https://github.com/pytorch-labs/tritonparse/wiki/05.-Code-Formatting" && exit 1)
 
-    test:
+    test-from-source:
         runs-on: 4-core-ubuntu-gpu-t4
         timeout-minutes: 120
         steps:
@@ -156,3 +156,58 @@ jobs:
                   COVERAGE: ${{ github.event.inputs.coverage || 'false' }}
               run: |
                   bash .ci/run-tests.sh
+
+    test-from-pip:
+        runs-on: 4-core-ubuntu-gpu-t4
+        timeout-minutes: 120
+        steps:
+            - uses: actions/checkout@v4
+
+            - name: Set up Python 3.11
+              uses: actions/setup-python@v4
+              with:
+                  python-version: "3.11"
+
+            - name: Get daily cache timestamp
+              id: daily-cache
+              run: |
+                  # Calculate date (e.g., 2024-01-15) for daily cache expiration
+                  DATE_STAMP=$(date +"%Y-%m-%d")
+                  echo "date=$DATE_STAMP" >> $GITHUB_OUTPUT
+                  echo "Using daily cache stamp: $DATE_STAMP"
+
+            - name: Cache pip dependencies
+              uses: actions/cache@v3
+              with:
+                  path: ~/.cache/pip
+                  key: ${{ runner.os }}-pip-3.11-${{ steps.daily-cache.outputs.date }}
+                  restore-keys: |
+                      ${{ runner.os }}-pip-3.11-
+
+            - name: Setup environment
+              env:
+                  CONDA_ENV: tritonparse-pip
+                  PYTHON_VERSION: "3.11"
+                  CUDA_VERSION: "12.8"
+              run: |
+                  bash .ci/setup.sh
+
+            - name: Install Triton from pip
+              env:
+                  CONDA_ENV: tritonparse-pip
+              run: |
+                  bash .ci/install-triton-pip.sh
+
+            - name: Install project dependencies
+              env:
+                  CONDA_ENV: tritonparse-pip
+              run: |
+                  bash .ci/install-project.sh
+
+            - name: Run tests
+              env:
+                  CONDA_ENV: tritonparse-pip
+                  TEST_TYPE: ${{ github.event.inputs.test-type || 'all' }}
+                  COVERAGE: ${{ github.event.inputs.coverage || 'false' }}
+              run: |
+                  bash .ci/run-tests.sh

README.md

@@ -69,9 +69,12 @@ cd tritonparse
 pip install -e .
 ```
 
-**Prerequisites:** Python ≥ 3.10, Triton > 3.3.1 ([install from source](https://github.com/triton-lang/triton)), GPU required (NVIDIA/AMD)
+**Prerequisites:** Python ≥ 3.10, Triton ≥ 3.4.0, GPU required (NVIDIA/AMD)
 
-TritonParse relies on new features in Triton > 3.3.1. Please install Triton from source for now.
+TritonParse relies on new features in Triton. Please install the latest version of Triton:
+```bash
+pip install triton
+```
 
 ## 📚 Complete Documentation
 

tritonparse/structured_logging.py

@@ -351,7 +351,11 @@ def extract_python_source_info(trace_data: Dict[str, Any], source):
         return
 
     # Get the original Python source code for the kernel
-    if isinstance(fn := source.fn, JITFunction) and hasattr(fn, "starting_line_number"):
+    if (
+        isinstance(fn := source.fn, JITFunction)
+        and hasattr(fn, "starting_line_number")
+        and hasattr(fn, "raw_src")
+    ):
         start_line_number = fn.starting_line_number
         source_lines = fn.raw_src
     else: