A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 19:45 UTC
[CI] Add TPU v1 test (#14834)
Signed-off-by: Richard Liu <ricliu@google.com>
Richard Liu committed
Mar 14, 2025 at 14:13 UTC
40677783aa1f59019424ff1828c54b696e4cfc3a
1 file changed
+27
.buildkite/run-tpu-v1-test.sh
new
+27
@@ -0,0 +1,27 @@
1
+#!/bin/bash
2
+
3
+set -e
4
+
5
+# Build the docker image.
6
+docker build -f Dockerfile.tpu -t vllm-tpu .
7
+
8
+# Set up cleanup.
9
+remove_docker_container() { docker rm -f tpu-test || true; }
10
+trap remove_docker_container EXIT
11
+# Remove the container that might not be cleaned up in the previous run.
12
+remove_docker_container
13
+
14
+# For HF_TOKEN.
15
+source /etc/environment
16
+# Run a simple end-to-end example.
17
+docker run --privileged --net host --shm-size=16G -it \
18
+ -e "HF_TOKEN=$HF_TOKEN" -e "VLLM_USE_V1=1" --name tpu-test \
19
+ vllm-tpu /bin/bash -c "python3 -m pip install git+https://github.com/thuml/depyf.git \
20
+ && python3 -m pip install pytest \
21
+ && python3 -m pip install lm_eval[api]==0.4.4 \
22
+ && pytest -v -s /workspace/vllm/tests/tpu/test_custom_dispatcher.py \
23
+ && pytest -v -s /workspace/vllm/tests/v1/tpu/test_basic.py \
24
+ && pytest -v -s /workspace/vllm/tests/entrypoints/llm/test_accuracy.py::test_lm_eval_accuracy_v1_engine \
25
+ && python3 /workspace/vllm/tests/tpu/test_compilation.py \
26
+ && python3 /workspace/vllm/tests/tpu/test_quantization_accuracy.py \
27
+ && python3 /workspace/vllm/examples/offline_inference/tpu.py"