From 137a7cf427eb6e1913db0e31b42b315dadccf56c Mon Sep 17 00:00:00 2001 From: Yong Wu Date: Wed, 14 Jan 2026 09:51:52 -0500 Subject: [PATCH] [CI] Add support for testing dependency commits before release --- ci/setup_python.env | 17 +++++++++++++ scripts/setup_test_env.sh | 25 +++++++++++++++++++ scripts/task_jit_run_tests_part1.sh | 3 +++ scripts/task_jit_run_tests_part2.sh | 3 +++ scripts/task_jit_run_tests_part3.sh | 3 +++ scripts/task_jit_run_tests_part4.sh | 3 +++ scripts/task_jit_run_tests_part5.sh | 3 +++ scripts/task_test_blackwell_kernels.sh | 3 +++ ...ask_test_jit_cache_package_build_import.sh | 3 +++ scripts/task_test_multi_node_comm_kernels.sh | 3 +++ scripts/task_test_nightly_build.sh | 3 +++ scripts/task_test_single_node_comm_kernels.sh | 3 +++ 12 files changed, 72 insertions(+) create mode 100644 ci/setup_python.env create mode 100755 scripts/setup_test_env.sh diff --git a/ci/setup_python.env b/ci/setup_python.env new file mode 100644 index 0000000000..ffd9c49ac2 --- /dev/null +++ b/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= diff --git a/scripts/setup_test_env.sh b/scripts/setup_test_env.sh new file mode 100755 index 0000000000..83480cbd6a --- /dev/null +++ b/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 diff --git a/scripts/task_jit_run_tests_part1.sh b/scripts/task_jit_run_tests_part1.sh index 60deeaac7b..824ed91b96 100755 --- a/scripts/task_jit_run_tests_part1.sh +++ b/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 diff --git a/scripts/task_jit_run_tests_part2.sh b/scripts/task_jit_run_tests_part2.sh index c85d9d7e28..919810e31e 100755 --- a/scripts/task_jit_run_tests_part2.sh +++ b/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 diff --git a/scripts/task_jit_run_tests_part3.sh b/scripts/task_jit_run_tests_part3.sh index 0f3df35b2c..03b7099329 100755 --- a/scripts/task_jit_run_tests_part3.sh +++ b/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 diff --git a/scripts/task_jit_run_tests_part4.sh b/scripts/task_jit_run_tests_part4.sh index 596decc174..5407a4d3c7 100755 --- a/scripts/task_jit_run_tests_part4.sh +++ b/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 diff --git a/scripts/task_jit_run_tests_part5.sh b/scripts/task_jit_run_tests_part5.sh index a4ada8334a..5848c29e01 100755 --- a/scripts/task_jit_run_tests_part5.sh +++ b/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 diff --git a/scripts/task_test_blackwell_kernels.sh b/scripts/task_test_blackwell_kernels.sh index c8a2e43691..130c5b0adc 100644 --- a/scripts/task_test_blackwell_kernels.sh +++ b/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} diff --git a/scripts/task_test_jit_cache_package_build_import.sh b/scripts/task_test_jit_cache_package_build_import.sh index 0627d7b82d..1a1dfa56f4 100755 --- a/scripts/task_test_jit_cache_package_build_import.sh +++ b/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 "========================================" diff --git a/scripts/task_test_multi_node_comm_kernels.sh b/scripts/task_test_multi_node_comm_kernels.sh index 0a7cad47ad..9eae39d705 100644 --- a/scripts/task_test_multi_node_comm_kernels.sh +++ b/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 diff --git a/scripts/task_test_nightly_build.sh b/scripts/task_test_nightly_build.sh index ad7773b5ab..01fd31443b 100755 --- a/scripts/task_test_nightly_build.sh +++ b/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 diff --git a/scripts/task_test_single_node_comm_kernels.sh b/scripts/task_test_single_node_comm_kernels.sh index 2bef87720b..fec4c0df3c 100644 --- a/scripts/task_test_single_node_comm_kernels.sh +++ b/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