| """ResourceBudget — psutil + rocm-smi-derived training-job sizing helper. |
| |
| Reads host RAM, CPU, GPU memory; emits a recommended `micro_batch_size` |
| that fits under the budget for a given `seq_len`. psutil + rocm-smi are |
| both lazy / optional — falls back to conservative defaults if neither is |
| available. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| import shutil |
| import subprocess |
| from typing import cast |
|
|
| from pydantic import BaseModel, ConfigDict, Field |
|
|
|
|
| class ResourceBudget(BaseModel): |
| model_config = ConfigDict(extra="forbid") |
|
|
| host_ram_gb: float = Field(ge=0.0) |
| cpu_cores: int = Field(ge=1) |
| gpu_count: int = Field(ge=0) |
| gpu_mem_gb_each: float = Field(ge=0.0) |
|
|
|
|
| def _host_ram_gb() -> float: |
| try: |
| import psutil |
|
|
| return float(psutil.virtual_memory().total) / (1024**3) |
| except ImportError: |
| return 0.0 |
|
|
|
|
| def _cpu_cores() -> int: |
| try: |
| import psutil |
|
|
| return int(psutil.cpu_count(logical=True) or 1) |
| except ImportError: |
| import os |
|
|
| return os.cpu_count() or 1 |
|
|
|
|
| def _gpu_mem_gb() -> tuple[int, float]: |
| """Return (gpu_count, gpu_mem_gb_each).""" |
| if shutil.which("rocm-smi") is None: |
| return (0, 0.0) |
| try: |
| out = subprocess.run( |
| ["rocm-smi", "--showmeminfo", "vram", "--json"], |
| capture_output=True, |
| text=True, |
| timeout=10.0, |
| check=False, |
| ) |
| except (subprocess.TimeoutExpired, OSError): |
| return (0, 0.0) |
| if out.returncode != 0: |
| return (0, 0.0) |
| try: |
| data = json.loads(out.stdout) |
| except json.JSONDecodeError: |
| return (0, 0.0) |
| cards = [k for k in data if k.startswith("card")] |
| if not cards: |
| return (0, 0.0) |
| total_bytes = 0.0 |
| for c in cards: |
| for k, v in (data[c] or {}).items(): |
| if "total" in k.lower(): |
| try: |
| total_bytes = max(total_bytes, float(v)) |
| break |
| except ValueError: |
| continue |
| if total_bytes == 0.0: |
| return (len(cards), 0.0) |
| return (len(cards), total_bytes / (1024**3)) |
|
|
|
|
| def detect() -> ResourceBudget: |
| """Return a `ResourceBudget` snapshot of the current host.""" |
| gpu_count, gpu_mem = _gpu_mem_gb() |
| return ResourceBudget( |
| host_ram_gb=_host_ram_gb(), |
| cpu_cores=_cpu_cores(), |
| gpu_count=gpu_count, |
| gpu_mem_gb_each=gpu_mem, |
| ) |
|
|
|
|
| def recommend_micro_batch(budget: ResourceBudget, seq_len: int, *, dtype_bytes: int = 2) -> int: |
| """Return a conservative `micro_batch_size` that fits under the budget. |
| |
| Heuristic: each token consumes ~`dtype_bytes` bytes for activations + a |
| factor for KV cache + gradients. We leave 30% headroom. |
| """ |
| if budget.gpu_mem_gb_each <= 0: |
| return 1 |
| bytes_per_token = dtype_bytes * 16 |
| available = budget.gpu_mem_gb_each * 0.7 * (1024**3) |
| bs = int(available / (seq_len * bytes_per_token)) |
| return cast(int, max(1, bs)) |
|
|
|
|
| __all__ = ["ResourceBudget", "detect", "recommend_micro_batch"] |
|
|
