Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 20 additions & 5 deletions codecarbon/core/cpu.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 "
Expand Down
24 changes: 23 additions & 1 deletion codecarbon/core/util.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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()
Expand Down
63 changes: 61 additions & 2 deletions tests/test_cpu.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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):
Expand Down
Loading