| """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 |
