| 1 | #!/bin/bash |
| 2 | |
| 3 | if [ ! -f "$1" ]; then |
| 4 | echo "Error: The env file '$1' does not exist." |
| 5 | exit 1 # Exit the script with a non-zero status to indicate an error |
| 6 | fi |
| 7 | |
| 8 | ENV_FILE=$1 |
| 9 | |
| 10 | # For testing on local vm, use `set -a` to export all variables |
| 11 | source /etc/environment |
| 12 | # shellcheck source=/dev/null |
| 13 | source "$ENV_FILE" |
| 14 | |
| 15 | remove_docker_container() { |
| 16 | docker rm -f "$CONTAINER_NAME" || true; |
| 17 | } |
| 18 | |
| 19 | trap remove_docker_container EXIT |
| 20 | |
| 21 | # Remove the container that might not be cleaned up in the previous run. |
| 22 | remove_docker_container |
| 23 | |
| 24 | LOG_ROOT=$(mktemp -d) |
| 25 | # If mktemp fails, set -e will cause the script to exit. |
| 26 | echo "Results will be stored in: $LOG_ROOT" |
| 27 | |
| 28 | if [ -z "$HF_TOKEN" ]; then |
| 29 | echo "Error: HF_TOKEN is not set or is empty." |
| 30 | exit 1 |
| 31 | fi |
| 32 | |
| 33 | # Make sure mounted disk or dir exists |
| 34 | if [ ! -d "$DOWNLOAD_DIR" ]; then |
| 35 | echo "Error: Folder $DOWNLOAD_DIR does not exist. This is useually a mounted drive. If no mounted drive, just create a folder." |
| 36 | exit 1 |
| 37 | fi |
| 38 | |
| 39 | echo "Run model $MODEL" |
| 40 | echo |
| 41 | |
| 42 | echo "starting docker...$CONTAINER_NAME" |
| 43 | echo |
| 44 | docker run \ |
| 45 | -v "$DOWNLOAD_DIR":"$DOWNLOAD_DIR" \ |
| 46 | --env-file "$ENV_FILE" \ |
| 47 | -e HF_TOKEN="$HF_TOKEN" \ |
| 48 | -e TARGET_COMMIT="$BUILDKITE_COMMIT" \ |
| 49 | -e MODEL="$MODEL" \ |
| 50 | -e WORKSPACE=/workspace \ |
| 51 | --name "$CONTAINER_NAME" \ |
| 52 | -d \ |
| 53 | --privileged \ |
| 54 | --network host \ |
| 55 | -v /dev/shm:/dev/shm \ |
| 56 | vllm/vllm-tpu-bm tail -f /dev/null |
| 57 | |
| 58 | echo "run script..." |
| 59 | echo |
| 60 | docker exec "$CONTAINER_NAME" /bin/bash -c ".buildkite/scripts/tpu/run_bm.sh" |
| 61 | |
| 62 | echo "copy result back..." |
| 63 | VLLM_LOG="$LOG_ROOT/$TEST_NAME"_vllm_log.txt |
| 64 | BM_LOG="$LOG_ROOT/$TEST_NAME"_bm_log.txt |
| 65 | docker cp "$CONTAINER_NAME:/workspace/vllm_log.txt" "$VLLM_LOG" |
| 66 | docker cp "$CONTAINER_NAME:/workspace/bm_log.txt" "$BM_LOG" |
| 67 | |
| 68 | throughput=$(grep "Request throughput (req/s):" "$BM_LOG" | sed 's/[^0-9.]//g') |
| 69 | echo "throughput for $TEST_NAME at $BUILDKITE_COMMIT: $throughput" |
| 70 | |
| 71 | if [ "$BUILDKITE" = "true" ]; then |
| 72 | echo "Running inside Buildkite" |
| 73 | buildkite-agent artifact upload "$VLLM_LOG" |
| 74 | buildkite-agent artifact upload "$BM_LOG" |
| 75 | else |
| 76 | echo "Not running inside Buildkite" |
| 77 | fi |
| 78 | |
| 79 | # |
| 80 | # compare the throughput with EXPECTED_THROUGHPUT |
| 81 | # and assert meeting the expectation |
| 82 | # |
| 83 | if [[ -z "$throughput" || ! "$throughput" =~ ^[0-9]+([.][0-9]+)?$ ]]; then |
| 84 | echo "Failed to get the throughput" |
| 85 | exit 1 |
| 86 | fi |
| 87 | |
| 88 | if (( $(echo "$throughput < $EXPECTED_THROUGHPUT" | bc -l) )); then |
| 89 | echo "Error: throughput($throughput) is less than expected($EXPECTED_THROUGHPUT)" |
| 90 | exit 1 |
| 91 | fi |