onw/dl.py

"""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()