Query Windows physical core count using ctypes and native Win32 API

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

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