onw/update.py

"""Updates of onw AMD itself from Hugging Face (REPO below).

  latest()        newest released version (from the repo's v* tags), or None when offline
  notes(since)    the CHANGELOG.md sections newer than `since`
  can_update()    only installs made by the one-line installer (%LOCALAPPDATA%\\onw-amd, ~/onw-amd, or a folder with the
                  .onw-installed marker) are updated in place; a development checkout is never touched
  apply(tag)      download that version into a staging folder first (a broken download leaves the current install as
                  it is), then replace the files, and run pip only when pyproject.toml changed

Restarting afterwards is the tray's job (onw.tray): it stops the server, starts a fresh tray through onw.relaunch and
reloads the model that was running.
"""
import json, os, re, shutil, subprocess, sys, time, urllib.request

from . import __version__
from .i18n import T

REPO = "ryugyosoft/onwAMD"


def vtuple(v):
    return tuple(int(x) for x in re.findall(r"\d+", v or "")[:4]) or (0,)


def _get(url, tries=4, timeout=30):
    req = urllib.request.Request(url, headers={"User-Agent": f"onw-amd/{__version__}"})
    for i in range(tries):
        try:
            with urllib.request.urlopen(req, timeout=timeout) as r:
                return r.read()
        except Exception:
            if i == tries - 1:
                raise
            time.sleep(2 * (i + 1))


def latest():
    """(version string like "0.8.1", tag) of the newest v* tag on the hub, or None."""
    try:
        refs = json.loads(_get(f"https://huggingface.co/api/models/{REPO}/refs", tries=2, timeout=15))
    except Exception:
        return None
    tags = [t["name"] for t in refs.get("tags", []) if re.fullmatch(r"v\d+(\.\d+)*", t.get("name", ""))]
    if not tags:
        return None
    tag = max(tags, key=vtuple)
    return tag[1:], tag


def newer(v, than=__version__):
    return vtuple(v) > vtuple(than)


def notes(since=__version__, tag="main", lang="ja"):
    """The changelog entries for versions newer than `since` (markdown), or ""."""
    name = "CHANGELOG.md" if lang == "ja" else "CHANGELOG_en.md"
    try:
        text = _get(f"https://huggingface.co/{REPO}/resolve/{tag}/{name}", tries=2, timeout=15).decode("utf8")
    except Exception:
        return ""
    out = []
    for sec in re.split(r"(?m)^## ", text)[1:]:
        m = re.match(r"v(\d+(?:\.\d+)*)", sec)
        if m and newer(m.group(1), since):
            out.append("## " + sec.strip())
    return "\n\n".join(out)


def install_dir():
    """The folder onw runs from (the editable install), or None if it is not an installer-made install."""
    d = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    home = os.path.join(os.environ.get("LOCALAPPDATA", ""), "onw-amd") if os.name == "nt" else os.path.expanduser("~/onw-amd")
    ok = os.path.exists(os.path.join(d, ".onw-installed")) or os.path.normcase(os.path.abspath(home)) == os.path.normcase(d)
    if not ok or os.path.isdir(os.path.join(d, ".git")) or not os.path.exists(os.path.join(d, "pyproject.toml")):
        return None
    return d


def can_update():
    return install_dir() is not None


def apply(tag, progress=lambda text: None):
    """Install version `tag` over the current install. Returns the new version string."""
    d = install_dir()
    if d is None:
        raise RuntimeError("this onw AMD is not an installer-made install (update it with git / hf instead)")
    info = json.loads(_get(f"https://huggingface.co/api/models/{REPO}/revision/{tag}"))
    files = [s["rfilename"] for s in info["siblings"] if s["rfilename"] != ".gitattributes"]
    stage = os.path.join(d, ".update")
    shutil.rmtree(stage, ignore_errors=True)

    def fetch(f):
        data = _get(f"https://huggingface.co/{REPO}/resolve/{tag}/{f}")
        if f.endswith((".sh", ".py", ".toml", ".txt", ".md", ".html")):
            data = data.replace(b"\r\n", b"\n")           # scripts keep Unix line endings on every platform
        if not data and f.endswith((".py", ".sh", ".bat", ".toml")) and not f.endswith("__init__.py"):
            raise RuntimeError(f"{f} came down empty; update cancelled")
        out = os.path.join(stage, f)
        os.makedirs(os.path.dirname(out), exist_ok=True)
        with open(out, "wb") as fh:
            fh.write(data)

    from concurrent.futures import ThreadPoolExecutor, as_completed
    with ThreadPoolExecutor(8) as ex:                 # 8 at a time: a slow connection costs far less than one by one
        jobs = [ex.submit(fetch, f) for f in files]
        for i, j in enumerate(as_completed(jobs)):
            j.result()
            progress(T("ダウンロード中 {0}/{1}", i + 1, len(files)))
    old_toml = open(os.path.join(d, "pyproject.toml"), "rb").read()
    new_toml = open(os.path.join(stage, "pyproject.toml"), "rb").read()
    progress(T("入れ替え中"))
    for f in files:                                   # replace file by file (a running Python keeps its loaded code)
        dst = os.path.join(d, f)
        os.makedirs(os.path.dirname(dst), exist_ok=True)
        for i in range(10):                           # Windows: a file being read right now (the page being served)
            try:                                      # cannot be replaced for that moment
                os.replace(os.path.join(stage, f), dst)
                break
            except PermissionError:
                if i == 9:
                    raise
                time.sleep(0.5)
    # modules the new version no longer has: stale code must not be imported by accident
    keep = {os.path.normcase(os.path.join(d, f)) for f in files}
    for root, dirs, names in os.walk(os.path.join(d, "onw")):
        dirs[:] = [x for x in dirs if x != "__pycache__"]
        for n in names:
            p = os.path.join(root, n)
            if n.endswith(".py") and os.path.normcase(p) not in keep:
                try:
                    os.remove(p)
                except OSError:
                    pass
    shutil.rmtree(stage, ignore_errors=True)
    open(os.path.join(d, ".onw-installed"), "w").write(tag + "\n")
    if new_toml != old_toml:                          # new or changed dependencies
        progress(T("パッケージを更新中（数分かかることがあります）"))
        kw = {"creationflags": subprocess.CREATE_NO_WINDOW} if os.name == "nt" else {}
        r = subprocess.run([sys.executable, "-m", "pip", "install", "-e", ".[tray]"],
                           cwd=d, capture_output=True, text=True, encoding="utf8", errors="replace", **kw)
        if r.returncode != 0:
            raise RuntimeError("pip failed: " + " ".join(r.stdout.split()[-40:]))
    return tag.lstrip("v")