1054 lines
37 KiB
Python
1054 lines
37 KiB
Python
"""CPU-side Hugging Face snapshot staging for Modal LLM profiles."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
import math
|
|
import os
|
|
import shutil
|
|
import socket
|
|
import threading
|
|
import time
|
|
from contextlib import contextmanager
|
|
from dataclasses import asdict, dataclass, replace
|
|
from pathlib import Path
|
|
from typing import Any, Callable, Iterator, Mapping
|
|
from uuid import uuid4
|
|
|
|
if __package__:
|
|
from .llm_profiles import LLMModelProfile, get_llm_profile, load_llm_profiles
|
|
else: # pragma: no cover - the stable cloud entrypoint imports top-level modules.
|
|
from llm_profiles import LLMModelProfile, get_llm_profile, load_llm_profiles
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_COMPLETE_MARKER_FILENAME = ".comfy-modal-llm-complete.json"
|
|
_DEFAULT_LEASE_TIMEOUT_SECONDS = 7200.0
|
|
_LEASE_POLL_SECONDS = 2.0
|
|
_LEASE_HEARTBEAT_SECONDS = 30.0
|
|
_DEFAULT_LEASE_HEARTBEAT_STALE_SECONDS = 300.0
|
|
_DEFAULT_MINIMUM_FREE_DISK_BYTES = 8 * 1024**3
|
|
_SNAPSHOT_ALLOW_PATTERNS = (
|
|
"*.safetensors",
|
|
"*.safetensors.index.json",
|
|
"*.model",
|
|
"*.tiktoken",
|
|
"added_tokens.json",
|
|
"chat_template*",
|
|
"config.json",
|
|
"generation_config.json",
|
|
"hf_quant_config.json",
|
|
"merges.txt",
|
|
"preprocessor_config.json",
|
|
"processor*",
|
|
"special_tokens_map.json",
|
|
"tokenizer*",
|
|
"video_preprocessor_config.json",
|
|
"vocab*",
|
|
)
|
|
_TOKENIZER_ALLOW_PATTERNS = tuple(
|
|
pattern for pattern in _SNAPSHOT_ALLOW_PATTERNS if "safetensors" not in pattern
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class StagedModelSnapshot:
|
|
"""Describe a model snapshot made visible on the shared Modal Volume."""
|
|
|
|
profile_id: str
|
|
repository: str
|
|
revision: str
|
|
path: str
|
|
downloaded: bool
|
|
elapsed_seconds: float
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class LLMStagingProgress:
|
|
"""Describe one numeric Hugging Face snapshot-staging update."""
|
|
|
|
stage: str
|
|
message: str
|
|
value: float | None = None
|
|
maximum: float | None = None
|
|
unit: str | None = None
|
|
indeterminate: bool = False
|
|
model_reference: str | None = None
|
|
|
|
|
|
StagingProgressCallback = Callable[[LLMStagingProgress], None]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ResolvedStagedModelProfile:
|
|
"""Describe one resolved immutable profile and its staged snapshot."""
|
|
|
|
requested_reference: str
|
|
profile_id: str
|
|
repository: str
|
|
revision: str
|
|
backend: str
|
|
quantization_method: str
|
|
artifact_bytes: int
|
|
manifest_path: str | None
|
|
manifest_created: bool
|
|
security_scan_complete: bool
|
|
path: str
|
|
downloaded: bool
|
|
resolve_elapsed_seconds: float
|
|
elapsed_seconds: float
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
"""Return JSON-compatible staging metadata."""
|
|
return asdict(self)
|
|
|
|
|
|
def _snapshot_tqdm_class(
|
|
progress_callback: StagingProgressCallback,
|
|
) -> type[Any]:
|
|
"""Build a tqdm subclass that forwards Hugging Face file progress."""
|
|
from tqdm.auto import tqdm
|
|
|
|
class SnapshotProgressTqdm(tqdm):
|
|
"""Mirror snapshot_download's aggregate file bar to the caller."""
|
|
|
|
def __init__(self, *args: Any, **kwargs: Any) -> None:
|
|
"""Initialize the Hugging Face bar and report its initial state."""
|
|
super().__init__(*args, **kwargs)
|
|
self._report_progress()
|
|
|
|
def update(self, n: float | None = 1) -> bool | None:
|
|
"""Forward every aggregate file-count update."""
|
|
updated = super().update(n)
|
|
self._report_progress()
|
|
return updated
|
|
|
|
def _report_progress(self) -> None:
|
|
"""Publish the current tqdm count with compact user-facing copy."""
|
|
maximum = float(self.total) if self.total is not None else None
|
|
value = float(self.n)
|
|
description = str(
|
|
getattr(self, "desc", None) or "Downloading model snapshot"
|
|
).strip()
|
|
progress_callback(
|
|
LLMStagingProgress(
|
|
stage="download",
|
|
message=description,
|
|
value=value,
|
|
maximum=maximum,
|
|
unit=str(getattr(self, "unit", None) or "files"),
|
|
indeterminate=maximum is None,
|
|
)
|
|
)
|
|
|
|
return SnapshotProgressTqdm
|
|
|
|
|
|
def model_snapshot_path(storage_root: str | Path, profile: LLMModelProfile) -> Path:
|
|
"""Return the shared path, reusing a valid legacy profile path in place."""
|
|
resolved_root = Path(storage_root).resolve()
|
|
canonical_path = resolved_root / profile.storage_relative_path()
|
|
if canonical_path.exists():
|
|
return canonical_path
|
|
model_root = resolved_root / "llm_models"
|
|
for legacy_path in sorted(model_root.glob(f"*/{profile.revision}")):
|
|
if legacy_path != canonical_path and _legacy_snapshot_matches(
|
|
legacy_path,
|
|
profile,
|
|
):
|
|
logger.info(
|
|
"Reusing legacy Modal LLM weight path repository=%s revision=%s "
|
|
"path=%s.",
|
|
profile.repository,
|
|
profile.revision,
|
|
legacy_path,
|
|
)
|
|
return legacy_path
|
|
return canonical_path
|
|
|
|
|
|
def _marker_path(snapshot_path: Path) -> Path:
|
|
"""Return the completion marker path for one staged snapshot."""
|
|
return snapshot_path / _COMPLETE_MARKER_FILENAME
|
|
|
|
|
|
def _required_model_paths(
|
|
snapshot_path: Path,
|
|
profile: LLMModelProfile,
|
|
) -> tuple[Path, ...]:
|
|
"""Return every required model artifact inside one snapshot."""
|
|
model_filename = profile.backend_option("model_filename")
|
|
if model_filename is None:
|
|
return (snapshot_path / "config.json",)
|
|
filenames = [str(model_filename)]
|
|
mmproj_filename = profile.backend_option("mmproj_filename")
|
|
if mmproj_filename is not None:
|
|
filenames.append(str(mmproj_filename))
|
|
return tuple(snapshot_path / filename for filename in filenames)
|
|
|
|
|
|
def _tokenizer_source(profile: LLMModelProfile) -> tuple[str, str] | None:
|
|
"""Return an optional separately pinned tokenizer snapshot source."""
|
|
repository = profile.backend_option("tokenizer_repository")
|
|
revision = profile.backend_option("tokenizer_revision")
|
|
if repository is None and revision is None:
|
|
return None
|
|
normalized_repository = str(repository or "").strip()
|
|
normalized_revision = str(revision or "").strip().lower()
|
|
if "/" not in normalized_repository or len(normalized_revision) != 40:
|
|
raise ValueError(
|
|
f"Modal LLM profile {profile.profile_id!r} has an invalid separately "
|
|
"pinned tokenizer source."
|
|
)
|
|
return normalized_repository, normalized_revision
|
|
|
|
|
|
def _snapshot_has_required_artifacts(
|
|
snapshot_path: Path,
|
|
profile: LLMModelProfile,
|
|
) -> bool:
|
|
"""Return whether model and optional tokenizer artifacts are complete."""
|
|
if not all(
|
|
required_path.is_file()
|
|
for required_path in _required_model_paths(snapshot_path, profile)
|
|
):
|
|
return False
|
|
return (
|
|
_tokenizer_source(profile) is None
|
|
or (snapshot_path / "tokenizer_config.json").is_file()
|
|
)
|
|
|
|
|
|
def _model_allow_patterns(profile: LLMModelProfile) -> tuple[str, ...]:
|
|
"""Return a bounded allowlist for this profile's primary repository."""
|
|
model_filename = profile.backend_option("model_filename")
|
|
if model_filename is None:
|
|
return _SNAPSHOT_ALLOW_PATTERNS
|
|
selected_filenames = [str(model_filename)]
|
|
mmproj_filename = profile.backend_option("mmproj_filename")
|
|
if mmproj_filename is not None:
|
|
selected_filenames.append(str(mmproj_filename))
|
|
return (*_TOKENIZER_ALLOW_PATTERNS, *selected_filenames)
|
|
|
|
|
|
def _read_marker(snapshot_path: Path) -> dict[str, Any] | None:
|
|
"""Return a valid completion marker, if present."""
|
|
marker_path = _marker_path(snapshot_path)
|
|
try:
|
|
payload = json.loads(marker_path.read_text(encoding="utf-8"))
|
|
except (FileNotFoundError, json.JSONDecodeError, OSError):
|
|
return None
|
|
return payload if isinstance(payload, dict) else None
|
|
|
|
|
|
def is_model_snapshot_staged(
|
|
storage_root: str | Path, profile: LLMModelProfile
|
|
) -> bool:
|
|
"""Return whether the requested immutable snapshot is completely staged."""
|
|
snapshot_path = model_snapshot_path(storage_root, profile)
|
|
marker = _read_marker(snapshot_path)
|
|
return bool(
|
|
marker
|
|
and marker.get("repository") == profile.repository
|
|
and marker.get("revision") == profile.revision
|
|
and _snapshot_has_required_artifacts(snapshot_path, profile)
|
|
)
|
|
|
|
|
|
def _write_marker(snapshot_path: Path, profile: LLMModelProfile) -> None:
|
|
"""Atomically record that an immutable model snapshot is complete."""
|
|
marker_path = _marker_path(snapshot_path)
|
|
temporary_path = marker_path.with_suffix(".tmp")
|
|
marker_payload = {
|
|
"marker_version": 2,
|
|
"profile_id": profile.profile_id,
|
|
"repository": profile.repository,
|
|
"revision": profile.revision,
|
|
"completed_at_unix": time.time(),
|
|
}
|
|
temporary_path.write_text(
|
|
json.dumps(marker_payload, sort_keys=True, separators=(",", ":")),
|
|
encoding="utf-8",
|
|
)
|
|
os.replace(temporary_path, marker_path)
|
|
|
|
|
|
def _legacy_snapshot_matches(
|
|
snapshot_path: Path,
|
|
profile: LLMModelProfile,
|
|
) -> bool:
|
|
"""Return whether an old profile-keyed directory has these exact weights."""
|
|
marker = _read_marker(snapshot_path)
|
|
return bool(
|
|
marker
|
|
and marker.get("repository") == profile.repository
|
|
and marker.get("revision") == profile.revision
|
|
and _snapshot_has_required_artifacts(snapshot_path, profile)
|
|
)
|
|
|
|
|
|
def _process_start_identity(pid: int) -> str | None:
|
|
"""Return the Linux process start tick used to reject PID reuse."""
|
|
try:
|
|
stat_record = Path(f"/proc/{pid}/stat").read_text(encoding="utf-8")
|
|
except (FileNotFoundError, OSError):
|
|
return None
|
|
fields_after_command = stat_record.rpartition(") ")[2].split()
|
|
return fields_after_command[19] if len(fields_after_command) > 19 else None
|
|
|
|
|
|
def _lease_owner_payload(owner_id: str) -> dict[str, Any]:
|
|
"""Return an owner record that can be validated on this worker."""
|
|
pid = os.getpid()
|
|
return {
|
|
"version": 1,
|
|
"owner_id": owner_id,
|
|
"host_id": socket.gethostname(),
|
|
"pid": pid,
|
|
"process_start": _process_start_identity(pid),
|
|
"token": uuid4().hex,
|
|
"acquired_at": time.time(),
|
|
}
|
|
|
|
|
|
def _read_lease_owner(lease_path: Path) -> dict[str, Any] | None:
|
|
"""Read a snapshot lease owner, tolerating legacy text records."""
|
|
try:
|
|
raw_value = lease_path.read_text(encoding="utf-8")
|
|
except (FileNotFoundError, OSError):
|
|
return None
|
|
try:
|
|
payload = json.loads(raw_value)
|
|
except json.JSONDecodeError:
|
|
return None
|
|
return dict(payload) if isinstance(payload, Mapping) else None
|
|
|
|
|
|
def _local_lease_owner_is_alive(owner: Mapping[str, Any]) -> bool | None:
|
|
"""Return local owner liveness, or None for another worker host."""
|
|
if str(owner.get("host_id") or "") != socket.gethostname():
|
|
return None
|
|
try:
|
|
pid = int(owner.get("pid"))
|
|
except (TypeError, ValueError):
|
|
return False
|
|
if pid <= 0:
|
|
return False
|
|
expected_start = str(owner.get("process_start") or "")
|
|
actual_start = _process_start_identity(pid)
|
|
if actual_start is None:
|
|
if expected_start:
|
|
return False
|
|
try:
|
|
os.kill(pid, 0)
|
|
except ProcessLookupError:
|
|
return False
|
|
except PermissionError:
|
|
return True
|
|
return True
|
|
return not expected_start or actual_start == expected_start
|
|
|
|
|
|
def _lease_record_matches(lease_path: Path, owner: Mapping[str, Any]) -> bool:
|
|
"""Return whether a path still contains the observed owner token."""
|
|
current = _read_lease_owner(lease_path)
|
|
return bool(
|
|
current
|
|
and current.get("token")
|
|
and current.get("token") == owner.get("token")
|
|
)
|
|
|
|
|
|
def _remove_owned_lease(lease_path: Path, owner: Mapping[str, Any]) -> None:
|
|
"""Remove a lock only while its ownership token still matches."""
|
|
if not _lease_record_matches(lease_path, owner):
|
|
return
|
|
try:
|
|
lease_path.unlink()
|
|
except FileNotFoundError:
|
|
pass
|
|
|
|
|
|
def _heartbeat_snapshot_lease(
|
|
lease_path: Path,
|
|
owner: Mapping[str, Any],
|
|
stop_event: threading.Event,
|
|
) -> None:
|
|
"""Refresh an active lease so another container never steals live work."""
|
|
while not stop_event.wait(_LEASE_HEARTBEAT_SECONDS):
|
|
if not _lease_record_matches(lease_path, owner):
|
|
return
|
|
try:
|
|
lease_path.touch()
|
|
except FileNotFoundError:
|
|
return
|
|
except OSError as exc:
|
|
logger.warning(
|
|
"Unable to heartbeat LLM staging lease %s: %s",
|
|
lease_path,
|
|
exc,
|
|
)
|
|
return
|
|
|
|
|
|
def _snapshot_lease_wait_timeout_seconds() -> float:
|
|
"""Return the configured maximum wait to acquire one snapshot lease."""
|
|
timeout_seconds = float(
|
|
os.getenv("COMFY_MODAL_LLM_STAGE_LEASE_TIMEOUT_SECONDS", "7200")
|
|
)
|
|
if not math.isfinite(timeout_seconds) or timeout_seconds <= 0:
|
|
raise ValueError(
|
|
"COMFY_MODAL_LLM_STAGE_LEASE_TIMEOUT_SECONDS must be positive."
|
|
)
|
|
return timeout_seconds
|
|
|
|
|
|
def _snapshot_lease_stale_seconds(
|
|
existing_owner: Mapping[str, Any] | None,
|
|
) -> float:
|
|
"""Return the age that proves a heartbeating or legacy lease abandoned."""
|
|
if existing_owner is None or not existing_owner.get("token"):
|
|
return _DEFAULT_LEASE_TIMEOUT_SECONDS
|
|
raw_value = os.getenv(
|
|
"COMFY_MODAL_LLM_STAGE_LEASE_HEARTBEAT_STALE_SECONDS",
|
|
str(_DEFAULT_LEASE_HEARTBEAT_STALE_SECONDS),
|
|
)
|
|
stale_seconds = float(raw_value)
|
|
if (
|
|
not math.isfinite(stale_seconds)
|
|
or stale_seconds <= _LEASE_HEARTBEAT_SECONDS * 2
|
|
):
|
|
raise ValueError(
|
|
"COMFY_MODAL_LLM_STAGE_LEASE_HEARTBEAT_STALE_SECONDS must exceed "
|
|
f"{_LEASE_HEARTBEAT_SECONDS * 2:.0f} seconds."
|
|
)
|
|
return stale_seconds
|
|
|
|
|
|
def _remove_expired_snapshot_lease(
|
|
lease_path: Path,
|
|
existing_owner: Mapping[str, Any] | None,
|
|
) -> None:
|
|
"""Remove one expired lease without deleting a replacement owner."""
|
|
logger.warning(
|
|
"Removing expired LLM staging lease %s owner=%s.",
|
|
lease_path,
|
|
(existing_owner or {}).get("owner_id"),
|
|
)
|
|
if existing_owner is not None:
|
|
_remove_owned_lease(lease_path, existing_owner)
|
|
return
|
|
try:
|
|
lease_path.unlink()
|
|
except FileNotFoundError:
|
|
pass
|
|
|
|
|
|
def _report_snapshot_lease_wait(
|
|
progress_callback: StagingProgressCallback | None,
|
|
model_label: str,
|
|
) -> None:
|
|
"""Publish the first wait state for a contended snapshot."""
|
|
if progress_callback is None:
|
|
return
|
|
progress_callback(
|
|
LLMStagingProgress(
|
|
stage="waiting_for_download",
|
|
message=f"Waiting for another download of {model_label} to finish",
|
|
indeterminate=True,
|
|
)
|
|
)
|
|
|
|
|
|
def _wait_for_existing_snapshot_lease(
|
|
lease_path: Path,
|
|
*,
|
|
started_at: float,
|
|
timeout_seconds: float,
|
|
progress_callback: StagingProgressCallback | None,
|
|
model_label: str,
|
|
report_wait: bool,
|
|
) -> bool:
|
|
"""Wait once or reclaim an abandoned existing snapshot lease."""
|
|
existing_owner = _read_lease_owner(lease_path)
|
|
if (
|
|
existing_owner is not None
|
|
and _local_lease_owner_is_alive(existing_owner) is False
|
|
):
|
|
logger.warning(
|
|
"Removing abandoned LLM staging lease %s owner=%s.",
|
|
lease_path,
|
|
existing_owner.get("owner_id"),
|
|
)
|
|
_remove_owned_lease(lease_path, existing_owner)
|
|
return False
|
|
try:
|
|
lease_age = time.time() - lease_path.stat().st_mtime
|
|
except FileNotFoundError:
|
|
return False
|
|
if lease_age >= _snapshot_lease_stale_seconds(existing_owner):
|
|
_remove_expired_snapshot_lease(lease_path, existing_owner)
|
|
return False
|
|
if time.monotonic() - started_at >= timeout_seconds:
|
|
raise TimeoutError(
|
|
f"Timed out waiting {timeout_seconds:.0f}s for model staging lease "
|
|
f"{lease_path}."
|
|
)
|
|
if report_wait:
|
|
_report_snapshot_lease_wait(progress_callback, model_label)
|
|
time.sleep(_LEASE_POLL_SECONDS)
|
|
return True
|
|
|
|
|
|
def _acquire_snapshot_lease(
|
|
lease_path: Path,
|
|
owner: Mapping[str, Any],
|
|
*,
|
|
progress_callback: StagingProgressCallback | None,
|
|
model_label: str,
|
|
) -> None:
|
|
"""Acquire one exclusive snapshot lease, waiting or reclaiming as needed."""
|
|
timeout_seconds = _snapshot_lease_wait_timeout_seconds()
|
|
started_at = time.monotonic()
|
|
waiting_reported = False
|
|
while True:
|
|
try:
|
|
descriptor = os.open(lease_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
|
|
except FileExistsError:
|
|
waited = _wait_for_existing_snapshot_lease(
|
|
lease_path,
|
|
started_at=started_at,
|
|
timeout_seconds=timeout_seconds,
|
|
progress_callback=progress_callback,
|
|
model_label=model_label,
|
|
report_wait=not waiting_reported,
|
|
)
|
|
waiting_reported = waiting_reported or waited
|
|
continue
|
|
with os.fdopen(descriptor, "w", encoding="utf-8") as lease_file:
|
|
json.dump(owner, lease_file, sort_keys=True, separators=(",", ":"))
|
|
return
|
|
|
|
|
|
@contextmanager
|
|
def _snapshot_lease(
|
|
snapshot_path: Path,
|
|
*,
|
|
progress_callback: StagingProgressCallback | None = None,
|
|
model_label: str,
|
|
owner_id: str,
|
|
) -> Iterator[None]:
|
|
"""Serialize downloads with liveness-checked, heartbeating ownership."""
|
|
snapshot_path.parent.mkdir(parents=True, exist_ok=True)
|
|
lease_path = snapshot_path.parent / f".{snapshot_path.name}.download.lock"
|
|
owner = _lease_owner_payload(owner_id)
|
|
_acquire_snapshot_lease(
|
|
lease_path,
|
|
owner,
|
|
progress_callback=progress_callback,
|
|
model_label=model_label,
|
|
)
|
|
heartbeat_stop = threading.Event()
|
|
heartbeat_thread = threading.Thread(
|
|
target=_heartbeat_snapshot_lease,
|
|
args=(lease_path, owner, heartbeat_stop),
|
|
name="llm-snapshot-lease-heartbeat",
|
|
daemon=True,
|
|
)
|
|
heartbeat_thread.start()
|
|
try:
|
|
yield
|
|
finally:
|
|
heartbeat_stop.set()
|
|
heartbeat_thread.join(timeout=1.0)
|
|
_remove_owned_lease(lease_path, owner)
|
|
|
|
|
|
def _snapshot_existing_bytes(snapshot_path: Path) -> int:
|
|
"""Return bytes already present for a resumable model snapshot."""
|
|
total = 0
|
|
if not snapshot_path.exists():
|
|
return total
|
|
for candidate in snapshot_path.rglob("*"):
|
|
try:
|
|
if candidate.is_file() and not candidate.name.endswith(".lock"):
|
|
total += candidate.stat().st_size
|
|
except (FileNotFoundError, OSError):
|
|
continue
|
|
return total
|
|
|
|
|
|
def _minimum_free_disk_bytes() -> int:
|
|
"""Return the configured free-space reserve retained after staging."""
|
|
raw_value = os.getenv("COMFY_MODAL_LLM_MIN_FREE_DISK_GB")
|
|
if raw_value is None:
|
|
return _DEFAULT_MINIMUM_FREE_DISK_BYTES
|
|
value_gb = float(raw_value)
|
|
if not math.isfinite(value_gb) or value_gb < 0:
|
|
raise ValueError("COMFY_MODAL_LLM_MIN_FREE_DISK_GB must not be negative.")
|
|
return int(value_gb * 1024**3)
|
|
|
|
|
|
def _preflight_snapshot_capacity(
|
|
snapshot_path: Path,
|
|
profile: LLMModelProfile,
|
|
) -> tuple[int, int, int]:
|
|
"""Require room for missing artifacts plus a post-download safety reserve."""
|
|
snapshot_path.parent.mkdir(parents=True, exist_ok=True)
|
|
free_bytes = shutil.disk_usage(snapshot_path.parent).free
|
|
existing_bytes = _snapshot_existing_bytes(snapshot_path)
|
|
remaining_bytes = max(0, profile.artifact_bytes - existing_bytes)
|
|
required_bytes = remaining_bytes + _minimum_free_disk_bytes()
|
|
if free_bytes < required_bytes:
|
|
raise RuntimeError(
|
|
"Insufficient disk space to stage LLM profile "
|
|
f"{profile.profile_id!r}: {free_bytes / 1024**3:.2f} GiB free, "
|
|
f"{remaining_bytes / 1024**3:.2f} GiB still required, and "
|
|
f"{_minimum_free_disk_bytes() / 1024**3:.2f} GiB must remain free."
|
|
)
|
|
return free_bytes, remaining_bytes, required_bytes
|
|
|
|
|
|
def stage_model_profile(
|
|
profile_id: str,
|
|
storage_root: str | Path,
|
|
*,
|
|
profile: LLMModelProfile | None = None,
|
|
snapshot_download: Callable[..., str] | None = None,
|
|
progress_callback: StagingProgressCallback | None = None,
|
|
model_reference: str | None = None,
|
|
owner_id: str | None = None,
|
|
) -> StagedModelSnapshot:
|
|
"""Download one pinned curated or generated snapshot on a CPU worker."""
|
|
profile = profile or get_llm_profile(profile_id, storage_root=storage_root)
|
|
if profile.profile_id != profile_id:
|
|
raise ValueError(
|
|
f"Stage request id {profile_id!r} does not match profile "
|
|
f"{profile.profile_id!r}."
|
|
)
|
|
model_label = model_reference or profile.profile_id
|
|
|
|
def report(progress: LLMStagingProgress) -> None:
|
|
"""Attach the user-facing model reference to one staging event."""
|
|
if progress_callback is not None:
|
|
progress_callback(replace(progress, model_reference=model_label))
|
|
|
|
snapshot_path = model_snapshot_path(storage_root, profile)
|
|
started_at = time.perf_counter()
|
|
report(
|
|
LLMStagingProgress(
|
|
stage="snapshot_check",
|
|
message=f"Checking staged snapshot for {model_label}",
|
|
indeterminate=True,
|
|
)
|
|
)
|
|
if is_model_snapshot_staged(storage_root, profile):
|
|
logger.info(
|
|
"Reusing staged Modal LLM profile=%s revision=%s path=%s.",
|
|
profile.profile_id,
|
|
profile.revision,
|
|
snapshot_path,
|
|
)
|
|
report(
|
|
LLMStagingProgress(
|
|
stage="cached",
|
|
message=f"Using staged snapshot for {model_label}",
|
|
value=1,
|
|
maximum=1,
|
|
unit="model",
|
|
)
|
|
)
|
|
return StagedModelSnapshot(
|
|
profile_id=profile.profile_id,
|
|
repository=profile.repository,
|
|
revision=profile.revision,
|
|
path=str(snapshot_path),
|
|
downloaded=False,
|
|
elapsed_seconds=time.perf_counter() - started_at,
|
|
)
|
|
|
|
with _snapshot_lease(
|
|
snapshot_path,
|
|
progress_callback=report,
|
|
model_label=model_label,
|
|
owner_id=owner_id or f"pid-{os.getpid()}",
|
|
):
|
|
if is_model_snapshot_staged(storage_root, profile):
|
|
report(
|
|
LLMStagingProgress(
|
|
stage="cached",
|
|
message=(
|
|
f"Using snapshot staged by another worker for {model_label}"
|
|
),
|
|
value=1,
|
|
maximum=1,
|
|
unit="model",
|
|
)
|
|
)
|
|
return StagedModelSnapshot(
|
|
profile_id=profile.profile_id,
|
|
repository=profile.repository,
|
|
revision=profile.revision,
|
|
path=str(snapshot_path),
|
|
downloaded=False,
|
|
elapsed_seconds=time.perf_counter() - started_at,
|
|
)
|
|
free_bytes, remaining_bytes, _required_bytes = _preflight_snapshot_capacity(
|
|
snapshot_path,
|
|
profile,
|
|
)
|
|
report(
|
|
LLMStagingProgress(
|
|
stage="disk_check",
|
|
message=(
|
|
f"Storage ready for {model_label}: "
|
|
f"{free_bytes / 1024**3:.1f} GiB free, "
|
|
f"{remaining_bytes / 1024**3:.1f} GiB remaining"
|
|
),
|
|
value=free_bytes,
|
|
maximum=free_bytes,
|
|
unit="bytes",
|
|
)
|
|
)
|
|
if snapshot_download is None:
|
|
from huggingface_hub import (
|
|
snapshot_download as huggingface_snapshot_download,
|
|
)
|
|
|
|
snapshot_download = huggingface_snapshot_download
|
|
report(
|
|
LLMStagingProgress(
|
|
stage="download_preparing",
|
|
message=f"Preparing download for {model_label}",
|
|
indeterminate=True,
|
|
)
|
|
)
|
|
snapshot_path.mkdir(parents=True, exist_ok=True)
|
|
logger.info(
|
|
"Staging Modal LLM profile=%s repository=%s revision=%s to %s.",
|
|
profile.profile_id,
|
|
profile.repository,
|
|
profile.revision,
|
|
snapshot_path,
|
|
)
|
|
download_options: dict[str, Any] = {
|
|
"repo_id": profile.repository,
|
|
"revision": profile.revision,
|
|
"local_dir": str(snapshot_path),
|
|
"token": os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_HUB_TOKEN"),
|
|
"allow_patterns": _model_allow_patterns(profile),
|
|
}
|
|
if progress_callback is not None:
|
|
download_options["tqdm_class"] = _snapshot_tqdm_class(report)
|
|
resolved_path = Path(snapshot_download(**download_options)).resolve()
|
|
if resolved_path != snapshot_path.resolve():
|
|
raise RuntimeError(
|
|
f"Hugging Face staged {profile.profile_id!r} at unexpected path "
|
|
f"{resolved_path}; expected {snapshot_path.resolve()}."
|
|
)
|
|
tokenizer_source = _tokenizer_source(profile)
|
|
if tokenizer_source is not None:
|
|
tokenizer_repository, tokenizer_revision = tokenizer_source
|
|
tokenizer_options: dict[str, Any] = {
|
|
"repo_id": tokenizer_repository,
|
|
"revision": tokenizer_revision,
|
|
"local_dir": str(snapshot_path),
|
|
"token": download_options["token"],
|
|
"allow_patterns": _TOKENIZER_ALLOW_PATTERNS,
|
|
}
|
|
if progress_callback is not None:
|
|
tokenizer_options["tqdm_class"] = _snapshot_tqdm_class(
|
|
report
|
|
)
|
|
tokenizer_path = Path(snapshot_download(**tokenizer_options)).resolve()
|
|
if tokenizer_path != snapshot_path.resolve():
|
|
raise RuntimeError(
|
|
f"Hugging Face staged tokenizer for {profile.profile_id!r} at "
|
|
f"unexpected path {tokenizer_path}; expected "
|
|
f"{snapshot_path.resolve()}."
|
|
)
|
|
if not _snapshot_has_required_artifacts(snapshot_path, profile):
|
|
raise RuntimeError(
|
|
f"Staged Modal LLM profile {profile.profile_id!r} is missing its "
|
|
"required model or tokenizer artifacts."
|
|
)
|
|
_write_marker(snapshot_path, profile)
|
|
report(
|
|
LLMStagingProgress(
|
|
stage="staged",
|
|
message=f"Model snapshot ready for {model_label}",
|
|
value=1,
|
|
maximum=1,
|
|
unit="model",
|
|
)
|
|
)
|
|
result = StagedModelSnapshot(
|
|
profile_id=profile.profile_id,
|
|
repository=profile.repository,
|
|
revision=profile.revision,
|
|
path=str(snapshot_path),
|
|
downloaded=True,
|
|
elapsed_seconds=time.perf_counter() - started_at,
|
|
)
|
|
logger.info("Completed Modal LLM staging: %s", asdict(result))
|
|
return result
|
|
|
|
|
|
def _persist_supplied_profile(
|
|
model_reference: str,
|
|
storage_root: str | Path,
|
|
supplied: Mapping[str, Any],
|
|
) -> tuple[LLMModelProfile, str | None, bool, bool]:
|
|
"""Validate planner metadata and persist a generated manifest remotely."""
|
|
raw_profile = supplied.get("profile", supplied)
|
|
if not isinstance(raw_profile, Mapping):
|
|
raise ValueError(
|
|
f"Resolved LLM profile for {model_reference!r} must contain an object."
|
|
)
|
|
profile = LLMModelProfile.from_mapping(raw_profile)
|
|
_validate_supplied_profile_reference(model_reference, profile)
|
|
if profile.source != "generated":
|
|
return profile, None, False, True
|
|
if __package__:
|
|
from .llm_profiles import generated_profile_manifest_path
|
|
else: # pragma: no cover - stable cloud entrypoint imports top-level modules.
|
|
from llm_profiles import generated_profile_manifest_path
|
|
|
|
manifest_path = generated_profile_manifest_path(storage_root, profile.profile_id)
|
|
manifest_payload = _supplied_profile_manifest_payload(
|
|
model_reference,
|
|
profile,
|
|
supplied,
|
|
)
|
|
existing_scan = _existing_supplied_profile_scan(
|
|
manifest_path,
|
|
raw_profile,
|
|
)
|
|
if existing_scan is not None:
|
|
return profile, str(manifest_path), False, existing_scan
|
|
manifest_path.parent.mkdir(parents=True, exist_ok=True)
|
|
temporary_path = manifest_path.with_suffix(f".{uuid4().hex}.tmp")
|
|
temporary_path.write_text(
|
|
json.dumps(manifest_payload, sort_keys=True, separators=(",", ":")),
|
|
encoding="utf-8",
|
|
)
|
|
os.replace(temporary_path, manifest_path)
|
|
return profile, str(manifest_path), True, bool(
|
|
manifest_payload["security_scan_complete"]
|
|
)
|
|
|
|
|
|
def _supplied_profile_manifest_payload(
|
|
model_reference: str,
|
|
profile: LLMModelProfile,
|
|
supplied: Mapping[str, Any],
|
|
) -> dict[str, Any]:
|
|
"""Build the generated manifest persisted from trusted planner metadata."""
|
|
security_scan_complete = supplied.get("security_scan_complete", False)
|
|
if not isinstance(security_scan_complete, bool):
|
|
raise ValueError(
|
|
f"Resolved LLM profile for {model_reference!r} has an invalid "
|
|
"security scan state."
|
|
)
|
|
return {
|
|
"schema_version": profile.schema_version,
|
|
"compatibility_policy_version": profile.compatibility_policy_version,
|
|
"requested_reference": model_reference,
|
|
"resolved_at_unix": time.time(),
|
|
"security_scan_complete": security_scan_complete,
|
|
"profile": profile.to_mapping(),
|
|
}
|
|
|
|
|
|
def _existing_supplied_profile_scan(
|
|
manifest_path: Path,
|
|
raw_profile: Mapping[str, Any],
|
|
) -> bool | None:
|
|
"""Return an existing matching manifest's scan state, or None if absent."""
|
|
if manifest_path.is_file():
|
|
try:
|
|
existing = json.loads(manifest_path.read_text(encoding="utf-8"))
|
|
except (json.JSONDecodeError, OSError) as exc:
|
|
raise RuntimeError(
|
|
f"Resolved LLM manifest {manifest_path} is unreadable: {exc}"
|
|
) from exc
|
|
if not isinstance(existing, Mapping) or existing.get("profile") != raw_profile:
|
|
raise RuntimeError(
|
|
f"Resolved LLM profile collision at {manifest_path}."
|
|
)
|
|
return bool(existing.get("security_scan_complete", False))
|
|
return None
|
|
|
|
|
|
def _validate_supplied_profile_reference(
|
|
model_reference: str,
|
|
profile: LLMModelProfile,
|
|
) -> None:
|
|
"""Reject planner metadata that does not describe its requested reference."""
|
|
curated_profile = load_llm_profiles().get(model_reference)
|
|
if curated_profile is not None:
|
|
if profile != curated_profile:
|
|
raise ValueError(
|
|
f"Resolved metadata for curated profile {model_reference!r} does "
|
|
"not match the checked-in profile."
|
|
)
|
|
return
|
|
if model_reference.startswith("hf-"):
|
|
if profile.profile_id != model_reference:
|
|
raise ValueError(
|
|
f"Resolved metadata profile id {profile.profile_id!r} does not "
|
|
f"match requested id {model_reference!r}."
|
|
)
|
|
return
|
|
if profile.source != "generated":
|
|
raise ValueError(
|
|
f"Resolved metadata for Hugging Face reference {model_reference!r} "
|
|
"must contain a generated profile."
|
|
)
|
|
if __package__:
|
|
from .llm_resolver import HuggingFaceModelReference
|
|
else: # pragma: no cover - stable cloud entrypoint imports top-level modules.
|
|
from llm_resolver import HuggingFaceModelReference
|
|
|
|
parsed = HuggingFaceModelReference.parse(model_reference)
|
|
if profile.repository != parsed.repository:
|
|
raise ValueError(
|
|
f"Resolved metadata repository {profile.repository!r} does not match "
|
|
f"requested repository {parsed.repository!r}."
|
|
)
|
|
requested_revision = (parsed.requested_revision or "").lower()
|
|
if len(requested_revision) == 40 and requested_revision != profile.revision:
|
|
raise ValueError(
|
|
f"Resolved metadata revision {profile.revision!r} does not match "
|
|
f"requested revision {requested_revision!r}."
|
|
)
|
|
|
|
|
|
def _resolve_profile_for_staging(
|
|
model_reference: str,
|
|
storage_root: str | Path,
|
|
supplied_profile: Mapping[str, Any] | None = None,
|
|
) -> tuple[LLMModelProfile, str | None, bool, bool]:
|
|
"""Resolve one curated, generated, or Hugging Face model reference."""
|
|
if supplied_profile is not None:
|
|
return _persist_supplied_profile(
|
|
model_reference,
|
|
storage_root,
|
|
supplied_profile,
|
|
)
|
|
curated_profiles = load_llm_profiles()
|
|
if model_reference in curated_profiles:
|
|
return curated_profiles[model_reference], None, False, True
|
|
if model_reference.startswith("hf-"):
|
|
return (
|
|
get_llm_profile(model_reference, storage_root=storage_root),
|
|
None,
|
|
False,
|
|
True,
|
|
)
|
|
if __package__:
|
|
from .llm_resolver import resolve_model_profile
|
|
else: # pragma: no cover - stable cloud entrypoint imports top-level modules.
|
|
from llm_resolver import resolve_model_profile
|
|
|
|
resolved = resolve_model_profile(model_reference, storage_root)
|
|
return (
|
|
resolved.profile,
|
|
resolved.manifest_path,
|
|
resolved.manifest_created,
|
|
resolved.security_scan_complete,
|
|
)
|
|
|
|
|
|
def resolve_and_stage_model_references(
|
|
model_references: list[str],
|
|
storage_root: str | Path,
|
|
*,
|
|
progress_callback: StagingProgressCallback | None = None,
|
|
resolved_profiles: Mapping[str, Mapping[str, Any]] | None = None,
|
|
owner_id: str | None = None,
|
|
) -> list[ResolvedStagedModelProfile]:
|
|
"""Resolve and stage model references on any CPU-backed remote worker."""
|
|
results: list[ResolvedStagedModelProfile] = []
|
|
for model_reference in model_references:
|
|
supplied_profile = (resolved_profiles or {}).get(model_reference)
|
|
if progress_callback is not None:
|
|
is_hugging_face_reference = (
|
|
"/" in model_reference and not model_reference.startswith("hf-")
|
|
)
|
|
progress_callback(
|
|
LLMStagingProgress(
|
|
stage=(
|
|
"resolved_metadata"
|
|
if supplied_profile is not None
|
|
else "metadata"
|
|
if is_hugging_face_reference
|
|
else "profile_resolution"
|
|
),
|
|
message=(
|
|
f"Using planner-resolved metadata for {model_reference}"
|
|
if supplied_profile is not None
|
|
else f"Inspecting Hugging Face metadata for {model_reference}"
|
|
if is_hugging_face_reference
|
|
else f"Resolving model profile {model_reference}"
|
|
),
|
|
indeterminate=True,
|
|
model_reference=model_reference,
|
|
)
|
|
)
|
|
resolve_started_at = time.perf_counter()
|
|
(
|
|
profile,
|
|
manifest_path,
|
|
manifest_created,
|
|
scan_complete,
|
|
) = (
|
|
_resolve_profile_for_staging(
|
|
model_reference,
|
|
storage_root,
|
|
supplied_profile,
|
|
)
|
|
if supplied_profile is not None
|
|
else _resolve_profile_for_staging(model_reference, storage_root)
|
|
)
|
|
resolve_elapsed_seconds = time.perf_counter() - resolve_started_at
|
|
staged = stage_model_profile(
|
|
profile.profile_id,
|
|
storage_root,
|
|
profile=profile,
|
|
progress_callback=progress_callback,
|
|
model_reference=model_reference,
|
|
owner_id=owner_id,
|
|
)
|
|
results.append(
|
|
ResolvedStagedModelProfile(
|
|
requested_reference=model_reference,
|
|
profile_id=staged.profile_id,
|
|
repository=staged.repository,
|
|
revision=staged.revision,
|
|
backend=profile.backend,
|
|
quantization_method=profile.quantization_method,
|
|
artifact_bytes=profile.artifact_bytes,
|
|
manifest_path=manifest_path,
|
|
manifest_created=manifest_created,
|
|
security_scan_complete=scan_complete,
|
|
path=staged.path,
|
|
downloaded=staged.downloaded,
|
|
resolve_elapsed_seconds=resolve_elapsed_seconds,
|
|
elapsed_seconds=staged.elapsed_seconds,
|
|
)
|
|
)
|
|
return results
|
|
|
|
|
|
__all__ = [
|
|
"LLMStagingProgress",
|
|
"ResolvedStagedModelProfile",
|
|
"StagedModelSnapshot",
|
|
"is_model_snapshot_staged",
|
|
"model_snapshot_path",
|
|
"resolve_and_stage_model_references",
|
|
"stage_model_profile",
|
|
]
|