Improve ONNX setup: proactive VC++/AVX2 checks, friendlier errors, post-install verification

web_interface/app.py

@@ -77,6 +77,180 @@ REQUIREMENTS = [
 ]
 
 
+# ---------------------------------------------------------------------------
+# Setup helpers (pre-flight checks, friendlier install, post-install verify)
+#
+# These are called from _bootstrap / _bootstrap_portable below. They depend
+# on the platform/runtime checks defined later in this file (_has_vcruntime,
+# _has_avx2, _install_vcredist_with_uac) — Python resolves those at call
+# time, so definition order doesn't matter.
+# ---------------------------------------------------------------------------
+
+def _check_internet(timeout: float = 4.0) -> bool:
+    """Cheap reachability probe for pypi.org.
+
+    Conservative on unknown errors: only a clean URLError (no DNS, refused,
+    timeout) returns False. We don't want a flaky probe to falsely block
+    setup when the real install would have worked."""
+    try:
+        import urllib.error
+        import urllib.request
+        urllib.request.urlopen("https://pypi.org/simple/", timeout=timeout)
+        return True
+    except urllib.error.URLError:
+        return False
+    except Exception:  # noqa: BLE001
+        return True
+
+
+def _friendly_pip_install(py: Path, packages: list[str], action: str) -> None:
+    """Run `pip install` and translate non-zero exits into actionable hints.
+
+    pip prints its own error output, so we add a one-line summary plus a
+    short list of common fixes — enough for a non-tech user to unblock
+    themselves without reading a stack trace."""
+    cmd = [str(py), "-m", "pip", "install", *packages]
+    try:
+        rc = subprocess.call(cmd)
+    except KeyboardInterrupt:
+        print("\n[setup] Cancelled by user.")
+        sys.exit(130)
+    if rc == 0:
+        return
+    print()
+    print(f"[setup] Failed to {action} (pip exit code {rc}).")
+    print("Common fixes:")
+    print("  - Check your internet connection (corporate proxy/VPN can block pypi.org)")
+    print("  - Make sure you have at least ~600 MB of free disk space")
+    print("  - Run this script from a folder you own (Desktop/Documents) so pip can write")
+    print("  - Once the issue is fixed, re-run with --reinstall to retry")
+    sys.exit(1)
+
+
+def _requirements_for_this_cpu() -> list[str]:
+    """Return REQUIREMENTS adjusted for this CPU.
+
+    On non-AVX2 CPUs, pin onnxruntime up front to the last AVX-only release
+    so the install matches what the loader will accept — instead of pulling
+    a modern wheel pip is happy with but the import will reject."""
+    if not _has_avx2():
+        print(
+            "[setup] This CPU lacks AVX2 — pinning onnxruntime to "
+            f"{ORT_NO_AVX2_FALLBACK_VERSION} (last release with AVX-only wheels)."
+        )
+        return [
+            f"onnxruntime=={ORT_NO_AVX2_FALLBACK_VERSION}" if r.startswith("onnxruntime") else r
+            for r in REQUIREMENTS
+        ]
+    return list(REQUIREMENTS)
+
+
+def _pre_install_environment_setup() -> None:
+    """Catch known-broken environments before we spend ~2 minutes on pip.
+
+    Verifies pypi.org is reachable and (on Windows) installs the VC++
+    runtime up front when missing. Doing this here saves the user from a
+    confusing install -> launch -> crash -> install -> launch dance, since
+    a missing VC++ runtime would otherwise only surface when onnxruntime is
+    first imported."""
+    if not _check_internet():
+        print("[setup] Could not reach https://pypi.org/.")
+        print("        Check your internet connection, then re-run this script.")
+        print("        If you're behind a corporate proxy, set HTTPS_PROXY first, e.g.")
+        print("            set HTTPS_PROXY=http://proxy.example.com:8080")
+        sys.exit(1)
+
+    if os.name == "nt" and not _has_vcruntime():
+        print("[setup] Microsoft Visual C++ runtime is missing — installing it now.")
+        print("        (one-time, ~25 MB; please accept the UAC prompt that appears)")
+        if not _install_vcredist_with_uac():
+            print("[setup] Could not install the VC++ runtime automatically.")
+            print(f"        Download it from {VC_REDIST_URL}, run it manually,")
+            print("        then re-run this script.")
+            sys.exit(1)
+
+
+def _verify_onnxruntime_loads(py: Path) -> tuple[bool, str]:
+    """Spawn `<py> -c 'import onnxruntime'` and report whether it succeeded.
+
+    A clean pip install does not guarantee the wheel will load: missing
+    vcruntime140_1.dll, antivirus quarantine of onnxruntime_pybind11_state.pyd,
+    and 32/64-bit Python mismatches all manifest only at first import.
+    Probing here lets the bootstrap recover before writing the success
+    marker, so the user only sees one set of progress messages."""
+    code = (
+        "import sys\n"
+        "try:\n"
+        "    import onnxruntime  # noqa: F401\n"
+        "except BaseException as e:\n"
+        "    sys.stderr.write(repr(e))\n"
+        "    sys.exit(1)\n"
+    )
+    try:
+        proc = subprocess.run(
+            [str(py), "-c", code], check=False, capture_output=True, text=True
+        )
+    except OSError as e:
+        return False, str(e)
+    if proc.returncode == 0:
+        return True, ""
+    return False, (proc.stderr or "").strip() or "unknown import failure"
+
+
+def _bootstrap_recover(py: Path, err_repr: str) -> bool:
+    """Single-shot recovery for a failing onnxruntime import during setup.
+
+    Mirrors the in-app recovery in _recover_after_onnxruntime_load_failure
+    but stays in the setup stage (no re-exec), so the caller can re-run the
+    verification probe and only mark the install good once it actually works.
+    Returns True if a recovery action was attempted."""
+    if os.name != "nt":
+        return False
+    if "DLL load failed" not in err_repr and "ImportError" not in err_repr:
+        return False
+    if not _has_vcruntime():
+        print("[setup] VC++ runtime still missing — retrying installer.")
+        return _install_vcredist_with_uac()
+    if not _has_avx2():
+        print(
+            "[setup] CPU lacks AVX2 — reinstalling onnxruntime "
+            f"{ORT_NO_AVX2_FALLBACK_VERSION} ..."
+        )
+        try:
+            subprocess.check_call(
+                [str(py), "-m", "pip", "install", "--force-reinstall",
+                 f"onnxruntime=={ORT_NO_AVX2_FALLBACK_VERSION}"]
+            )
+            return True
+        except subprocess.CalledProcessError:
+            return False
+    return False
+
+
+def _finalize_install(py: Path, marker_path: Path) -> None:
+    """Verify onnxruntime imports, run one round of recovery if needed, then
+    write the success marker. Exits with a tailored diagnostic if recovery
+    can't make the import work."""
+    print("[setup] Verifying onnxruntime can load ...")
+    ok, err = _verify_onnxruntime_loads(py)
+    if not ok:
+        print(f"[setup] Initial import failed: {err}")
+        if _bootstrap_recover(py, err):
+            ok, err = _verify_onnxruntime_loads(py)
+    if not ok:
+        print()
+        print("[setup] onnxruntime still fails to load.")
+        print(f"        Last error: {err}")
+        print("Likely causes:")
+        print("  - Antivirus quarantined onnxruntime_pybind11_state.pyd")
+        print("    (open Windows Defender's Protection History and restore it)")
+        print("  - Corrupt install — re-run this script with --reinstall")
+        print("  - 32-bit Python loading 64-bit wheels (or vice versa)")
+        sys.exit(1)
+    marker_path.write_text("ok\n", encoding="utf-8")
+    print("[setup] Done.")
+
+
 # ---------------------------------------------------------------------------
 # Bootstrap
 # ---------------------------------------------------------------------------
@@ -115,23 +289,32 @@ def _in_target_venv() -> bool:
 
 
 def _bootstrap(force_reinstall: bool) -> None:
+    needs_install = force_reinstall or not VENV_DIR.exists() or not MARKER.exists()
+    if needs_install:
+        _pre_install_environment_setup()
+
     if not VENV_DIR.exists():
-        print(f"[bootstrap] Creating virtualenv at {VENV_DIR} ...")
-        venv.EnvBuilder(with_pip=True, clear=False, upgrade_deps=False).create(VENV_DIR)
+        print(f"[setup] Creating Python environment in {VENV_DIR.name}/ ...")
+        try:
+            venv.EnvBuilder(with_pip=True, clear=False, upgrade_deps=False).create(VENV_DIR)
+        except Exception as e:  # noqa: BLE001
+            print(f"[setup] Could not create virtualenv: {e}")
+            print("        Try running this script from a folder you own (Desktop/Documents).")
+            sys.exit(1)
 
     py = _venv_python()
-    needs_install = force_reinstall or not MARKER.exists()
     if needs_install:
-        print("[bootstrap] Upgrading pip ...")
-        subprocess.check_call([str(py), "-m", "pip", "install", "--upgrade", "pip"])
-        print(f"[bootstrap] Installing: {', '.join(REQUIREMENTS)}")
-        subprocess.check_call([str(py), "-m", "pip", "install", *REQUIREMENTS])
-        MARKER.write_text("ok\n", encoding="utf-8")
+        print("[setup] Upgrading pip ...")
+        _friendly_pip_install(py, ["--upgrade", "pip"], "upgrade pip")
+        reqs = _requirements_for_this_cpu()
+        print("[setup] Installing dependencies (~500 MB, takes ~1-3 minutes) ...")
+        _friendly_pip_install(py, reqs, "install required packages")
+        _finalize_install(py, MARKER)
+        print("[setup] Launching the app.\n")
     else:
-        print("[bootstrap] Requirements already installed (delete .venv/.bootstrapped to redo).")
+        print("[setup] Environment ready (delete .venv/.bootstrapped to re-run setup).")
 
     args = [a for a in sys.argv[1:] if a != "--reinstall"]
-    print("[bootstrap] Re-launching inside venv ...\n")
     sys.exit(subprocess.call([str(py), str(Path(__file__).resolve()), *args]))
 
 
@@ -155,6 +338,10 @@ def _bootstrap_portable(force_reinstall: bool) -> None:
     we're already running inside the target interpreter.
     """
     py = Path(sys.executable)
+    needs_install = force_reinstall or not PYENV_MARKER.exists()
+
+    if needs_install:
+        _pre_install_environment_setup()
 
     # 1) Enable site-packages by uncommenting `import site` in pythonXY._pth.
     pth_files = sorted(py.parent.glob("python*._pth"))
@@ -170,9 +357,9 @@ def _bootstrap_portable(force_reinstall: bool) -> None:
         if new_text != text:
             try:
                 pth.write_text(new_text, encoding="utf-8")
-                print(f"[bootstrap] Enabled site-packages in {pth.name}.")
+                print(f"[setup] Enabled site-packages in {pth.name}.")
             except OSError as e:
-                print(f"[bootstrap] Could not edit {pth}: {e}")
+                print(f"[setup] Could not edit {pth}: {e}")
 
     # 2) Install pip if missing.
     has_pip = subprocess.call(
@@ -181,12 +368,21 @@ def _bootstrap_portable(force_reinstall: bool) -> None:
         stderr=subprocess.DEVNULL,
     ) == 0
     if not has_pip:
-        print("[bootstrap] Installing pip into portable Python ...")
+        print("[setup] Installing pip into portable Python ...")
         import urllib.request
         getpip = py.parent / "_get-pip.py"
         try:
-            urllib.request.urlretrieve("https://bootstrap.pypa.io/get-pip.py", getpip)
-            subprocess.check_call([str(py), str(getpip), "--no-warn-script-location"])
+            try:
+                urllib.request.urlretrieve("https://bootstrap.pypa.io/get-pip.py", getpip)
+            except Exception as e:  # noqa: BLE001
+                print(f"[setup] Could not download get-pip.py: {e}")
+                print("        Check your internet connection, then re-run this script.")
+                sys.exit(1)
+            try:
+                subprocess.check_call([str(py), str(getpip), "--no-warn-script-location"])
+            except subprocess.CalledProcessError:
+                print("[setup] get-pip.py failed. Try deleting .pyenv\\ and re-running.")
+                sys.exit(1)
         finally:
             try:
                 getpip.unlink()
@@ -194,15 +390,15 @@ def _bootstrap_portable(force_reinstall: bool) -> None:
                 pass
 
     # 3) Install requirements (skipped if marker says we already did it).
-    needs_install = force_reinstall or not PYENV_MARKER.exists()
     if needs_install:
-        print("[bootstrap] Upgrading pip ...")
-        subprocess.check_call([str(py), "-m", "pip", "install", "--upgrade", "pip"])
-        print(f"[bootstrap] Installing: {', '.join(REQUIREMENTS)}")
-        subprocess.check_call([str(py), "-m", "pip", "install", *REQUIREMENTS])
-        PYENV_MARKER.write_text("ok\n", encoding="utf-8")
+        print("[setup] Upgrading pip ...")
+        _friendly_pip_install(py, ["--upgrade", "pip"], "upgrade pip")
+        reqs = _requirements_for_this_cpu()
+        print("[setup] Installing dependencies (~500 MB, takes ~1-3 minutes) ...")
+        _friendly_pip_install(py, reqs, "install required packages")
+        _finalize_install(py, PYENV_MARKER)
     else:
-        print("[bootstrap] Requirements already installed (delete .pyenv/.bootstrapped to redo).")
+        print("[setup] Environment ready (delete .pyenv\\.bootstrapped to re-run setup).")
 
 
 # ---------------------------------------------------------------------------
@@ -228,12 +424,19 @@ def _looks_like_native_load_failure(err: BaseException) -> bool:
 
 def _has_vcruntime() -> bool:
     """True if the VC++ 2015-2022 runtime DLLs are present in System32. We
-    look at vcruntime140.dll and msvcp140.dll because onnxruntime's native
-    code links against both."""
+    look at vcruntime140.dll, vcruntime140_1.dll, and msvcp140.dll because
+    onnxruntime's native code links against all three — vcruntime140_1.dll
+    in particular ships the SEH unwinding helpers used by C++ exceptions on
+    x64, and a host with the older 14.0 runtime but no _1.dll still fails to
+    load with the same generic "DLL load failed" message."""
     if os.name != "nt":
         return True
     sys32 = Path(os.environ.get("WINDIR", r"C:\Windows")) / "System32"
-    return (sys32 / "vcruntime140.dll").exists() and (sys32 / "msvcp140.dll").exists()
+    return (
+        (sys32 / "vcruntime140.dll").exists()
+        and (sys32 / "vcruntime140_1.dll").exists()
+        and (sys32 / "msvcp140.dll").exists()
+    )
 
 
 def _has_avx2() -> bool:
@@ -949,17 +1152,21 @@ def main() -> None:
 
     force = "--reinstall" in sys.argv[1:]
 
-    if _running_portable_python():
-        # run.bat downloaded an embeddable Python because the user had none.
-        # Install requirements directly into it; no venv re-exec needed.
-        _bootstrap_portable(force_reinstall=force)
-        _run_app()
-        return
+    try:
+        if _running_portable_python():
+            # run.bat downloaded an embeddable Python because the user had none.
+            # Install requirements directly into it; no venv re-exec needed.
+            _bootstrap_portable(force_reinstall=force)
+            _run_app()
+            return
 
-    if not _in_target_venv():
-        _bootstrap(force_reinstall=force)
-        return  # _bootstrap re-execs and exits
-    _run_app()
+        if not _in_target_venv():
+            _bootstrap(force_reinstall=force)
+            return  # _bootstrap re-execs and exits
+        _run_app()
+    except KeyboardInterrupt:
+        print("\n[setup] Cancelled by user.")
+        sys.exit(130)
 
 
 if __name__ == "__main__":