mirror of
https://github.com/unslothai/unsloth.git
synced 2026-08-24 16:23:51 +00:00
* fix(studio): harden training setup and model loading Scan every model load root before approving remote code, and pin third-party codec sources to verified revisions. Align dataset option validation across the UI and backend, preserve manual drafts, and include edits in training start identity. Require job-scoped stop requests and retain bounded early-cancel tombstones without unsafe eviction. * fix(studio): harden training lifecycle and audio loading Preserve early start cancellations with bounded tombstones and explicit capacity handling for concurrent requests. Pin and verify third-party audio sources and codec artifacts with safe archive extraction and offline cache migration. Keep automatic evaluation data separate from every split included in a combined training instruction. Keep training summaries and GPU progress state aligned with the active configuration through public feature exports. Make route and lifecycle tests deterministic by isolating Hub reachability and executor behavior. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(studio): correct completion masking and offline training Apply Alpaca response markers safely when completion-only training is enabled. Skip optional runtime dependency installs while Studio is offline. Limit causal-conv1d hooks to recognized model families without dropping supported architectures. Disable known-broken TileLang dispatch when offline repair is unavailable. Use cached GGUF and model-size metadata without offline Hub retries. * fix(studio): resolve causal conv kernels from model configs Detect causal-conv1d requirements from resolved model architectures before loading model code. Keep name matching as a fallback while excluding unrelated renamed checkpoints. * fix(studio): preserve registered training cancellation Allow registered start cancellations to reclaim the oldest expiring tombstone when unknown cancellation capacity is full. Preserve the hard capacity limit and 429 response for unregistered request IDs. * Fix reset job scoping, DAC fast path, scan target and subset splits for PR #8103 - /api/train/reset: an unscoped reset could force-terminate a run mid-cancel. The guard now refuses to touch a live run it cannot prove it owns. The field stays optional so pre-rework clients, which POST /reset with no body, keep working. - ensure_dac_speech_weights: install the download into the pinned destination, so later loads hit the fast path instead of re-downloading and re-hashing 295 MB under the install lock. - _requires_security_review_for_model: apply the same load_scan_target alias normalization the sibling remote-code check gained, else the Spark-TTS alias 404s and fails open to "no review needed". - _purge_package_bytecode: best effort. It runs without the install lock over a cache shared by the inference and training workers; 7 of 8 concurrent imports died on it. Also moved inside the try so a failure cannot strand the cache dir on sys.path, and snapshot sys.modules before the origin audit. - commitSubset: clear the backing splits too, else the render-phase draft sync reads the previous subset's split back into the boxes it just reset. - Fix two tests that fail on the branch: the DAC assertion pinned the old return value, and the causal-conv1d assertion matched call formatting. * Make the SSM runtime tests Windows aware ensure_ssm_runtime deliberately skips causal-conv1d on win32 (no prebuilt wheel), so the two install-order assertions only hold off Windows. Caught on a real windows-latest runner. * Update two frontend source contracts the branch moved - captureTrainingStartInputs now delegates to createTrainingStartInputIdentity, so the normalize/flags assertions belong against training-start-inputs.ts. - resetTraining takes a RequiredTrainingJobScope and always sends the body, which is stronger than the hasScope branch the contract pinned. Both fail on the branch today; caught by tests/studio, which the studio backend job does not cover. * Fix unscoped reset compat, pyc purge fail-open and DAC fallback for PR #8103 Corrects four things in my earlier commits on this branch. Unscoped /api/train/reset returned "superseded" (HTTP 200) for a live run. The pre-rework cancel dialog chains stopTrainingRun then a bodyless reset, so an older client read that 200 as success and cleared its UI while training kept running. Return "active" (409) instead: same answer a live run already gives, and one those clients already handle. It still never force-terminates, so a bodyless reset landing between current_job_id being set and _cancel_requested being cleared cannot kill the run that just started. The bytecode purge was made best-effort, but it is the only thing stopping a stale or planted .pyc shadowing a verified .py: the manifest skips __pycache__ and the origin audit reads __file__, which still names the .py. Tolerate only FileNotFoundError, the real concurrent-purge race, and let PermissionError fail the load again. The DAC fast path copies 295 MB inside the hub cache with only Timeout caught, so a full disk turned a hash-verified download into a hard failure. Fall back to the verified hub path on OSError. commitSubset cleared both splits, but setDatasetSubset already does that; the extra setDatasetSplit(null) only cost a runDatasetCheck against an assumed "train" split. * Restore the eval split reset and the cancelled-run dismiss for PR #8103 Two corrections to 4e5389c21. commitSubset: I removed setDatasetEvalSplit(null) because setDatasetSubset already nulls datasetEvalSplit. It does, but it never resets evalSteps, and setDatasetEvalSplit is not a plain setter: it zeroes evalSteps and runs streamingCompatiblePatch. Without it, changing the subset left evaluation armed with no split, which routes/training.py rejects with 422 once streaming is on, and which silently auto-detects an eval split otherwise. Restored. setDatasetSplit(null) stays out, since its only unique effect was a runDatasetCheck against an assumed "train". Unscoped reset: returning "active" for every live run was too broad. The pre-rework cancel dialog only dismisses after stopTrainingRun succeeded, so _cancel_requested is already set and clearing the UI is right; 409 there just wedged the overlay behind a "Training still active" toast. Now 409 only when no stop was requested, which is the stale-tab case the change was for. Still no force_terminate on an unscoped reset. Also covers the __pycache__ branch of the purge, which is the route a planted .pyc actually takes; the existing test only reached the top-level .pyc loop. * Keep live start cancellations and fall back on a full disk for PR #8103 Two fixes plus the regression tests the earlier lifecycle fixes shipped without. Cancelling the active start at tombstone capacity reclaimed a slot by deleting the soonest-expiring entry. Expired ones are already pruned a few lines above, so that entry was always live, and dropping it let a delayed /start spawn the job it had cancelled. Reserve capacity instead: only the owner of the active start reaches that branch and there is at most one, so the table lands at cap + 1 rather than forgetting a cancellation. Unregistered ids still hit the hard cap. ensure_dac_speech_weights migrated a pre-existing legacy file with an unguarded copy, so a hub cache that cannot absorb a second 295 MB copy failed weights that had already passed size and sha256. Same fallback the download branch below it already uses. Tests: renamed and local SSM checkpoints resolving from config rather than name, the owner of an active start staying cancellable at capacity, live cancellations surviving that cancel, and the full-disk legacy DAC fallback. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Opt the pinned source checkouts into Git long paths for PR #8103 Cross-OS CI caught 14 failures on windows-latest, all of them the pinned checkout dying with "error: unable to write file ...: Filename too long". Git for Windows still enforces MAX_PATH unless told otherwise, and the cache nests a 40-char revision, a staging dir and .git/objects under the studio home. A venv-inferred home already measures about 253 of the 260 characters, so a slightly longer user or install path fails on a normal Windows machine, not just under the deeper pytest tmp dir. Passed per invocation with -c so no user or system Git config is touched, and it is a no-op off Windows. * Bound pending cancels, reach legacy DAC weights, and delete read-only checkouts for PR #8103 Three fixes, two of them on my own previous commit. Moving the owner cancel from evict-oldest to a one-slot overshoot also changed what the hardcoded reclaim_capacity=True on the pending non-owner branch did: it used to evict, so the table stayed at the cap, and it started overshooting instead. Start plus cancel could then be repeated to grow it without bound (1224 entries against a cap of 1024). That branch now takes the plain reservation, so the extra slot belongs to the owner of the active run alone. The DAC legacy fallback sat behind destination.parent.mkdir() and the install lock, both of which need a writable cache, so a read-only or full hub cache raised before weights we can already verify were ever looked at. Fall back at both points, and only to an artifact that passes the same size and sha256 check. Cross-OS CI then caught replacing a pinned checkout failing on windows-latest with WinError 5: Git marks .git/objects read-only and Windows will not delete a read-only file, so any repair or revision change died there. Clear the attribute and retry, only when the path is genuinely not writable, so an open handle still surfaces. * Delete the cached checkout the Windows-safe way in the migration test for PR #8103 --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: danielhanchen <moonshotaisubstack@gmail.com>
1107 lines
44 KiB
Python
1107 lines
44 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import gzip
|
|
import hashlib
|
|
import importlib
|
|
import json
|
|
import os
|
|
import re
|
|
import shutil
|
|
import stat
|
|
import subprocess
|
|
import sys
|
|
import tarfile
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
import urllib.error
|
|
import urllib.request
|
|
import uuid
|
|
from dataclasses import dataclass
|
|
from pathlib import Path, PureWindowsPath
|
|
from types import ModuleType
|
|
|
|
from filelock import FileLock, Timeout
|
|
|
|
from utils.native_path_leases import child_env_without_native_path_secret
|
|
from utils.paths.storage_roots import cache_root
|
|
from utils.subprocess_compat import (
|
|
windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs,
|
|
)
|
|
|
|
|
|
@dataclass(frozen = True)
|
|
class PinnedSource:
|
|
name: str
|
|
package: str
|
|
repository: str
|
|
revision: str
|
|
required_files: tuple[str, ...]
|
|
omitted_files: tuple[str, ...] = ()
|
|
generated_files: tuple[tuple[str, str], ...] = ()
|
|
source_tree_digest: str | None = None
|
|
runtime_tree_digest: str | None = None
|
|
archive_url: str | None = None
|
|
|
|
|
|
SPARK_TTS_SOURCE = PinnedSource(
|
|
name = "Spark-TTS",
|
|
package = "sparktts",
|
|
repository = "https://github.com/SparkAudio/Spark-TTS",
|
|
revision = "2f1ea9082400547242641f5271b6f941c9f439d1",
|
|
required_files = (
|
|
"sparktts/models/audio_tokenizer.py",
|
|
"sparktts/utils/audio.py",
|
|
),
|
|
generated_files = (("sparktts/__init__.py", ""),),
|
|
source_tree_digest = "20ff9f4c9e380b89248b828e9f39ec14572c43ff4a8d87b76190dbb3214b1b27",
|
|
runtime_tree_digest = "f14510e491a87ab287910e1d3f80e6b3d1bcea91b7f91f3baa45a3181a6993ba",
|
|
archive_url = (
|
|
"https://github.com/SparkAudio/Spark-TTS/archive/"
|
|
"2f1ea9082400547242641f5271b6f941c9f439d1.tar.gz"
|
|
),
|
|
)
|
|
|
|
OUTETTS_SOURCE = PinnedSource(
|
|
name = "OuteTTS",
|
|
package = "outetts",
|
|
repository = "https://github.com/edwko/OuteTTS",
|
|
revision = "f5eac6e70d792844c6a6959d900a47af2c061a5b",
|
|
required_files = (
|
|
"outetts/models/config.py",
|
|
"outetts/utils/preprocessing.py",
|
|
"outetts/version/v3/audio_processor.py",
|
|
"outetts/version/v3/prompt_processor.py",
|
|
),
|
|
omitted_files = (
|
|
"outetts/interface.py",
|
|
"outetts/models/gguf_model.py",
|
|
),
|
|
generated_files = (("outetts/__init__.py", ""),),
|
|
source_tree_digest = "817299085cb018839d37bf43505c9a742188bdb0f6ead8e1ea19a8643f0bb49f",
|
|
runtime_tree_digest = "b9f878aeb2de4d3ab0a5b1f75a5d04f2a137e4369143099bb24f6d6a41301fab",
|
|
archive_url = (
|
|
"https://github.com/edwko/OuteTTS/archive/"
|
|
"f5eac6e70d792844c6a6959d900a47af2c061a5b.tar.gz"
|
|
),
|
|
)
|
|
|
|
_REVISION_PATTERN = re.compile(r"[0-9a-f]{40}")
|
|
_SHA256_PATTERN = re.compile(r"[0-9a-f]{64}")
|
|
_IMPORT_LOCK = threading.RLock()
|
|
|
|
_DAC_REPOSITORY = "ibm-research/DAC.speech.v1.0"
|
|
_DAC_REVISION = "1ea7f64cd0678415e2d8c32d67b190722cb9b149"
|
|
_DAC_FILENAME = "weights_24khz_1.5kbps_v1.0.pth"
|
|
_DAC_SIZE = 295731578
|
|
_DAC_SHA256 = "d77ca0b04df942ec64e6a7a162bcac093b1127700acdaec0079f40d32c4405fb"
|
|
|
|
_ARCHIVE_MAX_DOWNLOAD_BYTES = 32 * 1024 * 1024
|
|
_ARCHIVE_MAX_MEMBERS = 10_000
|
|
_ARCHIVE_MAX_UNCOMPRESSED_BYTES = 128 * 1024 * 1024
|
|
_ARCHIVE_MAX_TAR_BYTES = 160 * 1024 * 1024
|
|
_ARCHIVE_SOCKET_TIMEOUT_SECONDS = 15
|
|
_ARCHIVE_DOWNLOAD_DEADLINE_SECONDS = 300
|
|
|
|
# Git for Windows still enforces MAX_PATH (260) unless told otherwise, and the pinned cache
|
|
# nests a 40-char revision, a staging dir and .git/objects under the studio home; a
|
|
# venv-inferred home already reaches ~253 chars, so a slightly longer one fails the
|
|
# checkout with "Filename too long". Passed per-invocation so no user config is touched.
|
|
_GIT_LONG_PATHS = ["-c", "core.longpaths=true"]
|
|
|
|
|
|
def _git(arguments: list[str], *, source_name: str) -> subprocess.CompletedProcess:
|
|
env = child_env_without_native_path_secret()
|
|
env["GIT_TERMINAL_PROMPT"] = "0"
|
|
env["GIT_LFS_SKIP_SMUDGE"] = "1"
|
|
env["GIT_NO_REPLACE_OBJECTS"] = "1"
|
|
try:
|
|
return subprocess.run(
|
|
["git", *_GIT_LONG_PATHS, *arguments],
|
|
check = True,
|
|
capture_output = True,
|
|
text = True,
|
|
encoding = "utf-8",
|
|
errors = "replace",
|
|
timeout = 300,
|
|
env = env,
|
|
**_windows_hidden_subprocess_kwargs(),
|
|
)
|
|
except FileNotFoundError as error:
|
|
raise RuntimeError(f"Git is required to install the pinned {source_name} source") from error
|
|
except subprocess.TimeoutExpired as error:
|
|
raise RuntimeError(f"Timed out while installing the pinned {source_name} source") from error
|
|
except subprocess.CalledProcessError as error:
|
|
detail = (error.stderr or error.stdout or "").strip()
|
|
message = f"Could not install the pinned {source_name} source"
|
|
raise RuntimeError(f"{message}: {detail}" if detail else message) from error
|
|
|
|
|
|
def _git_bytes(
|
|
arguments: list[str], *, source_name: str, input_data: bytes
|
|
) -> subprocess.CompletedProcess:
|
|
env = child_env_without_native_path_secret()
|
|
env["GIT_TERMINAL_PROMPT"] = "0"
|
|
env["GIT_LFS_SKIP_SMUDGE"] = "1"
|
|
env["GIT_NO_REPLACE_OBJECTS"] = "1"
|
|
try:
|
|
return subprocess.run(
|
|
["git", *_GIT_LONG_PATHS, *arguments],
|
|
check = True,
|
|
capture_output = True,
|
|
input = input_data,
|
|
timeout = 300,
|
|
env = env,
|
|
**_windows_hidden_subprocess_kwargs(),
|
|
)
|
|
except FileNotFoundError as error:
|
|
raise RuntimeError(f"Git is required to install the pinned {source_name} source") from error
|
|
except subprocess.TimeoutExpired as error:
|
|
raise RuntimeError(f"Timed out while installing the pinned {source_name} source") from error
|
|
except subprocess.CalledProcessError as error:
|
|
detail = (error.stderr or b"").decode("utf-8", errors = "replace").strip()
|
|
message = f"Could not install the pinned {source_name} source"
|
|
raise RuntimeError(f"{message}: {detail}" if detail else message) from error
|
|
|
|
|
|
def _generated_cache_path(relative: str) -> bool:
|
|
normalized = relative.replace("\\", "/")
|
|
return "/__pycache__/" in f"/{normalized}" and normalized.endswith((".pyc", ".pyo"))
|
|
|
|
|
|
def _package_path_parts(relative: str, spec: PinnedSource, *, kind: str) -> tuple[str, ...]:
|
|
normalized = relative.replace("\\", "/")
|
|
parts = tuple(normalized.split("/"))
|
|
if (
|
|
normalized != relative
|
|
or not normalized
|
|
or normalized.startswith("/")
|
|
or any(part in ("", ".", "..") for part in parts)
|
|
or any(PureWindowsPath(part).drive for part in parts)
|
|
or parts[0] != spec.package
|
|
):
|
|
raise ValueError(f"Invalid {kind} path for {spec.name}: {relative}")
|
|
return parts
|
|
|
|
|
|
def _configured_package_paths(
|
|
relatives: tuple[str, ...], spec: PinnedSource, *, kind: str
|
|
) -> tuple[str, ...]:
|
|
validated = []
|
|
seen = set()
|
|
for relative in relatives:
|
|
_package_path_parts(relative, spec, kind = kind)
|
|
if relative in seen:
|
|
raise ValueError(f"Invalid {kind} path for {spec.name}: {relative}")
|
|
seen.add(relative)
|
|
validated.append(relative)
|
|
return tuple(validated)
|
|
|
|
|
|
def _generated_file_contents(spec: PinnedSource) -> dict[str, bytes]:
|
|
generated = {}
|
|
for relative, content in spec.generated_files:
|
|
_package_path_parts(relative, spec, kind = "generated")
|
|
if relative in generated:
|
|
raise ValueError(f"Invalid generated path for {spec.name}: {relative}")
|
|
generated[relative] = content.encode("utf-8")
|
|
return generated
|
|
|
|
|
|
def _tracked_package_blobs(checkout: Path, spec: PinnedSource) -> dict[str, str]:
|
|
output = _git(
|
|
[
|
|
"-C",
|
|
str(checkout),
|
|
"ls-tree",
|
|
"-r",
|
|
"-z",
|
|
spec.revision,
|
|
"--",
|
|
spec.package,
|
|
],
|
|
source_name = spec.name,
|
|
).stdout
|
|
blobs = {}
|
|
for record in (record for record in output.split("\0") if record):
|
|
metadata, separator, relative = record.partition("\t")
|
|
fields = metadata.split(" ")
|
|
if separator != "\t" or len(fields) != 3:
|
|
raise ValueError(f"Invalid tracked tree entry for {spec.name}")
|
|
mode, object_type, object_id = fields
|
|
_package_path_parts(relative, spec, kind = "tracked")
|
|
if (
|
|
mode not in ("100644", "100755")
|
|
or object_type != "blob"
|
|
or _REVISION_PATTERN.fullmatch(object_id) is None
|
|
or relative in blobs
|
|
):
|
|
raise ValueError(f"Invalid tracked tree entry for {spec.name}: {relative}")
|
|
blobs[relative] = object_id
|
|
return blobs
|
|
|
|
|
|
def _pinned_blob_digests(
|
|
checkout: Path, object_ids: tuple[str, ...], spec: PinnedSource
|
|
) -> dict[str, str]:
|
|
unique_object_ids = tuple(dict.fromkeys(object_ids))
|
|
if not unique_object_ids:
|
|
return {}
|
|
result = _git_bytes(
|
|
["-C", str(checkout), "cat-file", "--batch"],
|
|
source_name = spec.name,
|
|
input_data = "".join(f"{object_id}\n" for object_id in unique_object_ids).encode("ascii"),
|
|
).stdout
|
|
digests = {}
|
|
offset = 0
|
|
for expected_object_id in unique_object_ids:
|
|
header_end = result.find(b"\n", offset)
|
|
if header_end < 0:
|
|
raise ValueError(f"Invalid pinned blob data for {spec.name}")
|
|
fields = result[offset:header_end].split(b" ")
|
|
if len(fields) != 3:
|
|
raise ValueError(f"Invalid pinned blob data for {spec.name}")
|
|
object_id, object_type, size_value = fields
|
|
try:
|
|
size = int(size_value)
|
|
except ValueError as error:
|
|
raise ValueError(f"Invalid pinned blob data for {spec.name}") from error
|
|
content_start = header_end + 1
|
|
content_end = content_start + size
|
|
if (
|
|
object_id.decode("ascii", errors = "replace") != expected_object_id
|
|
or object_type != b"blob"
|
|
or size < 0
|
|
or content_end >= len(result)
|
|
or result[content_end : content_end + 1] != b"\n"
|
|
):
|
|
raise ValueError(f"Invalid pinned blob data for {spec.name}")
|
|
digests[expected_object_id] = hashlib.sha256(result[content_start:content_end]).hexdigest()
|
|
offset = content_end + 1
|
|
if offset != len(result):
|
|
raise ValueError(f"Invalid pinned blob data for {spec.name}")
|
|
return digests
|
|
|
|
|
|
def _package_file(root: Path, relative: str, spec: PinnedSource) -> Path:
|
|
parts = _package_path_parts(relative, spec, kind = "tracked")
|
|
path = root.joinpath(*parts)
|
|
current = root
|
|
for part in parts:
|
|
current = current / part
|
|
if current.is_symlink():
|
|
raise ValueError(f"Symlinks are not allowed in {spec.name} source")
|
|
if not path.is_file():
|
|
raise ValueError(f"Missing tracked file in {spec.name} source: {relative}")
|
|
return path
|
|
|
|
|
|
def _checkout_manifest(checkout: Path, spec: PinnedSource) -> dict[str, str]:
|
|
package_root = checkout / spec.package
|
|
if package_root.is_symlink() or not package_root.is_dir():
|
|
raise ValueError(f"Missing {spec.package} package")
|
|
omitted = _configured_package_paths(spec.omitted_files, spec, kind = "omitted")
|
|
excluded = set(omitted) | set(_generated_file_contents(spec))
|
|
tracked_blobs = _tracked_package_blobs(checkout, spec)
|
|
pinned_digests = _pinned_blob_digests(
|
|
checkout,
|
|
tuple(tracked_blobs.values()),
|
|
spec,
|
|
)
|
|
manifest = {}
|
|
for relative, object_id in tracked_blobs.items():
|
|
path = _package_file(checkout, relative, spec)
|
|
digest = hashlib.sha256(path.read_bytes()).hexdigest()
|
|
if digest != pinned_digests[object_id]:
|
|
raise ValueError(f"Tracked file does not match the pinned {spec.name} blob: {relative}")
|
|
if relative not in excluded:
|
|
manifest[relative] = digest
|
|
return manifest
|
|
|
|
|
|
def _manifest_digest(manifest: dict[str, str]) -> str:
|
|
payload = json.dumps(
|
|
manifest,
|
|
sort_keys = True,
|
|
separators = (",", ":"),
|
|
).encode("utf-8")
|
|
return hashlib.sha256(payload).hexdigest()
|
|
|
|
|
|
def _filesystem_source_manifest(source: Path, spec: PinnedSource) -> dict[str, str]:
|
|
if source.is_symlink() or not source.is_dir():
|
|
raise ValueError(f"Missing {spec.name} source")
|
|
package_root = source / spec.package
|
|
if package_root.is_symlink() or not package_root.is_dir():
|
|
raise ValueError(f"Missing {spec.package} package")
|
|
excluded = set(_configured_package_paths(spec.omitted_files, spec, kind = "omitted")) | set(
|
|
_generated_file_contents(spec)
|
|
)
|
|
manifest = {}
|
|
for path in sorted(package_root.rglob("*")):
|
|
relative = path.relative_to(source).as_posix()
|
|
if path.is_symlink():
|
|
raise ValueError(f"Symlinks are not allowed in {spec.name} source")
|
|
if path.is_dir() or _generated_cache_path(relative):
|
|
continue
|
|
if not path.is_file():
|
|
raise ValueError(f"Special files are not allowed in {spec.name} source")
|
|
if relative not in excluded:
|
|
manifest[relative] = hashlib.sha256(path.read_bytes()).hexdigest()
|
|
return manifest
|
|
|
|
|
|
def _sealed_source_manifest(source: Path, spec: PinnedSource) -> dict[str, str] | None:
|
|
if spec.source_tree_digest is None:
|
|
return None
|
|
try:
|
|
manifest = _filesystem_source_manifest(source, spec)
|
|
except (OSError, ValueError):
|
|
return None
|
|
if _manifest_digest(manifest) != spec.source_tree_digest:
|
|
return None
|
|
return manifest
|
|
|
|
|
|
def _runtime_manifest(runtime: Path, spec: PinnedSource) -> dict[str, str]:
|
|
package_root = runtime / spec.package
|
|
if package_root.is_symlink() or not package_root.is_dir():
|
|
raise ValueError(f"Missing {spec.package} package")
|
|
manifest = {}
|
|
for path in sorted(package_root.rglob("*")):
|
|
relative = path.relative_to(runtime).as_posix()
|
|
if path.is_symlink():
|
|
raise ValueError(f"Symlinks are not allowed in {spec.name} runtime source")
|
|
if path.is_dir() or _generated_cache_path(relative):
|
|
continue
|
|
if not path.is_file():
|
|
raise ValueError(f"Special files are not allowed in {spec.name} runtime source")
|
|
manifest[relative] = hashlib.sha256(path.read_bytes()).hexdigest()
|
|
return manifest
|
|
|
|
|
|
def _expected_runtime_manifest(checkout: Path, spec: PinnedSource) -> dict[str, str]:
|
|
manifest = _checkout_manifest(checkout, spec)
|
|
for relative, content in _generated_file_contents(spec).items():
|
|
manifest[relative] = hashlib.sha256(content).hexdigest()
|
|
return manifest
|
|
|
|
|
|
def _valid_checkout(path: Path, spec: PinnedSource) -> bool:
|
|
if path.is_symlink() or not path.is_dir():
|
|
return False
|
|
try:
|
|
required_files = _configured_package_paths(
|
|
spec.required_files,
|
|
spec,
|
|
kind = "required",
|
|
)
|
|
for relative in required_files:
|
|
required = path.joinpath(*_package_path_parts(relative, spec, kind = "required"))
|
|
if required.is_symlink() or not required.is_file():
|
|
return False
|
|
head = (
|
|
_git(
|
|
["-C", str(path), "rev-parse", "HEAD"],
|
|
source_name = spec.name,
|
|
)
|
|
.stdout.strip()
|
|
.lower()
|
|
)
|
|
branch = _git(
|
|
["-C", str(path), "rev-parse", "--abbrev-ref", "HEAD"],
|
|
source_name = spec.name,
|
|
).stdout.strip()
|
|
origin = _git(
|
|
["-C", str(path), "remote", "get-url", "origin"],
|
|
source_name = spec.name,
|
|
).stdout.strip()
|
|
status = _git(
|
|
["-C", str(path), "status", "--porcelain=v1", "--untracked-files=all"],
|
|
source_name = spec.name,
|
|
).stdout
|
|
ignored = _git(
|
|
["-C", str(path), "ls-files", "--others", "--ignored", "--exclude-standard", "-z"],
|
|
source_name = spec.name,
|
|
).stdout
|
|
_checkout_manifest(path, spec)
|
|
except (OSError, RuntimeError, ValueError):
|
|
return False
|
|
return (
|
|
head == spec.revision
|
|
and branch == "HEAD"
|
|
and origin.rstrip("/").removesuffix(".git")
|
|
== spec.repository.rstrip("/").removesuffix(".git")
|
|
and not status
|
|
and not ignored
|
|
)
|
|
|
|
|
|
def _clear_read_only(function, path, _error) -> None:
|
|
# Git marks .git/objects read-only, and Windows refuses to delete a read-only file, so
|
|
# replacing a checkout dies with WinError 5. Only retry when the path really is not
|
|
# writable: an open handle (WinError 32) must still surface rather than spin.
|
|
if os.access(path, os.W_OK):
|
|
raise
|
|
os.chmod(path, os.stat(path).st_mode | stat.S_IWRITE)
|
|
function(path)
|
|
|
|
|
|
def _remove_owned_path(path: Path) -> None:
|
|
if path.is_symlink() or path.is_file():
|
|
path.unlink(missing_ok = True)
|
|
elif path.is_dir():
|
|
# onexc replaced onerror in 3.12; the handler signature is the same either way.
|
|
handler = (
|
|
{"onexc": _clear_read_only}
|
|
if sys.version_info >= (3, 12)
|
|
else {"onerror": _clear_read_only}
|
|
)
|
|
shutil.rmtree(path, **handler)
|
|
|
|
|
|
def _replace_owned_directory(staging: Path, destination: Path) -> None:
|
|
displaced = None
|
|
if destination.exists() or destination.is_symlink():
|
|
displaced = destination.with_name(f".{destination.name}.invalid-{uuid.uuid4().hex}")
|
|
os.replace(destination, displaced)
|
|
try:
|
|
os.replace(staging, destination)
|
|
except Exception:
|
|
if displaced is not None and not destination.exists():
|
|
os.replace(displaced, destination)
|
|
displaced = None
|
|
raise
|
|
finally:
|
|
if displaced is not None:
|
|
_remove_owned_path(displaced)
|
|
|
|
|
|
def _install_checkout(destination: Path, spec: PinnedSource) -> None:
|
|
workspace = Path(tempfile.mkdtemp(prefix = ".source-", dir = destination.parent))
|
|
checkout = workspace / "checkout"
|
|
hooks = workspace / "hooks"
|
|
hooks.mkdir()
|
|
hook_config = f"core.hooksPath={hooks}"
|
|
try:
|
|
_git(["init", "--quiet", str(checkout)], source_name = spec.name)
|
|
_git(
|
|
["-C", str(checkout), "config", "core.autocrlf", "false"],
|
|
source_name = spec.name,
|
|
)
|
|
_git(
|
|
["-C", str(checkout), "remote", "add", "origin", spec.repository],
|
|
source_name = spec.name,
|
|
)
|
|
_git(
|
|
[
|
|
"-c",
|
|
hook_config,
|
|
"-C",
|
|
str(checkout),
|
|
"fetch",
|
|
"--quiet",
|
|
"--depth=1",
|
|
"--no-tags",
|
|
"origin",
|
|
spec.revision,
|
|
],
|
|
source_name = spec.name,
|
|
)
|
|
fetched = (
|
|
_git(
|
|
["-C", str(checkout), "rev-parse", "FETCH_HEAD^{commit}"],
|
|
source_name = spec.name,
|
|
)
|
|
.stdout.strip()
|
|
.lower()
|
|
)
|
|
if fetched != spec.revision:
|
|
raise RuntimeError(f"{spec.name} returned a different revision than the pinned source")
|
|
_git(
|
|
[
|
|
"-c",
|
|
hook_config,
|
|
"-C",
|
|
str(checkout),
|
|
"checkout",
|
|
"--quiet",
|
|
"--detach",
|
|
spec.revision,
|
|
],
|
|
source_name = spec.name,
|
|
)
|
|
if not _valid_checkout(checkout, spec):
|
|
raise RuntimeError(f"The downloaded {spec.name} source failed integrity validation")
|
|
_replace_owned_directory(checkout, destination)
|
|
finally:
|
|
_remove_owned_path(workspace)
|
|
|
|
|
|
def _archive_root_name(spec: PinnedSource) -> str:
|
|
repository_name = spec.repository.rstrip("/").rsplit("/", 1)[-1].removesuffix(".git")
|
|
if not repository_name:
|
|
raise RuntimeError(f"Invalid pinned {spec.name} repository")
|
|
return f"{repository_name}-{spec.revision}"
|
|
|
|
|
|
def _download_archive(url: str, destination: Path, spec: PinnedSource) -> None:
|
|
request = urllib.request.Request(url, headers = {"User-Agent": "Unsloth-Studio"})
|
|
deadline = time.monotonic() + _ARCHIVE_DOWNLOAD_DEADLINE_SECONDS
|
|
try:
|
|
if time.monotonic() >= deadline:
|
|
raise RuntimeError(f"Timed out downloading the pinned {spec.name} source archive")
|
|
with urllib.request.urlopen(
|
|
request,
|
|
timeout = _ARCHIVE_SOCKET_TIMEOUT_SECONDS,
|
|
) as response:
|
|
if time.monotonic() >= deadline:
|
|
raise RuntimeError(f"Timed out downloading the pinned {spec.name} source archive")
|
|
content_length = response.headers.get("Content-Length")
|
|
if content_length is not None:
|
|
try:
|
|
advertised_size = int(content_length)
|
|
except ValueError as error:
|
|
raise RuntimeError(f"Invalid {spec.name} archive response size") from error
|
|
if advertised_size < 0 or advertised_size > _ARCHIVE_MAX_DOWNLOAD_BYTES:
|
|
raise RuntimeError(f"The pinned {spec.name} archive is too large")
|
|
total = 0
|
|
read_chunk = getattr(response, "read1", None)
|
|
if not callable(read_chunk):
|
|
read_chunk = response.read
|
|
with destination.open("wb") as handle:
|
|
while True:
|
|
if time.monotonic() >= deadline:
|
|
raise RuntimeError(
|
|
f"Timed out downloading the pinned {spec.name} source archive"
|
|
)
|
|
chunk = read_chunk(1024 * 1024)
|
|
if time.monotonic() >= deadline:
|
|
raise RuntimeError(
|
|
f"Timed out downloading the pinned {spec.name} source archive"
|
|
)
|
|
if not chunk:
|
|
break
|
|
total += len(chunk)
|
|
if total > _ARCHIVE_MAX_DOWNLOAD_BYTES:
|
|
raise RuntimeError(f"The pinned {spec.name} archive is too large")
|
|
handle.write(chunk)
|
|
except RuntimeError:
|
|
raise
|
|
except (OSError, urllib.error.URLError) as error:
|
|
raise RuntimeError(f"Could not download the pinned {spec.name} source archive") from error
|
|
|
|
|
|
def _archive_member_parts(member: tarfile.TarInfo, spec: PinnedSource) -> tuple[str, ...]:
|
|
name = member.name[:-1] if member.isdir() and member.name.endswith("/") else member.name
|
|
parts = tuple(name.split("/"))
|
|
if (
|
|
not name
|
|
or name.startswith("/")
|
|
or "\\" in name
|
|
or any(part in ("", ".", "..") for part in parts)
|
|
or any(PureWindowsPath(part).drive for part in parts)
|
|
or parts[0] != _archive_root_name(spec)
|
|
):
|
|
raise RuntimeError(f"Invalid path in the pinned {spec.name} source archive")
|
|
return parts
|
|
|
|
|
|
class _BoundedArchiveReader:
|
|
def __init__(self, handle, limit: int):
|
|
self._handle = handle
|
|
self._limit = limit
|
|
self._read = 0
|
|
|
|
def read(self, size: int = -1) -> bytes:
|
|
remaining = self._limit - self._read
|
|
requested = remaining + 1 if size < 0 else min(size, remaining + 1)
|
|
data = self._handle.read(requested)
|
|
self._read += len(data)
|
|
if self._read > self._limit:
|
|
raise RuntimeError("The pinned source archive expands too large")
|
|
return data
|
|
|
|
|
|
def _install_archive_source(destination: Path, spec: PinnedSource) -> None:
|
|
if spec.archive_url is None or spec.source_tree_digest is None:
|
|
raise RuntimeError(f"The pinned {spec.name} source archive is not configured")
|
|
workspace = Path(tempfile.mkdtemp(prefix = ".archive-", dir = destination.parent))
|
|
archive = workspace / "source.tar.gz"
|
|
staging = workspace / "source"
|
|
staging.mkdir()
|
|
try:
|
|
_download_archive(spec.archive_url, archive, spec)
|
|
member_count = 0
|
|
uncompressed_bytes = 0
|
|
extracted = set()
|
|
try:
|
|
with archive.open("rb") as compressed:
|
|
with gzip.GzipFile(fileobj = compressed, mode = "rb") as decompressed:
|
|
reader = _BoundedArchiveReader(decompressed, _ARCHIVE_MAX_TAR_BYTES)
|
|
with tarfile.open(fileobj = reader, mode = "r|") as bundle:
|
|
for member in bundle:
|
|
member_count += 1
|
|
if member_count > _ARCHIVE_MAX_MEMBERS:
|
|
raise RuntimeError(
|
|
f"The pinned {spec.name} archive has too many entries"
|
|
)
|
|
parts = _archive_member_parts(member, spec)
|
|
if member.isdir():
|
|
continue
|
|
if not member.isfile() or member.size < 0:
|
|
raise RuntimeError(
|
|
f"The pinned {spec.name} archive contains a non-regular file"
|
|
)
|
|
uncompressed_bytes += member.size
|
|
if uncompressed_bytes > _ARCHIVE_MAX_UNCOMPRESSED_BYTES:
|
|
raise RuntimeError(
|
|
f"The pinned {spec.name} archive expands too large"
|
|
)
|
|
if len(parts) < 3 or parts[1] != spec.package:
|
|
continue
|
|
relative = "/".join(parts[1:])
|
|
_package_path_parts(relative, spec, kind = "archive")
|
|
if relative in extracted:
|
|
raise RuntimeError(
|
|
f"The pinned {spec.name} archive contains duplicate files"
|
|
)
|
|
extracted.add(relative)
|
|
source_file = bundle.extractfile(member)
|
|
if source_file is None:
|
|
raise RuntimeError(
|
|
f"The pinned {spec.name} archive contains an unreadable file"
|
|
)
|
|
destination_file = staging.joinpath(*parts[1:])
|
|
destination_file.parent.mkdir(parents = True, exist_ok = True)
|
|
remaining = member.size
|
|
with source_file, destination_file.open("wb") as handle:
|
|
while remaining:
|
|
chunk = source_file.read(min(1024 * 1024, remaining))
|
|
if not chunk:
|
|
raise RuntimeError(
|
|
f"The pinned {spec.name} archive ended unexpectedly"
|
|
)
|
|
handle.write(chunk)
|
|
remaining -= len(chunk)
|
|
except (tarfile.TarError, EOFError, OSError) as error:
|
|
raise RuntimeError(f"The pinned {spec.name} source archive is invalid") from error
|
|
if _sealed_source_manifest(staging, spec) is None:
|
|
raise RuntimeError(f"The pinned {spec.name} source archive failed integrity validation")
|
|
_replace_owned_directory(staging, destination)
|
|
finally:
|
|
_remove_owned_path(workspace)
|
|
|
|
|
|
def _valid_runtime(
|
|
runtime: Path,
|
|
spec: PinnedSource,
|
|
checkout: Path | None = None,
|
|
) -> bool:
|
|
if runtime.is_symlink() or not runtime.is_dir():
|
|
return False
|
|
try:
|
|
required_files = _configured_package_paths(
|
|
spec.required_files,
|
|
spec,
|
|
kind = "required",
|
|
)
|
|
omitted_files = _configured_package_paths(
|
|
spec.omitted_files,
|
|
spec,
|
|
kind = "omitted",
|
|
)
|
|
top_level = {path.name for path in runtime.iterdir() if path.name != "__pycache__"}
|
|
if top_level != {spec.package}:
|
|
return False
|
|
for relative in required_files:
|
|
required = runtime.joinpath(*_package_path_parts(relative, spec, kind = "required"))
|
|
if required.is_symlink() or not required.is_file():
|
|
return False
|
|
for relative in omitted_files:
|
|
omitted = runtime.joinpath(*_package_path_parts(relative, spec, kind = "omitted"))
|
|
if omitted.exists() or omitted.is_symlink():
|
|
return False
|
|
for relative, content in _generated_file_contents(spec).items():
|
|
generated = runtime / relative
|
|
if generated.is_symlink() or not generated.is_file():
|
|
return False
|
|
if generated.read_bytes() != content:
|
|
return False
|
|
manifest = _runtime_manifest(runtime, spec)
|
|
if spec.runtime_tree_digest is not None:
|
|
return _manifest_digest(manifest) == spec.runtime_tree_digest
|
|
return checkout is not None and manifest == _expected_runtime_manifest(checkout, spec)
|
|
except (OSError, RuntimeError, ValueError):
|
|
return False
|
|
|
|
|
|
def _install_runtime(runtime: Path, checkout: Path, spec: PinnedSource) -> None:
|
|
workspace = Path(tempfile.mkdtemp(prefix = ".runtime-", dir = runtime.parent))
|
|
staging = workspace / "runtime"
|
|
staging.mkdir()
|
|
try:
|
|
if spec.source_tree_digest is not None:
|
|
source_manifest = _sealed_source_manifest(checkout, spec)
|
|
if source_manifest is None:
|
|
raise RuntimeError(f"The cached {spec.name} source failed integrity validation")
|
|
else:
|
|
source_manifest = _checkout_manifest(checkout, spec)
|
|
for relative, expected_digest in source_manifest.items():
|
|
source_file = checkout / relative
|
|
destination_file = staging / relative
|
|
destination_file.parent.mkdir(parents = True, exist_ok = True)
|
|
shutil.copy2(source_file, destination_file)
|
|
if hashlib.sha256(destination_file.read_bytes()).hexdigest() != expected_digest:
|
|
raise RuntimeError(f"{spec.name} source changed while preparing its runtime")
|
|
for relative, content in _generated_file_contents(spec).items():
|
|
destination_file = staging / relative
|
|
destination_file.parent.mkdir(parents = True, exist_ok = True)
|
|
destination_file.write_bytes(content)
|
|
if not _valid_runtime(staging, spec, checkout):
|
|
raise RuntimeError(f"The prepared {spec.name} runtime failed integrity validation")
|
|
_replace_owned_directory(staging, runtime)
|
|
finally:
|
|
_remove_owned_path(workspace)
|
|
|
|
|
|
def ensure_pinned_source(
|
|
spec: PinnedSource, *, legacy_sources: tuple[Path | str, ...] = ()
|
|
) -> Path:
|
|
revision = spec.revision.lower()
|
|
if _REVISION_PATTERN.fullmatch(revision) is None or revision != spec.revision:
|
|
raise RuntimeError(f"{spec.name} source revision must be a lowercase full Git commit")
|
|
for digest in (spec.source_tree_digest, spec.runtime_tree_digest):
|
|
if digest is not None and _SHA256_PATTERN.fullmatch(digest) is None:
|
|
raise RuntimeError(f"{spec.name} source digest must be a lowercase SHA-256")
|
|
if (spec.source_tree_digest is None) != (spec.runtime_tree_digest is None):
|
|
raise RuntimeError(f"{spec.name} source and runtime digests must be configured together")
|
|
|
|
parent = cache_root() / "third-party-sources" / spec.name
|
|
version_root = parent / revision
|
|
checkout = version_root / "source"
|
|
runtime = version_root / "runtime-v1"
|
|
if _valid_runtime(runtime, spec):
|
|
return runtime.resolve()
|
|
|
|
version_root.mkdir(parents = True, exist_ok = True)
|
|
try:
|
|
with FileLock(str(parent / ".install.lock"), timeout = 300):
|
|
if _valid_runtime(runtime, spec):
|
|
return runtime.resolve()
|
|
|
|
source = None
|
|
if spec.source_tree_digest is not None:
|
|
for candidate in (checkout, *(Path(value) for value in legacy_sources)):
|
|
if _sealed_source_manifest(candidate, spec) is not None:
|
|
source = candidate
|
|
break
|
|
elif _valid_checkout(checkout, spec):
|
|
source = checkout
|
|
|
|
if source is not None and _valid_runtime(runtime, spec, source):
|
|
return runtime.resolve()
|
|
|
|
if source is None:
|
|
from utils.utils import hf_env_offline
|
|
|
|
if hf_env_offline():
|
|
raise RuntimeError(
|
|
f"The pinned {spec.name} source is not cached and Studio is offline"
|
|
)
|
|
if spec.archive_url is not None:
|
|
_install_archive_source(checkout, spec)
|
|
else:
|
|
_install_checkout(checkout, spec)
|
|
source = checkout
|
|
_install_runtime(runtime, source, spec)
|
|
except Timeout as error:
|
|
raise RuntimeError(f"Timed out waiting for another {spec.name} installation") from error
|
|
|
|
if not _valid_runtime(runtime, spec, checkout):
|
|
raise RuntimeError(f"The installed {spec.name} source failed integrity validation")
|
|
return runtime.resolve()
|
|
|
|
|
|
def ensure_spark_tts_source(model_repo_path: Path | str | None = None) -> Path:
|
|
legacy_parent = Path(model_repo_path).parent if model_repo_path is not None else Path.cwd()
|
|
legacy_sources = (legacy_parent / "Spark-TTS",)
|
|
return ensure_pinned_source(SPARK_TTS_SOURCE, legacy_sources = legacy_sources)
|
|
|
|
|
|
def ensure_outetts_source() -> Path:
|
|
backend_root = Path(__file__).resolve().parents[1]
|
|
return ensure_pinned_source(
|
|
OUTETTS_SOURCE,
|
|
legacy_sources = (
|
|
backend_root / "core" / "inference" / "OuteTTS",
|
|
backend_root / "core" / "training" / "inference" / "OuteTTS",
|
|
),
|
|
)
|
|
|
|
|
|
def _artifact_matches(path: Path, *, expected_size: int, expected_sha256: str) -> bool:
|
|
try:
|
|
if not path.is_file() or path.stat().st_size != expected_size:
|
|
return False
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as handle:
|
|
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest() == expected_sha256
|
|
except OSError:
|
|
return False
|
|
|
|
|
|
def _install_verified_artifact(source: Path, destination: Path) -> None:
|
|
workspace = Path(tempfile.mkdtemp(prefix = ".artifact-", dir = destination.parent))
|
|
staging = workspace / destination.name
|
|
try:
|
|
shutil.copyfile(source, staging)
|
|
if not _artifact_matches(
|
|
staging,
|
|
expected_size = _DAC_SIZE,
|
|
expected_sha256 = _DAC_SHA256,
|
|
):
|
|
raise RuntimeError("The cached DAC speech weights changed during migration")
|
|
os.replace(staging, destination)
|
|
finally:
|
|
_remove_owned_path(workspace)
|
|
|
|
|
|
def _default_legacy_dac_weights_path() -> Path | None:
|
|
if sys.platform == "win32":
|
|
appdata = (os.environ.get("APPDATA") or "").strip()
|
|
if not appdata:
|
|
return None
|
|
return Path(appdata) / "outeai" / "dac" / _DAC_FILENAME
|
|
return Path.home() / ".cache" / "outeai" / "dac" / _DAC_FILENAME
|
|
|
|
|
|
def ensure_dac_speech_weights(legacy_path: Path | str | None = None) -> Path:
|
|
from huggingface_hub import hf_hub_download
|
|
from utils.hf_cache_settings import active_hf_hub_cache
|
|
from utils.utils import hf_env_offline
|
|
|
|
hub_cache = Path(active_hf_hub_cache())
|
|
destination = (
|
|
hub_cache
|
|
/ "studio-pinned-artifacts"
|
|
/ _DAC_REPOSITORY.replace("/", "--")
|
|
/ _DAC_REVISION
|
|
/ _DAC_FILENAME
|
|
)
|
|
if _artifact_matches(
|
|
destination,
|
|
expected_size = _DAC_SIZE,
|
|
expected_sha256 = _DAC_SHA256,
|
|
):
|
|
return destination.resolve()
|
|
|
|
def _verified_legacy() -> Path | None:
|
|
candidate = (
|
|
Path(legacy_path) if legacy_path is not None else _default_legacy_dac_weights_path()
|
|
)
|
|
if candidate is not None and _artifact_matches(
|
|
candidate,
|
|
expected_size = _DAC_SIZE,
|
|
expected_sha256 = _DAC_SHA256,
|
|
):
|
|
return candidate.resolve()
|
|
return None
|
|
|
|
try:
|
|
destination.parent.mkdir(parents = True, exist_ok = True)
|
|
except OSError:
|
|
# A read-only or full hub cache must not hide weights we can already verify.
|
|
fallback = _verified_legacy()
|
|
if fallback is None:
|
|
raise
|
|
return fallback
|
|
try:
|
|
with FileLock(str(destination.parent / ".install.lock"), timeout = 300):
|
|
if _artifact_matches(
|
|
destination,
|
|
expected_size = _DAC_SIZE,
|
|
expected_sha256 = _DAC_SHA256,
|
|
):
|
|
return destination.resolve()
|
|
|
|
legacy = (
|
|
Path(legacy_path) if legacy_path is not None else _default_legacy_dac_weights_path()
|
|
)
|
|
if legacy is not None and _artifact_matches(
|
|
legacy,
|
|
expected_size = _DAC_SIZE,
|
|
expected_sha256 = _DAC_SHA256,
|
|
):
|
|
# Same as the download branch below: the copy is an optimisation, so a full
|
|
# disk must not reject weights that already passed the size and sha256 check.
|
|
try:
|
|
_install_verified_artifact(legacy, destination)
|
|
except OSError:
|
|
return legacy.resolve()
|
|
return destination.resolve()
|
|
|
|
offline = hf_env_offline()
|
|
download_error = None
|
|
downloaded = None
|
|
try:
|
|
downloaded = Path(
|
|
hf_hub_download(
|
|
repo_id = _DAC_REPOSITORY,
|
|
filename = _DAC_FILENAME,
|
|
revision = _DAC_REVISION,
|
|
cache_dir = str(hub_cache),
|
|
local_files_only = offline,
|
|
)
|
|
)
|
|
except Exception as error:
|
|
download_error = error
|
|
|
|
if downloaded is not None and _artifact_matches(
|
|
downloaded,
|
|
expected_size = _DAC_SIZE,
|
|
expected_sha256 = _DAC_SHA256,
|
|
):
|
|
# Populate the pinned destination so later loads hit the fast path above
|
|
# instead of re-downloading and re-hashing 295 MB under the install lock.
|
|
# The copy is an optimisation, so a full disk must not fail a verified
|
|
# download; fall back to the hub path the caller used before.
|
|
try:
|
|
_install_verified_artifact(downloaded, destination)
|
|
except OSError:
|
|
return downloaded.resolve()
|
|
return destination.resolve()
|
|
|
|
if download_error is not None:
|
|
raise RuntimeError(
|
|
"The pinned DAC speech weights are unavailable in the active Hugging Face cache"
|
|
) from download_error
|
|
raise RuntimeError("The downloaded DAC speech weights failed integrity validation")
|
|
except OSError:
|
|
# Same reasoning as the mkdir above: taking the lock needs a writable cache, and
|
|
# verified weights we already hold are a better answer than failing the load.
|
|
fallback = _verified_legacy()
|
|
if fallback is None:
|
|
raise
|
|
return fallback
|
|
except Timeout as error:
|
|
raise RuntimeError("Timed out waiting for the DAC speech weights installation") from error
|
|
|
|
|
|
def _module_is_inside(module: ModuleType, package_root: Path) -> bool:
|
|
origins = []
|
|
origin = getattr(module, "__file__", None)
|
|
if origin:
|
|
origins.append(origin)
|
|
origins.extend(getattr(module, "__path__", ()) or ())
|
|
if not origins:
|
|
return False
|
|
for value in origins:
|
|
try:
|
|
if not Path(value).resolve().is_relative_to(package_root):
|
|
return False
|
|
except (OSError, ValueError):
|
|
return False
|
|
return True
|
|
|
|
|
|
def _purge_package_bytecode(package_root: Path) -> None:
|
|
# This is the only thing stopping a stale or planted .pyc from shadowing a verified .py:
|
|
# the manifest skips __pycache__ entirely, and the origin audit reads __file__, which
|
|
# still names the .py. So only the concurrent-purge race is tolerated (another worker
|
|
# deleting the same tree without the install lock); a PermissionError must stay fatal.
|
|
for directory, child_directories, files in os.walk(package_root, topdown = True):
|
|
directory_path = Path(directory)
|
|
for name in tuple(child_directories):
|
|
path = directory_path / name
|
|
if path.is_symlink():
|
|
child_directories.remove(name)
|
|
if name == "__pycache__":
|
|
with contextlib.suppress(FileNotFoundError):
|
|
path.unlink()
|
|
elif name == "__pycache__":
|
|
child_directories.remove(name)
|
|
with contextlib.suppress(FileNotFoundError):
|
|
shutil.rmtree(path)
|
|
for name in files:
|
|
if name.endswith((".pyc", ".pyo")):
|
|
(directory_path / name).unlink(missing_ok = True)
|
|
|
|
|
|
def _remove_package_modules(package: str) -> None:
|
|
for name in list(sys.modules):
|
|
if name == package or name.startswith(f"{package}."):
|
|
sys.modules.pop(name, None)
|
|
|
|
|
|
def import_pinned_module(module_name: str, *, package: str, source: Path | str) -> ModuleType:
|
|
if module_name != package and not module_name.startswith(f"{package}."):
|
|
raise ValueError(f"Only {package} modules can be imported from this pinned source")
|
|
source_root = Path(source).resolve()
|
|
unresolved_package_root = source_root / package
|
|
if unresolved_package_root.is_symlink() or not unresolved_package_root.is_dir():
|
|
raise RuntimeError(f"The pinned {package} package is missing")
|
|
package_root = unresolved_package_root.resolve()
|
|
package_init = package_root / "__init__.py"
|
|
if package_init.is_symlink() or not package_init.is_file():
|
|
raise RuntimeError(f"The pinned {package} package is not sealed")
|
|
|
|
with _IMPORT_LOCK:
|
|
for name, loaded_module in list(sys.modules.items()):
|
|
if name != package and not name.startswith(f"{package}."):
|
|
continue
|
|
if not _module_is_inside(loaded_module, package_root):
|
|
sys.modules.pop(name, None)
|
|
|
|
source_value = str(source_root)
|
|
while source_value in sys.path:
|
|
sys.path.remove(source_value)
|
|
sys.path.insert(0, source_value)
|
|
try:
|
|
# Inside the try: anything raising here would otherwise strand the cache dir at
|
|
# sys.path[0] for the process lifetime, with nothing imported and no rollback.
|
|
_purge_package_bytecode(package_root)
|
|
importlib.invalidate_caches()
|
|
module = importlib.import_module(module_name)
|
|
invalid_modules = sorted(
|
|
name
|
|
# Snapshot: another thread importing here would otherwise raise
|
|
# "dictionary changed size during iteration" out of a good codec load.
|
|
for name, loaded_module in list(sys.modules.items())
|
|
if (name == package or name.startswith(f"{package}."))
|
|
and not _module_is_inside(loaded_module, package_root)
|
|
)
|
|
if invalid_modules:
|
|
names = ", ".join(invalid_modules)
|
|
raise RuntimeError(
|
|
f"{package} loaded package modules from outside the pinned source: {names}"
|
|
)
|
|
return module
|
|
except BaseException:
|
|
while source_value in sys.path:
|
|
sys.path.remove(source_value)
|
|
_remove_package_modules(package)
|
|
raise
|
|
|
|
|
|
def deactivate_pinned_package(package: str, source: Path | str | None) -> None:
|
|
with _IMPORT_LOCK:
|
|
if source is not None:
|
|
source_value = str(Path(source).resolve())
|
|
while source_value in sys.path:
|
|
sys.path.remove(source_value)
|
|
_remove_package_modules(package)
|
|
|
|
|
|
def import_sparktts_module(module_name: str, source: Path | str) -> ModuleType:
|
|
return import_pinned_module(module_name, package = "sparktts", source = source)
|
|
|
|
|
|
def import_outetts_module(module_name: str, source: Path | str) -> ModuleType:
|
|
return import_pinned_module(module_name, package = "outetts", source = source)
|