diff --git a/scripts/upload.py b/scripts/upload.py index 441585c940..127e615779 100644 --- a/scripts/upload.py +++ b/scripts/upload.py @@ -1,17 +1,31 @@ import boto3 from pathlib import Path import mimetypes +import hashlib import os +import json from multiprocessing.pool import ThreadPool AKEY = os.environ["AWS_S3_ACCESS_KEY"] SKEY = os.environ["AWS_S3_SECRET_KEY"] + BUCKET = "keras.io" USE_THREADING = True +HASH_CACHE = "contents_hashes.json" s3 = boto3.client("s3", aws_access_key_id=AKEY, aws_secret_access_key=SKEY) +def hash_file(fpath): + h = hashlib.sha256() + b = bytearray(128 * 1024) + mv = memoryview(b) + with open(fpath, "rb", buffering=0) as f: + while n := f.readinto(mv): + h.update(mv[:n]) + return h.hexdigest()[:8] + + def upload_file(bucket, fpath, key_name, redirect=None): print(f"...Upload to {bucket}:{key_name}") mime = mimetypes.guess_type(fpath)[0] @@ -23,6 +37,19 @@ def upload_file(bucket, fpath, key_name, redirect=None): ) +def load_hash_cache(): + s3.download_file(BUCKET, HASH_CACHE, HASH_CACHE) + with open(HASH_CACHE) as f: + contents = f.read() + return json.loads(contents) + + +def save_hash_cache(hash_cache): + with open(HASH_CACHE, "w") as f: + f.write(json.dumps(hash_cache)) + upload_file(BUCKET, HASH_CACHE, HASH_CACHE) + + def wrapped_upload_file(args): bucket, fpath, key_name = args upload_file(bucket, fpath, key_name) @@ -44,10 +71,10 @@ def cleanup(site_directory, redirect_directory): s3.delete_object(Bucket=BUCKET, Key=key) -def upload_dir(directory, include_img=True): +def upload_dir(directory, include_img=True, hash_cache=None): print(f"Uploading files from '{directory}'...") all_targets = [] - for dp, dn, fn in os.walk(directory): + for dp, _, fn in os.walk(directory): if fn: for f in fn: fpath = os.path.join(dp, f) @@ -60,6 +87,18 @@ def upload_dir(directory, include_img=True): print("> " + fpath) print(">>>>>> " + key_name) all_targets.append((BUCKET, fpath, key_name)) + + if hash_cache is not None: + filtered_targets = [] + new_hash_cache = {} + for bucket, fpath, key_name in all_targets: + new_hash = hash_file(fpath) + old_hash = hash_cache.get(key_name) + if new_hash != old_hash: + filtered_targets.append((bucket, fpath, key_name)) + new_hash_cache[key_name] = new_hash + all_targets = filtered_targets + if USE_THREADING: pool = ThreadPool(processes=8) pool.map(wrapped_upload_file, all_targets) @@ -67,10 +106,13 @@ def upload_dir(directory, include_img=True): for args in all_targets: wrapped_upload_file(args) + if hash_cache is not None: + return new_hash_cache + def upload_redirects(directory): print("Uploading redirects...") - for dp, dn, fn in os.walk(directory): + for dp, _, fn in os.walk(directory): if fn: for f in fn: fpath = os.path.join(dp, f) @@ -87,4 +129,8 @@ def upload_redirects(directory): if __name__ == "__main__": root = Path(__file__).parent.parent.resolve() - upload_dir(os.path.join(root, "site"), include_img=True) + hash_cache = load_hash_cache() + hash_cache = upload_dir( + os.path.join(root, "site"), include_img=True, hash_cache=hash_cache + ) + save_hash_cache(hash_cache)