mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
fix: reliable, self-healing AI feature-bundle installs (#472)
Make on-demand AI feature-bundle installs reliable and self-healing, closing the failure modes behind most "some tool doesn't work" reports. Multi-bundle installs: tools needing more than one bundle (Passport Photo, Enhance Faces) install every required bundle from one action and stay not-installed until all are present. Verified across all 19 AI tools. Downloads: self-heal the accelerated Hugging Face (Xet) client so an upgraded venv no longer silently falls back to slow urllib; restart instead of corrupting a resumed partial when a proxy ignores Range and returns 200; verify the completed size; fail fast on disk-full and HTTP 4xx; retry transient errors five times; add hf_transfer fallback and document Xet egress. Install integrity: crash-atomic venv writes so a killed or out-of-space install can no longer tear the shared venv and break other tools; a boot breadcrumb reseeds a torn venv to a clean state automatically; a post-install smoke import test refuses to record a bundle whose libraries cannot load; an install watchdog stops a wedged installer that would otherwise hold the venv writer lock forever. Adds unit and end-to-end tests for every failure mode above.
This commit is contained in:
SnapOtter
@@ -15,6 +15,7 @@ Final result is a JSON object on stdout.
|
||||
import errno
|
||||
import glob
|
||||
import hashlib
|
||||
import importlib
|
||||
import json
|
||||
import os
|
||||
import platform
|
||||
@@ -27,6 +28,9 @@ import urllib.error
|
||||
import urllib.request
|
||||
from datetime import datetime, timezone
|
||||
|
||||
DOWNLOAD_CHUNK_BYTES = 4 * 1024 * 1024
|
||||
DOWNLOAD_META_BYTES = 64 * 1024 * 1024
|
||||
|
||||
|
||||
# -- Helpers --
|
||||
|
||||
@@ -149,6 +153,166 @@ def verify_sha256(filepath: str, expected: str) -> bool:
|
||||
|
||||
# -- Download with resume --
|
||||
|
||||
def _set_env_temporarily(key: str, value: str):
|
||||
previous = os.environ.get(key)
|
||||
os.environ[key] = value
|
||||
return previous
|
||||
|
||||
|
||||
def _restore_env(key: str, previous) -> None:
|
||||
if previous is None:
|
||||
os.environ.pop(key, None)
|
||||
else:
|
||||
os.environ[key] = previous
|
||||
|
||||
|
||||
def _cleanup_hf_local_dir(local_dir: str, archive_file: str) -> None:
|
||||
if "/" in archive_file:
|
||||
top_level = archive_file.split("/", 1)[0]
|
||||
shutil.rmtree(os.path.join(local_dir, top_level), ignore_errors=True)
|
||||
shutil.rmtree(os.path.join(local_dir, ".cache"), ignore_errors=True)
|
||||
|
||||
|
||||
def ensure_hf_hub(venv_path: str) -> None:
|
||||
"""Guarantee the accelerated Hugging Face client is importable before the
|
||||
download so the multi-GB bundle transfer takes the fast Xet path.
|
||||
|
||||
The installer runs under the on-disk venv (PYTHON_VENV_PATH, i.e.
|
||||
/data/ai/venv in Docker). That venv is normally seeded from the image's
|
||||
/opt/venv, which bakes huggingface-hub[hf_xet]. But an install whose venv
|
||||
predates the base package (an upgrade where the reseed stamp didn't move, a
|
||||
hand-copied or offline-imported venv) would import-fail in
|
||||
download_with_hf_hub and silently fall back to the slow single-stream urllib
|
||||
downloader. Self-heal by pip-installing the client into this same venv.
|
||||
|
||||
A bundle install already requires network and lifts the offline guard (see
|
||||
main()), so this adds no new offline dependency; if the pip install fails we
|
||||
fall through to the resumable urllib downloader, the correct degraded path.
|
||||
"""
|
||||
try:
|
||||
import huggingface_hub # noqa: F401
|
||||
|
||||
return
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
python_path = os.path.join(venv_path, "bin", "python3")
|
||||
if not os.path.exists(python_path):
|
||||
return
|
||||
|
||||
emit_progress(1, "Preparing accelerated download client...")
|
||||
try:
|
||||
subprocess.run(
|
||||
[
|
||||
python_path, "-m", "pip", "install", "--quiet",
|
||||
"huggingface-hub[hf_xet,hf_transfer]==0.36.2",
|
||||
],
|
||||
capture_output=True, text=True, timeout=300, check=True,
|
||||
)
|
||||
# The finder caches the venv's site-packages listing; drop it so the
|
||||
# just-installed package is visible to the import in download_with_hf_hub.
|
||||
importlib.invalidate_caches()
|
||||
except Exception as e:
|
||||
emit_progress(1, f"Accelerated client unavailable ({e}); using resumable download.")
|
||||
|
||||
|
||||
def download_with_hf_hub(
|
||||
bundle_repo: str,
|
||||
archive_file: str,
|
||||
dest: str,
|
||||
expected_size: int,
|
||||
progress_start: int,
|
||||
progress_end: int,
|
||||
force_download: bool = False,
|
||||
) -> bool:
|
||||
"""Download through huggingface_hub when available.
|
||||
|
||||
huggingface_hub 0.32+ can use hf_xet for faster large-file transfers and
|
||||
manages retries/resume internally. Return False when the client is missing
|
||||
or fails so callers can fall back to the manual urllib downloader.
|
||||
"""
|
||||
try:
|
||||
from huggingface_hub import hf_hub_download
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
# Enable hf_transfer (Rust multi-connection downloader) ONLY when the
|
||||
# package is actually importable. For a Xet-backed repo hf_xet takes
|
||||
# precedence and this is a no-op, but if the Xet CAS endpoint is unreachable
|
||||
# (e.g. a firewall that allows huggingface.co but blocks transfer.xethub.hf.co)
|
||||
# hf_hub_download falls back to plain HTTP, and hf_transfer makes that
|
||||
# fallback multi-connection instead of single-stream. Gating on the import
|
||||
# avoids the "HF_HUB_ENABLE_HF_TRANSFER set but package missing" hard error
|
||||
# on a venv that only has hf_xet.
|
||||
try:
|
||||
import hf_transfer # noqa: F401
|
||||
|
||||
os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
local_dir = os.path.dirname(dest)
|
||||
os.makedirs(local_dir, exist_ok=True)
|
||||
emit_progress(progress_start, "Downloading with accelerated Hugging Face client...")
|
||||
|
||||
previous_progress = _set_env_temporarily("HF_HUB_DISABLE_PROGRESS_BARS", "1")
|
||||
try:
|
||||
downloaded_path = hf_hub_download(
|
||||
repo_id=bundle_repo,
|
||||
filename=archive_file,
|
||||
repo_type="model",
|
||||
local_dir=local_dir,
|
||||
force_download=force_download,
|
||||
)
|
||||
except Exception as e:
|
||||
emit_progress(
|
||||
progress_start,
|
||||
f"Accelerated download unavailable, using resumable fallback: {e}",
|
||||
)
|
||||
# Reclaim any partial blob/metadata hf_hub_download staged under
|
||||
# local_dir/.cache so the urllib fallback starts clean and disk is freed.
|
||||
_cleanup_hf_local_dir(local_dir, archive_file)
|
||||
return False
|
||||
finally:
|
||||
_restore_env("HF_HUB_DISABLE_PROGRESS_BARS", previous_progress)
|
||||
|
||||
try:
|
||||
if not os.path.exists(downloaded_path):
|
||||
emit_progress(
|
||||
progress_start,
|
||||
"Accelerated download did not produce an archive, using resumable fallback...",
|
||||
)
|
||||
return False
|
||||
|
||||
if os.path.abspath(downloaded_path) != os.path.abspath(dest):
|
||||
if os.path.exists(dest):
|
||||
os.unlink(dest)
|
||||
os.replace(downloaded_path, dest)
|
||||
|
||||
size = os.path.getsize(dest)
|
||||
if expected_size > 0:
|
||||
pct = min(size / expected_size, 1.0)
|
||||
progress = int(progress_start + pct * (progress_end - progress_start))
|
||||
else:
|
||||
progress = progress_end
|
||||
emit_progress(
|
||||
min(progress, progress_end),
|
||||
f"Downloaded with accelerated client ({size / (1024**3):.1f} GB)",
|
||||
)
|
||||
return True
|
||||
except Exception as e:
|
||||
emit_progress(
|
||||
progress_start,
|
||||
f"Accelerated download post-processing failed, using resumable fallback: {e}",
|
||||
)
|
||||
return False
|
||||
finally:
|
||||
# Always drop the hf staging tree (local_dir/<top>, local_dir/.cache).
|
||||
# On success the archive is already moved to dest; on any failure this
|
||||
# stops the transient hf cache copy from leaking across the fallback.
|
||||
_cleanup_hf_local_dir(local_dir, archive_file)
|
||||
|
||||
|
||||
def download_with_resume(
|
||||
url: str,
|
||||
dest: str,
|
||||
@@ -180,7 +344,15 @@ def download_with_resume(
|
||||
if bytes_downloaded == 0 and os.path.exists(partial_path):
|
||||
os.unlink(partial_path)
|
||||
|
||||
max_retries = 3
|
||||
def _cleanup_partial() -> None:
|
||||
for p in (partial_path, meta_path):
|
||||
if os.path.exists(p):
|
||||
try:
|
||||
os.unlink(p)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
max_retries = 5
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
headers = {"User-Agent": "snapotter-installer/2.0"}
|
||||
@@ -193,10 +365,19 @@ def download_with_resume(
|
||||
|
||||
req = urllib.request.Request(url, headers=headers)
|
||||
with urllib.request.urlopen(req, timeout=300) as resp:
|
||||
status = getattr(resp, "status", None) or resp.getcode()
|
||||
# If we asked to resume (sent a Range) but the server sent the
|
||||
# whole file back (200 instead of 206 Partial Content -- a proxy
|
||||
# or CDN that ignores Range), restart from byte 0. Appending a
|
||||
# full body onto the existing partial would corrupt the archive
|
||||
# and fail the checksum on every retry.
|
||||
if bytes_downloaded > 0 and status != 206:
|
||||
bytes_downloaded = 0
|
||||
mode = "ab" if bytes_downloaded > 0 else "wb"
|
||||
next_meta_at = bytes_downloaded + DOWNLOAD_META_BYTES
|
||||
with open(partial_path, mode) as f:
|
||||
while True:
|
||||
chunk = resp.read(65536)
|
||||
chunk = resp.read(DOWNLOAD_CHUNK_BYTES)
|
||||
if not chunk:
|
||||
break
|
||||
f.write(chunk)
|
||||
@@ -212,10 +393,20 @@ def download_with_resume(
|
||||
stage = f"Downloading... {bytes_downloaded / (1024**3):.1f} GB"
|
||||
emit_progress(progress, stage)
|
||||
|
||||
# Write meta periodically (every 10 MB)
|
||||
if bytes_downloaded % (10 * 1024 * 1024) < 65536:
|
||||
# Write meta periodically so a crash can resume.
|
||||
if bytes_downloaded >= next_meta_at:
|
||||
with open(meta_path, "w") as mf:
|
||||
json.dump({"bytesDownloaded": bytes_downloaded}, mf)
|
||||
next_meta_at = bytes_downloaded + DOWNLOAD_META_BYTES
|
||||
|
||||
# Guard against a truncated body or an error page served as the
|
||||
# archive: the completed size must match what the manifest expects.
|
||||
# A mismatch is retryable (transient truncation / a stale CDN edge).
|
||||
if expected_size > 0 and bytes_downloaded != expected_size:
|
||||
raise RuntimeError(
|
||||
f"incomplete download: got {bytes_downloaded} bytes, "
|
||||
f"expected {expected_size} (truncated response or error page)"
|
||||
)
|
||||
|
||||
# Download complete
|
||||
os.rename(partial_path, dest)
|
||||
@@ -223,27 +414,58 @@ def download_with_resume(
|
||||
os.unlink(meta_path)
|
||||
return
|
||||
|
||||
except Exception as e:
|
||||
# Write meta for resume on next attempt
|
||||
with open(meta_path, "w") as mf:
|
||||
json.dump({"bytesDownloaded": bytes_downloaded}, mf)
|
||||
|
||||
if attempt < max_retries - 1:
|
||||
delay = 10 * (2 ** attempt)
|
||||
emit_progress(
|
||||
progress_start,
|
||||
f"Download failed (attempt {attempt + 1}/{max_retries}), "
|
||||
f"retrying in {delay}s: {e}",
|
||||
)
|
||||
time.sleep(delay)
|
||||
else:
|
||||
# Clean up on final failure
|
||||
for p in (partial_path, meta_path):
|
||||
if os.path.exists(p):
|
||||
os.unlink(p)
|
||||
except urllib.error.HTTPError as e:
|
||||
# HTTPError subclasses OSError, so it MUST be caught before the
|
||||
# OSError clause below. 4xx (except 408 Timeout / 429 Too Many
|
||||
# Requests) won't fix on retry -- a wrong URL, a private repo, or a
|
||||
# removed archive -- so fail fast with the manual-download hint.
|
||||
if 400 <= e.code < 500 and e.code not in (408, 429):
|
||||
_cleanup_partial()
|
||||
raise RuntimeError(
|
||||
f"Failed to download after {max_retries} attempts: {e}"
|
||||
f"Download failed with HTTP {e.code} ({e.reason}). The archive "
|
||||
f"URL may be wrong or access-restricted."
|
||||
)
|
||||
_retry_or_raise(e, attempt, max_retries, bytes_downloaded, meta_path,
|
||||
progress_start, _cleanup_partial)
|
||||
except OSError as e:
|
||||
# Disk full is not transient: retrying can't create space. Fail fast
|
||||
# with an actionable message instead of burning the backoff budget.
|
||||
# (URLError/connection errors also land here; their errno is None, so
|
||||
# they fall through to the retry path.)
|
||||
if getattr(e, "errno", None) == errno.ENOSPC:
|
||||
_cleanup_partial()
|
||||
raise RuntimeError(
|
||||
"Ran out of disk space while downloading the bundle. "
|
||||
"Free up space and retry."
|
||||
)
|
||||
_retry_or_raise(e, attempt, max_retries, bytes_downloaded, meta_path,
|
||||
progress_start, _cleanup_partial)
|
||||
except Exception as e:
|
||||
_retry_or_raise(e, attempt, max_retries, bytes_downloaded, meta_path,
|
||||
progress_start, _cleanup_partial)
|
||||
|
||||
|
||||
def _retry_or_raise(err, attempt, max_retries, bytes_downloaded, meta_path,
|
||||
progress_start, cleanup) -> None:
|
||||
"""Shared transient-failure handler for download_with_resume: persist resume
|
||||
metadata and back off, or clean up and raise on the final attempt."""
|
||||
try:
|
||||
with open(meta_path, "w") as mf:
|
||||
json.dump({"bytesDownloaded": bytes_downloaded}, mf)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
if attempt < max_retries - 1:
|
||||
delay = min(60, 5 * (2 ** attempt))
|
||||
emit_progress(
|
||||
progress_start,
|
||||
f"Download failed (attempt {attempt + 1}/{max_retries}), "
|
||||
f"retrying in {delay}s: {err}",
|
||||
)
|
||||
time.sleep(delay)
|
||||
else:
|
||||
cleanup()
|
||||
raise RuntimeError(f"Failed to download after {max_retries} attempts: {err}")
|
||||
|
||||
|
||||
# -- Safe tar extraction --
|
||||
@@ -270,11 +492,15 @@ def safe_extract(tar_path: str, staging_dir: str) -> None:
|
||||
# -- File move --
|
||||
|
||||
def move_tree(src: str, dst: str) -> None:
|
||||
"""Merge src into dst, overwriting existing files. Renames entries where
|
||||
possible so that on the same filesystem no copy (and thus no transient
|
||||
doubling of the payload on disk) occurs; falls back to a copy only across
|
||||
filesystems. The old copytree+rmtree approach duplicated the whole tree on
|
||||
disk during the move, which could exhaust the host on a tight-disk node."""
|
||||
"""Merge src into dst, replacing entries crash-atomically where possible.
|
||||
|
||||
This writes into the SHARED /data/ai/venv site-packages, so a crash mid-move
|
||||
must never leave a package in a half-replaced state (that tears the venv and
|
||||
breaks every other AI tool). For a file replacing a file, os.replace swaps in
|
||||
place with NO delete-then-write window, so an interruption leaves either the
|
||||
old or the new file intact, never a missing one. Cross-filesystem copies go
|
||||
through a temp sibling then an atomic rename for the same reason. Renames
|
||||
(vs copytree) also avoid transiently doubling the payload on disk."""
|
||||
if not os.path.isdir(src):
|
||||
return
|
||||
os.makedirs(dst, exist_ok=True)
|
||||
@@ -285,22 +511,31 @@ def move_tree(src: str, dst: str) -> None:
|
||||
# Both dirs exist: merge recursively rather than replace.
|
||||
move_tree(s, d)
|
||||
continue
|
||||
if os.path.exists(d):
|
||||
if os.path.isdir(d):
|
||||
shutil.rmtree(d)
|
||||
else:
|
||||
os.remove(d)
|
||||
try:
|
||||
os.rename(s, d)
|
||||
except OSError as e:
|
||||
if getattr(e, "errno", None) == errno.EXDEV:
|
||||
# Cross-filesystem: rename isn't allowed, fall back to a copy.
|
||||
if os.path.isdir(s):
|
||||
shutil.copytree(s, d)
|
||||
# A type mismatch (dir<->file) can't be atomically swapped by rename,
|
||||
# so clear the destination first. A file-over-file or new entry needs
|
||||
# no pre-delete: os.replace is atomic and leaves no torn window.
|
||||
if os.path.exists(d) and os.path.isdir(d) != os.path.isdir(s):
|
||||
if os.path.isdir(d):
|
||||
shutil.rmtree(d)
|
||||
else:
|
||||
shutil.copy2(s, d)
|
||||
else:
|
||||
os.remove(d)
|
||||
os.replace(s, d)
|
||||
continue
|
||||
except OSError as e:
|
||||
if getattr(e, "errno", None) != errno.EXDEV:
|
||||
raise
|
||||
# Cross-filesystem: rename isn't allowed. Copy to a temp sibling and then
|
||||
# atomically replace, so a mid-copy ENOSPC never leaves a truncated file
|
||||
# where a working one used to be.
|
||||
if os.path.isdir(s):
|
||||
if os.path.exists(d):
|
||||
shutil.rmtree(d)
|
||||
shutil.copytree(s, d)
|
||||
else:
|
||||
tmp = d + ".part"
|
||||
shutil.copy2(s, tmp)
|
||||
os.replace(tmp, d)
|
||||
# Remove whatever remains of src (emptied by renames, or copied originals).
|
||||
shutil.rmtree(src, ignore_errors=True)
|
||||
|
||||
@@ -461,8 +696,20 @@ def _install() -> None:
|
||||
|
||||
emit_progress(2, f"Downloading {bundle.get('name', bundle_id)} bundle...")
|
||||
|
||||
# Make sure the accelerated Xet client is importable in this venv, so a
|
||||
# drifted/upgraded venv doesn't silently fall back to slow urllib.
|
||||
ensure_hf_hub(venv_path)
|
||||
|
||||
try:
|
||||
download_with_resume(url, tar_path, compressed_size, 2, 85)
|
||||
if not download_with_hf_hub(
|
||||
bundle_repo,
|
||||
archive_file,
|
||||
tar_path,
|
||||
compressed_size,
|
||||
2,
|
||||
85,
|
||||
):
|
||||
download_with_resume(url, tar_path, compressed_size, 2, 85)
|
||||
except RuntimeError as e:
|
||||
fail(
|
||||
f"{e}\n\n"
|
||||
@@ -478,7 +725,16 @@ def _install() -> None:
|
||||
os.unlink(tar_path)
|
||||
emit_progress(86, "Checksum mismatch, retrying download...")
|
||||
try:
|
||||
download_with_resume(url, tar_path, compressed_size, 2, 85)
|
||||
if not download_with_hf_hub(
|
||||
bundle_repo,
|
||||
archive_file,
|
||||
tar_path,
|
||||
compressed_size,
|
||||
2,
|
||||
85,
|
||||
force_download=True,
|
||||
):
|
||||
download_with_resume(url, tar_path, compressed_size, 2, 85)
|
||||
except RuntimeError as e:
|
||||
fail(str(e))
|
||||
|
||||
@@ -501,6 +757,11 @@ def _install() -> None:
|
||||
except Exception as e:
|
||||
if os.path.exists(staging_dir):
|
||||
shutil.rmtree(staging_dir, ignore_errors=True)
|
||||
if isinstance(e, OSError) and getattr(e, "errno", None) == errno.ENOSPC:
|
||||
fail(
|
||||
"Ran out of disk space while extracting the bundle. "
|
||||
"Free up space and retry."
|
||||
)
|
||||
fail(f"Failed to extract archive: {e}")
|
||||
|
||||
# -- Read bundle.json from tar --
|
||||
@@ -539,10 +800,28 @@ def _install() -> None:
|
||||
# -- Move site-packages --
|
||||
emit_progress(92, "Installing packages...")
|
||||
site_packages_dir = get_site_packages_dir(venv_path)
|
||||
venv_writing_marker = os.path.join(ai_dir, "venv.writing")
|
||||
|
||||
try:
|
||||
if os.path.isdir(staging_sp) and site_packages_dir:
|
||||
# Breadcrumb the destructive shared-venv write. If the process is
|
||||
# killed mid-move (OOM/SIGKILL/power loss), move_tree can leave the
|
||||
# venv torn, which breaks OTHER installed tools. The marker survives
|
||||
# the crash; on next boot recoverInterruptedInstalls sees it and
|
||||
# reseeds the venv back to a known-good base. We clear it the instant
|
||||
# the site-packages move completes, since the venv is consistent
|
||||
# again then (a later models-move failure can't tear the venv).
|
||||
with open(venv_writing_marker, "w") as mf:
|
||||
json.dump(
|
||||
{
|
||||
"bundleId": bundle_id,
|
||||
"startedAt": datetime.now(timezone.utc).isoformat(),
|
||||
},
|
||||
mf,
|
||||
)
|
||||
move_tree(staging_sp, site_packages_dir)
|
||||
if os.path.exists(venv_writing_marker):
|
||||
os.unlink(venv_writing_marker)
|
||||
|
||||
# -- Move models --
|
||||
emit_progress(95, "Installing models...")
|
||||
@@ -560,6 +839,39 @@ def _install() -> None:
|
||||
emit_progress(97, "Finalizing...")
|
||||
apply_fixups(staging_dir, venv_path)
|
||||
|
||||
# -- Verify the bundle actually imports --
|
||||
# File-copy completion does NOT prove the bundle works: an incomplete
|
||||
# extraction or an ABI mismatch (e.g. a numpy/torch/protobuf skew) can leave
|
||||
# every file present yet the module unimportable, so the tool "installs" but
|
||||
# fails at first use. Import the bundle's key native libraries in the venv
|
||||
# now; if that fails, refuse to mark the bundle installed so the user gets a
|
||||
# clear retry instead of a silently broken tool.
|
||||
smoke_imports = bundle.get("smokeImports") or []
|
||||
if smoke_imports and os.environ.get("SNAPOTTER_SKIP_INSTALL_SMOKE") != "1":
|
||||
emit_progress(99, "Verifying installation...")
|
||||
venv_python = os.path.join(venv_path, "bin", "python3")
|
||||
if os.path.exists(venv_python):
|
||||
import_stmt = "\n".join(f"import {mod}" for mod in smoke_imports)
|
||||
try:
|
||||
proc = subprocess.run(
|
||||
[venv_python, "-c", import_stmt],
|
||||
capture_output=True, text=True, timeout=300,
|
||||
)
|
||||
except subprocess.TimeoutExpired:
|
||||
shutil.rmtree(staging_dir, ignore_errors=True)
|
||||
fail("Installation verification timed out. Please retry the install.")
|
||||
if proc.returncode != 0:
|
||||
shutil.rmtree(staging_dir, ignore_errors=True)
|
||||
tail = "\n".join((proc.stderr or "").strip().splitlines()[-6:])
|
||||
fail(
|
||||
"Installation verification failed: the bundle installed but its "
|
||||
"libraries could not be loaded, so the tool would not work.\n"
|
||||
f"{tail}\n\n"
|
||||
"This usually means an interrupted or corrupted install. Retry the "
|
||||
"install; if it keeps failing, use Settings > AI Features > Reset AI "
|
||||
"Environment, then reinstall."
|
||||
)
|
||||
|
||||
# -- Write installed.json --
|
||||
emit_progress(98, "Recording installation...")
|
||||
installed = read_installed(ai_dir)
|
||||
|
||||
@@ -0,0 +1,151 @@
|
||||
import importlib.util
|
||||
import os
|
||||
import sys
|
||||
import types
|
||||
|
||||
|
||||
def load_installer():
|
||||
script_path = os.path.join(os.path.dirname(__file__), "..", "install_feature.py")
|
||||
spec = importlib.util.spec_from_file_location("install_feature_under_test", script_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def test_download_with_hf_hub_uses_accelerated_client(monkeypatch, tmp_path):
|
||||
installer = load_installer()
|
||||
downloaded = tmp_path / "hf-cache" / "bundle.tar.gz"
|
||||
downloaded.parent.mkdir()
|
||||
|
||||
calls = {}
|
||||
|
||||
def fake_hf_hub_download(**kwargs):
|
||||
calls.update(kwargs)
|
||||
downloaded.write_bytes(b"archive")
|
||||
return str(downloaded)
|
||||
|
||||
fake_module = types.ModuleType("huggingface_hub")
|
||||
fake_module.hf_hub_download = fake_hf_hub_download
|
||||
monkeypatch.setitem(sys.modules, "huggingface_hub", fake_module)
|
||||
|
||||
progress = []
|
||||
monkeypatch.setattr(installer, "emit_progress", lambda p, s: progress.append((p, s)))
|
||||
|
||||
dest = tmp_path / "staging" / "object-eraser-colorize-amd64-gpu.tar.gz"
|
||||
dest.parent.mkdir()
|
||||
|
||||
assert (
|
||||
installer.download_with_hf_hub(
|
||||
"snapotter/feature-bundles",
|
||||
"v2.0.0/object-eraser-colorize-amd64-gpu.tar.gz",
|
||||
str(dest),
|
||||
100,
|
||||
2,
|
||||
85,
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
assert dest.read_bytes() == b"archive"
|
||||
assert calls["repo_id"] == "snapotter/feature-bundles"
|
||||
assert calls["repo_type"] == "model"
|
||||
assert calls["filename"] == "v2.0.0/object-eraser-colorize-amd64-gpu.tar.gz"
|
||||
assert any("accelerated" in stage.lower() for _, stage in progress)
|
||||
|
||||
|
||||
def test_download_with_hf_hub_cleans_cache_when_download_raises(monkeypatch, tmp_path):
|
||||
"""A failed accelerated download must not leak its .cache staging tree onto
|
||||
disk before the urllib fallback runs."""
|
||||
installer = load_installer()
|
||||
staging = tmp_path / "staging"
|
||||
staging.mkdir()
|
||||
# Simulate a partial hf cache tree left behind by a failed transfer.
|
||||
leaked_cache = staging / ".cache" / "huggingface" / "download"
|
||||
leaked_cache.mkdir(parents=True)
|
||||
(leaked_cache / "blob.incomplete").write_bytes(b"partial")
|
||||
leaked_nested = staging / "v2.0.0"
|
||||
leaked_nested.mkdir()
|
||||
|
||||
def fake_hf_hub_download(**_kwargs):
|
||||
raise RuntimeError("xet CAS unreachable")
|
||||
|
||||
fake_module = types.ModuleType("huggingface_hub")
|
||||
fake_module.hf_hub_download = fake_hf_hub_download
|
||||
monkeypatch.setitem(sys.modules, "huggingface_hub", fake_module)
|
||||
monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
|
||||
|
||||
dest = staging / "object-eraser-colorize-amd64-gpu.tar.gz"
|
||||
|
||||
assert (
|
||||
installer.download_with_hf_hub(
|
||||
"deepsafe/feature-bundles",
|
||||
"v2.0.0/object-eraser-colorize-amd64-gpu.tar.gz",
|
||||
str(dest),
|
||||
100,
|
||||
2,
|
||||
85,
|
||||
)
|
||||
is False
|
||||
)
|
||||
# Both the .cache tree and the nested archive dir are reclaimed.
|
||||
assert not (staging / ".cache").exists()
|
||||
assert not (staging / "v2.0.0").exists()
|
||||
|
||||
|
||||
def test_ensure_hf_hub_noops_when_client_already_importable(monkeypatch, tmp_path):
|
||||
installer = load_installer()
|
||||
fake_module = types.ModuleType("huggingface_hub")
|
||||
monkeypatch.setitem(sys.modules, "huggingface_hub", fake_module)
|
||||
|
||||
ran = {"pip": False}
|
||||
monkeypatch.setattr(
|
||||
installer.subprocess, "run", lambda *a, **k: ran.__setitem__("pip", True)
|
||||
)
|
||||
monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
|
||||
|
||||
installer.ensure_hf_hub(str(tmp_path))
|
||||
assert ran["pip"] is False
|
||||
|
||||
|
||||
def test_ensure_hf_hub_self_heals_missing_client(monkeypatch, tmp_path):
|
||||
"""On a drifted venv where huggingface_hub is missing, ensure_hf_hub must
|
||||
pip-install it into that venv rather than let the caller fall back to the
|
||||
slow single-stream urllib downloader silently."""
|
||||
installer = load_installer()
|
||||
monkeypatch.delitem(sys.modules, "huggingface_hub", raising=False)
|
||||
|
||||
# Make the huggingface_hub import fail deterministically so ensure_hf_hub
|
||||
# takes its self-heal branch.
|
||||
import builtins
|
||||
|
||||
real_import = builtins.__import__
|
||||
|
||||
def blocked_import(name, *a, **k):
|
||||
if name == "huggingface_hub":
|
||||
raise ImportError("No module named 'huggingface_hub'")
|
||||
return real_import(name, *a, **k)
|
||||
|
||||
monkeypatch.setattr(builtins, "__import__", blocked_import)
|
||||
|
||||
venv = tmp_path / "venv"
|
||||
(venv / "bin").mkdir(parents=True)
|
||||
(venv / "bin" / "python3").write_text("")
|
||||
|
||||
pip_calls = []
|
||||
monkeypatch.setattr(
|
||||
installer.subprocess,
|
||||
"run",
|
||||
lambda cmd, **k: pip_calls.append(cmd) or types.SimpleNamespace(returncode=0),
|
||||
)
|
||||
monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
|
||||
|
||||
installer.ensure_hf_hub(str(venv))
|
||||
|
||||
assert len(pip_calls) == 1
|
||||
cmd = pip_calls[0]
|
||||
assert cmd[0] == str(venv / "bin" / "python3")
|
||||
assert "install" in cmd
|
||||
spec = next(part for part in cmd if part.startswith("huggingface-hub["))
|
||||
assert "hf_xet" in spec
|
||||
assert "hf_transfer" in spec
|
||||
@@ -0,0 +1,134 @@
|
||||
"""Crash-atomicity tests for move_tree, which writes into the SHARED venv.
|
||||
|
||||
The invariant under test: a crash (or ENOSPC) partway through move_tree must
|
||||
never leave a destination entry missing. Because the venv is shared by every AI
|
||||
tool, a half-replaced package is what tears the venv and breaks unrelated tools.
|
||||
"""
|
||||
|
||||
import errno
|
||||
import importlib.util
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def load_installer():
|
||||
script_path = os.path.join(os.path.dirname(__file__), "..", "install_feature.py")
|
||||
spec = importlib.util.spec_from_file_location("install_feature_move_under_test", script_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def test_file_over_file_is_replaced_without_a_delete_window(monkeypatch, tmp_path):
|
||||
installer = load_installer()
|
||||
src = tmp_path / "src"
|
||||
dst = tmp_path / "dst"
|
||||
src.mkdir()
|
||||
dst.mkdir()
|
||||
(dst / "pkg.py").write_text("OLD")
|
||||
(src / "pkg.py").write_text("NEW")
|
||||
|
||||
removed = []
|
||||
real_remove = os.remove
|
||||
monkeypatch.setattr(installer.os, "remove", lambda p: removed.append(p) or real_remove(p))
|
||||
|
||||
installer.move_tree(str(src), str(dst))
|
||||
|
||||
assert (dst / "pkg.py").read_text() == "NEW"
|
||||
# The old file was atomically replaced, never deleted-then-rewritten.
|
||||
assert str(dst / "pkg.py") not in removed
|
||||
|
||||
|
||||
def test_crash_mid_move_leaves_every_dest_old_or_new_never_missing(monkeypatch, tmp_path):
|
||||
installer = load_installer()
|
||||
src = tmp_path / "src"
|
||||
dst = tmp_path / "dst"
|
||||
src.mkdir()
|
||||
dst.mkdir()
|
||||
for name in ("a.py", "b.py", "c.py"):
|
||||
(dst / name).write_text(f"OLD_{name}")
|
||||
(src / name).write_text(f"NEW_{name}")
|
||||
|
||||
real_replace = os.replace
|
||||
state = {"n": 0}
|
||||
|
||||
def flaky_replace(s, d):
|
||||
state["n"] += 1
|
||||
if state["n"] == 2:
|
||||
raise OSError("simulated crash mid-move")
|
||||
return real_replace(s, d)
|
||||
|
||||
monkeypatch.setattr(installer.os, "replace", flaky_replace)
|
||||
|
||||
with pytest.raises(OSError):
|
||||
installer.move_tree(str(src), str(dst))
|
||||
|
||||
# Regardless of listdir order, every destination file must still exist and
|
||||
# hold either its old or its new content -- never a torn/missing entry.
|
||||
for name in ("a.py", "b.py", "c.py"):
|
||||
assert (dst / name).exists()
|
||||
assert (dst / name).read_text() in (f"OLD_{name}", f"NEW_{name}")
|
||||
|
||||
|
||||
def test_merges_directories_and_overwrites_files(tmp_path):
|
||||
installer = load_installer()
|
||||
src = tmp_path / "src"
|
||||
dst = tmp_path / "dst"
|
||||
(src / "pkg").mkdir(parents=True)
|
||||
(dst / "pkg").mkdir(parents=True)
|
||||
(dst / "pkg" / "keep.py").write_text("KEEP")
|
||||
(src / "pkg" / "keep.py").write_text("UPDATED")
|
||||
(src / "pkg" / "new.py").write_text("NEW")
|
||||
|
||||
installer.move_tree(str(src), str(dst))
|
||||
|
||||
assert (dst / "pkg" / "keep.py").read_text() == "UPDATED"
|
||||
assert (dst / "pkg" / "new.py").read_text() == "NEW"
|
||||
|
||||
|
||||
def test_type_mismatch_dir_replaces_file(tmp_path):
|
||||
installer = load_installer()
|
||||
src = tmp_path / "src"
|
||||
dst = tmp_path / "dst"
|
||||
src.mkdir()
|
||||
dst.mkdir()
|
||||
(dst / "x").write_text("i-am-a-file")
|
||||
(src / "x").mkdir()
|
||||
(src / "x" / "inner.py").write_text("dir-content")
|
||||
|
||||
installer.move_tree(str(src), str(dst))
|
||||
|
||||
assert (dst / "x").is_dir()
|
||||
assert (dst / "x" / "inner.py").read_text() == "dir-content"
|
||||
|
||||
|
||||
def test_exdev_file_copies_via_temp_then_atomic_replace(monkeypatch, tmp_path):
|
||||
installer = load_installer()
|
||||
src = tmp_path / "src"
|
||||
dst = tmp_path / "dst"
|
||||
src.mkdir()
|
||||
dst.mkdir()
|
||||
(dst / "f.py").write_text("OLD")
|
||||
(src / "f.py").write_text("NEW")
|
||||
|
||||
real_replace = os.replace
|
||||
seen = {"exdev": False, "part": False}
|
||||
|
||||
def exdev_for_direct_move(s, d):
|
||||
if s == str(src / "f.py"):
|
||||
seen["exdev"] = True
|
||||
raise OSError(errno.EXDEV, "cross-device link")
|
||||
if s.endswith(".part"):
|
||||
seen["part"] = True
|
||||
return real_replace(s, d)
|
||||
|
||||
monkeypatch.setattr(installer.os, "replace", exdev_for_direct_move)
|
||||
|
||||
installer.move_tree(str(src), str(dst))
|
||||
|
||||
assert (dst / "f.py").read_text() == "NEW"
|
||||
assert seen["exdev"] and seen["part"]
|
||||
# No leftover temp file.
|
||||
assert not (dst / "f.py.part").exists()
|
||||
@@ -0,0 +1,238 @@
|
||||
"""Robustness tests for download_with_resume: the urllib fallback downloader.
|
||||
|
||||
These simulate the network/disk failure modes users hit (flaky connections,
|
||||
Range-ignoring proxies, truncated bodies, disk full, dead URLs) with no real
|
||||
network, and assert the downloader either self-recovers or fails fast with a
|
||||
clear, actionable error instead of corrupting the archive or hanging.
|
||||
"""
|
||||
|
||||
import builtins
|
||||
import errno
|
||||
import importlib.util
|
||||
import os
|
||||
import urllib.error
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def load_installer():
|
||||
script_path = os.path.join(os.path.dirname(__file__), "..", "install_feature.py")
|
||||
spec = importlib.util.spec_from_file_location("install_feature_resume_under_test", script_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
class FakeResp:
|
||||
"""Minimal stand-in for the urlopen response context manager."""
|
||||
|
||||
def __init__(self, data: bytes, status: int = 200):
|
||||
self._data = data
|
||||
self._pos = 0
|
||||
self.status = status
|
||||
|
||||
def read(self, n: int) -> bytes:
|
||||
chunk = self._data[self._pos : self._pos + n]
|
||||
self._pos += len(chunk)
|
||||
return chunk
|
||||
|
||||
def getcode(self) -> int:
|
||||
return self.status
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *_):
|
||||
return False
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def no_sleep(monkeypatch):
|
||||
"""Never actually sleep during backoff; record calls instead."""
|
||||
installer = load_installer()
|
||||
calls = []
|
||||
monkeypatch.setattr(installer.time, "sleep", lambda s: calls.append(s))
|
||||
return calls
|
||||
|
||||
|
||||
def _patch_urlopen(monkeypatch, installer, handler):
|
||||
"""handler(req, call_index) -> FakeResp or raises."""
|
||||
state = {"n": 0}
|
||||
|
||||
def fake_urlopen(req, timeout=None):
|
||||
idx = state["n"]
|
||||
state["n"] += 1
|
||||
return handler(req, idx)
|
||||
|
||||
monkeypatch.setattr(installer.urllib.request, "urlopen", fake_urlopen)
|
||||
return state
|
||||
|
||||
|
||||
def test_fresh_download_succeeds(monkeypatch, tmp_path):
|
||||
installer = load_installer()
|
||||
monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
|
||||
data = b"x" * 4096
|
||||
_patch_urlopen(monkeypatch, installer, lambda req, i: FakeResp(data, 200))
|
||||
|
||||
dest = tmp_path / "bundle.tar.gz"
|
||||
installer.download_with_resume("https://h/f", str(dest), len(data), 2, 85)
|
||||
|
||||
assert dest.read_bytes() == data
|
||||
assert not (tmp_path / "bundle.tar.gz.partial").exists()
|
||||
assert not (tmp_path / "bundle.tar.gz.meta").exists()
|
||||
|
||||
|
||||
def test_range_ignored_200_restarts_instead_of_corrupting(monkeypatch, tmp_path, no_sleep):
|
||||
"""A resumable partial exists, but the server ignores Range and returns 200
|
||||
with the full body. The downloader must restart (truncate) rather than
|
||||
append the full body onto the partial and corrupt the archive."""
|
||||
installer = load_installer()
|
||||
monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
|
||||
data = b"GOOD" * 1024
|
||||
|
||||
dest = tmp_path / "bundle.tar.gz"
|
||||
partial = tmp_path / "bundle.tar.gz.partial"
|
||||
meta = tmp_path / "bundle.tar.gz.meta"
|
||||
# A stale partial that would corrupt the file if appended to.
|
||||
partial.write_bytes(b"STALEPARTIAL")
|
||||
meta.write_text('{"bytesDownloaded": 12}')
|
||||
|
||||
# Server ignores the Range header and always sends the full body as 200.
|
||||
_patch_urlopen(monkeypatch, installer, lambda req, i: FakeResp(data, 200))
|
||||
|
||||
installer.download_with_resume("https://h/f", str(dest), len(data), 2, 85)
|
||||
|
||||
# Exactly the good bytes, not STALEPARTIAL + data.
|
||||
assert dest.read_bytes() == data
|
||||
assert no_sleep == [] # no retry needed; it restarted cleanly in one pass
|
||||
|
||||
|
||||
def test_proper_206_resume_appends(monkeypatch, tmp_path):
|
||||
"""When the server honors Range with 206, the partial is kept and only the
|
||||
remaining bytes are appended."""
|
||||
installer = load_installer()
|
||||
monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
|
||||
full = b"HEADER" + b"TAILDATA" * 512
|
||||
head_len = 6
|
||||
|
||||
dest = tmp_path / "bundle.tar.gz"
|
||||
partial = tmp_path / "bundle.tar.gz.partial"
|
||||
meta = tmp_path / "bundle.tar.gz.meta"
|
||||
partial.write_bytes(full[:head_len])
|
||||
meta.write_text(f'{{"bytesDownloaded": {head_len}}}')
|
||||
|
||||
def handler(req, i):
|
||||
assert req.get_header("Range") == f"bytes={head_len}-"
|
||||
return FakeResp(full[head_len:], 206)
|
||||
|
||||
_patch_urlopen(monkeypatch, installer, handler)
|
||||
installer.download_with_resume("https://h/f", str(dest), len(full), 2, 85)
|
||||
assert dest.read_bytes() == full
|
||||
|
||||
|
||||
def test_disk_full_fails_fast_without_retry(monkeypatch, tmp_path, no_sleep):
|
||||
installer = load_installer()
|
||||
monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
|
||||
data = b"y" * 4096
|
||||
dest = tmp_path / "bundle.tar.gz"
|
||||
partial_path = str(dest) + ".partial"
|
||||
|
||||
real_open = builtins.open
|
||||
|
||||
def fake_open(path, mode="r", *a, **k):
|
||||
if str(path) == partial_path and ("w" in mode or "a" in mode):
|
||||
real = real_open(path, mode, *a, **k)
|
||||
|
||||
class NoSpace:
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *_):
|
||||
real.close()
|
||||
return False
|
||||
|
||||
def write(self, _data):
|
||||
raise OSError(errno.ENOSPC, "No space left on device")
|
||||
|
||||
return NoSpace()
|
||||
return real_open(path, mode, *a, **k)
|
||||
|
||||
monkeypatch.setattr(builtins, "open", fake_open)
|
||||
_patch_urlopen(monkeypatch, installer, lambda req, i: FakeResp(data, 200))
|
||||
|
||||
with pytest.raises(RuntimeError, match="disk space"):
|
||||
installer.download_with_resume("https://h/f", str(dest), len(data), 2, 85)
|
||||
|
||||
assert no_sleep == [] # disk-full is not retried
|
||||
assert not os.path.exists(partial_path)
|
||||
|
||||
|
||||
def test_http_404_fails_fast_without_retry(monkeypatch, tmp_path, no_sleep):
|
||||
installer = load_installer()
|
||||
monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
|
||||
dest = tmp_path / "bundle.tar.gz"
|
||||
|
||||
def handler(req, i):
|
||||
raise urllib.error.HTTPError("https://h/f", 404, "Not Found", {}, None)
|
||||
|
||||
_patch_urlopen(monkeypatch, installer, handler)
|
||||
|
||||
with pytest.raises(RuntimeError, match="HTTP 404"):
|
||||
installer.download_with_resume("https://h/f", str(dest), 4096, 2, 85)
|
||||
|
||||
assert no_sleep == [] # a 404 won't fix on retry
|
||||
|
||||
|
||||
def test_429_is_retried(monkeypatch, tmp_path, no_sleep):
|
||||
installer = load_installer()
|
||||
monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
|
||||
data = b"z" * 2048
|
||||
dest = tmp_path / "bundle.tar.gz"
|
||||
|
||||
def handler(req, i):
|
||||
if i == 0:
|
||||
raise urllib.error.HTTPError("https://h/f", 429, "Too Many Requests", {}, None)
|
||||
return FakeResp(data, 200)
|
||||
|
||||
_patch_urlopen(monkeypatch, installer, handler)
|
||||
installer.download_with_resume("https://h/f", str(dest), len(data), 2, 85)
|
||||
|
||||
assert dest.read_bytes() == data
|
||||
assert len(no_sleep) == 1 # backed off once, then succeeded
|
||||
|
||||
|
||||
def test_truncated_body_is_retried_then_succeeds(monkeypatch, tmp_path, no_sleep):
|
||||
installer = load_installer()
|
||||
monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
|
||||
full = b"c" * 4096
|
||||
dest = tmp_path / "bundle.tar.gz"
|
||||
|
||||
def handler(req, i):
|
||||
if i == 0:
|
||||
return FakeResp(full[:-100], 200) # truncated / error page
|
||||
return FakeResp(full, 200)
|
||||
|
||||
_patch_urlopen(monkeypatch, installer, handler)
|
||||
installer.download_with_resume("https://h/f", str(dest), len(full), 2, 85)
|
||||
|
||||
assert dest.read_bytes() == full
|
||||
assert len(no_sleep) == 1
|
||||
|
||||
|
||||
def test_connection_error_retried_then_raises_after_max(monkeypatch, tmp_path, no_sleep):
|
||||
installer = load_installer()
|
||||
monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
|
||||
dest = tmp_path / "bundle.tar.gz"
|
||||
|
||||
def handler(req, i):
|
||||
raise urllib.error.URLError("connection reset by peer")
|
||||
|
||||
_patch_urlopen(monkeypatch, installer, handler)
|
||||
|
||||
with pytest.raises(RuntimeError, match="after 5 attempts"):
|
||||
installer.download_with_resume("https://h/f", str(dest), 4096, 2, 85)
|
||||
|
||||
assert len(no_sleep) == 4 # 5 attempts => 4 backoffs
|
||||
assert not os.path.exists(str(dest) + ".partial")
|
||||
assert not os.path.exists(str(dest) + ".meta")
|
||||
Reference in New Issue
Block a user