1#!/usr/bin/env python3 2# Copyright (c) Meta Platforms, Inc. and affiliates. 3# All rights reserved. 4# 5# This source code is licensed under the BSD-style license found in the 6# LICENSE file in the root directory of this source tree. 7 8# pyre-strict 9 10""" 11Compare binary file sizes. Used by the Skycastle workflow to ensure no change adds excessive size to executorch. 12 13Usage: file_size_compare.py [-h] --compare-file FILE [--base-file FILE] [-s, --max-size SIZE] [-e, --error-size SIZE] [-w, --warning-size SIZE] 14 15Exit Codes: 16 0 - OK 17 1 - Comparison yielded a warning 18 2 - Comparison yielded an error 19 3 - Script errored while executing 20""" 21 22import argparse 23import os 24import sys 25from pathlib import Path 26 27# Exit codes. 28EXIT_OK = 0 29EXIT_WARNING = 1 30EXIT_ERROR = 2 31EXIT_SCRIPT_ERROR = 3 32 33# TTY ANSI color codes. 34TTY_GREEN = "\033[0;32m" 35TTY_RED = "\033[0;31m" 36TTY_RESET = "\033[0m" 37 38# Error message printed if size is exceeded. 39SIZE_ERROR_MESSAGE = """This diff is increasing the binary size of ExecuTorch (the PyTorch Edge model executor) by a large amount. 40ExecuTorch has strict size requirements due to its embedded use case. Please follow these steps: 411. Check the output of the two steps (Build ... with the base commit/diff version) and compare their executable section sizes. 422. Contact a member of #pytorch_edge_portability so we can better help you. 43""" 44 45 46def create_file_path(file_name: str) -> Path: 47 """Create Path object from file name string.""" 48 file_path = Path(file_name) 49 if not file_path.is_file(): 50 print(f"{file_path} is not a valid file path") 51 sys.exit(EXIT_SCRIPT_ERROR) 52 return file_path 53 54 55def get_file_size(file_path: Path) -> int: 56 """Get the size of a file on disk.""" 57 return os.path.getsize(file_path) 58 59 60def print_ansi(ansi_code: str) -> None: 61 """Print an ANSI escape code.""" 62 if sys.stdout.isatty(): 63 print(ansi_code, end="") 64 65 66def print_size_diff(compare_file: str, base_file: str, delta: int) -> None: 67 """Print the size difference.""" 68 if delta > 0: 69 print(f"{compare_file} is {delta} bytes bigger than {base_file}.") 70 else: 71 print_ansi(TTY_GREEN) 72 print(f"{compare_file} is {abs(delta)} bytes SMALLER than {base_file}. Great!") 73 print_ansi(TTY_RESET) 74 75 76def print_size_error() -> None: 77 """Print an error message for excessive size.""" 78 print_ansi(TTY_RED) 79 print(SIZE_ERROR_MESSAGE) 80 print_ansi(TTY_RESET) 81 82 83def compare_against_base( 84 base_file: str, compare_file: str, warning_size: int, error_size: int 85) -> int: 86 """Compare test binary file size against base revision binary file size.""" 87 base_file = create_file_path(base_file) 88 compare_file = create_file_path(compare_file) 89 90 diff = get_file_size(compare_file) - get_file_size(base_file) 91 print_size_diff(compare_file.name, base_file.name, diff) 92 93 if diff >= error_size: 94 print_size_error() 95 return EXIT_ERROR 96 elif diff >= warning_size: 97 return EXIT_WARNING 98 else: 99 return EXIT_OK 100 101 102def compare_against_max(compare_file: str, max_size: int) -> int: 103 """Compare test binary file size against maximum value.""" 104 compare_file = create_file_path(compare_file) 105 106 diff = get_file_size(compare_file) - max_size 107 print_size_diff(compare_file.name, "specified max size", diff) 108 109 if diff > 0: 110 print_size_error() 111 return EXIT_ERROR 112 else: 113 return EXIT_OK 114 115 116def main() -> int: 117 # Parse arguments. 118 parser = argparse.ArgumentParser(description="Compare binary file size") 119 parser.add_argument( 120 "--compare-file", 121 metavar="FILE", 122 type=str, 123 required=True, 124 help="Binary to compare against size args or base revision binary", 125 ) 126 parser.add_argument( 127 "--base-file", 128 metavar="FILE", 129 type=str, 130 help="Base revision binary", 131 dest="base_file", 132 ) 133 parser.add_argument( 134 "-s, --max-size", 135 metavar="SIZE", 136 type=int, 137 help="Max size of the binary, in bytes", 138 dest="max_size", 139 ) 140 parser.add_argument( 141 "-e, --error-size", 142 metavar="SIZE", 143 type=int, 144 help="Size difference between binaries constituting an error, in bytes", 145 dest="error_size", 146 ) 147 parser.add_argument( 148 "-w, --warning-size", 149 metavar="SIZE", 150 type=int, 151 help="Size difference between binaries constituting a warning, in bytes", 152 dest="warning_size", 153 ) 154 155 args = parser.parse_args() 156 157 if args.base_file is not None: 158 if args.max_size is not None: 159 print("Cannot specify both base file and maximum size arguments.") 160 sys.exit(EXIT_SCRIPT_ERROR) 161 162 if args.error_size is None or args.warning_size is None: 163 print( 164 "When comparing against base revision, error and warning sizes must be specified." 165 ) 166 sys.exit(EXIT_SCRIPT_ERROR) 167 168 return compare_against_base( 169 args.base_file, args.compare_file, args.warning_size, args.error_size 170 ) 171 elif args.max_size is not None: 172 return compare_against_max(args.compare_file, args.max_size) 173 174 175if __name__ == "__main__": 176 sys.exit(main()) 177