File size: 5,979 Bytes
d0bf6e2 00a5163 d0bf6e2 00a5163 d0bf6e2 00a5163 d0bf6e2 00a5163 d0bf6e2 00a5163 d0bf6e2 00a5163 d0bf6e2 00a5163 d0bf6e2 00a5163 d0bf6e2 00a5163 d0bf6e2 | 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 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 | """Download one model (or just FastFlowLM) as its own process: the manager can cancel it by ending the process, and
a crash inside a download cannot take the tray down. Progress goes to a small JSON file the manager reads.
usage (internal): python -m onw.dl KEY STATUS_FILE a model ("...-FLM"): FastFlowLM first when missing,
`flm pull`, then every file checked against flm's hashes
python -m onw.dl --repair KEY STATUS_FILE the model's damaged files fetched again (after a check)
python -m onw.dl --verify KEY STATUS_FILE just the check; the result (damaged files) in STATUS_FILE
python -m onw.dl --engine STATUS_FILE just FastFlowLM
"""
import json, os, re, subprocess, sys, time
from . import app, lemonade
class Status:
def __init__(self, path):
self.path, self.last = path, 0.0
def __call__(self, frac, text, force=False, **extra):
if not force and time.time() - self.last < 0.5:
return
self.last = time.time()
with open(self.path + ".tmp", "w", encoding="utf8") as f:
json.dump({"frac": round(frac, 4), "text": text, **extra}, f, ensure_ascii=False)
os.replace(self.path + ".tmp", self.path)
def engine(status):
if not lemonade.up_to_date():
lemonade.install(status)
status(1.0, "FastFlowLM " + lemonade.VERSION, force=True)
def pull(m, status, cfg, force=False):
"""`flm pull <checkpoint>`; progress from its "[FLM] Downloading X/Y: file" lines (lemonade: flm_download)."""
p = subprocess.Popen(lemonade.flm_pull_command(m["checkpoint"], force), stdout=subprocess.PIPE,
stderr=subprocess.STDOUT, env=lemonade.flm_env(cfg), text=True, encoding="utf8",
errors="replace", **lemonade.NO_WINDOW)
# flm prints "Downloading 2/6: model.q4nx" per file and, with \r, "Downloading: 44.6% (85.7MB / 191.9MB)" while a
# file comes down: the bytes give the progress (one big weights file would otherwise sit at 0% for minutes)
file_pat = re.compile(r"Downloading\s+(\d+)/(\d+):\s*(.*)")
byte_pat = re.compile(r"Downloading:\s*([\d.]+)%\s*\(([\d.]+)\s*([KMGT]?B)\s*/\s*([\d.]+)\s*([KMGT]?B)\)")
unit = {"B": 1, "KB": 2**10, "MB": 2**20, "GB": 2**30, "TB": 2**40}
total = lemonade.flm_download_size(m["checkpoint"]) # the files' sizes (the footprint is memory, not download)
done = cur = cur_total = 0.0 # bytes of the finished files, of the current one
i = n = 0
name = ""
def report():
whole = max(total, done + cur_total, 1)
frac = min((done + cur) / whole, 0.99)
status(frac, f"{(done + cur) / 2**30:.1f} / {whole / 2**30:.1f} GB ({frac * 100:.0f}%)"
+ (f" · {i}/{n} {name[:40]}" if n else ""))
for line in p.stdout:
print(line, end="", flush=True)
line = re.sub(r"\x1b\[[0-9;?]*[A-Za-z]", "", line)
g = file_pat.search(line)
if g:
done, cur, cur_total = done + cur_total, 0.0, 0.0 # the previous file is complete
i, n, name = int(g.group(1)), int(g.group(2)), g.group(3).strip()
report()
continue
g = byte_pat.search(line)
if g:
cur = float(g.group(2)) * unit.get(g.group(3).upper(), 1)
cur_total = float(g.group(4)) * unit.get(g.group(5).upper(), 1)
report()
if p.wait() != 0:
sys.exit(f"flm pull failed (exit code {p.returncode})")
def verify(m, status, cfg):
"""The model's files against flm's hashes. -> damaged file names ([] intact, None: no hashes to check)."""
status(0.0, "checking…", force=True)
bad = lemonade.flm_verify(m["checkpoint"], cfg,
lambda frac, text: status(frac, f"checking {text} ({frac * 100:.0f}%)"))
print(f"[onw] verify {m['checkpoint']}: " + ("no hashes" if bad is None else ", ".join(bad) or "intact"), flush=True)
return bad
def repair(m, bad, status, cfg):
"""Fetch the damaged files again: remove them so that `flm pull` downloads them; if flm still counts the model
as complete, pull it whole with --force. Exits when the files are still damaged."""
for name, path, *_ in lemonade.flm_files(m["checkpoint"], cfg):
if name in bad and os.path.exists(path):
os.remove(path)
pull(m, status, cfg)
bad = verify(m, status, cfg)
if bad:
pull(m, status, cfg, force=True)
bad = verify(m, status, cfg)
if bad:
sys.exit("downloaded files are damaged (checked against FastFlowLM's hashes): " + ", ".join(bad))
def model_entry(key, cfg):
ok, msg = lemonade.flm_validate(cfg)
if not ok:
sys.exit(f"FastFlowLM NPU check failed: {msg}")
m = lemonade.flm_models(cfg).get(key)
if m is None:
sys.exit(f"FastFlowLM has no model {key}")
return m
def main():
args = sys.argv[1:]
if args[0] == "--engine":
engine(Status(args[1]))
return
if args[0] in ("--verify", "--repair"):
mode, key, status = args[0], args[1], Status(args[2])
else:
mode, key, status = "--download", args[0], Status(args[1])
cfg = app.load_config()
status(0.0, "…", force=True)
if mode == "--verify":
m = lemonade.flm_models(cfg).get(key)
if m is None:
sys.exit(f"FastFlowLM has no model {key}")
bad = verify(m, status, cfg)
status(1.0, "done", force=True, bad=bad)
return
engine(lambda frac, text, force=False: status(0.0, text, force))
m = model_entry(key, cfg)
if mode == "--download":
pull(m, status, cfg)
bad = verify(m, status, cfg) # (--repair: finds the damaged files again, here)
if bad:
repair(m, bad, status, cfg)
status(1.0, "done", force=True, bad=[])
if __name__ == "__main__":
main()
|