shared_metrics.py

# Shared metrics helper for all EvoNet mini nodes
import os
import platform

# HuggingFace Spaces hardware specs (chính xác theo docs HF)
HF_HARDWARE_SPECS = {
    "cpu-basic":    {"ram_gb": 16,  "cpu": 2,  "gpu": None},
    "cpu-upgrade":  {"ram_gb": 32,  "cpu": 8,  "gpu": None},
    "t4-small":     {"ram_gb": 16,  "cpu": 4,  "gpu": "T4-16GB"},
    "t4-medium":    {"ram_gb": 30,  "cpu": 8,  "gpu": "T4-16GB×2"},
    "a10g-small":   {"ram_gb": 24,  "cpu": 4,  "gpu": "A10G-24GB"},
    "a10g-large":   {"ram_gb": 48,  "cpu": 12, "gpu": "A10G-24GB×4"},
    "a100-large":   {"ram_gb": 142, "cpu": 12, "gpu": "A100-80GB"},
    "zero-a10g":    {"ram_gb": 24,  "cpu": 4,  "gpu": "A10G-24GB"},
    "zero-t4":      {"ram_gb": 16,  "cpu": 4,  "gpu": "T4-16GB"},
}

def get_hf_metrics():
    """
    Lấy thông số phần cứng dựa trên SPACE_HARDWARE env var của HuggingFace.
    Tránh đọc psutil vì nó trả về RAM của HOST máy chủ vật lý (~128GB),
    không phải RAM được cấp cho container Space.
    """
    hardware_tier = os.environ.get("SPACE_HARDWARE", "cpu-basic").strip()
    spec = HF_HARDWARE_SPECS.get(hardware_tier, {"ram_gb": 16, "cpu": 2, "gpu": None})
    
    ram_mb = spec["ram_gb"] * 1024
    cpu_count = spec["cpu"]
    
    # Ước tính RAM đã dùng (dùng psutil nếu có, giới hạn trong spec)
    used_mb = 0
    try:
        import psutil
        real_used = psutil.Process().memory_info().rss / 1024 / 1024
        # Cap ở spec của tier để tránh hiển thị sai
        used_mb = min(round(real_used), ram_mb)
    except Exception:
        used_mb = round(ram_mb * 0.25)  # Giả sử 25% nếu không có psutil
    
    # Max concurrent workers: mỗi worker ~2GB RAM
    available_mb = max(0, ram_mb - used_mb)
    max_workers = max(1, available_mb // 2048)
    
    return {
        "total_memory_mb": ram_mb,
        "used_memory_mb": used_mb,
        "cpu_count": cpu_count,
        "cpu_percent": 0,
        "platform": platform.system(),
        "max_workers": max_workers,
        "hardware_tier": hardware_tier,
        "gpu": spec.get("gpu"),
    }