1"""Uploads a new binary to s3 and updates its hash in the config file. 2 3You'll need to have appropriate credentials on the PyTorch AWS buckets, see: 4https://boto3.amazonaws.com/v1/documentation/api/latest/guide/quickstart.html#configuration 5for how to configure them. 6""" 7 8import argparse 9import hashlib 10import json 11import logging 12import os 13 14import boto3 # type: ignore[import] 15 16 17def compute_file_sha256(path: str) -> str: 18 """Compute the SHA256 hash of a file and return it as a hex string.""" 19 # If the file doesn't exist, return an empty string. 20 if not os.path.exists(path): 21 return "" 22 23 hash = hashlib.sha256() 24 25 # Open the file in binary mode and hash it. 26 with open(path, "rb") as f: 27 for b in f: 28 hash.update(b) 29 30 # Return the hash as a hexadecimal string. 31 return hash.hexdigest() 32 33 34def main() -> None: 35 parser = argparse.ArgumentParser( 36 description="s3 binary updater", 37 fromfile_prefix_chars="@", 38 ) 39 parser.add_argument( 40 "--config-json", 41 required=True, 42 help="path to config json that you are trying to update", 43 ) 44 parser.add_argument( 45 "--linter", 46 required=True, 47 help="name of linter you're trying to update", 48 ) 49 parser.add_argument( 50 "--platform", 51 required=True, 52 help="which platform you are uploading the binary for", 53 ) 54 parser.add_argument( 55 "--file", 56 required=True, 57 help="file to upload", 58 ) 59 parser.add_argument( 60 "--dry-run", 61 action="store_true", 62 help="if set, don't actually upload/write hash", 63 ) 64 args = parser.parse_args() 65 logging.basicConfig(level=logging.INFO) 66 67 config = json.load(open(args.config_json)) 68 linter_config = config[args.linter][args.platform] 69 bucket = linter_config["s3_bucket"] 70 object_name = linter_config["object_name"] 71 72 # Upload the file 73 logging.info( 74 "Uploading file %s to s3 bucket: %s, object name: %s", 75 args.file, 76 bucket, 77 object_name, 78 ) 79 if not args.dry_run: 80 s3_client = boto3.client("s3") 81 s3_client.upload_file(args.file, bucket, object_name) 82 83 # Update hash in repo 84 hash_of_new_binary = compute_file_sha256(args.file) 85 logging.info("Computed new hash for binary %s", hash_of_new_binary) 86 87 linter_config["hash"] = hash_of_new_binary 88 config_dump = json.dumps(config, indent=4, sort_keys=True) 89 90 logging.info("Writing out new config:") 91 logging.info(config_dump) 92 if not args.dry_run: 93 with open(args.config_json, "w") as f: 94 f.write(config_dump) 95 96 97if __name__ == "__main__": 98 main() 99