| """System profiler (PRIVATE to Hardware_Module): CPU, RAM (cgroup aware), GPUs, model size.""" |
| from __future__ import annotations |
|
|
| import json |
| import os |
| import re |
| from pathlib import Path |
| from typing import Dict, Optional |
|
|
| GIB = 1024 ** 3 |
|
|
|
|
| def _read_int(path: str) -> Optional[int]: |
| try: |
| raw = Path(path).read_text().strip() |
| if raw == "max": |
| return None |
| value = int(raw) |
| return value if value < (1 << 60) else None |
| except Exception: |
| return None |
|
|
|
|
| def cgroup_memory() -> Dict[str, Optional[int]]: |
| """Container limits (GitHub Codespaces / Docker / Kaggle) - psutil alone over-reports.""" |
| limit = _read_int("/sys/fs/cgroup/memory.max") or _read_int("/sys/fs/cgroup/memory/memory.limit_in_bytes") |
| used = _read_int("/sys/fs/cgroup/memory.current") or _read_int("/sys/fs/cgroup/memory/memory.usage_in_bytes") |
| return {"limit": limit, "used": used} |
|
|
|
|
| def estimate_model(model_path: str, token: Optional[str] = None) -> Dict: |
| """Parameter count (in billions) + layer geometry from config.json (local dir or hub repo).""" |
| cfg: Dict = {} |
| local = Path(str(model_path)) / "config.json" |
| try: |
| if local.exists(): |
| cfg = json.loads(local.read_text(encoding="utf-8")) |
| else: |
| from huggingface_hub import hf_hub_download |
| cfg = json.loads(Path(hf_hub_download(str(model_path), "config.json", token=token or None)).read_text()) |
| except Exception: |
| cfg = {} |
| hidden = int(cfg.get("hidden_size", 0)) |
| layers = int(cfg.get("num_hidden_layers", 0)) |
| vocab = int(cfg.get("vocab_size", 0)) |
| inter = int(cfg.get("intermediate_size", 4 * hidden if hidden else 0)) |
| heads = int(cfg.get("num_attention_heads", 1) or 1) |
| kv_heads = int(cfg.get("num_key_value_heads", heads) or heads) |
| head_dim = int(cfg.get("head_dim", hidden // heads if hidden else 0) or 0) |
| tied = bool(cfg.get("tie_word_embeddings", False)) |
| if hidden and layers and vocab: |
| attn = hidden * heads * head_dim + 2 * hidden * kv_heads * head_dim + heads * head_dim * hidden |
| mlp = 3 * hidden * inter |
| per_layer = attn + mlp + 2 * hidden |
| embed = vocab * hidden |
| total = embed * (1 if tied else 2) + layers * per_layer + hidden |
| else: |
| match = re.search(r"(\d+(?:\.\d+)?)\s*[bB]", str(model_path)) |
| total = int(float(match.group(1)) * 1e9) if match else int(1e9) |
| layers = layers or 24 |
| embed = int(total * 0.1) |
| per_layer = (total - embed) // layers |
| return {"params": int(total), "params_b": round(total / 1e9, 3), "num_layers": layers, |
| "layer_params": int(per_layer), "embed_params": int(embed), "tied_embeddings": tied, |
| "hidden_size": hidden, "vocab_size": vocab} |
|
|
|
|
| def profile(model_path: str, token: Optional[str] = None) -> Dict: |
| import psutil |
| vm = psutil.virtual_memory() |
| cg = cgroup_memory() |
| total = vm.total |
| avail = vm.available |
| if cg["limit"]: |
| total = min(total, cg["limit"]) |
| if cg["used"] is not None: |
| avail = min(avail, max(cg["limit"] - cg["used"], 0)) |
| gpus = [] |
| try: |
| import torch |
| if torch.cuda.is_available(): |
| for i in range(torch.cuda.device_count()): |
| free, tot = torch.cuda.mem_get_info(i) |
| major, _minor = torch.cuda.get_device_capability(i) |
| gpus.append({"index": i, "name": torch.cuda.get_device_name(i), "free_bytes": int(free), |
| "total_bytes": int(tot), "bf16": major >= 8}) |
| except Exception: |
| gpus = [] |
| model = estimate_model(model_path, token) |
| return { |
| "cpu_cores": os.cpu_count() or 1, |
| "ram_total_bytes": int(total), "ram_avail_bytes": int(avail), |
| "gpu_count": len(gpus), "gpus": gpus, |
| "gpu_free_bytes": sum(g["free_bytes"] for g in gpus), |
| "model": model, |
| } |
|
|
|
|
| def describe(prof: Dict) -> str: |
| m = prof["model"] |
| lines = [ |
| f" CPU cores : {prof['cpu_cores']}", |
| f" RAM total/free : {prof['ram_total_bytes'] / GIB:.2f} GiB / {prof['ram_avail_bytes'] / GIB:.2f} GiB", |
| f" GPUs : {prof['gpu_count']}" |
| + "".join(f"\n cuda:{g['index']} {g['name']} free {g['free_bytes'] / GIB:.1f}/{g['total_bytes'] / GIB:.1f} GiB" |
| for g in prof["gpus"]), |
| f" Model : ~{m['params_b']}B params, {m['num_layers']} layers, hidden {m['hidden_size']}, vocab {m['vocab_size']}", |
| ] |
| return "\n".join(lines) |
