cesnimda
46 lines
2.2 KiB
Python
46 lines
2.2 KiB
Python
"""Prepare the reviewed trained checkpoint in a NEW nonproduction directory; never reset caches."""
|
|
from pathlib import Path
|
|
import hashlib
|
|
import json
|
|
import shutil
|
|
import sys
|
|
import time
|
|
import urllib.parse
|
|
import urllib.request
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent/'model'))
|
|
from verify import verify_fixture
|
|
|
|
|
|
def prepare(destination):
|
|
destination=Path(destination)
|
|
manifest=json.loads(Path(__file__).with_name('ARTIFACTS.json').read_text())['weights-only']
|
|
revision=manifest['sourceRevision']
|
|
assert revision=='a4f8f3ea906ed274767e9906dbaede7531d660ff'
|
|
assert manifest['license'].startswith('Apache-2.0')
|
|
assert not destination.exists(), 'Use a new isolated directory; existing data is never replaced.'
|
|
assert shutil.disk_usage(destination.parent).free>8*1024**3, 'Require8GiB available for bounded preparation.'
|
|
destination.mkdir(mode=0o700)
|
|
deadline=time.monotonic()+1200
|
|
for entry in manifest['files']:
|
|
name=entry['file']; assert Path(name).name==name and name not in ('.','..')
|
|
url=entry.get('url') or entry['source']
|
|
parsed=urllib.parse.urlsplit(url)
|
|
assert parsed.scheme=='https' and parsed.netloc=='huggingface.co'
|
|
assert parsed.path.startswith('/sshleifer/distilbart-cnn-12-6/resolve/'+revision+'/')
|
|
temporary=destination/(name+'.partial'); count=0; digest=hashlib.sha256()
|
|
with urllib.request.urlopen(url,timeout=60) as response, temporary.open('xb') as stream:
|
|
while block:=response.read(1024*1024):
|
|
count+=len(block)
|
|
assert count<=entry['bytes'] and time.monotonic()<deadline, 'Artifact size/deadline exceeded.'
|
|
digest.update(block); stream.write(block)
|
|
assert count==entry['bytes'] and digest.hexdigest()==entry['sha256'], 'Artifact identity mismatch; do not load partial files.'
|
|
temporary.rename(destination/name)
|
|
verify_fixture(destination,manifest)
|
|
print(json.dumps({'status':'PASS','directory':str(destination.resolve()),'revision':revision,'files':len(manifest['files'])}))
|
|
|
|
|
|
if __name__=='__main__':
|
|
assert len(sys.argv)==2, 'Usage: python prepare.py NEW_NONPRODUCTION_CACHE_DIRECTORY'
|
|
prepare(sys.argv[1])
|