1#!/usr/bin/env python3 2# Tool quickly rebuild one or two files with debug info 3# Mimics following behavior: 4# - touch file 5# - ninja -j1 -v -n torch_python | sed -e 's/-O[23]/-g/g' -e 's#\[[0-9]\+\/[0-9]\+\] \+##' |sh 6# - Copy libs from build/lib to torch/lib folder 7 8from __future__ import annotations 9 10import subprocess 11import sys 12from pathlib import Path 13from typing import Any 14 15 16PYTORCH_ROOTDIR = Path(__file__).resolve().parent.parent 17TORCH_DIR = PYTORCH_ROOTDIR / "torch" 18TORCH_LIB_DIR = TORCH_DIR / "lib" 19BUILD_DIR = PYTORCH_ROOTDIR / "build" 20BUILD_LIB_DIR = BUILD_DIR / "lib" 21 22 23def check_output(args: list[str], cwd: str | None = None) -> str: 24 return subprocess.check_output(args, cwd=cwd).decode("utf-8") 25 26 27def parse_args() -> Any: 28 from argparse import ArgumentParser 29 30 parser = ArgumentParser(description="Incremental build PyTorch with debinfo") 31 parser.add_argument("--verbose", action="store_true") 32 parser.add_argument("files", nargs="*") 33 return parser.parse_args() 34 35 36def get_lib_extension() -> str: 37 if sys.platform == "linux": 38 return "so" 39 if sys.platform == "darwin": 40 return "dylib" 41 raise RuntimeError(f"Usupported platform {sys.platform}") 42 43 44def create_symlinks() -> None: 45 """Creates symlinks from build/lib to torch/lib""" 46 if not TORCH_LIB_DIR.exists(): 47 raise RuntimeError(f"Can't create symlinks as {TORCH_LIB_DIR} does not exist") 48 if not BUILD_LIB_DIR.exists(): 49 raise RuntimeError(f"Can't create symlinks as {BUILD_LIB_DIR} does not exist") 50 for torch_lib in TORCH_LIB_DIR.glob(f"*.{get_lib_extension()}"): 51 if torch_lib.is_symlink(): 52 continue 53 build_lib = BUILD_LIB_DIR / torch_lib.name 54 if not build_lib.exists(): 55 raise RuntimeError(f"Can't find {build_lib} corresponding to {torch_lib}") 56 torch_lib.unlink() 57 torch_lib.symlink_to(build_lib) 58 59 60def has_build_ninja() -> bool: 61 return (BUILD_DIR / "build.ninja").exists() 62 63 64def is_devel_setup() -> bool: 65 output = check_output([sys.executable, "-c", "import torch;print(torch.__file__)"]) 66 return output.strip() == str(TORCH_DIR / "__init__.py") 67 68 69def create_build_plan() -> list[tuple[str, str]]: 70 output = check_output( 71 ["ninja", "-j1", "-v", "-n", "torch_python"], cwd=str(BUILD_DIR) 72 ) 73 rc = [] 74 for line in output.split("\n"): 75 if not line.startswith("["): 76 continue 77 line = line.split("]", 1)[1].strip() 78 if line.startswith(": &&") and line.endswith("&& :"): 79 line = line[4:-4] 80 line = line.replace("-O2", "-g").replace("-O3", "-g") 81 name = line.split("-o ", 1)[1].split(" ")[0] 82 rc.append((name, line)) 83 return rc 84 85 86def main() -> None: 87 if sys.platform == "win32": 88 print("Not supported on Windows yet") 89 sys.exit(-95) 90 if not is_devel_setup(): 91 print( 92 "Not a devel setup of PyTorch, please run `python3 setup.py develop --user` first" 93 ) 94 sys.exit(-1) 95 if not has_build_ninja(): 96 print("Only ninja build system is supported at the moment") 97 sys.exit(-1) 98 args = parse_args() 99 for file in args.files: 100 if file is None: 101 continue 102 Path(file).touch() 103 build_plan = create_build_plan() 104 if len(build_plan) == 0: 105 return print("Nothing to do") 106 if len(build_plan) > 100: 107 print("More than 100 items needs to be rebuild, run `ninja torch_python` first") 108 sys.exit(-1) 109 for idx, (name, cmd) in enumerate(build_plan): 110 print(f"[{idx + 1 } / {len(build_plan)}] Building {name}") 111 if args.verbose: 112 print(cmd) 113 subprocess.check_call(["sh", "-c", cmd], cwd=BUILD_DIR) 114 create_symlinks() 115 116 117if __name__ == "__main__": 118 main() 119