[TPU] patch TPU wheel build script to resolve metadata issue #27279
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request correctly addresses a metadata issue in TPU wheel builds by properly using setuptools-scm for version overriding in setup.py. The accompanying changes in the build.sh script to patch hardcoded package names are also appropriate. My feedback focuses on improving the robustness of the patching mechanism within the build script to make it safer and more resilient to future code changes.
| sed -i \ | ||
| -e "s/importlib.metadata.version(\(['\"]\)vllm\1)/importlib.metadata.version(\1vllm-tpu\1)/" \ | ||
| -e "s/importlib.metadata.metadata(\(['\"]\)vllm\1)/importlib.metadata.metadata(\1vllm-tpu\1)/" \ | ||
| -e "s/version(\(['\"]\)vllm\1)/version(\1vllm-tpu\1)/" \ | ||
| "${CHANGE_FILE_LIST[@]}" |
There was a problem hiding this comment.
For improved robustness and safety, I recommend two changes to this sed command:
-
Global Replacement: Add the
gflag to each substitution. This ensures that all occurrences of the pattern on a single line are replaced, not just the first. While the current files may only have one occurrence per line, this makes the script more resilient to future changes. -
Safer Patching: Use
sed -i.bakto create backup files. This is more robust than relying on a secondsedcommand in thecleanupfunction to revert changes, especially if the script is terminated abruptly (e.g., withkill -9). Thecleanupfunction should then be updated to restore from these.bakfiles, which is consistent with howpyproject.tomlis handled in this script.
| sed -i \ | |
| -e "s/importlib.metadata.version(\(['\"]\)vllm\1)/importlib.metadata.version(\1vllm-tpu\1)/" \ | |
| -e "s/importlib.metadata.metadata(\(['\"]\)vllm\1)/importlib.metadata.metadata(\1vllm-tpu\1)/" \ | |
| -e "s/version(\(['\"]\)vllm\1)/version(\1vllm-tpu\1)/" \ | |
| "${CHANGE_FILE_LIST[@]}" | |
| sed -i.bak \ | |
| -e "s/importlib.metadata.version(\(['\"]\)vllm\1)/importlib.metadata.version(\1vllm-tpu\1)/g" \ | |
| -e "s/importlib.metadata.metadata(\(['\"]\)vllm\1)/importlib.metadata.metadata(\1vllm-tpu\1)/g" \ | |
| -e "s/version(\(['\"]\)vllm\1)/version(\1vllm-tpu\1)/g" \ | |
| "${CHANGE_FILE_LIST[@]}" |
There was a problem hiding this comment.
Can you at least leave a comment for what these edits are doing? I don't naturally read sed :)
There was a problem hiding this comment.
Done, thanks for the reminder :)
| echo "Restoring vllm code..." | ||
| sed -i \ | ||
| -e "s/importlib.metadata.version(\(['\"]\)vllm-tpu\1)/importlib.metadata.version(\1vllm\1)/" \ | ||
| -e "s/importlib.metadata.metadata(\(['\"]\)vllm-tpu\1)/importlib.metadata.metadata(\1vllm\1)/" \ | ||
| -e "s/version(\(['\"]\)vllm-tpu\1)/version(\1vllm\1)/" \ | ||
| "${CHANGE_FILE_LIST[@]}" |
There was a problem hiding this comment.
To complement the use of sed -i.bak for patching, the cleanup logic should be updated to restore files from their backups. This is a more robust and safer approach than running another sed command to revert the changes, as it protects against incomplete reversions if the script is forcefully terminated.
| echo "Restoring vllm code..." | |
| sed -i \ | |
| -e "s/importlib.metadata.version(\(['\"]\)vllm-tpu\1)/importlib.metadata.version(\1vllm\1)/" \ | |
| -e "s/importlib.metadata.metadata(\(['\"]\)vllm-tpu\1)/importlib.metadata.metadata(\1vllm\1)/" \ | |
| -e "s/version(\(['\"]\)vllm-tpu\1)/version(\1vllm\1)/" \ | |
| "${CHANGE_FILE_LIST[@]}" | |
| echo "Restoring vllm code..." | |
| for file in "${CHANGE_FILE_LIST[@]}"; do | |
| if [ -f "${file}.bak" ]; then | |
| mv "${file}.bak" "${file}" | |
| fi | |
| done |
setup.py
Outdated
| # wheels (e.g. CPU, TPU) without modifying the source. | ||
| if env_version := os.getenv("VLLM_VERSION_OVERRIDE"): | ||
| return env_version | ||
| print(f"Overriding VLLM-TPU version with: {env_version}") |
There was a problem hiding this comment.
| print(f"Overriding VLLM-TPU version with: {env_version}") | |
| print(f"Overriding VLLM version with {env_version} from VLLM_VERSION_OVERRIDE") |
| sed -i \ | ||
| -e "s/importlib.metadata.version(\(['\"]\)vllm\1)/importlib.metadata.version(\1vllm-tpu\1)/" \ | ||
| -e "s/importlib.metadata.metadata(\(['\"]\)vllm\1)/importlib.metadata.metadata(\1vllm-tpu\1)/" \ | ||
| -e "s/version(\(['\"]\)vllm\1)/version(\1vllm-tpu\1)/" \ | ||
| "${CHANGE_FILE_LIST[@]}" |
There was a problem hiding this comment.
Can you at least leave a comment for what these edits are doing? I don't naturally read sed :)
Signed-off-by: Johnny Yang <johnnyyang@google.com>
mgoin
added
tpu
…oject#27279) Signed-off-by: Johnny Yang <johnnyyang@google.com> Signed-off-by: Elizabeth Thomas <email2eliza@gmail.com>
…oject#27279) Signed-off-by: Johnny Yang <johnnyyang@google.com> Signed-off-by: George D. Torres <gdavtor@gmail.com>
…oject#27279) Signed-off-by: Johnny Yang <johnnyyang@google.com> Signed-off-by: Bram Wasti <bwasti@meta.com>
3 participants
Purpose
Patch tpu wheel build script to resolve metadata issue described here
Test Plan
E2E test to verify built python wheel.
Test Result
Successful.
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.