File size: 1,785 Bytes
7be70d4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 | #!/usr/bin/env python3
"""Verify the exported model bundle using only the standard library, on CPU."""
import argparse
import hashlib
import json
from pathlib import Path
import sys
def verify(root):
root = Path(root).resolve()
manifest = json.loads((root / 'bundle_manifest.json').read_text())
if manifest.get('status') != 'complete':
raise ValueError('bundle is not marked complete')
limit = manifest['max_file_bytes']
total = 0
for row in manifest['files']:
relative = Path(row['path'])
if relative.is_absolute() or '..' in relative.parts:
raise ValueError('unsafe manifest path')
path = root / relative
if path.is_symlink() or not path.is_file() or path.stat().st_size != row['bytes']:
raise ValueError('missing/changed file: ' + str(relative))
if path.stat().st_size > limit:
raise ValueError('file exceeds upload shard limit: ' + str(relative))
digest = hashlib.sha256()
with path.open('rb') as handle:
for block in iter(lambda: handle.read(8*1024*1024), b''):
digest.update(block)
if digest.hexdigest() != row['sha256']:
raise ValueError('SHA256 differs: ' + str(relative))
total += row['bytes']
print(json.dumps({'status': 'verified', 'files': len(manifest['files']),
'file_bytes': total, 'max_file_bytes': limit,
'GPU_used': False, 'model_loaded': False}, indent=2))
return manifest
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--bundle', type=Path, default=Path(__file__).resolve().parents[1])
args = parser.parse_args()
verify(args.bundle)
if __name__ == '__main__':
main()
|