| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import os |
| 5 | import sys |
| 6 | import zipfile |
| 7 | |
| 8 | # Read the VLLM_MAX_SIZE_MB environment variable, defaulting to 500 MiB |
| 9 | # Note that we have 800 MiB quota, please use it wisely. |
| 10 | # See https://github.com/pypi/support/issues/6326 . |
| 11 | # Please also sync the value with the one in Dockerfile. |
| 12 | VLLM_MAX_SIZE_MB = int(os.environ.get("VLLM_MAX_SIZE_MB", 500)) |
| 13 | |
| 14 | |
| 15 | def print_top_10_largest_files(zip_file): |
| 16 | """Print the top 10 largest files in the given zip file.""" |
| 17 | with zipfile.ZipFile(zip_file, "r") as z: |
| 18 | file_sizes = [(f, z.getinfo(f).file_size) for f in z.namelist()] |
| 19 | file_sizes.sort(key=lambda x: x[1], reverse=True) |
| 20 | for f, size in file_sizes[:10]: |
| 21 | print(f"{f}: {size / (1024 * 1024):.2f} MBs uncompressed.") |
| 22 | |
| 23 | |
| 24 | def check_wheel_size(directory): |
| 25 | """Check the size of .whl files in the given directory.""" |
| 26 | for root, _, files in os.walk(directory): |
| 27 | for file_name in files: |
| 28 | if file_name.endswith(".whl"): |
| 29 | wheel_path = os.path.join(root, file_name) |
| 30 | wheel_size_mb = os.path.getsize(wheel_path) / (1024 * 1024) |
| 31 | if wheel_size_mb > VLLM_MAX_SIZE_MB: |
| 32 | print( |
| 33 | f"Not allowed: Wheel {wheel_path} is larger " |
| 34 | f"({wheel_size_mb:.2f} MB) than the limit " |
| 35 | f"({VLLM_MAX_SIZE_MB} MB)." |
| 36 | ) |
| 37 | print_top_10_largest_files(wheel_path) |
| 38 | return 1 |
| 39 | else: |
| 40 | print( |
| 41 | f"Wheel {wheel_path} is within the allowed size " |
| 42 | f"({wheel_size_mb:.2f} MB)." |
| 43 | ) |
| 44 | return 0 |
| 45 | |
| 46 | |
| 47 | if __name__ == "__main__": |
| 48 | if len(sys.argv) < 2: |
| 49 | print("Usage: python check-wheel-size.py <directory>") |
| 50 | sys.exit(1) |
| 51 | |
| 52 | directory = sys.argv[1] |
| 53 | sys.exit(check_wheel_size(directory)) |