diff --git a/codecarbon/core/resource_tracker.py b/codecarbon/core/resource_tracker.py index c1ecc6b02..adef7b947 100644 --- a/codecarbon/core/resource_tracker.py +++ b/codecarbon/core/resource_tracker.py @@ -3,7 +3,13 @@ from codecarbon.core import cpu, gpu, powermetrics from codecarbon.core.config import parse_gpu_ids -from codecarbon.core.util import detect_cpu_model, is_linux_os, is_mac_os, is_windows_os +from codecarbon.core.util import ( + detect_cpu_model, + is_linux_os, + is_mac_arm, + is_mac_os, + is_windows_os, +) from codecarbon.external.hardware import CPU, GPU, MODE_CPU_LOAD, AppleSiliconChip from codecarbon.external.logger import logger from codecarbon.external.ram import RAM @@ -99,7 +105,7 @@ def _get_install_instructions(self): """Get CPU tracking installation instructions for the current OS.""" if is_mac_os(): cpu_model = detect_cpu_model() - if "M1" in cpu_model or "M2" in cpu_model or "M3" in cpu_model: + if cpu_model and is_mac_arm(cpu_model): return "Mac OS and ARM processor detected: Please enable PowerMetrics sudo to measure CPU" else: return "Mac OS detected: Please install Intel Power Gadget or enable PowerMetrics sudo to measure CPU" diff --git a/codecarbon/core/util.py b/codecarbon/core/util.py index b9ec93b7b..da13dd301 100644 --- a/codecarbon/core/util.py +++ b/codecarbon/core/util.py @@ -88,6 +88,10 @@ def is_mac_os() -> bool: return system.startswith("dar") +def is_mac_arm(cpu_model: str) -> bool: + return bool(re.search(r"\bM\d{1,2}\b", cpu_model)) + + def is_windows_os() -> bool: system = sys.platform.lower() return system.startswith("win") diff --git a/docs/introduction/methodology.md b/docs/introduction/methodology.md index d57cd9453..31769d3e6 100644 --- a/docs/introduction/methodology.md +++ b/docs/introduction/methodology.md @@ -186,7 +186,7 @@ Tracks Intel processors energy consumption using the . But has been discontinued. There is a discussion about it on [github issues #457](https://github.com/mlco2/codecarbon/issues/457). -- **Apple Silicon Chips (M1, M2)** +- **Apple Silicon Chips (M1, M2, M3, ...)** Apple Silicon Chips contain both the CPU and the GPU. diff --git a/tests/test_core_util.py b/tests/test_core_util.py index f22d87262..6c1ba6f14 100644 --- a/tests/test_core_util.py +++ b/tests/test_core_util.py @@ -1,7 +1,9 @@ import shutil import tempfile -from codecarbon.core.util import backup, detect_cpu_model, resolve_path +import pytest + +from codecarbon.core.util import backup, detect_cpu_model, is_mac_arm, resolve_path def test_detect_cpu_model_caching(): @@ -42,3 +44,31 @@ def test_backup(): backup(first_file.name) backup_of_backup_path = resolve_path(f"{first_file.name}_0.bak") assert backup_of_backup_path.exists() + + +@pytest.mark.parametrize( + "cpu_model, expected", + [ + # Apple Silicon chips that should match + ("Apple M1", True), + ("Apple M2", True), + ("Apple M3", True), + ("Apple M4", True), + ("Apple M1 Pro", True), + ("Apple M2 Max", True), + ("Apple M3 Ultra", True), + ("Apple M4 Pro", True), + ("Apple M10", True), + # Non-Apple ARM or unrelated chips that should NOT match + ("Intel Core i7-9750H", False), + ("AMD Ryzen 9 5900X", False), + ("Qualcomm Snapdragon 8cx Gen 3", False), + # Partial matches that should NOT match (no word boundary) + ("SuperM2000 Processor", False), + ("M2fast chip", False), + # Empty string + ("", False), + ], +) +def test_is_mac_arm(cpu_model, expected): + assert is_mac_arm(cpu_model) == expected diff --git a/tests/test_resource_tracker.py b/tests/test_resource_tracker.py new file mode 100644 index 000000000..f6c67df5e --- /dev/null +++ b/tests/test_resource_tracker.py @@ -0,0 +1,43 @@ +from unittest.mock import MagicMock, patch + +import pytest + +from codecarbon.core.resource_tracker import ResourceTracker + + +@pytest.mark.parametrize( + "is_mac, is_windows, is_linux, cpu_model, expected_fragment", + [ + # Mac + ARM chip + (True, False, False, "Apple M4", "PowerMetrics sudo"), + # Mac + Intel chip + (True, False, False, "Intel Core i7", "Intel Power Gadget"), + # Mac + cpu_model is None + (True, False, False, None, "Intel Power Gadget"), + # Windows + (False, True, False, "Intel Core i7", "Intel Power Gadget"), + # Linux + (False, False, True, "Intel Core i7", "RAPL"), + # Unknown OS + (False, False, False, "Intel Core i7", ""), + ], +) +def test_get_install_instructions( + is_mac, is_windows, is_linux, cpu_model, expected_fragment +): + tracker = MagicMock() + resource_tracker = ResourceTracker(tracker) + + with ( + patch("codecarbon.core.resource_tracker.is_mac_os", return_value=is_mac), + patch( + "codecarbon.core.resource_tracker.is_windows_os", return_value=is_windows + ), + patch("codecarbon.core.resource_tracker.is_linux_os", return_value=is_linux), + patch( + "codecarbon.core.resource_tracker.detect_cpu_model", return_value=cpu_model + ), + ): + result = resource_tracker._get_install_instructions() + + assert expected_fragment in result