diff --git a/.ci/install-triton-pip.sh b/.ci/install-triton-pip.sh new file mode 100644 index 0000000..4e98f6a --- /dev/null +++ b/.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 diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index e216ee2..9466e5c 100644 --- a/.github/workflows/test.yml +++ b/.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 diff --git a/README.md b/README.md index 2ab1165..da51c50 100644 --- a/README.md +++ b/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 diff --git a/tritonparse/structured_logging.py b/tritonparse/structured_logging.py index f234a3b..64d4630 100644 --- a/tritonparse/structured_logging.py +++ b/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: