mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
SnapOtterandGitHub
a731c3d1fe
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.
906 lines
35 KiB
Python
906 lines
35 KiB
Python
"""Pre-built AI bundle installer for SnapOtter.
|
|
|
|
Downloads a pre-built tar.gz archive (or uses a local file), verifies its
|
|
SHA256 checksum, extracts site-packages and models, and writes installed.json.
|
|
|
|
Invoked by the Node.js backend as a subprocess.
|
|
|
|
Usage:
|
|
python3 install_feature.py <bundleId> <manifestPath> <modelsDir>
|
|
|
|
Progress is reported via JSON lines on stderr (parsed by the Node bridge).
|
|
Final result is a JSON object on stdout.
|
|
"""
|
|
|
|
import errno
|
|
import glob
|
|
import hashlib
|
|
import importlib
|
|
import json
|
|
import os
|
|
import platform
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import tarfile
|
|
import time
|
|
import urllib.error
|
|
import urllib.request
|
|
from datetime import datetime, timezone
|
|
|
|
DOWNLOAD_CHUNK_BYTES = 4 * 1024 * 1024
|
|
DOWNLOAD_META_BYTES = 64 * 1024 * 1024
|
|
|
|
|
|
# -- Helpers --
|
|
|
|
def emit_progress(percent: int, stage: str) -> None:
|
|
"""Emit a progress update via stderr JSON line."""
|
|
sys.stderr.write(json.dumps({"progress": percent, "stage": stage}) + "\n")
|
|
sys.stderr.flush()
|
|
|
|
|
|
def fail(message: str) -> None:
|
|
"""Print error to stderr and exit non-zero."""
|
|
sys.stderr.write(json.dumps({"error": message}) + "\n")
|
|
sys.stderr.flush()
|
|
sys.exit(1)
|
|
|
|
|
|
# -- Architecture detection --
|
|
|
|
def detect_arch() -> str:
|
|
"""Return the bundle archive key for this host.
|
|
|
|
Only two archive variants are currently published to the bundle repo:
|
|
'amd64-gpu' and 'arm64-cpu' (see deepsafe/feature-bundles). There is no
|
|
CPU-only amd64 variant yet, so amd64 hosts always resolve to 'amd64-gpu'
|
|
even when no GPU is present: this downloads working CUDA-capable
|
|
packages, just larger than a CPU-only host strictly needs. Do not change
|
|
this to branch on GPU presence without first publishing an 'amd64-cpu'
|
|
archive for every bundle; requesting a key that doesn't exist in the
|
|
manifest fails the install outright (see the archives.get(arch) lookup
|
|
below), which would be worse than the current oversized-but-working
|
|
download.
|
|
"""
|
|
machine = platform.machine().lower()
|
|
if machine in ("aarch64", "arm64"):
|
|
return "arm64-cpu"
|
|
return "amd64-gpu"
|
|
|
|
|
|
# -- Disk space --
|
|
|
|
def _existing_ancestor(path: str) -> str:
|
|
"""Nearest existing ancestor of path (so disk_usage never raises on a
|
|
not-yet-created dir like the venv)."""
|
|
p = os.path.abspath(path)
|
|
while p and not os.path.exists(p):
|
|
parent = os.path.dirname(p)
|
|
if parent == p:
|
|
break
|
|
p = parent
|
|
return p or "/"
|
|
|
|
|
|
def check_disk_space(path: str, needed_bytes: int) -> None:
|
|
"""Fail if insufficient disk space on the filesystem holding path."""
|
|
usage = shutil.disk_usage(_existing_ancestor(path))
|
|
if usage.free < needed_bytes:
|
|
free_gb = usage.free / (1024 ** 3)
|
|
need_gb = needed_bytes / (1024 ** 3)
|
|
fail(
|
|
f"Insufficient disk space: need {need_gb:.1f} GB, "
|
|
f"have {free_gb:.1f} GB free. "
|
|
f"Free up space and retry."
|
|
)
|
|
|
|
|
|
def estimate_extracted(compressed: int, extracted: int) -> int:
|
|
"""Extracted-size estimate for the disk preflight. When the manifest omits
|
|
extractedSize (0), the budget would otherwise collapse to just the
|
|
compressed size and under-reserve for the extracted payload; fall back to a
|
|
conservative 3x of compressed (measured extracted/compressed ratios reach
|
|
~3x). This is only the early sanity bail; the accurate guard is the
|
|
real-on-disk re-check just before the destructive venv write."""
|
|
return extracted if extracted > 0 else compressed * 3
|
|
|
|
|
|
def dir_size(path: str) -> int:
|
|
"""Total size in bytes of all files under path (best-effort)."""
|
|
total = 0
|
|
for root, _dirs, files in os.walk(path):
|
|
for f in files:
|
|
try:
|
|
total += os.path.getsize(os.path.join(root, f))
|
|
except OSError:
|
|
pass
|
|
return total
|
|
|
|
|
|
def same_filesystem(a: str, b: str) -> bool:
|
|
"""True if paths a and b live on the same filesystem (so a rename between
|
|
them is a cheap metadata op rather than a full copy)."""
|
|
try:
|
|
return os.stat(_existing_ancestor(a)).st_dev == os.stat(_existing_ancestor(b)).st_dev
|
|
except OSError:
|
|
return False
|
|
|
|
|
|
# -- Venv site-packages discovery --
|
|
|
|
def get_site_packages_dir(venv_path: str) -> str:
|
|
"""Find the site-packages directory inside a Python venv."""
|
|
matches = glob.glob(os.path.join(venv_path, "lib", "python*", "site-packages"))
|
|
if matches:
|
|
return matches[0]
|
|
return ""
|
|
|
|
|
|
# -- SHA256 verification --
|
|
|
|
def verify_sha256(filepath: str, expected: str) -> bool:
|
|
"""Stream-hash a file and compare to expected hex digest."""
|
|
h = hashlib.sha256()
|
|
with open(filepath, "rb") as f:
|
|
while True:
|
|
chunk = f.read(8192)
|
|
if not chunk:
|
|
break
|
|
h.update(chunk)
|
|
return h.hexdigest() == expected
|
|
|
|
|
|
# -- 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,
|
|
expected_size: int,
|
|
progress_start: int,
|
|
progress_end: int,
|
|
) -> None:
|
|
"""Download a file with resume support via Range headers.
|
|
|
|
Uses .partial and .meta sidecar files for crash recovery.
|
|
"""
|
|
partial_path = dest + ".partial"
|
|
meta_path = dest + ".meta"
|
|
|
|
# Check for existing partial download
|
|
bytes_downloaded = 0
|
|
if os.path.exists(partial_path) and os.path.exists(meta_path):
|
|
try:
|
|
with open(meta_path, "r") as f:
|
|
meta = json.load(f)
|
|
bytes_downloaded = meta.get("bytesDownloaded", 0)
|
|
if bytes_downloaded > 0:
|
|
actual_size = os.path.getsize(partial_path)
|
|
if actual_size != bytes_downloaded:
|
|
bytes_downloaded = 0 # Mismatch, restart
|
|
except (json.JSONDecodeError, OSError):
|
|
bytes_downloaded = 0
|
|
|
|
if bytes_downloaded == 0 and os.path.exists(partial_path):
|
|
os.unlink(partial_path)
|
|
|
|
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"}
|
|
if bytes_downloaded > 0:
|
|
headers["Range"] = f"bytes={bytes_downloaded}-"
|
|
emit_progress(
|
|
progress_start,
|
|
f"Resuming download from {bytes_downloaded / (1024**3):.1f} GB...",
|
|
)
|
|
|
|
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(DOWNLOAD_CHUNK_BYTES)
|
|
if not chunk:
|
|
break
|
|
f.write(chunk)
|
|
bytes_downloaded += len(chunk)
|
|
|
|
# Update progress
|
|
if expected_size > 0:
|
|
pct = bytes_downloaded / expected_size
|
|
progress = int(
|
|
progress_start + pct * (progress_end - progress_start)
|
|
)
|
|
progress = min(progress, progress_end)
|
|
stage = f"Downloading... {bytes_downloaded / (1024**3):.1f} GB"
|
|
emit_progress(progress, stage)
|
|
|
|
# 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)
|
|
if os.path.exists(meta_path):
|
|
os.unlink(meta_path)
|
|
return
|
|
|
|
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"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 --
|
|
|
|
def safe_extract(tar_path: str, staging_dir: str) -> None:
|
|
"""Extract a tar.gz with security guards."""
|
|
os.makedirs(staging_dir, exist_ok=True)
|
|
with tarfile.open(tar_path, "r:gz") as tf:
|
|
for member in tf.getmembers():
|
|
# Block symlinks, hardlinks, devices
|
|
if not member.isfile() and not member.isdir():
|
|
raise RuntimeError(f"Blocked unsafe tar entry type: {member.name}")
|
|
# Block absolute paths and traversal
|
|
if member.name.startswith("/") or ".." in member.name.split("/"):
|
|
raise RuntimeError(f"Blocked unsafe tar path: {member.name}")
|
|
# The filter= kwarg was added in Python 3.12; the manual guards above
|
|
# already block unsafe entries on older interpreters (e.g. 3.11).
|
|
if sys.version_info >= (3, 12):
|
|
tf.extractall(staging_dir, filter="data")
|
|
else:
|
|
tf.extractall(staging_dir)
|
|
|
|
|
|
# -- File move --
|
|
|
|
def move_tree(src: str, dst: str) -> None:
|
|
"""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)
|
|
for name in os.listdir(src):
|
|
s = os.path.join(src, name)
|
|
d = os.path.join(dst, name)
|
|
if os.path.isdir(s) and os.path.isdir(d):
|
|
# Both dirs exist: merge recursively rather than replace.
|
|
move_tree(s, d)
|
|
continue
|
|
try:
|
|
# 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:
|
|
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)
|
|
|
|
|
|
# -- Fixups (NCCL wheel) --
|
|
|
|
def apply_fixups(staging_dir: str, venv_path: str) -> None:
|
|
"""Install any wheels from fixups/ directory (local only, no network)."""
|
|
fixups_dir = os.path.join(staging_dir, "fixups")
|
|
if not os.path.isdir(fixups_dir):
|
|
return
|
|
wheels = [f for f in os.listdir(fixups_dir) if f.endswith(".whl")]
|
|
if not wheels:
|
|
return
|
|
python_path = os.path.join(venv_path, "bin", "python3")
|
|
if not os.path.exists(python_path):
|
|
return
|
|
for wheel in wheels:
|
|
pkg_name = wheel.split("-")[0]
|
|
try:
|
|
subprocess.run(
|
|
[python_path, "-m", "pip", "install", "--no-index",
|
|
f"--find-links={fixups_dir}", pkg_name],
|
|
capture_output=True, text=True, timeout=60,
|
|
)
|
|
except Exception:
|
|
pass # Non-fatal
|
|
|
|
|
|
# -- installed.json management --
|
|
|
|
def read_installed(ai_dir: str) -> dict:
|
|
"""Read the current installed.json, returning empty structure if missing."""
|
|
path = os.path.join(ai_dir, "installed.json")
|
|
if not os.path.exists(path):
|
|
return {"bundles": {}}
|
|
try:
|
|
with open(path, "r") as f:
|
|
return json.load(f)
|
|
except (json.JSONDecodeError, OSError):
|
|
return {"bundles": {}}
|
|
|
|
|
|
def write_installed_atomic(ai_dir: str, data: dict) -> None:
|
|
"""Write installed.json atomically (write .tmp then rename)."""
|
|
path = os.path.join(ai_dir, "installed.json")
|
|
tmp_path = path + ".tmp"
|
|
with open(tmp_path, "w") as f:
|
|
json.dump(data, f, indent=2)
|
|
f.write("\n")
|
|
os.rename(tmp_path, path)
|
|
|
|
|
|
# -- Main --
|
|
|
|
def main() -> None:
|
|
"""Run the install with runtime-download restrictions lifted.
|
|
|
|
In strict offline mode (SNAPOTTER_ALLOW_MODEL_DOWNLOAD=0) the sidecar
|
|
runs with HF_HUB_OFFLINE/TRANSFORMERS_OFFLINE=1; a bundle install is an
|
|
explicitly user-initiated download, so those flags are lifted here
|
|
regardless. The previous values are restored in the finally block because
|
|
this script can run in-process inside the long-lived dispatcher, where
|
|
os.environ changes would otherwise leak into every later request.
|
|
"""
|
|
saved = {key: os.environ.get(key) for key in ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE")}
|
|
os.environ["HF_HUB_OFFLINE"] = "0"
|
|
os.environ["TRANSFORMERS_OFFLINE"] = "0"
|
|
try:
|
|
_install()
|
|
finally:
|
|
for key, value in saved.items():
|
|
if value is None:
|
|
os.environ.pop(key, None)
|
|
else:
|
|
os.environ[key] = value
|
|
|
|
|
|
def _install() -> None:
|
|
if len(sys.argv) < 4:
|
|
fail(
|
|
f"Usage: {sys.argv[0]} <bundleId> <manifestPath> <modelsDir>\n"
|
|
f"Got {len(sys.argv) - 1} argument(s)"
|
|
)
|
|
|
|
bundle_id = sys.argv[1]
|
|
manifest_path = sys.argv[2]
|
|
models_dir = sys.argv[3]
|
|
ai_dir = os.path.dirname(models_dir)
|
|
staging_base = os.path.join(ai_dir, "staging")
|
|
venv_path = os.environ.get("PYTHON_VENV_PATH", os.path.join(ai_dir, "venv"))
|
|
|
|
# -- Load manifest --
|
|
emit_progress(0, "Reading manifest...")
|
|
|
|
try:
|
|
with open(manifest_path, "r") as f:
|
|
manifest = json.load(f)
|
|
except Exception as e:
|
|
fail(f"Failed to read manifest: {e}")
|
|
|
|
bundles = manifest.get("bundles", {})
|
|
if bundle_id not in bundles:
|
|
fail(f"Unknown bundle: {bundle_id}")
|
|
|
|
bundle = bundles[bundle_id]
|
|
archives = bundle.get("archives")
|
|
if not archives:
|
|
fail(f"Bundle '{bundle_id}' has no archives in manifest (v2 required)")
|
|
|
|
# -- Detect architecture --
|
|
arch = detect_arch()
|
|
archive = archives.get(arch)
|
|
if not archive:
|
|
fail(f"No archive for architecture '{arch}' in bundle '{bundle_id}'")
|
|
|
|
archive_file = archive["file"]
|
|
expected_sha256 = archive["sha256"]
|
|
compressed_size = archive.get("compressedSize", 0)
|
|
extracted_size = archive.get("extractedSize", 0)
|
|
|
|
# -- Check for local file override (testing / offline) --
|
|
local_path = os.environ.get("SNAPOTTER_BUNDLE_LOCAL_PATH")
|
|
|
|
if local_path:
|
|
# Local mode: use the file directly, verify checksum
|
|
emit_progress(5, "Using local bundle archive...")
|
|
tar_path = local_path
|
|
|
|
if not os.path.exists(tar_path):
|
|
fail(f"Local bundle file not found: {tar_path}")
|
|
|
|
# Verify checksum
|
|
emit_progress(10, "Verifying checksum...")
|
|
if not verify_sha256(tar_path, expected_sha256):
|
|
fail(
|
|
f"SHA256 checksum mismatch for local file.\n"
|
|
f"Expected: {expected_sha256}\n"
|
|
f"This usually means the manifest and archive are out of sync."
|
|
)
|
|
else:
|
|
# Remote mode: download from HuggingFace
|
|
bundle_repo = manifest.get("bundleRepo", "deepsafe/feature-bundles")
|
|
url = f"https://huggingface.co/{bundle_repo}/resolve/main/{archive_file}"
|
|
|
|
# Disk space check (early sanity bail before a multi-GB download).
|
|
# estimate_extracted covers the extractedSize:0 case so the budget can't
|
|
# collapse to just the compressed size; the accurate guard is the
|
|
# real-on-disk re-check just before the destructive venv write below.
|
|
needed = compressed_size + estimate_extracted(compressed_size, extracted_size)
|
|
needed += 500 * 1024 * 1024 # 500 MB buffer
|
|
if needed > 0:
|
|
check_disk_space(ai_dir, needed)
|
|
|
|
# Download
|
|
os.makedirs(staging_base, exist_ok=True)
|
|
tar_path = os.path.join(staging_base, f"{bundle_id}-{arch}.tar.gz")
|
|
|
|
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:
|
|
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"
|
|
f"You can download the bundle manually from:\n"
|
|
f" {url}\n"
|
|
f"Then upload it via Settings > AI Features > Offline Import."
|
|
)
|
|
|
|
# Verify checksum
|
|
emit_progress(86, "Verifying integrity...")
|
|
if not verify_sha256(tar_path, expected_sha256):
|
|
# Delete and retry once from scratch
|
|
os.unlink(tar_path)
|
|
emit_progress(86, "Checksum mismatch, retrying download...")
|
|
try:
|
|
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))
|
|
|
|
if not verify_sha256(tar_path, expected_sha256):
|
|
os.unlink(tar_path)
|
|
fail(
|
|
f"SHA256 checksum mismatch after re-download.\n"
|
|
f"Expected: {expected_sha256}\n"
|
|
f"The archive may be corrupted. Try again later."
|
|
)
|
|
|
|
# -- Extract to staging --
|
|
staging_dir = os.path.join(ai_dir, f"staging-{bundle_id}")
|
|
emit_progress(88, "Extracting packages and models...")
|
|
|
|
try:
|
|
if os.path.exists(staging_dir):
|
|
shutil.rmtree(staging_dir)
|
|
safe_extract(tar_path, staging_dir)
|
|
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 --
|
|
bundle_json_path = os.path.join(staging_dir, "bundle.json")
|
|
if not os.path.exists(bundle_json_path):
|
|
shutil.rmtree(staging_dir, ignore_errors=True)
|
|
fail("Archive is missing bundle.json")
|
|
|
|
try:
|
|
with open(bundle_json_path, "r") as f:
|
|
bundle_meta = json.load(f)
|
|
except Exception as e:
|
|
shutil.rmtree(staging_dir, ignore_errors=True)
|
|
fail(f"Invalid bundle.json: {e}")
|
|
|
|
version = bundle_meta.get("version", manifest.get("imageVersion", "unknown"))
|
|
model_ids = bundle_meta.get("models", [])
|
|
|
|
# -- Disk re-check before the first destructive venv write --
|
|
# The upfront check ran before the download and used an estimate; now the
|
|
# payload is really on disk, so measure it and verify there's room to place
|
|
# it before we start writing into the venv. Running here (after the
|
|
# local/remote branches merge) also covers the offline-import path, which
|
|
# skipped the upfront check entirely. Each budget is checked against the
|
|
# filesystem the bytes actually land on: when the venv lives on a
|
|
# different filesystem than staging, the site-packages payload is COPIED
|
|
# onto the venv's disk, so that disk (not ai_dir's) must hold it. Models
|
|
# stay under ai_dir either way, moving by rename.
|
|
disk_floor = 1024 ** 3 # 1 GB for fixups / installed.json / slack
|
|
staging_sp = os.path.join(staging_dir, "site-packages")
|
|
if not same_filesystem(staging_dir, venv_path):
|
|
sp_bytes = dir_size(staging_sp) if os.path.isdir(staging_sp) else 0
|
|
check_disk_space(venv_path, sp_bytes + disk_floor)
|
|
check_disk_space(ai_dir, disk_floor)
|
|
|
|
# -- 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...")
|
|
staging_models = os.path.join(staging_dir, "models")
|
|
if os.path.isdir(staging_models):
|
|
os.makedirs(models_dir, exist_ok=True)
|
|
move_tree(staging_models, models_dir)
|
|
except OSError as e:
|
|
shutil.rmtree(staging_dir, ignore_errors=True)
|
|
if getattr(e, "errno", None) == errno.ENOSPC:
|
|
fail("Ran out of disk space while installing the bundle. Free up space and retry.")
|
|
fail(f"Failed to install bundle files: {e}")
|
|
|
|
# -- Apply fixups --
|
|
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)
|
|
installed["bundles"][bundle_id] = {
|
|
"version": version,
|
|
"installedAt": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
|
|
"models": model_ids,
|
|
}
|
|
write_installed_atomic(ai_dir, installed)
|
|
|
|
# -- Cleanup --
|
|
if os.path.exists(staging_dir):
|
|
shutil.rmtree(staging_dir, ignore_errors=True)
|
|
# Clean up downloaded tar (but not if local override)
|
|
if not local_path and os.path.exists(tar_path):
|
|
os.unlink(tar_path)
|
|
|
|
# -- Done --
|
|
emit_progress(100, "Complete")
|
|
|
|
result = {
|
|
"success": True,
|
|
"bundleId": bundle_id,
|
|
"version": version,
|
|
"models": model_ids,
|
|
}
|
|
print(json.dumps(result))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|