diff --git a/src/parallax/server/server_info.py b/src/parallax/server/server_info.py index 0e2234e9..bb8f31e1 100644 --- a/src/parallax/server/server_info.py +++ b/src/parallax/server/server_info.py @@ -4,7 +4,9 @@ We haven't used other info, will wait until DHT implemented. """ +import logging import platform +import re import subprocess from dataclasses import asdict, dataclass from typing import Any, ClassVar, Dict, Optional @@ -24,6 +26,8 @@ except ImportError: psutil = None +logger = logging.getLogger(__name__) + @dataclass class HardwareInfo: @@ -60,7 +64,12 @@ def detect() -> "HardwareInfo": class AppleSiliconHardwareInfo(HardwareInfo): """HardwareInfo specialised for Apple silicon (M-series).""" - # From cpu-monkey.com + # Peak GPU-ALU FP16 TFLOPS per chip (cpu-monkey.com for M1-M4). M5 values are + # GPU-ALU-only for consistency with M1-M4: max-bin cores (10/20/40) * 0.808, + # M5's per-core ALU rate (~flat vs M4's 0.852; M5's gains are in its new + # Neural Accelerators, not the ALUs). With those accelerators engaged (what + # MLX uses for matmul) effective FP16 is roughly 2x these values (~70 for M5 + # Max) — a separate metric from the ALU figures tabulated here. _APPLE_PEAK_FP16: ClassVar[Dict[str, float]] = { "M1": 4.58, "M1 Pro": 10.6, @@ -75,8 +84,32 @@ class AppleSiliconHardwareInfo(HardwareInfo): "M4": 8.52, "M4 Pro": 17.04, "M4 Max": 34.08, + "M5": 8.08, + "M5 Pro": 16.16, + "M5 Max": 32.32, } + # Fallback FP16 for a chip not in the table (e.g. a new generation): scale by + # GPU core count at the latest known per-core rate, else a base-chip default. + _FALLBACK_FP16_PER_CORE: ClassVar[float] = 0.808 # latest known (M5-era) ALU rate + _FALLBACK_FP16_DEFAULT: ClassVar[float] = 8.08 # base-chip default if cores unknown + + @classmethod + def _gpu_core_count(cls) -> Optional[int]: + """Best-effort Apple GPU core count via system_profiler; None on failure.""" + try: + out = subprocess.check_output( + ["system_profiler", "SPDisplaysDataType"], text=True, timeout=10 + ) + match = re.search(r"Total Number of Cores:\s*(\d+)", out) + if match: + return int(match.group(1)) + except Exception: + # system_profiler may be unavailable/slow (e.g. CI, virtualized hosts); + # the caller falls back to a conservative default. + pass + return None + @classmethod def detect(cls) -> "AppleSiliconHardwareInfo": if psutil: @@ -92,13 +125,28 @@ def detect(cls) -> "AppleSiliconHardwareInfo": # For github action, we need to remove the "(Virtual)" suffix if short_name.endswith(" (Virtual)"): short_name = short_name.rsplit(" (Virtual)", maxsplit=1)[0] - try: - flops = cls._APPLE_PEAK_FP16[short_name] - except KeyError as e: - raise RuntimeError( - f"Unknown Apple silicon chip '{short_name}' detected. " - "Please add it to the _APPLE_PEAK_FP16 dictionary." - ) from e + flops = cls._APPLE_PEAK_FP16.get(short_name) + if flops is None: + # Unknown chip (likely a newer generation). Estimate from GPU core + # count rather than crashing, and warn so it can be added explicitly. + cores = cls._gpu_core_count() + if cores: + flops = round(cores * cls._FALLBACK_FP16_PER_CORE, 2) + logger.warning( + "Unknown Apple silicon chip '%s'; estimating %.2f TFLOPS FP16 " + "from %d GPU cores. Add it to _APPLE_PEAK_FP16 for an exact value.", + short_name, + flops, + cores, + ) + else: + flops = cls._FALLBACK_FP16_DEFAULT + logger.warning( + "Unknown Apple silicon chip '%s' and GPU core count unavailable; " + "falling back to %.2f TFLOPS FP16. Add it to _APPLE_PEAK_FP16.", + short_name, + flops, + ) return cls(num_gpus=1, total_ram_gb=round(total_gb, 1), chip=chip, tflops_fp16=flops) diff --git a/tests/test_server_info.py b/tests/test_server_info.py new file mode 100644 index 00000000..f0225964 --- /dev/null +++ b/tests/test_server_info.py @@ -0,0 +1,86 @@ +""" +Tests for Apple Silicon hardware detection in server_info. + +Covers the M5-family table entries and the graceful fallback for unknown +Apple silicon chips (regression for the crash reported in issue #439). +""" + +from unittest.mock import patch + +from parallax.server.server_info import AppleSiliconHardwareInfo + + +def _check_output_for(brand: str, gpu_cores=None): + """Build a subprocess.check_output side_effect simulating a given Mac. + + Args: + brand: chip name returned for the cpu brand_string query, e.g. "M5 Max". + gpu_cores: GPU core count returned by system_profiler, or None to + simulate it being unavailable (e.g. CI / virtualized hosts). + """ + + def _side_effect(cmd, *args, **kwargs): + if "machdep.cpu.brand_string" in cmd: + return f"Apple {brand}" + if "hw.memsize" in cmd: + return str(64 * 2**30) # 64 GB + if "SPDisplaysDataType" in cmd: + if gpu_cores is None: + raise FileNotFoundError("system_profiler unavailable") + return f"Graphics/Displays:\n Total Number of Cores: {gpu_cores}\n" + raise ValueError(f"unexpected command: {cmd}") + + return _side_effect + + +class TestAppleSiliconDetect: + """AppleSiliconHardwareInfo.detect() chip resolution and fallback.""" + + def test_known_chip_uses_table(self): + """A chip in the table resolves to its tabulated FP16 value.""" + with ( + patch("parallax.server.server_info.psutil", None), + patch("subprocess.check_output", side_effect=_check_output_for("M5 Max")), + ): + info = AppleSiliconHardwareInfo.detect() + assert info.chip == "Apple M5 Max" + assert info.tflops_fp16 == AppleSiliconHardwareInfo._APPLE_PEAK_FP16["M5 Max"] + + def test_unknown_chip_estimates_from_gpu_cores(self): + """An unknown chip is estimated from GPU core count, not crashed on.""" + with ( + patch("parallax.server.server_info.psutil", None), + patch("subprocess.check_output", side_effect=_check_output_for("M6 Max", gpu_cores=40)), + ): + info = AppleSiliconHardwareInfo.detect() + expected = round(40 * AppleSiliconHardwareInfo._FALLBACK_FP16_PER_CORE, 2) + assert info.tflops_fp16 == expected + + def test_unknown_chip_without_core_count_uses_default(self): + """When GPU core count can't be read, a conservative default is used.""" + with ( + patch("parallax.server.server_info.psutil", None), + patch( + "subprocess.check_output", side_effect=_check_output_for("M6 Ultra", gpu_cores=None) + ), + ): + info = AppleSiliconHardwareInfo.detect() + assert info.tflops_fp16 == AppleSiliconHardwareInfo._FALLBACK_FP16_DEFAULT + + def test_unknown_chip_does_not_raise(self): + """Regression for #439: an unknown chip must not raise.""" + with ( + patch("parallax.server.server_info.psutil", None), + patch( + "subprocess.check_output", side_effect=_check_output_for("M99 Ultra", gpu_cores=24) + ), + ): + info = AppleSiliconHardwareInfo.detect() # must not raise + assert info.tflops_fp16 > 0 + + def test_m5_family_entries_present(self): + """The M5 family is tabulated (issue #439).""" + table = AppleSiliconHardwareInfo._APPLE_PEAK_FP16 + assert table["M5"] == 8.08 + assert table["M5 Pro"] == 16.16 + assert table["M5 Max"] == 32.32