141 lines
6.2 KiB
Python
141 lines
6.2 KiB
Python
"""Download pinned official safetensors, no framework or shared cache."""
|
|
import argparse
|
|
import concurrent.futures
|
|
import hashlib
|
|
import json
|
|
import pathlib
|
|
import os
|
|
import shutil
|
|
import time
|
|
import threading
|
|
import urllib.parse
|
|
import urllib.request
|
|
|
|
|
|
def fetch_json(url):
|
|
with urllib.request.urlopen(url, timeout=90) as response:
|
|
return json.load(response)
|
|
|
|
|
|
def ranged_download(url, part, total, workers=8):
|
|
"""Resume a sequential prefix; journal completed byte ranges before final hashing."""
|
|
journal = part.with_suffix(part.suffix + '.ranges.json')
|
|
if journal.exists():
|
|
state = json.loads(journal.read_text())
|
|
if state['total'] != total:
|
|
raise RuntimeError('Range journal size mismatch')
|
|
else:
|
|
state = {'total': total, 'prefix': part.stat().st_size if part.exists() else 0, 'done': []}
|
|
journal.write_text(json.dumps(state))
|
|
lock = threading.Lock()
|
|
fd = os.open(part, os.O_CREAT | os.O_RDWR, 0o600)
|
|
chunk_size = 64 * 1024 * 1024
|
|
ranges = [(start, min(start + chunk_size, total) - 1)
|
|
for start in range(state['prefix'], total, chunk_size)]
|
|
|
|
def one(bounds):
|
|
start, end = bounds
|
|
if start in state['done']:
|
|
return
|
|
for attempt in range(3):
|
|
try:
|
|
req = urllib.request.Request(url, headers={'Range': f'bytes={start}-{end}'})
|
|
with urllib.request.urlopen(req, timeout=180) as response:
|
|
expected = f'bytes {start}-{end}/{total}'
|
|
if response.status != 206 or response.headers.get('Content-Range') != expected:
|
|
raise RuntimeError('Server did not honor exact requested range')
|
|
offset = start
|
|
while chunk := response.read(min(4 * 1024 * 1024, end - offset + 1)):
|
|
written = os.pwrite(fd, chunk, offset)
|
|
if written != len(chunk):
|
|
raise RuntimeError('Short positional write')
|
|
offset += written
|
|
if offset > end:
|
|
break
|
|
if offset != end + 1:
|
|
raise RuntimeError('Incomplete range')
|
|
with lock:
|
|
os.fsync(fd)
|
|
state['done'].append(start)
|
|
tmp = journal.with_suffix('.tmp')
|
|
tmp.write_text(json.dumps(state))
|
|
tmp.replace(journal)
|
|
done = state['prefix'] + sum(b - a + 1 for a, b in ranges if a in state['done'])
|
|
print(f'Weight download {done}/{total} bytes ({100 * done / total:.1f}%)', flush=True)
|
|
return
|
|
except Exception:
|
|
if attempt == 2:
|
|
raise
|
|
time.sleep(2 ** attempt)
|
|
try:
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as pool:
|
|
list(pool.map(one, ranges))
|
|
finally:
|
|
os.close(fd)
|
|
journal.unlink()
|
|
|
|
|
|
def main():
|
|
p = argparse.ArgumentParser()
|
|
p.add_argument('--root', type=pathlib.Path, required=True)
|
|
p.add_argument('--repo', default='Qwen/Qwen3.5-2B')
|
|
args = p.parse_args()
|
|
root = args.root.resolve()
|
|
root.mkdir(parents=True, exist_ok=True)
|
|
weights = root / 'weights'
|
|
weights.mkdir(exist_ok=True)
|
|
marker = weights / '.pilot-owned'
|
|
marker.write_text(args.repo)
|
|
old_manifest = root / 'manifest.json'
|
|
pinned = json.loads(old_manifest.read_text())['revision'] if old_manifest.exists() else None
|
|
endpoint = f'https://huggingface.co/api/models/{args.repo}'
|
|
if pinned:
|
|
endpoint += f'/revision/{pinned}'
|
|
info = fetch_json(endpoint + '?blobs=true')
|
|
revision = info['sha']
|
|
files = [f for f in info['siblings'] if f['rfilename'].endswith('.safetensors')
|
|
or f['rfilename'] in ('config.json', 'model.safetensors.index.json', 'README.md')]
|
|
total = sum(f.get('size', 0) for f in files)
|
|
if shutil.disk_usage(root).free < total + 3_000_000_000:
|
|
raise RuntimeError('Insufficient space for checkpoint plus 3 GB margin')
|
|
manifest = {'repo': args.repo, 'revision': revision, 'expected_bytes': total,
|
|
'started_unix': time.time(), 'files': [], 'license': info.get('cardData', {}).get('license')}
|
|
(root / 'source_metadata.json').write_text(json.dumps(info, indent=2))
|
|
for entry in files:
|
|
name = entry['rfilename']
|
|
if pathlib.PurePosixPath(name).name != name:
|
|
raise RuntimeError(f'Unexpected nested path: {name}')
|
|
dest = weights / name
|
|
part = dest.with_suffix(dest.suffix + '.partial')
|
|
digest = hashlib.sha256()
|
|
url = f'https://huggingface.co/{args.repo}/resolve/{revision}/{urllib.parse.quote(name)}?download=true'
|
|
print(f'Downloading {name}, expected {entry.get("size")} bytes', flush=True)
|
|
if entry.get('size', 0) > 128 * 1024 * 1024:
|
|
ranged_download(url, part, entry['size'])
|
|
with part.open('rb') as stream:
|
|
while chunk := stream.read(4 * 1024 * 1024):
|
|
digest.update(chunk)
|
|
else:
|
|
with urllib.request.urlopen(url, timeout=180) as response, part.open('wb') as out:
|
|
while chunk := response.read(4 * 1024 * 1024):
|
|
out.write(chunk)
|
|
digest.update(chunk)
|
|
sha = digest.hexdigest()
|
|
if entry.get('size') and part.stat().st_size != entry['size']:
|
|
raise RuntimeError(f'Size mismatch: {name}')
|
|
expected_sha = entry.get('lfs', {}).get('sha256')
|
|
if expected_sha and sha != expected_sha:
|
|
raise RuntimeError(f'SHA256 mismatch: {name}')
|
|
part.rename(dest)
|
|
manifest['files'].append({'name': name, 'bytes': dest.stat().st_size,
|
|
'sha256': sha, 'upstream_sha256': expected_sha})
|
|
(root / 'manifest.json').write_text(json.dumps(manifest, indent=2))
|
|
manifest['completed_unix'] = time.time()
|
|
(root / 'manifest.json').write_text(json.dumps(manifest, indent=2))
|
|
shutil.copy2(weights / 'config.json', root / 'model_config.json')
|
|
print(json.dumps({'revision': revision, 'bytes': sum(f['bytes'] for f in manifest['files'])}), flush=True)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|