| 1 | #!/bin/bash |
| 2 | |
| 3 | # This script build the GH200 docker image and run the offline inference inside the container. |
| 4 | # It serves a sanity check for compilation and basic model usage. |
| 5 | set -ex |
| 6 | |
| 7 | # Skip the new torch installation during build since we are using the specified version for arm64 in the Dockerfile |
| 8 | python3 tools/use_existing_torch.py |
| 9 | |
| 10 | # Try building the docker image |
| 11 | DOCKER_BUILDKIT=1 docker build . \ |
| 12 | --file docker/Dockerfile \ |
| 13 | --target vllm-openai \ |
| 14 | --platform "linux/arm64" \ |
| 15 | -t gh200-test \ |
| 16 | --build-arg max_jobs=66 \ |
| 17 | --build-arg nvcc_threads=2 \ |
| 18 | --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0-b8b5f17a7d9ccfc25bbc5cf17b3fcea12964a042 \ |
| 19 | --build-arg torch_cuda_arch_list="9.0+PTX" |
| 20 | |
| 21 | # Setup cleanup |
| 22 | remove_docker_container() { docker rm -f gh200-test || true; } |
| 23 | trap remove_docker_container EXIT |
| 24 | remove_docker_container |
| 25 | |
| 26 | # Run the image and test offline inference |
| 27 | docker run -e HF_TOKEN -e VLLM_WORKER_MULTIPROC_METHOD=spawn -v /root/.cache/huggingface:/root/.cache/huggingface --name gh200-test --gpus=all --entrypoint="" gh200-test bash -c ' |
| 28 | python3 examples/basic/offline_inference/generate.py --model meta-llama/Llama-3.2-1B |
| 29 | ' |