Merge pull request #1054 from arno-hc/experiment

Fix issues for Windows hosts

Author: benoit-cty

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 "

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()

tests/test_cpu.py

@@ -11,13 +11,45 @@
     IntelPowerGadget,
     IntelRAPL,
     is_powergadget_available,
+    is_psutil_available,
 )
 from codecarbon.core.units import Energy, Power, Time
 from codecarbon.core.util import count_physical_cpus
 from codecarbon.external.hardware import CPU
 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):