[CI] Add support for testing dependency commits before release

ci/setup_python.env

@@ -0,0 +1,17 @@
+# CI Python Environment Overrides
+# ================================
+# This file allows overriding Python package versions for CI testing.
+# Useful for testing specific commits of dependencies before they are released.
+#
+# Usage:
+#   1. In a test PR, uncomment and set the variable(s) below
+#   2. CI will install the specified version(s) before running tests
+#   3. Do NOT merge these changes to main - close the PR after testing
+#
+# Example:
+#   TVM_FFI_REF=abc123def456  # Test a specific TVM-FFI commit
+#   TVM_FFI_REF=main          # Test TVM-FFI main branch
+#   TVM_FFI_REF=v0.2.0        # Test a specific tag
+
+# Uncomment to override TVM-FFI version:
+# TVM_FFI_REF=

scripts/setup_test_env.sh

@@ -0,0 +1,25 @@
+#!/bin/bash
+# Setup test environment with optional package overrides
+# This script should be sourced at the beginning of CI test scripts.
+#
+# It reads ci/setup_python.env and installs any overridden package versions.
+# This is useful for testing specific commits of dependencies (e.g., TVM-FFI)
+# before they are officially released.
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
+
+# Source the environment override file if it exists
+if [ -f "${REPO_ROOT}/ci/setup_python.env" ]; then
+  source "${REPO_ROOT}/ci/setup_python.env"
+fi
+
+# Override TVM-FFI if specified
+if [ -n "${TVM_FFI_REF:-}" ]; then
+  echo "========================================"
+  echo "Overriding TVM-FFI with ref: ${TVM_FFI_REF}"
+  echo "========================================"
+  pip install --force-reinstall "git+https://github.com/apache/tvm-ffi.git@${TVM_FFI_REF}"
+  echo "TVM-FFI override complete."
+  echo ""
+fi

scripts/task_jit_run_tests_part1.sh

@@ -6,6 +6,9 @@ set -x
 : ${CUDA_VISIBLE_DEVICES:=0}
 : ${SKIP_INSTALL:=0}
 
+# Source test environment setup (handles package overrides like TVM-FFI)
+source "$(dirname "${BASH_SOURCE[0]}")/setup_test_env.sh"
+
 if [ "$SKIP_INSTALL" = "0" ]; then
   pip install -e . -v
 fi

scripts/task_jit_run_tests_part2.sh

@@ -6,6 +6,9 @@ set -x
 : ${CUDA_VISIBLE_DEVICES:=0}
 : ${SKIP_INSTALL:=0}
 
+# Source test environment setup (handles package overrides like TVM-FFI)
+source "$(dirname "${BASH_SOURCE[0]}")/setup_test_env.sh"
+
 if [ "$SKIP_INSTALL" = "0" ]; then
   pip install -e . -v
 fi

scripts/task_jit_run_tests_part3.sh

@@ -6,6 +6,9 @@ set -x
 : ${CUDA_VISIBLE_DEVICES:=0}
 : ${SKIP_INSTALL:=0}
 
+# Source test environment setup (handles package overrides like TVM-FFI)
+source "$(dirname "${BASH_SOURCE[0]}")/setup_test_env.sh"
+
 if [ "$SKIP_INSTALL" = "0" ]; then
   pip install -e . -v
 fi

scripts/task_jit_run_tests_part4.sh

@@ -6,6 +6,9 @@ set -x
 : ${CUDA_VISIBLE_DEVICES:=0}
 : ${SKIP_INSTALL:=0}
 
+# Source test environment setup (handles package overrides like TVM-FFI)
+source "$(dirname "${BASH_SOURCE[0]}")/setup_test_env.sh"
+
 if [ "$SKIP_INSTALL" = "0" ]; then
   pip install -e . -v
 fi

scripts/task_jit_run_tests_part5.sh

@@ -6,6 +6,9 @@ set -x
 : ${CUDA_VISIBLE_DEVICES:=0}
 : ${SKIP_INSTALL:=0}
 
+# Source test environment setup (handles package overrides like TVM-FFI)
+source "$(dirname "${BASH_SOURCE[0]}")/setup_test_env.sh"
+
 if [ "$SKIP_INSTALL" = "0" ]; then
   pip install -e . -v
 fi

scripts/task_test_blackwell_kernels.sh

@@ -2,6 +2,9 @@
 
 set -eo pipefail
 
+# Source test environment setup (handles package overrides like TVM-FFI)
+source "$(dirname "${BASH_SOURCE[0]}")/setup_test_env.sh"
+
 : ${JUNIT_DIR:=$(realpath ./junit)}
 : ${MAX_JOBS:=$(nproc)}
 : ${CUDA_VISIBLE_DEVICES:=0}

scripts/task_test_jit_cache_package_build_import.sh

@@ -3,6 +3,9 @@
 set -eo pipefail
 set -x
 
+# Source test environment setup (handles package overrides like TVM-FFI)
+source "$(dirname "${BASH_SOURCE[0]}")/setup_test_env.sh"
+
 echo "========================================"
 echo "Starting flashinfer-jit-cache test script"
 echo "========================================"

scripts/task_test_multi_node_comm_kernels.sh

@@ -5,6 +5,9 @@ set -x
 : ${MAX_JOBS:=$(nproc)}
 : ${CUDA_VISIBLE_DEVICES:=0}
 
+# Source test environment setup (handles package overrides like TVM-FFI)
+source "$(dirname "${BASH_SOURCE[0]}")/setup_test_env.sh"
+
 # Clean Python bytecode cache to avoid stale imports (e.g., after module refactoring)
 # echo "Cleaning Python bytecode cache..."
 # find . -type d -name __pycache__ -exec rm -rf {} + 2>/dev/null || true

scripts/task_test_nightly_build.sh

@@ -3,6 +3,9 @@
 set -eo pipefail
 set -x
 
+# Source test environment setup (handles package overrides like TVM-FFI)
+source "$(dirname "${BASH_SOURCE[0]}")/setup_test_env.sh"
+
 # This script installs nightly build packages and runs tests
 # Expected dist directories to be in current directory or specified via env vars
 

scripts/task_test_single_node_comm_kernels.sh

@@ -5,6 +5,9 @@ set -x
 : ${MAX_JOBS:=$(nproc)}
 : ${CUDA_VISIBLE_DEVICES:=0}
 
+# Source test environment setup (handles package overrides like TVM-FFI)
+source "$(dirname "${BASH_SOURCE[0]}")/setup_test_env.sh"
+
 # Clean Python bytecode cache to avoid stale imports (e.g., after module refactoring)
 echo "Cleaning Python bytecode cache..."
 find . -type d -name __pycache__ -exec rm -rf {} + 2>/dev/null || true