From 749f73d7d906c4184996445b50e07f426adaf65b Mon Sep 17 00:00:00 2001 From: Nathan Rusch Date: Fri, 14 Aug 2026 20:51:43 -0700 Subject: [PATCH] Add Windows core count query method using `ctypes` and native Win32 API Signed-off-by: Nathan Rusch --- src/rez/tests/test_utils.py | 46 ++++++++++++--- src/rez/utils/platform_.py | 112 +++++++++++++++++++++++++++++++++++- 2 files changed, 150 insertions(+), 8 deletions(-) diff --git a/src/rez/tests/test_utils.py b/src/rez/tests/test_utils.py index fb16a3f88d..2a2967f3f9 100644 --- a/src/rez/tests/test_utils.py +++ b/src/rez/tests/test_utils.py @@ -127,33 +127,65 @@ def test_wmic_no_match_returns_none(self): # -- _physical_cores fallback chain --------------------------------------- - def test_physical_cores_prefers_powershell(self): - """powershell succeeds → wmic is never called.""" - with unittest.mock.patch.object(platform_, '_physical_cores_from_powershell', + def test_physical_cores_prefers_native(self): + """If native succeeds, powershell and wmic are never called.""" + with unittest.mock.patch.object(platform_, '_physical_cores_native', + return_value=4) as mock_native, \ + unittest.mock.patch.object(platform_, '_physical_cores_from_powershell', return_value=4) as mock_ps, \ unittest.mock.patch.object(platform_, '_physical_cores_from_wmic', return_value=4) as mock_wmic: result = platform_._physical_cores() + mock_native.assert_called_once() + mock_ps.assert_not_called() + mock_wmic.assert_not_called() + self.assertEqual(result, 4) + + def test_physical_cores_powershell_fallback(self): + """If native fails and powershell succeeds, wmic is never called.""" + with unittest.mock.patch.object(platform_, '_physical_cores_native', + return_value=None) as mock_native, \ + unittest.mock.patch.object(platform_, '_physical_cores_from_powershell', + return_value=4) as mock_ps, \ + unittest.mock.patch.object(platform_, '_physical_cores_from_wmic', + return_value=4) as mock_wmic: + result = platform_._physical_cores() + mock_native.assert_called_once() mock_ps.assert_called_once() mock_wmic.assert_not_called() self.assertEqual(result, 4) def test_physical_cores_falls_back_to_wmic(self): """powershell returns None (e.g. older Windows) → wmic result used.""" - with unittest.mock.patch.object(platform_, '_physical_cores_from_powershell', + with unittest.mock.patch.object(platform_, '_physical_cores_native', + return_value=None), \ + unittest.mock.patch.object(platform_, '_physical_cores_from_powershell', return_value=None), \ unittest.mock.patch.object(platform_, '_physical_cores_from_wmic', return_value=4): self.assertEqual(platform_._physical_cores(), 4) - def test_physical_cores_both_fail_returns_none(self): - """Both helpers return None → _physical_cores returns None.""" - with unittest.mock.patch.object(platform_, '_physical_cores_from_powershell', + def test_physical_cores_all_fail_returns_none(self): + """All helpers return None → _physical_cores returns None.""" + with unittest.mock.patch.object(platform_, '_physical_cores_native', + return_value=None), \ + unittest.mock.patch.object(platform_, '_physical_cores_from_powershell', return_value=None), \ unittest.mock.patch.object(platform_, '_physical_cores_from_wmic', return_value=None): self.assertIsNone(platform_._physical_cores()) + # -- _physical_cores internal consistency --------------------------------- + + def test_physical_cores_methods_same_result(self): + """All helpers that succeed return the same result.""" + native_cores = platform_._physical_cores_native() + pwsh_cores = platform_._physical_cores_from_powershell() + wmic_cores = platform_._physical_cores_from_wmic() + values = {native_cores, pwsh_cores, wmic_cores} + values.discard(None) + assert len(values) < 2 + # -- integration smoke-tests ---------------------------------------------- def test_logical_cores_is_positive(self): diff --git a/src/rez/utils/platform_.py b/src/rez/utils/platform_.py index bc3af7eb5d..53afdb10a7 100644 --- a/src/rez/utils/platform_.py +++ b/src/rez/utils/platform_.py @@ -516,6 +516,114 @@ def symlink(self, source: str, link_name: str): def _terminal_emulator_command(self) -> str: return "START" + def _physical_cores_native(self) -> int | None: + """Query the physical CPU count using `ctypes` wrappers for the Win32 API.""" + # try-except everything in case rez is running in an interpreter without `ctypes`. + try: + import ctypes + from ctypes import wintypes + + kernel32 = ctypes.WinDLL("kernel32", use_last_error=True) + + # Taken from the `GetLogicalProcessorInformationEx` API reference: + # https://learn.microsoft.com/en-us/windows/win32/api/sysinfoapi/nf-sysinfoapi-getlogicalprocessorinformationex + # > When this function is called with a relationship type of RelationProcessorCore, it + # > returns a PROCESSOR_RELATIONSHIP structure for every active processor core in every + # > processor group in the system. + RelationProcessorCore = 0 + + # Taken from the "System Error Codes" list: + # https://learn.microsoft.com/en-us/windows/win32/debug/system-error-codes--0-499- + ERROR_INSUFFICIENT_BUFFER = 0x7A + + class SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX(ctypes.Structure): + """ + typedef struct _SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX { + LOGICAL_PROCESSOR_RELATIONSHIP Relationship; + DWORD Size; + union { + PROCESSOR_RELATIONSHIP Processor; + NUMA_NODE_RELATIONSHIP NumaNode; + CACHE_RELATIONSHIP Cache; + GROUP_RELATIONSHIP Group; + SHARED_COMPUTE_UNIT_RELATIONSHIP SharedComputeUnit; + } DUMMYUNIONNAME; + } SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX; + """ + + _fields_ = [ + ("Relationship", ctypes.c_int), # LOGICAL_PROCESSOR_RELATIONSHIP enum + ("Size", wintypes.DWORD), + # NOTE: We only ever instantiate this struct from an offset into an existing + # buffer. We need the `Relationship` field for counting cores, and the `Size` + # field for computing buffer strides, but we can omit any remaining fields. + ] + + # BOOL GetLogicalProcessorInformationEx( + # [in] LOGICAL_PROCESSOR_RELATIONSHIP RelationshipType, + # [out, optional] PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX Buffer, + # [in, out] PDWORD ReturnedLength + # ); + GetLogicalProcessorInformationEx = kernel32.GetLogicalProcessorInformationEx + GetLogicalProcessorInformationEx.argtypes = [ + ctypes.c_int, # RelationshipType + ctypes.c_void_p, # Buffer + ctypes.POINTER(wintypes.DWORD), # ReturnedLength + ] + GetLogicalProcessorInformationEx.restype = wintypes.BOOL + + buffer_size = wintypes.DWORD(0) + + # Call `GetLogicalProcessorInformationEx` once with a null buffer, which will "fail" + # with a known error, but will set the `ReturnedLength` pointer's value to the required + # buffer size. + success = GetLogicalProcessorInformationEx( + RelationProcessorCore, + None, + ctypes.byref(buffer_size), + ) + if success: + # XXX: This should never happen, but we'll be paranoid. + return None + + if ctypes.get_last_error() != ERROR_INSUFFICIENT_BUFFER: + # The call failed for an unknown reason. + return None + + if not buffer_size.value: + # XXX: Again, I don't think this can ever happen. + return None + + # Now we can allocate a buffer of the required size and call the function again. + buffer = (ctypes.c_byte * buffer_size.value)() + success = GetLogicalProcessorInformationEx( + RelationProcessorCore, + buffer, + ctypes.byref(buffer_size), + ) + if not success: + return None + + # Walk the buffer of `SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX` struct instances and + # count the number of matching entries. + offset = 0 + num_cpus = 0 + buffer_start = ctypes.addressof(buffer) + + while offset < buffer_size.value: + proc_info = SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX.from_address(buffer_start + offset) + + # This check should be redundant given that we're already passing this relationship + # type to `GetLogicalProcessorInformationEx`, but just to be safe... + if proc_info.Relationship == RelationProcessorCore: + num_cpus += 1 + + offset += proc_info.Size + + return num_cpus if num_cpus else None + except Exception: + return None + def _physical_cores_from_powershell(self) -> int | None: # wmic was removed in Windows 11 24H2; use PowerShell/CimInstance instead. # powershell.exe (Windows PowerShell 5.1) ships with all Windows 10/11 installs. @@ -570,7 +678,9 @@ def _physical_cores_from_wmic(self) -> int | None: return sum(map(int, result)) def _physical_cores(self) -> int | None: - return self._physical_cores_from_powershell() or self._physical_cores_from_wmic() + return ( + self._physical_cores_native() or self._physical_cores_from_powershell() or self._physical_cores_from_wmic() + ) def _difftool(self): # although meld would be preferred, fc ships with all Windows versions back to DOS