From eb494e700591247a473f05564b6f81f1047e2f87 Mon Sep 17 00:00:00 2001
From: Anton Brand <anton.brand@ubisoft.com>
Date: Wed, 4 Dec 2024 14:06:24 -0500
Subject: [PATCH] [RQD] Fix core detection on Windows platforms (#1468)

**Link the Issue(s) this Pull Request is related to.**
#1469

**Summarize your change.**
Fix an error on Windows platforms where a submitted job could not be
picked up properly due to available processors incorrect detection.

The `reserveHT` method of the RQD `Machine` class relies on a
`__procs_by_physid_and_coreid` attribute that is not correctly filled
when on a Windows platform.

See:
- `reserveHT` relying on `__procs_by_physid_and_coreid`:
https://github.com/AcademySoftwareFoundation/OpenCue/blob/master/rqd/rqd/rqmachine.py#L842
- `__procs_by_physid_and_coreid` being filled only if on a Linux
platform:
https://github.com/AcademySoftwareFoundation/OpenCue/blob/master/rqd/rqd/rqmachine.py#L613

---------

Signed-off-by: Diego Tavares <dtavares@imageworks.com>
Co-authored-by: Diego Tavares <dtavares@imageworks.com>
Co-authored-by: Kern Attila GERMAIN <5556461+KernAttila@users.noreply.github.com>
---
 rqd/rqd/rqmachine.py | 75 +++++++++++++++++++++++++++++++++++---------
 rqd/setup.py         |  3 +-
 2 files changed, 63 insertions(+), 15 deletions(-)

diff --git a/rqd/rqd/rqmachine.py b/rqd/rqd/rqmachine.py
index a6d9a6317..8694afee4 100644
--- a/rqd/rqd/rqmachine.py
+++ b/rqd/rqd/rqmachine.py
@@ -670,20 +670,8 @@ def __initMachineStats(self, pathCpuInfo=None):
             hyperthreadingMultiplier = 1
 
         if platform.system() == 'Windows':
-            # Windows memory information
-            stat = self.getWindowsMemory()
-            TEMP_DEFAULT = 1048576
-            self.__renderHost.total_mcp = TEMP_DEFAULT
-            self.__renderHost.total_mem = int(stat.ullTotalPhys / 1024)
-            self.__renderHost.total_swap = int(stat.ullTotalPageFile / 1024)
-
-            # Windows CPU information
-            logical_core_count = psutil.cpu_count(logical=True)
-            actual_core_count = psutil.cpu_count(logical=False)
-            hyperthreadingMultiplier = logical_core_count // actual_core_count
-
-            __totalCores = logical_core_count * rqd.rqconstants.CORE_VALUE
-            __numProcs = 1  # TODO: figure out how to count sockets in Python
+            logicalCoreCount, __numProcs, hyperthreadingMultiplier = self.__initStatsFromWindows()
+            __totalCores = logicalCoreCount * rqd.rqconstants.CORE_VALUE
 
         # All other systems will just have one proc/core
         if not __numProcs or not __totalCores:
@@ -713,6 +701,65 @@ def __initMachineStats(self, pathCpuInfo=None):
         if hyperthreadingMultiplier >= 1:
             self.__renderHost.attributes['hyperthreadingMultiplier'] = str(hyperthreadingMultiplier)
 
+    def __initStatsFromWindows(self):
+        """Init machine stats for Windows platforms.
+
+        @rtype:  tuple
+        @return: A 3-items tuple containing:
+            - the number of logical cores
+            - the number of physical processors
+            - the hyper-threading multiplier
+        """
+        # Windows memory information
+        stat = self.getWindowsMemory()
+        TEMP_DEFAULT = 1048576
+        self.__renderHost.total_mcp = TEMP_DEFAULT
+        self.__renderHost.total_mem = int(stat.ullTotalPhys / 1024)
+        self.__renderHost.total_swap = int(stat.ullTotalPageFile / 1024)
+
+        # Windows CPU information
+        self.__updateProcsMappingsFromWindows()
+
+        logicalCoreCount = psutil.cpu_count(logical=True)
+        actualCoreCount = psutil.cpu_count(logical=False)
+        hyperThreadingMultiplier = logicalCoreCount // actualCoreCount
+
+        physicalProcessorCount = len(self.__procs_by_physid_and_coreid)
+
+        return logicalCoreCount, physicalProcessorCount, hyperThreadingMultiplier
+
+    def __updateProcsMappingsFromWindows(self):
+        """
+        Update `__procs_by_physid_and_coreid` and `__physid_and_coreid_by_proc` mappings
+        for Windows platforms.
+        """
+        # Windows-specific
+        import wmi  # pylint:disable=import-outside-toplevel,import-error
+
+        # Reset mappings
+        self.__procs_by_physid_and_coreid = {}
+        self.__physid_and_coreid_by_proc = {}
+
+        # Connect to the Windows Management Instrumentation (WMI) interface
+        wmiInstance = wmi.WMI()
+
+        # Retrieve CPU information using WMI
+        for physicalId, processor in enumerate(wmiInstance.Win32_Processor()):
+
+            threadPerCore = processor.NumberOfLogicalProcessors // processor.NumberOfCores
+            procId = 0
+
+            for coreId in range(processor.NumberOfCores):
+                for _ in range(threadPerCore):
+                    self.__procs_by_physid_and_coreid.setdefault(
+                        str(physicalId), {}
+                    ).setdefault(str(coreId), set()).add(str(procId))
+                    self.__physid_and_coreid_by_proc[str(procId)] = (
+                        str(physicalId),
+                        str(coreId),
+                    )
+                    procId += 1
+
     def getWindowsMemory(self):
         """Gets information on system memory, Windows compatible version."""
         # From
diff --git a/rqd/setup.py b/rqd/setup.py
index a9735b142..69c88c706 100644
--- a/rqd/setup.py
+++ b/rqd/setup.py
@@ -59,7 +59,8 @@
         'grpcio',
         'grpcio-tools',
         'psutil',
-        'pywin32; platform_system == "Windows"'
+        'pywin32==224; platform_system == "Windows"',
+        'wmi==1.5.1; platform_system == "Windows"'
     ]
 )