1import argparse 2import os 3 4from yaml import dump, Dumper 5 6from torch._export.serde import schema_check 7 8 9if __name__ == "__main__": 10 parser = argparse.ArgumentParser(prog="update_schema") 11 parser.add_argument( 12 "--prefix", type=str, required=True, help="The root of pytorch directory." 13 ) 14 parser.add_argument( 15 "--dry-run", 16 action="store_true", 17 help="Print the schema instead of writing it to file.", 18 ) 19 parser.add_argument( 20 "--force-unsafe", 21 action="store_true", 22 help="!!! Only use this option when you are a chad. !!! Force to write the schema even if schema validation doesn't pass.", 23 ) 24 args = parser.parse_args() 25 26 assert os.path.exists( 27 args.prefix 28 ), f"Assuming path {args.prefix} is the root of pytorch directory, but it doesn't exist." 29 30 commit = schema_check.update_schema() 31 32 if os.path.exists(args.prefix + commit.path): 33 if commit.result["SCHEMA_VERSION"] < commit.base["SCHEMA_VERSION"]: 34 raise RuntimeError( 35 f"Schema version downgraded from {commit.base['SCHEMA_VERSION']} to {commit.result['SCHEMA_VERSION']}." 36 ) 37 38 if commit.result["TREESPEC_VERSION"] < commit.base["TREESPEC_VERSION"]: 39 raise RuntimeError( 40 f"Treespec version downgraded from {commit.base['TREESPEC_VERSION']} to {commit.result['TREESPEC_VERSION']}." 41 ) 42 else: 43 assert args.force_unsafe, "Existing schema yaml file not found, please use --force-unsafe to try again." 44 45 next_version, reason = schema_check.check(commit, args.force_unsafe) 46 47 if next_version is not None and next_version != commit.result["SCHEMA_VERSION"]: 48 raise RuntimeError( 49 f"Schema version is not updated from {commit.base['SCHEMA_VERSION']} to {next_version}.\n" 50 + "Please either:\n" 51 + " 1. update schema.py to not break compatibility.\n" 52 + " or 2. bump the schema version to the expected value.\n" 53 + " or 3. use --force-unsafe to override schema.yaml (not recommended).\n " 54 + "and try again.\n" 55 + f"Reason: {reason}" 56 ) 57 58 header = ( 59 "# @" + "generated by " + os.path.basename(__file__).rsplit(".", 1)[0] + ".py" 60 ) 61 header += f"\n# checksum<<{commit.checksum_result}>>" 62 payload = dump(commit.result, Dumper=Dumper, sort_keys=False) 63 64 content = header + "\n" + payload 65 66 if args.dry_run: 67 print(content) 68 print("\nWill write the above schema to" + args.prefix + commit.path) 69 else: 70 with open(args.prefix + commit.path, "w") as f: 71 f.write(content) 72