diff --git a/codecarbon/core/cpu.py b/codecarbon/core/cpu.py index 3ab71cfab..9ed09d20a 100644 --- a/codecarbon/core/cpu.py +++ b/codecarbon/core/cpu.py @@ -208,14 +208,29 @@ def is_rapl_available(rapl_dir: Optional[str] = None) -> bool: def is_psutil_available(): try: - nice = psutil.cpu_times().nice - if nice > 0.0001: - return True + cpu_times = psutil.cpu_times() + + # platforms like Windows do not have 'nice' attribute + if hasattr(cpu_times, "nice"): + nice = cpu_times.nice + if nice > 0.0001: + return True + else: + logger.debug( + f"is_psutil_available(): psutil.cpu_times().nice is too small: {nice}" + ) + return False + else: + # Fallback: check if psutil works by calling cpu_percent logger.debug( - f"is_psutil_available() : psutil.cpu_times().nice is too small : {nice} !" + "is_psutil_available(): no 'nice' attribute, using fallback check." ) - return False + + # check CPU utilization usable + psutil.cpu_percent(interval=0.0, percpu=False) + return True + except Exception as e: logger.debug( "Not using the psutil interface, an exception occurred while instantiating " diff --git a/codecarbon/core/util.py b/codecarbon/core/util.py index facb79ab5..b9ec93b7b 100644 --- a/codecarbon/core/util.py +++ b/codecarbon/core/util.py @@ -103,7 +103,7 @@ def count_physical_cpus(): import subprocess if platform.system() == "Windows": - return int(os.environ.get("NUMBER_OF_PROCESSORS", 1)) + return _windows_get_physical_sockets() else: try: output = subprocess.check_output(["lscpu"], text=True) @@ -119,6 +119,28 @@ def count_physical_cpus(): return 1 +def _windows_get_physical_sockets(): + try: + # use PowerShell to count number of objects of class Win32_Processor + cmd = [ + "powershell", + "-NoProfile", + "-Command", + "(Get-CimInstance -ClassName Win32_Processor).Count", + ] + result = subprocess.run( + cmd, capture_output=True, text=True, timeout=10, check=True + ) + + output = result.stdout.strip() or "1" + logger.debug(f"Detected {output} physical sockets on Windows.") + return int(output) + + except (subprocess.CalledProcessError, subprocess.TimeoutExpired, ValueError) as e: + logger.error(f"Error detecting physical sockets on Windows: {e}") + return 1 # Fallback:at least one socket + + def count_cpus() -> int: if SLURM_JOB_ID is None: return psutil.cpu_count() diff --git a/tests/test_cpu.py b/tests/test_cpu.py index 5453ac4ef..f7d34aca7 100644 --- a/tests/test_cpu.py +++ b/tests/test_cpu.py @@ -11,6 +11,7 @@ IntelPowerGadget, IntelRAPL, is_powergadget_available, + is_psutil_available, ) from codecarbon.core.units import Energy, Power, Time from codecarbon.core.util import count_physical_cpus @@ -18,6 +19,37 @@ from codecarbon.input import DataSource +class TestCPU(unittest.TestCase): + @mock.patch("psutil.cpu_times") + def test_is_psutil_available_with_nice(self, mock_cpu_times): + # Create a mock with 'nice' attribute + mock_times = mock.Mock() + mock_times.nice = 0.1 + mock_cpu_times.return_value = mock_times + self.assertTrue(is_psutil_available()) + + @mock.patch("psutil.cpu_times") + def test_is_psutil_available_with_small_nice(self, mock_cpu_times): + # Test when nice attribute is too small + mock_times = mock.Mock() + mock_times.nice = 0.00001 + mock_cpu_times.return_value = mock_times + self.assertFalse(is_psutil_available()) + + @mock.patch("psutil.cpu_times") + def test_is_psutil_available_without_nice(self, mock_cpu_times): + # Create a mock without 'nice' attribute (like Windows) + mock_times = mock.Mock(spec=[]) # Empty spec = no attributes + mock_cpu_times.return_value = mock_times + with mock.patch("psutil.cpu_percent") as mock_cpu_percent: + self.assertTrue(is_psutil_available()) + mock_cpu_percent.assert_called_once_with(interval=0.0, percpu=False) + + @mock.patch("psutil.cpu_times", side_effect=Exception("Test error")) + def test_is_psutil_not_available_on_exception(self, mock_cpu_times): + self.assertFalse(is_psutil_available()) + + class TestIntelPowerGadget(unittest.TestCase): @pytest.mark.integ_test def test_intel_power_gadget(self): @@ -306,10 +338,37 @@ def test_get_matching_cpu(self): class TestPhysicalCPU(unittest.TestCase): def test_count_physical_cpus_windows(self): with mock.patch("platform.system", return_value="Windows"): - with mock.patch.dict(os.environ, {"NUMBER_OF_PROCESSORS": "4"}): + + with mock.patch( + "subprocess.run", return_value=mock.Mock(returncode=0, stdout="4") + ): assert count_physical_cpus() == 4 - with mock.patch.dict(os.environ, {}, clear=True): + with mock.patch( + "subprocess.run", return_value=mock.Mock(returncode=0, stdout="") + ): + assert count_physical_cpus() == 1 + + def test_count_physical_cpus_windows_with_error(self): + with mock.patch("platform.system", return_value="Windows"): + # Test CalledProcessError + with mock.patch( + "subprocess.run", + side_effect=subprocess.CalledProcessError(1, "powershell"), + ): + assert count_physical_cpus() == 1 + + # Test TimeoutExpired + with mock.patch( + "subprocess.run", + side_effect=subprocess.TimeoutExpired("powershell", 10), + ): + assert count_physical_cpus() == 1 + + # Test ValueError when converting invalid output + with mock.patch( + "subprocess.run", return_value=mock.Mock(returncode=0, stdout="invalid") + ): assert count_physical_cpus() == 1 def test_count_physical_cpus_linux(self):