File size: 2,946 Bytes
9d614f1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | """PUBLIC INTERFACE of Hardware_Module. Other modules may import ONLY this file.
prepare_base_model(cfg, fresh) -> purge (optional) + authenticated fresh HF download -> local folder
profile_system(model_dir, cfg) -> CPU/RAM/GPU/model-size profile
plan_execution(profile, cfg) -> tier 1/2/3 plan with exact RAM-vs-Disk split
load_model(model_dir, plan, cfg, row_recipe=None) -> model (tier 3 models carry a `_rss_ring` handle)
release_memory() -> gc + CUDA cache flush
"""
from __future__ import annotations
import gc
from pathlib import Path
from typing import Dict, Optional
from .internal import cache_cleaner_and_downloader as cleaner
from .internal import system_profiler as profiler
from .internal import tier_allocation_planner as planner
from .internal.sequential_ring_loader import SequentialRingLoader
def purge_old_artifacts(cfg: Dict) -> list:
print("[Hardware] purging old model leftovers ...")
return cleaner.purge_old_artifacts(Path(cfg["root"]), cfg["base_model_repo"])
def prepare_base_model(cfg: Dict, fresh: bool = False) -> str:
"""Return the local folder that holds BASE_MODEL_REPO (downloads it if needed)."""
root = Path(cfg["root"])
if fresh:
cleaner.purge_old_artifacts(root, cfg["base_model_repo"])
dest = root / "base_model"
path = cleaner.download_base_model(cfg["base_model_repo"], dest, cfg.get("hf_token"), force=fresh)
return str(path)
def profile_system(model_dir: str, cfg: Dict) -> Dict:
return profiler.profile(model_dir, cfg.get("hf_token") if cleaner.token_is_valid(cfg.get("hf_token")) else None)
def plan_execution(prof: Dict, cfg: Dict) -> Dict:
force = cfg.get("force_tier")
plan = planner.make_plan(prof, int(force) if force else None)
print("[Hardware] system profile:\n" + profiler.describe(prof))
print("[Hardware] execution plan:\n" + planner.describe_plan(plan, prof))
return plan
def load_model(model_dir: str, plan: Dict, cfg: Dict, row_recipe: Optional[Dict] = None):
import torch
root = Path(cfg["root"])
torch.set_num_threads(max(1, int(cfg.get("cpu_threads") or (torch.get_num_threads()))))
print(f"[Hardware] loading model from {model_dir} (tier {plan['tier']}) ...")
if plan["tier"] == 3:
loader = SequentialRingLoader(
model_dir, dtype=getattr(torch, plan["dtype"]), window=plan.get("window") or 2,
evict_ends=plan.get("evict_ends", True), overlay_dir=root / "offload_disk" / "overlay",
row_recipe=row_recipe)
model = loader.build()
model._rss_ring = loader
else:
model = planner.materialize_model(model_dir, plan, root / "offload_disk")
model._rss_ring = None
model._rss_plan = plan
return model
def release_memory() -> None:
gc.collect()
try:
import torch
if torch.cuda.is_available():
torch.cuda.empty_cache()
except Exception:
pass |