| 1 | #!/bin/bash |
| 2 | |
| 3 | set -xu |
| 4 | |
| 5 | |
| 6 | remove_docker_container() { |
| 7 | docker rm -f tpu-test || true; |
| 8 | } |
| 9 | |
| 10 | trap remove_docker_container EXIT |
| 11 | |
| 12 | # Remove the container that might not be cleaned up in the previous run. |
| 13 | remove_docker_container |
| 14 | |
| 15 | # Build the docker image. |
| 16 | docker build -f docker/Dockerfile.tpu -t vllm-tpu . |
| 17 | |
| 18 | # Set up cleanup. |
| 19 | cleanup_docker() { |
| 20 | # Get Docker's root directory |
| 21 | docker_root=$(docker info -f '{{.DockerRootDir}}') |
| 22 | if [ -z "$docker_root" ]; then |
| 23 | echo "Failed to determine Docker root directory." |
| 24 | exit 1 |
| 25 | fi |
| 26 | echo "Docker root directory: $docker_root" |
| 27 | # Check disk usage of the filesystem where Docker's root directory is located |
| 28 | disk_usage=$(df "$docker_root" | tail -1 | awk '{print $5}' | sed 's/%//') |
| 29 | # Define the threshold |
| 30 | threshold=70 |
| 31 | if [ "$disk_usage" -gt "$threshold" ]; then |
| 32 | echo "Disk usage is above $threshold%. Cleaning up Docker images and volumes..." |
| 33 | # Remove dangling images (those that are not tagged and not used by any container) |
| 34 | docker image prune -f |
| 35 | # Remove unused volumes / force the system prune for old images as well. |
| 36 | docker volume prune -f && docker system prune --force --filter "until=72h" --all |
| 37 | echo "Docker images and volumes cleanup completed." |
| 38 | else |
| 39 | echo "Disk usage is below $threshold%. No cleanup needed." |
| 40 | fi |
| 41 | } |
| 42 | cleanup_docker |
| 43 | |
| 44 | # For HF_TOKEN. |
| 45 | source /etc/environment |
| 46 | |
| 47 | docker run --privileged --net host --shm-size=16G -it \ |
| 48 | -e "HF_TOKEN=$HF_TOKEN" --name tpu-test \ |
| 49 | vllm-tpu /bin/bash -c ' |
| 50 | set -e # Exit immediately if a command exits with a non-zero status. |
| 51 | set -u # Treat unset variables as an error. |
| 52 | |
| 53 | echo "--- Starting script inside Docker container ---" |
| 54 | |
| 55 | # Create results directory |
| 56 | RESULTS_DIR=$(mktemp -d) |
| 57 | # If mktemp fails, set -e will cause the script to exit. |
| 58 | echo "Results will be stored in: $RESULTS_DIR" |
| 59 | |
| 60 | # Install dependencies |
| 61 | echo "--- Installing Python dependencies ---" |
| 62 | python3 -m pip install --progress-bar off git+https://github.com/thuml/depyf.git \ |
| 63 | && python3 -m pip install --progress-bar off pytest pytest-asyncio tpu-info \ |
| 64 | && python3 -m pip install --progress-bar off "lm-eval[api]>=0.4.12" \ |
| 65 | && python3 -m pip install --progress-bar off hf-transfer tblib==3.1.0 |
| 66 | echo "--- Python dependencies installed ---" |
| 67 | |
| 68 | export VLLM_XLA_CHECK_RECOMPILATION=1 |
| 69 | export VLLM_XLA_CACHE_PATH= |
| 70 | |
| 71 | echo "--- Hardware Information ---" |
| 72 | # tpu-info |
| 73 | echo "--- Starting Tests ---" |
| 74 | set +e |
| 75 | overall_script_exit_code=0 |
| 76 | |
| 77 | # --- Test Definitions --- |
| 78 | # If a test fails, this function will print logs and will not cause the main script to exit. |
| 79 | run_test() { |
| 80 | local test_num=$1 |
| 81 | local test_name=$2 |
| 82 | local test_command=$3 |
| 83 | local log_file="$RESULTS_DIR/test_${test_num}.log" |
| 84 | local actual_exit_code |
| 85 | |
| 86 | echo "--- TEST_$test_num: Running $test_name ---" |
| 87 | |
| 88 | # Execute the test command. |
| 89 | eval "$test_command" > >(tee -a "$log_file") 2> >(tee -a "$log_file" >&2) |
| 90 | actual_exit_code=$? |
| 91 | |
| 92 | echo "TEST_${test_num}_COMMAND_EXIT_CODE: $actual_exit_code" # This goes to main log |
| 93 | echo "TEST_${test_num}_COMMAND_EXIT_CODE: $actual_exit_code" >> "$log_file" # Also to per-test log |
| 94 | |
| 95 | if [ "$actual_exit_code" -ne 0 ]; then |
| 96 | echo "TEST_$test_num ($test_name) FAILED with exit code $actual_exit_code." >&2 |
| 97 | echo "--- Log for failed TEST_$test_num ($test_name) ---" >&2 |
| 98 | if [ -f "$log_file" ]; then |
| 99 | cat "$log_file" >&2 |
| 100 | else |
| 101 | echo "Log file $log_file not found for TEST_$test_num ($test_name)." >&2 |
| 102 | fi |
| 103 | echo "--- End of log for TEST_$test_num ($test_name) ---" >&2 |
| 104 | return "$actual_exit_code" # Return the failure code |
| 105 | else |
| 106 | echo "TEST_$test_num ($test_name) PASSED." |
| 107 | return 0 # Return success |
| 108 | fi |
| 109 | } |
| 110 | |
| 111 | # Helper function to call run_test and update the overall script exit code |
| 112 | run_and_track_test() { |
| 113 | local test_num_arg="$1" |
| 114 | local test_name_arg="$2" |
| 115 | local test_command_arg="$3" |
| 116 | |
| 117 | # Run the test |
| 118 | run_test "$test_num_arg" "$test_name_arg" "$test_command_arg" |
| 119 | local test_specific_exit_code=$? |
| 120 | |
| 121 | # If the test failed, set the overall script exit code to 1 |
| 122 | if [ "$test_specific_exit_code" -ne 0 ]; then |
| 123 | # No need for extra echo here, run_test already logged the failure. |
| 124 | overall_script_exit_code=1 |
| 125 | fi |
| 126 | } |
| 127 | |
| 128 | # --- Actual Test Execution --- |
| 129 | run_and_track_test 1 "test_struct_output_generate.py" \ |
| 130 | "python3 -m pytest -s -v /workspace/vllm/tests/entrypoints/llm/test_struct_output_generate.py -k \"not test_structured_output_with_reasoning_matrices\"" |
| 131 | run_and_track_test 2 "test_moe_pallas.py" \ |
| 132 | "python3 -m pytest -s -v /workspace/vllm/tests/tpu/test_moe_pallas.py" |
| 133 | run_and_track_test 3 "test_lora.py" \ |
| 134 | "VLLM_XLA_CHECK_RECOMPILATION=0 python3 -m pytest -s -v /workspace/vllm/tests/tpu/lora/test_lora.py" |
| 135 | run_and_track_test 4 "test_tpu_qkv_linear.py" \ |
| 136 | "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_tpu_qkv_linear.py" |
| 137 | run_and_track_test 5 "test_spmd_model_weight_loading.py" \ |
| 138 | "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_spmd_model_weight_loading.py" |
| 139 | run_and_track_test 6 "test_kv_cache_update_kernel.py" \ |
| 140 | "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_kv_cache_update_kernel.py" |
| 141 | run_and_track_test 7 "test_tpu_int8.py" \ |
| 142 | "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_tpu_int8.py" |
| 143 | |
| 144 | # After all tests have been attempted, exit with the overall status. |
| 145 | if [ "$overall_script_exit_code" -ne 0 ]; then |
| 146 | echo "--- One or more tests FAILED. Overall script exiting with failure code 1. ---" |
| 147 | else |
| 148 | echo "--- All tests have completed and PASSED. Overall script exiting with success code 0. ---" |
| 149 | fi |
| 150 | exit "$overall_script_exit_code" |
| 151 | ' # IMPORTANT: This is the closing single quote for the bash -c "..." command. Ensure it is present and correct. |
| 152 | |
| 153 | # Capture the exit code of the docker run command |
| 154 | DOCKER_RUN_EXIT_CODE=$? |
| 155 | |
| 156 | # The trap will run for cleanup. |
| 157 | # Exit the main script with the Docker run command's exit code. |
| 158 | if [ "$DOCKER_RUN_EXIT_CODE" -ne 0 ]; then |
| 159 | echo "Docker run command failed with exit code $DOCKER_RUN_EXIT_CODE." |
| 160 | exit "$DOCKER_RUN_EXIT_CODE" |
| 161 | else |
| 162 | echo "Docker run command completed successfully." |
| 163 | exit 0 |
| 164 | fi |
| 165 | # TODO: This test fails because it uses RANDOM_SEED sampling |
| 166 | # pytest -v -s /workspace/vllm/tests/tpu/test_custom_dispatcher.py \ |