dekun
104 lines
3.0 KiB
Python
104 lines
3.0 KiB
Python
"""根据心跳上报的显卡型号和显存，给出节点算力权重。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
from dataclasses import dataclass
|
|
from decimal import Decimal
|
|
|
|
HIGH = Decimal("1")
|
|
MID = Decimal("0.5")
|
|
LOW = Decimal("0.05")
|
|
NONE = Decimal("0.02")
|
|
DEFAULT_WEIGHT = LOW
|
|
GENERATION_40 = "40"
|
|
GENERATION_30 = "30"
|
|
GENERATION_10 = "10"
|
|
RECOVERY_40_DAYS = 210
|
|
RECOVERY_30_DAYS = 365
|
|
RECOVERY_10_DAYS = 540
|
|
NODE_ID_BASE = 1000
|
|
GENESIS_COUNT = 20
|
|
GENESIS_RECOVERY_DAYS = 30
|
|
|
|
_MODEL_NUMBER = re.compile(r"(RTX|GTX|GT)\s*(\d{3,4})")
|
|
_DATACENTER = ("H100", "H200", "A100", "A800", "A6000", "A40", "A30", "L40")
|
|
_INTEGRATED = ("INTEL", "UHD", "IRIS", "BASIC DISPLAY", "MICROSOFT")
|
|
|
|
|
|
@dataclass
|
|
class HardwareReport:
|
|
gpu_model: str | None
|
|
vram_gb: float | None
|
|
cpu_count: int | None
|
|
memory_gb: float | None
|
|
|
|
|
|
def display_node_id(node_index: int) -> int:
|
|
"""第 1 台显示 1001，第 2 台显示 1002。"""
|
|
return NODE_ID_BASE + node_index
|
|
|
|
|
|
def is_genesis(node_index: int) -> bool:
|
|
return 1 <= node_index <= GENESIS_COUNT
|
|
|
|
|
|
def displayed_population(actual: int) -> int:
|
|
"""全网台数展示时加上 1000 的基数。3 台在线显示为 1003。"""
|
|
return NODE_ID_BASE + actual
|
|
|
|
|
|
def score_hardware(gpu_model: str | None, vram_gb: float | None) -> Decimal:
|
|
"""高端主流卡 1.0，中端 0.5，旧卡或低端独显 0.05，没有独显 0.02。"""
|
|
name = (gpu_model or "").strip()
|
|
upper = name.upper()
|
|
if not name or _is_integrated(upper):
|
|
return NONE
|
|
if any(token in upper for token in _DATACENTER):
|
|
return HIGH
|
|
match = _MODEL_NUMBER.search(upper)
|
|
if match is not None:
|
|
return _from_model(match.group(1), int(match.group(2)))
|
|
if vram_gb is not None and vram_gb >= 8:
|
|
return HIGH
|
|
if vram_gb is not None and vram_gb >= 4:
|
|
return MID
|
|
return LOW
|
|
|
|
|
|
def classify_generation(gpu_model: str | None) -> tuple[str, int]:
|
|
"""40 系及以上 210 天，30 系 365 天，其余包括 10 系和核显 540 天。"""
|
|
upper = (gpu_model or "").strip().upper()
|
|
if not upper or _is_integrated(upper):
|
|
return GENERATION_10, RECOVERY_10_DAYS
|
|
if any(token in upper for token in _DATACENTER):
|
|
return GENERATION_40, RECOVERY_40_DAYS
|
|
match = _MODEL_NUMBER.search(upper)
|
|
if match is not None and match.group(1) == "RTX":
|
|
number = int(match.group(2))
|
|
if number >= 4000:
|
|
return GENERATION_40, RECOVERY_40_DAYS
|
|
if number >= 3000:
|
|
return GENERATION_30, RECOVERY_30_DAYS
|
|
return GENERATION_10, RECOVERY_10_DAYS
|
|
|
|
|
|
def _is_integrated(upper: str) -> bool:
|
|
if any(token in upper for token in ("NVIDIA", "GEFORCE", "QUADRO")):
|
|
return False
|
|
if any(token in upper for token in _INTEGRATED):
|
|
return True
|
|
return "RADEON" in upper and "RX" not in upper
|
|
|
|
|
|
def _from_model(family: str, number: int) -> Decimal:
|
|
if family == "RTX":
|
|
if number >= 3060:
|
|
return HIGH
|
|
return MID
|
|
if family == "GTX":
|
|
if number >= 1060:
|
|
return MID
|
|
return LOW
|
|
return LOW
|