From bcce6ecbaa3b72f519d10458b4d025bf380d2798 Mon Sep 17 00:00:00 2001 From: Alexandra Bara Date: Wed, 8 Oct 2025 19:24:03 -0500 Subject: [PATCH 1/6] added plugin from internal --- .../inband/device_enumeration/__init__.py | 25 ++++ .../device_enumeration/analyzer_args.py | 50 +++++++ .../device_enumeration_analyzer.py | 84 ++++++++++++ .../device_enumeration_collector.py | 124 ++++++++++++++++++ .../device_enumeration_plugin.py | 45 +++++++ .../device_enumeration/deviceenumdata.py | 32 +++++ 6 files changed, 360 insertions(+) create mode 100644 nodescraper/plugins/inband/device_enumeration/__init__.py create mode 100644 nodescraper/plugins/inband/device_enumeration/analyzer_args.py create mode 100644 nodescraper/plugins/inband/device_enumeration/device_enumeration_analyzer.py create mode 100644 nodescraper/plugins/inband/device_enumeration/device_enumeration_collector.py create mode 100644 nodescraper/plugins/inband/device_enumeration/device_enumeration_plugin.py create mode 100644 nodescraper/plugins/inband/device_enumeration/deviceenumdata.py diff --git a/nodescraper/plugins/inband/device_enumeration/__init__.py b/nodescraper/plugins/inband/device_enumeration/__init__.py new file mode 100644 index 00000000..fdb6d9c8 --- /dev/null +++ b/nodescraper/plugins/inband/device_enumeration/__init__.py @@ -0,0 +1,25 @@ +############################################################################### +# +# MIT License +# +# Copyright (c) 2025 Advanced Micro Devices, Inc. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +############################################################################### diff --git a/nodescraper/plugins/inband/device_enumeration/analyzer_args.py b/nodescraper/plugins/inband/device_enumeration/analyzer_args.py new file mode 100644 index 00000000..77913fe8 --- /dev/null +++ b/nodescraper/plugins/inband/device_enumeration/analyzer_args.py @@ -0,0 +1,50 @@ +############################################################################### +# +# MIT License +# +# Copyright (c) 2025 Advanced Micro Devices, Inc. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +############################################################################### + +from nodescraper.models import AnalyzerArgs +from nodescraper.plugins.inband.device_enumeration.deviceenumdata import ( + DeviceEnumerationDataModel, +) + + +class DeviceEnumerationAnalyzerArgs(AnalyzerArgs): + cpu_count: list | int = (None,) + gpu_count: list | int = (None,) + vf_count: list | int = (None,) + + @classmethod + def build_from_model( + cls, datamodel: DeviceEnumerationDataModel + ) -> "DeviceEnumerationAnalyzerArgs": + """build analyzer args from data model + + Args: + datamodel (DeviceEnumerationDataModel): data model for plugin + + Returns: + DeviceEnumerationAnalyzerArgs: instance of analyzer args class + """ + return cls(exp_bios_version=datamodel.bios_version) diff --git a/nodescraper/plugins/inband/device_enumeration/device_enumeration_analyzer.py b/nodescraper/plugins/inband/device_enumeration/device_enumeration_analyzer.py new file mode 100644 index 00000000..173c7f4d --- /dev/null +++ b/nodescraper/plugins/inband/device_enumeration/device_enumeration_analyzer.py @@ -0,0 +1,84 @@ +############################################################################### +# +# MIT License +# +# Copyright (c) 2025 Advanced Micro Devices, Inc. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +########################### + +from errorscraper.event import EventCategory, EventPriority + +from nodescraper.interfaces import DataAnalyzer +from nodescraper.models import TaskResult, TaskStatus +from nodescraper.plugins.inband.deviceenumdata import DeviceEnumerationDataModel + + +class DeviceEnumerationAnalyzer( + DataAnalyzer[DeviceEnumerationDataModel, DeviceEnumerationAnalyzerArgs] +): + """Check Device Enumeration matches expected cpu and gpu count + supported by all OSs, SKUs, and platforms.""" + + DATA_MODEL = DeviceEnumerationDataModel + + def analyze_data( + self, data: DeviceEnumerationDataModel, args: Optional[DeviceEnumerationAnalyzerArgs] = None + ) -> TaskResult: + + if not args: + self.result.message = "Expected Device Enumeration expected data not provided" + self.result.status = TaskStatus.NOT_RAN + return self.result + + if isinstance(args.cpu_count, int): + cpu_count = [args.cpu_count] + if isinstance(args.gpu_count, int): + gpu_count = [args.gpu_count] + if isinstance(args.vf_count, int): + vf_count = [args.vf_count] + + checks = {} + if cpu_count not in [None, []]: + checks["cpu_count"] = cpu_count + if gpu_count not in [None, []]: + checks["gpu_count"] = gpu_count + if vf_count not in [None, []]: + checks["vf_count"] = vf_count + + self.result.message = "" + for check, accepted_counts in checks.items(): + actual_count = getattr(data, check) + if actual_count not in accepted_counts: + message = f"Expected {check} in {accepted_counts}, but got {actual_count}. " + self.result.message += message + self.result.status = TaskStatus.ERRORS_DETECTED + self._log_event( + category=EventCategory.PLATFORM, + description=message, + data={check: actual_count}, + priority=EventPriority.CRITICAL, + console_log=True, + ) + if self.result.message == "": + self.result.status = TaskStatus.OK + self.result.message = f"Device Enumeration validated on {checks.keys()}." + + return self.result diff --git a/nodescraper/plugins/inband/device_enumeration/device_enumeration_collector.py b/nodescraper/plugins/inband/device_enumeration/device_enumeration_collector.py new file mode 100644 index 00000000..f387851f --- /dev/null +++ b/nodescraper/plugins/inband/device_enumeration/device_enumeration_collector.py @@ -0,0 +1,124 @@ +############################################################################### +# +# MIT License +# +# Copyright (c) 2025 Advanced Micro Devices, Inc. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +############################################################################### + + +from nodescraper.base import InBandDataCollector +from nodescraper.connection.inband.inband import CommandArtifact +from nodescraper.enums import EventCategory, EventPriority, OSFamily +from nodescraper.models import TaskResult, TaskStatus +from nodescraper.plugins.inband.devenumdata import DeviceEnumerationDataModel + + +class DeviceEnumerationCollector(InBandDataCollector[DeviceEnumerationDataModel]): + """Collect CPU and GPU count""" + + DATA_MODEL = DeviceEnumerationDataModel + + def _warning( + self, + description: str, + command: CommandArtifact, + category: EventCategory = EventCategory.PLATFORM, + ): + self._log_event( + category=category, + description=description, + data={ + "command": command.command, + "stdout": command.stdout, + "stderr": command.stderr, + "exit_code": command.exit_code, + }, + priority=EventPriority.WARNING, + ) + + def collect_data(self, args=None) -> tuple[TaskResult, DeviceEnumerationDataModel | None]: + """ + Read CPU and GPU count + On Linux, use lscpu and lspci + On Windows, use WMI and hyper-v cmdlets + """ + device_enum = None + if self.system_info.os_family == OSFamily.LINUX: + baremetal_device_id = self.system_info.devid_ep + sriov_device_id = self.system_info.devid_ep_vf + + cpu_count_res = self._run_system_command("lscpu | grep Socket | awk '{ print $2 }'") + gpu_count_res = self._run_system_command(f"lspci -d :{baremetal_device_id} | wc -l") + vf_count_res = self._run_system_command(f"lspci -d :{sriov_device_id} | wc -l") + else: + cpu_count_res = self._run_system_command( + 'powershell -Command "(Get-WmiObject -Class Win32_Processor | Measure-Object).Count"' + ) + gpu_count_res = self._run_system_command( + 'powershell -Command "(wmic path win32_VideoController get name | findstr AMD | Measure-Object).Count"' + ) + vf_count_res = self._run_system_command( + 'powershell -Command "(Get-VMHostPartitionableGpu | Measure-Object).Count"' + ) + cpu_count, gpu_count, vf_count = [None, None, None] + + if cpu_count_res.exit_code == 0: + cpu_count = int(cpu_count_res.stdout) + else: + self._warning(description="Cannot determine CPU count", command=cpu_count_res) + + if gpu_count_res.exit_code == 0: + gpu_count = int(gpu_count_res.stdout) + else: + self._warning(description="Cannot determine GPU count", command=gpu_count_res) + + if vf_count_res.exit_code == 0: + vf_count = int(vf_count_res.stdout) + else: + self._warning( + description="Cannot determine VF count", + command=vf_count_res, + category=EventCategory.SW_DRIVER, + ) + + if cpu_count or gpu_count or vf_count: + device_enum = DeviceEnumerationDataModel( + cpu_count=cpu_count, gpu_count=gpu_count, vf_count=vf_count + ) + self._log_event( + category=EventCategory.PLATFORM, + description=f"Counted {device_enum.cpu_count} CPUs, {device_enum.gpu_count} GPUs, {device_enum.vf_count} VFs", + data=device_enum.model_dump(exclude_none=True), + priority=EventPriority.INFO, + ) + self.result.message = f"Device Enumeration: {device_enum.model_dump(exclude_none=True)}" + self.result.status = TaskStatus.OK + else: + self.result.message = "Device Enumeration info not found" + self.result.status = TaskStatus.EXECUTION_FAILURE + self._log_event( + category=EventCategory.SW_DRIVER, + description=self.result.message, + priority=EventPriority.CRITICAL, + ) + + return self.result, device_enum diff --git a/nodescraper/plugins/inband/device_enumeration/device_enumeration_plugin.py b/nodescraper/plugins/inband/device_enumeration/device_enumeration_plugin.py new file mode 100644 index 00000000..3bb14467 --- /dev/null +++ b/nodescraper/plugins/inband/device_enumeration/device_enumeration_plugin.py @@ -0,0 +1,45 @@ +############################################################################### +# +# MIT License +# +# Copyright (c) 2025 Advanced Micro Devices, Inc. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +############################################################################### +from nodescraper.base import InBandDataPlugin + +from .analyzer_args import DeviceEnumerationAnalyzerArgs +from .device_enumeration_analyzer import DeviceEnumerationAnalyzer +from .device_enumeration_collector import DeviceEnumerationCollector +from .device_enumerationdata import DeviceEnumerationDataModel + + +class DeviceEnumerationPlugin( + InBandDataPlugin[DeviceEnumerationDataModel, None, DeviceEnumerationAnalyzerArgs] +): + """Plugin for collection and analysis of BIOS data""" + + DATA_MODEL = DeviceEnumerationDataModel + + COLLECTOR = DeviceEnumerationCollector + + ANALYZER = DeviceEnumerationAnalyzer + + ANALYZER_ARGS = DeviceEnumerationAnalyzerArgs diff --git a/nodescraper/plugins/inband/device_enumeration/deviceenumdata.py b/nodescraper/plugins/inband/device_enumeration/deviceenumdata.py new file mode 100644 index 00000000..3a6082df --- /dev/null +++ b/nodescraper/plugins/inband/device_enumeration/deviceenumdata.py @@ -0,0 +1,32 @@ +############################################################################### +# +# MIT License +# +# Copyright (c) 2025 Advanced Micro Devices, Inc. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +############################################################################### +from nodescraper.models import DataModel + + +class DeviceEnumerationDataModel(DataModel): + cpu_count: int | None = None + gpu_count: int | None = None + vf_count: int | None = None From de1d824becd71de324694eaef084e9bd316de6ee Mon Sep 17 00:00:00 2001 From: Alexandra Bara Date: Mon, 17 Nov 2025 18:06:37 -0600 Subject: [PATCH 2/6] utest + picking up info through lspci --- .../inband/device_enumeration/__init__.py | 4 + .../device_enumeration/analyzer_args.py | 18 ++- .../device_enumeration_analyzer.py | 34 ++--- .../device_enumeration_collector.py | 55 ++++--- .../device_enumeration_plugin.py | 2 +- .../device_enumeration/deviceenumdata.py | 8 +- .../test_device_enumeration_analyzer.py | 127 ++++++++++++++++ .../test_device_enumeration_collector.py | 143 ++++++++++++++++++ 8 files changed, 341 insertions(+), 50 deletions(-) create mode 100644 test/unit/plugin/test_device_enumeration_analyzer.py create mode 100644 test/unit/plugin/test_device_enumeration_collector.py diff --git a/nodescraper/plugins/inband/device_enumeration/__init__.py b/nodescraper/plugins/inband/device_enumeration/__init__.py index fdb6d9c8..a5073399 100644 --- a/nodescraper/plugins/inband/device_enumeration/__init__.py +++ b/nodescraper/plugins/inband/device_enumeration/__init__.py @@ -23,3 +23,7 @@ # SOFTWARE. # ############################################################################### +from .analyzer_args import DeviceEnumerationAnalyzerArgs +from .device_enumeration_plugin import DeviceEnumerationPlugin + +__all__ = ["DeviceEnumerationPlugin", "DeviceEnumerationAnalyzerArgs"] diff --git a/nodescraper/plugins/inband/device_enumeration/analyzer_args.py b/nodescraper/plugins/inband/device_enumeration/analyzer_args.py index 77913fe8..8c3303cb 100644 --- a/nodescraper/plugins/inband/device_enumeration/analyzer_args.py +++ b/nodescraper/plugins/inband/device_enumeration/analyzer_args.py @@ -23,17 +23,17 @@ # SOFTWARE. # ############################################################################### +from typing import Optional, Union from nodescraper.models import AnalyzerArgs -from nodescraper.plugins.inband.device_enumeration.deviceenumdata import ( - DeviceEnumerationDataModel, -) + +from .deviceenumdata import DeviceEnumerationDataModel class DeviceEnumerationAnalyzerArgs(AnalyzerArgs): - cpu_count: list | int = (None,) - gpu_count: list | int = (None,) - vf_count: list | int = (None,) + cpu_count: Optional[Union[list[int], int]] = None + gpu_count: Optional[Union[list[int], int]] = None + vf_count: Optional[Union[list[int], int]] = None @classmethod def build_from_model( @@ -47,4 +47,8 @@ def build_from_model( Returns: DeviceEnumerationAnalyzerArgs: instance of analyzer args class """ - return cls(exp_bios_version=datamodel.bios_version) + return cls( + cpu_count=datamodel.cpu_count, + gpu_count=datamodel.gpu_count, + vf_count=datamodel.vf_count, + ) diff --git a/nodescraper/plugins/inband/device_enumeration/device_enumeration_analyzer.py b/nodescraper/plugins/inband/device_enumeration/device_enumeration_analyzer.py index 173c7f4d..3cdfa802 100644 --- a/nodescraper/plugins/inband/device_enumeration/device_enumeration_analyzer.py +++ b/nodescraper/plugins/inband/device_enumeration/device_enumeration_analyzer.py @@ -22,13 +22,15 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. # -########################### - -from errorscraper.event import EventCategory, EventPriority +############################################################################### +from typing import Optional +from nodescraper.enums import EventCategory, EventPriority, ExecutionStatus from nodescraper.interfaces import DataAnalyzer -from nodescraper.models import TaskResult, TaskStatus -from nodescraper.plugins.inband.deviceenumdata import DeviceEnumerationDataModel +from nodescraper.models import TaskResult + +from .analyzer_args import DeviceEnumerationAnalyzerArgs +from .deviceenumdata import DeviceEnumerationDataModel class DeviceEnumerationAnalyzer( @@ -45,22 +47,20 @@ def analyze_data( if not args: self.result.message = "Expected Device Enumeration expected data not provided" - self.result.status = TaskStatus.NOT_RAN + self.result.status = ExecutionStatus.NOT_RAN return self.result - if isinstance(args.cpu_count, int): - cpu_count = [args.cpu_count] - if isinstance(args.gpu_count, int): - gpu_count = [args.gpu_count] - if isinstance(args.vf_count, int): - vf_count = [args.vf_count] + # Convert to lists if integers, otherwise use as-is + cpu_count = [args.cpu_count] if isinstance(args.cpu_count, int) else args.cpu_count + gpu_count = [args.gpu_count] if isinstance(args.gpu_count, int) else args.gpu_count + vf_count = [args.vf_count] if isinstance(args.vf_count, int) else args.vf_count checks = {} - if cpu_count not in [None, []]: + if cpu_count is not None and cpu_count != []: checks["cpu_count"] = cpu_count - if gpu_count not in [None, []]: + if gpu_count is not None and gpu_count != []: checks["gpu_count"] = gpu_count - if vf_count not in [None, []]: + if vf_count is not None and vf_count != []: checks["vf_count"] = vf_count self.result.message = "" @@ -69,7 +69,7 @@ def analyze_data( if actual_count not in accepted_counts: message = f"Expected {check} in {accepted_counts}, but got {actual_count}. " self.result.message += message - self.result.status = TaskStatus.ERRORS_DETECTED + self.result.status = ExecutionStatus.ERROR self._log_event( category=EventCategory.PLATFORM, description=message, @@ -78,7 +78,7 @@ def analyze_data( console_log=True, ) if self.result.message == "": - self.result.status = TaskStatus.OK + self.result.status = ExecutionStatus.OK self.result.message = f"Device Enumeration validated on {checks.keys()}." return self.result diff --git a/nodescraper/plugins/inband/device_enumeration/device_enumeration_collector.py b/nodescraper/plugins/inband/device_enumeration/device_enumeration_collector.py index f387851f..42e5e42e 100644 --- a/nodescraper/plugins/inband/device_enumeration/device_enumeration_collector.py +++ b/nodescraper/plugins/inband/device_enumeration/device_enumeration_collector.py @@ -23,20 +23,35 @@ # SOFTWARE. # ############################################################################### - +from typing import Optional from nodescraper.base import InBandDataCollector from nodescraper.connection.inband.inband import CommandArtifact -from nodescraper.enums import EventCategory, EventPriority, OSFamily -from nodescraper.models import TaskResult, TaskStatus -from nodescraper.plugins.inband.devenumdata import DeviceEnumerationDataModel +from nodescraper.enums import EventCategory, EventPriority, ExecutionStatus, OSFamily +from nodescraper.models import TaskResult + +from .deviceenumdata import DeviceEnumerationDataModel -class DeviceEnumerationCollector(InBandDataCollector[DeviceEnumerationDataModel]): +class DeviceEnumerationCollector(InBandDataCollector[DeviceEnumerationDataModel, None]): """Collect CPU and GPU count""" DATA_MODEL = DeviceEnumerationDataModel + # Linux commands + CMD_CPU_COUNT_LINUX = "lscpu | grep Socket | awk '{ print $2 }'" + CMD_GPU_COUNT_LINUX = "lspci -d 1002: | grep -i 'VGA\\|Display\\|3D' | wc -l" + CMD_VF_COUNT_LINUX = "lspci -d 1002: | grep -i 'Virtual Function' | wc -l" + + # Windows commands + CMD_CPU_COUNT_WINDOWS = ( + 'powershell -Command "(Get-WmiObject -Class Win32_Processor | Measure-Object).Count"' + ) + CMD_GPU_COUNT_WINDOWS = 'powershell -Command "(wmic path win32_VideoController get name | findstr AMD | Measure-Object).Count"' + CMD_VF_COUNT_WINDOWS = ( + 'powershell -Command "(Get-VMHostPartitionableGpu | Measure-Object).Count"' + ) + def _warning( self, description: str, @@ -55,7 +70,7 @@ def _warning( priority=EventPriority.WARNING, ) - def collect_data(self, args=None) -> tuple[TaskResult, DeviceEnumerationDataModel | None]: + def collect_data(self, args=None) -> tuple[TaskResult, Optional[DeviceEnumerationDataModel]]: """ Read CPU and GPU count On Linux, use lscpu and lspci @@ -63,22 +78,18 @@ def collect_data(self, args=None) -> tuple[TaskResult, DeviceEnumerationDataMode """ device_enum = None if self.system_info.os_family == OSFamily.LINUX: - baremetal_device_id = self.system_info.devid_ep - sriov_device_id = self.system_info.devid_ep_vf + # Count CPU sockets + cpu_count_res = self._run_sut_cmd(self.CMD_CPU_COUNT_LINUX) - cpu_count_res = self._run_system_command("lscpu | grep Socket | awk '{ print $2 }'") - gpu_count_res = self._run_system_command(f"lspci -d :{baremetal_device_id} | wc -l") - vf_count_res = self._run_system_command(f"lspci -d :{sriov_device_id} | wc -l") + # Count all AMD GPUs (vendor ID 1002) + gpu_count_res = self._run_sut_cmd(self.CMD_GPU_COUNT_LINUX) + + # Count AMD Virtual Functions + vf_count_res = self._run_sut_cmd(self.CMD_VF_COUNT_LINUX) else: - cpu_count_res = self._run_system_command( - 'powershell -Command "(Get-WmiObject -Class Win32_Processor | Measure-Object).Count"' - ) - gpu_count_res = self._run_system_command( - 'powershell -Command "(wmic path win32_VideoController get name | findstr AMD | Measure-Object).Count"' - ) - vf_count_res = self._run_system_command( - 'powershell -Command "(Get-VMHostPartitionableGpu | Measure-Object).Count"' - ) + cpu_count_res = self._run_sut_cmd(self.CMD_CPU_COUNT_WINDOWS) + gpu_count_res = self._run_sut_cmd(self.CMD_GPU_COUNT_WINDOWS) + vf_count_res = self._run_sut_cmd(self.CMD_VF_COUNT_WINDOWS) cpu_count, gpu_count, vf_count = [None, None, None] if cpu_count_res.exit_code == 0: @@ -111,10 +122,10 @@ def collect_data(self, args=None) -> tuple[TaskResult, DeviceEnumerationDataMode priority=EventPriority.INFO, ) self.result.message = f"Device Enumeration: {device_enum.model_dump(exclude_none=True)}" - self.result.status = TaskStatus.OK + self.result.status = ExecutionStatus.OK else: self.result.message = "Device Enumeration info not found" - self.result.status = TaskStatus.EXECUTION_FAILURE + self.result.status = ExecutionStatus.EXECUTION_FAILURE self._log_event( category=EventCategory.SW_DRIVER, description=self.result.message, diff --git a/nodescraper/plugins/inband/device_enumeration/device_enumeration_plugin.py b/nodescraper/plugins/inband/device_enumeration/device_enumeration_plugin.py index 3bb14467..baf2aa2d 100644 --- a/nodescraper/plugins/inband/device_enumeration/device_enumeration_plugin.py +++ b/nodescraper/plugins/inband/device_enumeration/device_enumeration_plugin.py @@ -28,7 +28,7 @@ from .analyzer_args import DeviceEnumerationAnalyzerArgs from .device_enumeration_analyzer import DeviceEnumerationAnalyzer from .device_enumeration_collector import DeviceEnumerationCollector -from .device_enumerationdata import DeviceEnumerationDataModel +from .deviceenumdata import DeviceEnumerationDataModel class DeviceEnumerationPlugin( diff --git a/nodescraper/plugins/inband/device_enumeration/deviceenumdata.py b/nodescraper/plugins/inband/device_enumeration/deviceenumdata.py index 3a6082df..74939209 100644 --- a/nodescraper/plugins/inband/device_enumeration/deviceenumdata.py +++ b/nodescraper/plugins/inband/device_enumeration/deviceenumdata.py @@ -23,10 +23,12 @@ # SOFTWARE. # ############################################################################### +from typing import Optional + from nodescraper.models import DataModel class DeviceEnumerationDataModel(DataModel): - cpu_count: int | None = None - gpu_count: int | None = None - vf_count: int | None = None + cpu_count: Optional[int] = None + gpu_count: Optional[int] = None + vf_count: Optional[int] = None diff --git a/test/unit/plugin/test_device_enumeration_analyzer.py b/test/unit/plugin/test_device_enumeration_analyzer.py new file mode 100644 index 00000000..8f91fadc --- /dev/null +++ b/test/unit/plugin/test_device_enumeration_analyzer.py @@ -0,0 +1,127 @@ +############################################################################### +# +# MIT License +# +# Copyright (c) 2025 Advanced Micro Devices, Inc. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +############################################################################### +import pytest + +from nodescraper.enums.eventcategory import EventCategory +from nodescraper.enums.eventpriority import EventPriority +from nodescraper.enums.executionstatus import ExecutionStatus +from nodescraper.models.systeminfo import OSFamily +from nodescraper.plugins.inband.device_enumeration.analyzer_args import ( + DeviceEnumerationAnalyzerArgs, +) +from nodescraper.plugins.inband.device_enumeration.device_enumeration_analyzer import ( + DeviceEnumerationAnalyzer, +) +from nodescraper.plugins.inband.device_enumeration.deviceenumdata import ( + DeviceEnumerationDataModel, +) + + +@pytest.fixture +def device_enumeration_analyzer(system_info): + return DeviceEnumerationAnalyzer(system_info=system_info) + + +@pytest.fixture +def device_enumeration_data(): + return DeviceEnumerationDataModel(cpu_count=4, gpu_count=4, vf_count=8) + + +def test_analyze_passing_linux(system_info, device_enumeration_analyzer, device_enumeration_data): + """Test a normal passing case with matching config""" + system_info.os_family = OSFamily.LINUX + + args = DeviceEnumerationAnalyzerArgs(cpu_count=4, gpu_count=4, vf_count=8) + + result = device_enumeration_analyzer.analyze_data(data=device_enumeration_data, args=args) + + assert result.status == ExecutionStatus.OK + assert len(result.events) == 0 + + +def test_analyze_passing_windows(system_info, device_enumeration_analyzer, device_enumeration_data): + """Test a normal passing case on Windows""" + system_info.os_family = OSFamily.WINDOWS + + args = DeviceEnumerationAnalyzerArgs(gpu_count=4, vf_count=8) + + result = device_enumeration_analyzer.analyze_data(data=device_enumeration_data, args=args) + + assert result.status == ExecutionStatus.OK + assert len(result.events) == 0 + + +def test_analyze_no_args(device_enumeration_analyzer, device_enumeration_data): + """Test with no analyzer args provided""" + + result = device_enumeration_analyzer.analyze_data(data=device_enumeration_data, args=None) + + assert result.status == ExecutionStatus.NOT_RAN + assert "Expected Device Enumeration expected data not provided" in result.message + assert len(result.events) == 0 + + +def test_analyze_unexpected_counts(device_enumeration_analyzer, device_enumeration_data): + """Test with config specifying different device counts""" + + args = DeviceEnumerationAnalyzerArgs(cpu_count=1, gpu_count=10) + + result = device_enumeration_analyzer.analyze_data(data=device_enumeration_data, args=args) + + assert result.status == ExecutionStatus.ERROR + assert "but got" in result.message + + for event in result.events: + assert event.priority == EventPriority.CRITICAL + assert event.category == EventCategory.PLATFORM.value + + +def test_analyze_mismatched_cpu_count(device_enumeration_analyzer): + """Test with invalid device enumeration on SUT""" + + data = DeviceEnumerationDataModel(cpu_count=5, gpu_count=4, vf_count=8) + args = DeviceEnumerationAnalyzerArgs(cpu_count=4, gpu_count=4) + + result = device_enumeration_analyzer.analyze_data(data=data, args=args) + + assert result.status == ExecutionStatus.ERROR + assert "but got" in result.message + + for event in result.events: + assert event.priority == EventPriority.CRITICAL + assert event.category == EventCategory.PLATFORM.value + + +def test_analyze_list_of_accepted_counts(device_enumeration_analyzer): + """Test with a list of acceptable counts""" + + data = DeviceEnumerationDataModel(cpu_count=4, gpu_count=4, vf_count=8) + args = DeviceEnumerationAnalyzerArgs(cpu_count=[2, 4, 8], gpu_count=[4, 8]) + + result = device_enumeration_analyzer.analyze_data(data=data, args=args) + + assert result.status == ExecutionStatus.OK + assert len(result.events) == 0 diff --git a/test/unit/plugin/test_device_enumeration_collector.py b/test/unit/plugin/test_device_enumeration_collector.py new file mode 100644 index 00000000..a5a1ef30 --- /dev/null +++ b/test/unit/plugin/test_device_enumeration_collector.py @@ -0,0 +1,143 @@ +############################################################################### +# +# MIT License +# +# Copyright (c) 2025 Advanced Micro Devices, Inc. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +############################################################################### +from unittest.mock import MagicMock + +import pytest + +from nodescraper.enums.executionstatus import ExecutionStatus +from nodescraper.enums.systeminteraction import SystemInteractionLevel +from nodescraper.models.systeminfo import OSFamily +from nodescraper.plugins.inband.device_enumeration.device_enumeration_collector import ( + DeviceEnumerationCollector, +) +from nodescraper.plugins.inband.device_enumeration.deviceenumdata import ( + DeviceEnumerationDataModel, +) + + +@pytest.fixture +def device_enumeration_collector(system_info, conn_mock): + return DeviceEnumerationCollector( + system_info=system_info, + system_interaction_level=SystemInteractionLevel.PASSIVE, + connection=conn_mock, + ) + + +def test_collect_linux(system_info, device_enumeration_collector): + """Test linux typical output""" + system_info.os_family = OSFamily.LINUX + + device_enumeration_collector._run_sut_cmd = MagicMock( + side_effect=[ + MagicMock( + exit_code=0, + stdout="2", + stderr="", + command="lscpu | grep Socket | awk '{ print $2 }'", + ), + MagicMock( + exit_code=0, + stdout="8", + stderr="", + command="lspci -d 1002: | grep -i 'VGA\\|Display\\|3D' | wc -l", + ), + MagicMock( + exit_code=0, + stdout="0", + stderr="", + command="lspci -d 1002: | grep -i 'Virtual Function' | wc -l", + ), + ] + ) + + result, data = device_enumeration_collector.collect_data() + assert result.status == ExecutionStatus.OK + assert data == DeviceEnumerationDataModel(cpu_count=2, gpu_count=8, vf_count=0) + + +def test_collect_windows(system_info, device_enumeration_collector): + """Test windows typical output""" + system_info.os_family = OSFamily.WINDOWS + + device_enumeration_collector._run_sut_cmd = MagicMock( + side_effect=[ + MagicMock( + exit_code=0, + stdout="2", + stderr="", + command='powershell -Command "(Get-WmiObject -Class Win32_Processor | Measure-Object).Count"', + ), + MagicMock( + exit_code=0, + stdout="8", + stderr="", + command='powershell -Command "(wmic path win32_VideoController get name | findstr AMD | Measure-Object).Count"', + ), + MagicMock( + exit_code=0, + stdout="8", + stderr="", + command='powershell -Command "(Get-VMHostPartitionableGpu | Measure-Object).Count"', + ), + ] + ) + + result, data = device_enumeration_collector.collect_data() + assert result.status == ExecutionStatus.OK + assert data == DeviceEnumerationDataModel(cpu_count=2, gpu_count=8, vf_count=8) + + +def test_collect_error(system_info, device_enumeration_collector): + """Test with bad exit code""" + system_info.os_family = OSFamily.LINUX + + device_enumeration_collector._run_sut_cmd = MagicMock( + side_effect=[ + MagicMock( + exit_code=1, + stdout="some output", + stderr="command failed", + command="lscpu | grep Socket | awk '{ print $2 }'", + ), + MagicMock( + exit_code=1, + stdout="some output", + stderr="command failed", + command="lspci -d 1002: | grep -i 'VGA\\|Display\\|3D' | wc -l", + ), + MagicMock( + exit_code=1, + stdout="some output", + stderr="command failed", + command="lspci -d 1002: | grep -i 'Virtual Function' | wc -l", + ), + ] + ) + + result, data = device_enumeration_collector.collect_data() + assert result.status == ExecutionStatus.EXECUTION_FAILURE + assert data is None From 7b06590d8cdf513a27301ae2dff8cdf2fd458b7d Mon Sep 17 00:00:00 2001 From: Alexandra Bara Date: Wed, 19 Nov 2025 15:15:00 -0600 Subject: [PATCH 3/6] moved vendor id into systeminfo class --- nodescraper/models/systeminfo.py | 1 + .../device_enumeration_collector.py | 15 ++++++++------- 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/nodescraper/models/systeminfo.py b/nodescraper/models/systeminfo.py index d2be4ae4..e82d6212 100644 --- a/nodescraper/models/systeminfo.py +++ b/nodescraper/models/systeminfo.py @@ -41,3 +41,4 @@ class SystemInfo(BaseModel): platform: Optional[str] = None metadata: Optional[dict] = Field(default_factory=dict) location: Optional[SystemLocation] = SystemLocation.LOCAL + vendorid_ep: int = 0x1002 diff --git a/nodescraper/plugins/inband/device_enumeration/device_enumeration_collector.py b/nodescraper/plugins/inband/device_enumeration/device_enumeration_collector.py index 42e5e42e..92430d5c 100644 --- a/nodescraper/plugins/inband/device_enumeration/device_enumeration_collector.py +++ b/nodescraper/plugins/inband/device_enumeration/device_enumeration_collector.py @@ -38,12 +38,10 @@ class DeviceEnumerationCollector(InBandDataCollector[DeviceEnumerationDataModel, DATA_MODEL = DeviceEnumerationDataModel - # Linux commands CMD_CPU_COUNT_LINUX = "lscpu | grep Socket | awk '{ print $2 }'" - CMD_GPU_COUNT_LINUX = "lspci -d 1002: | grep -i 'VGA\\|Display\\|3D' | wc -l" - CMD_VF_COUNT_LINUX = "lspci -d 1002: | grep -i 'Virtual Function' | wc -l" + CMD_GPU_COUNT_LINUX = "lspci -d {vendorid_ep}: | grep -i 'VGA\\|Display\\|3D' | wc -l" + CMD_VF_COUNT_LINUX = "lspci -d {vendorid_ep}: | grep -i 'Virtual Function' | wc -l" - # Windows commands CMD_CPU_COUNT_WINDOWS = ( 'powershell -Command "(Get-WmiObject -Class Win32_Processor | Measure-Object).Count"' ) @@ -81,11 +79,14 @@ def collect_data(self, args=None) -> tuple[TaskResult, Optional[DeviceEnumeratio # Count CPU sockets cpu_count_res = self._run_sut_cmd(self.CMD_CPU_COUNT_LINUX) - # Count all AMD GPUs (vendor ID 1002) - gpu_count_res = self._run_sut_cmd(self.CMD_GPU_COUNT_LINUX) + # Count all AMD GPUs + vendor_id = format(self.system_info.vendorid_ep, "x") + gpu_count_res = self._run_sut_cmd( + self.CMD_GPU_COUNT_LINUX.format(vendorid_ep=vendor_id) + ) # Count AMD Virtual Functions - vf_count_res = self._run_sut_cmd(self.CMD_VF_COUNT_LINUX) + vf_count_res = self._run_sut_cmd(self.CMD_VF_COUNT_LINUX.format(vendorid_ep=vendor_id)) else: cpu_count_res = self._run_sut_cmd(self.CMD_CPU_COUNT_WINDOWS) gpu_count_res = self._run_sut_cmd(self.CMD_GPU_COUNT_WINDOWS) From 957326210ec748e54a72ac977693cb2c67ffa170 Mon Sep 17 00:00:00 2001 From: Alexandra Bara Date: Tue, 25 Nov 2025 12:45:32 -0600 Subject: [PATCH 4/6] fix for checking analyzer args --- .../device_enumeration/device_enumeration_analyzer.py | 6 ++---- test/unit/plugin/test_device_enumeration_analyzer.py | 6 +++--- 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/nodescraper/plugins/inband/device_enumeration/device_enumeration_analyzer.py b/nodescraper/plugins/inband/device_enumeration/device_enumeration_analyzer.py index 3cdfa802..3903a1de 100644 --- a/nodescraper/plugins/inband/device_enumeration/device_enumeration_analyzer.py +++ b/nodescraper/plugins/inband/device_enumeration/device_enumeration_analyzer.py @@ -45,10 +45,8 @@ def analyze_data( self, data: DeviceEnumerationDataModel, args: Optional[DeviceEnumerationAnalyzerArgs] = None ) -> TaskResult: - if not args: - self.result.message = "Expected Device Enumeration expected data not provided" - self.result.status = ExecutionStatus.NOT_RAN - return self.result + if args is None: + args = DeviceEnumerationAnalyzerArgs() # Convert to lists if integers, otherwise use as-is cpu_count = [args.cpu_count] if isinstance(args.cpu_count, int) else args.cpu_count diff --git a/test/unit/plugin/test_device_enumeration_analyzer.py b/test/unit/plugin/test_device_enumeration_analyzer.py index 8f91fadc..28c3b666 100644 --- a/test/unit/plugin/test_device_enumeration_analyzer.py +++ b/test/unit/plugin/test_device_enumeration_analyzer.py @@ -75,12 +75,12 @@ def test_analyze_passing_windows(system_info, device_enumeration_analyzer, devic def test_analyze_no_args(device_enumeration_analyzer, device_enumeration_data): - """Test with no analyzer args provided""" + """Test with no analyzer args provided - should skip all checks""" result = device_enumeration_analyzer.analyze_data(data=device_enumeration_data, args=None) - assert result.status == ExecutionStatus.NOT_RAN - assert "Expected Device Enumeration expected data not provided" in result.message + assert result.status == ExecutionStatus.OK + assert "Device Enumeration validated" in result.message assert len(result.events) == 0 From 5c7cf60089f86eb1afff9c39d098d6a332a57ed4 Mon Sep 17 00:00:00 2001 From: Alexandra Bara Date: Wed, 26 Nov 2025 10:11:37 -0600 Subject: [PATCH 5/6] fix for default args --- .../device_enumeration/device_enumeration_analyzer.py | 6 +++++- test/unit/plugin/test_device_enumeration_analyzer.py | 6 +++--- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/nodescraper/plugins/inband/device_enumeration/device_enumeration_analyzer.py b/nodescraper/plugins/inband/device_enumeration/device_enumeration_analyzer.py index 3903a1de..d006bc4d 100644 --- a/nodescraper/plugins/inband/device_enumeration/device_enumeration_analyzer.py +++ b/nodescraper/plugins/inband/device_enumeration/device_enumeration_analyzer.py @@ -46,7 +46,11 @@ def analyze_data( ) -> TaskResult: if args is None: - args = DeviceEnumerationAnalyzerArgs() + self.result.status = ExecutionStatus.NOT_RAN + self.result.message = ( + "Expected Device Enumeration data not provided, skipping analysis." + ) + return self.result # Convert to lists if integers, otherwise use as-is cpu_count = [args.cpu_count] if isinstance(args.cpu_count, int) else args.cpu_count diff --git a/test/unit/plugin/test_device_enumeration_analyzer.py b/test/unit/plugin/test_device_enumeration_analyzer.py index 28c3b666..c58c9ea0 100644 --- a/test/unit/plugin/test_device_enumeration_analyzer.py +++ b/test/unit/plugin/test_device_enumeration_analyzer.py @@ -75,12 +75,12 @@ def test_analyze_passing_windows(system_info, device_enumeration_analyzer, devic def test_analyze_no_args(device_enumeration_analyzer, device_enumeration_data): - """Test with no analyzer args provided - should skip all checks""" + """Test with no analyzer args provided - should skip analysis""" result = device_enumeration_analyzer.analyze_data(data=device_enumeration_data, args=None) - assert result.status == ExecutionStatus.OK - assert "Device Enumeration validated" in result.message + assert result.status == ExecutionStatus.NOT_RAN + assert "Expected Device Enumeration data not provided, skipping analysis." in result.message assert len(result.events) == 0 From fa929a104b155e9f07e8bb66110c93fcf269b008 Mon Sep 17 00:00:00 2001 From: Alexandra Bara Date: Mon, 1 Dec 2025 11:05:05 -0600 Subject: [PATCH 6/6] addressed reviews --- .../device_enumeration/analyzer_args.py | 33 +++++++++++++++---- .../device_enumeration_analyzer.py | 17 ++++------ .../device_enumeration_collector.py | 19 +++++------ 3 files changed, 40 insertions(+), 29 deletions(-) diff --git a/nodescraper/plugins/inband/device_enumeration/analyzer_args.py b/nodescraper/plugins/inband/device_enumeration/analyzer_args.py index 8c3303cb..8f74ed00 100644 --- a/nodescraper/plugins/inband/device_enumeration/analyzer_args.py +++ b/nodescraper/plugins/inband/device_enumeration/analyzer_args.py @@ -23,7 +23,9 @@ # SOFTWARE. # ############################################################################### -from typing import Optional, Union +from typing import Any, Optional + +from pydantic import field_validator from nodescraper.models import AnalyzerArgs @@ -31,9 +33,26 @@ class DeviceEnumerationAnalyzerArgs(AnalyzerArgs): - cpu_count: Optional[Union[list[int], int]] = None - gpu_count: Optional[Union[list[int], int]] = None - vf_count: Optional[Union[list[int], int]] = None + cpu_count: Optional[list[int]] = None + gpu_count: Optional[list[int]] = None + vf_count: Optional[list[int]] = None + + @field_validator("cpu_count", "gpu_count", "vf_count", mode="before") + @classmethod + def normalize_to_list(cls, v: Any) -> Optional[list[int]]: + """Convert single integer values to lists for consistent handling. + + Args: + v: The input value (can be int, list[int], or None). + + Returns: + Optional[list[int]]: The normalized list value or None. + """ + if v is None: + return None + if isinstance(v, int): + return [v] + return v @classmethod def build_from_model( @@ -48,7 +67,7 @@ def build_from_model( DeviceEnumerationAnalyzerArgs: instance of analyzer args class """ return cls( - cpu_count=datamodel.cpu_count, - gpu_count=datamodel.gpu_count, - vf_count=datamodel.vf_count, + cpu_count=[datamodel.cpu_count] if datamodel.cpu_count is not None else None, + gpu_count=[datamodel.gpu_count] if datamodel.gpu_count is not None else None, + vf_count=[datamodel.vf_count] if datamodel.vf_count is not None else None, ) diff --git a/nodescraper/plugins/inband/device_enumeration/device_enumeration_analyzer.py b/nodescraper/plugins/inband/device_enumeration/device_enumeration_analyzer.py index d006bc4d..7cf39335 100644 --- a/nodescraper/plugins/inband/device_enumeration/device_enumeration_analyzer.py +++ b/nodescraper/plugins/inband/device_enumeration/device_enumeration_analyzer.py @@ -52,18 +52,13 @@ def analyze_data( ) return self.result - # Convert to lists if integers, otherwise use as-is - cpu_count = [args.cpu_count] if isinstance(args.cpu_count, int) else args.cpu_count - gpu_count = [args.gpu_count] if isinstance(args.gpu_count, int) else args.gpu_count - vf_count = [args.vf_count] if isinstance(args.vf_count, int) else args.vf_count - checks = {} - if cpu_count is not None and cpu_count != []: - checks["cpu_count"] = cpu_count - if gpu_count is not None and gpu_count != []: - checks["gpu_count"] = gpu_count - if vf_count is not None and vf_count != []: - checks["vf_count"] = vf_count + if args.cpu_count is not None and args.cpu_count != []: + checks["cpu_count"] = args.cpu_count + if args.gpu_count is not None and args.gpu_count != []: + checks["gpu_count"] = args.gpu_count + if args.vf_count is not None and args.vf_count != []: + checks["vf_count"] = args.vf_count self.result.message = "" for check, accepted_counts in checks.items(): diff --git a/nodescraper/plugins/inband/device_enumeration/device_enumeration_collector.py b/nodescraper/plugins/inband/device_enumeration/device_enumeration_collector.py index 92430d5c..88506f6e 100644 --- a/nodescraper/plugins/inband/device_enumeration/device_enumeration_collector.py +++ b/nodescraper/plugins/inband/device_enumeration/device_enumeration_collector.py @@ -74,7 +74,6 @@ def collect_data(self, args=None) -> tuple[TaskResult, Optional[DeviceEnumeratio On Linux, use lscpu and lspci On Windows, use WMI and hyper-v cmdlets """ - device_enum = None if self.system_info.os_family == OSFamily.LINUX: # Count CPU sockets cpu_count_res = self._run_sut_cmd(self.CMD_CPU_COUNT_LINUX) @@ -91,20 +90,21 @@ def collect_data(self, args=None) -> tuple[TaskResult, Optional[DeviceEnumeratio cpu_count_res = self._run_sut_cmd(self.CMD_CPU_COUNT_WINDOWS) gpu_count_res = self._run_sut_cmd(self.CMD_GPU_COUNT_WINDOWS) vf_count_res = self._run_sut_cmd(self.CMD_VF_COUNT_WINDOWS) - cpu_count, gpu_count, vf_count = [None, None, None] + + device_enum = DeviceEnumerationDataModel() if cpu_count_res.exit_code == 0: - cpu_count = int(cpu_count_res.stdout) + device_enum.cpu_count = int(cpu_count_res.stdout) else: self._warning(description="Cannot determine CPU count", command=cpu_count_res) if gpu_count_res.exit_code == 0: - gpu_count = int(gpu_count_res.stdout) + device_enum.gpu_count = int(gpu_count_res.stdout) else: self._warning(description="Cannot determine GPU count", command=gpu_count_res) if vf_count_res.exit_code == 0: - vf_count = int(vf_count_res.stdout) + device_enum.vf_count = int(vf_count_res.stdout) else: self._warning( description="Cannot determine VF count", @@ -112,10 +112,7 @@ def collect_data(self, args=None) -> tuple[TaskResult, Optional[DeviceEnumeratio category=EventCategory.SW_DRIVER, ) - if cpu_count or gpu_count or vf_count: - device_enum = DeviceEnumerationDataModel( - cpu_count=cpu_count, gpu_count=gpu_count, vf_count=vf_count - ) + if device_enum.cpu_count or device_enum.gpu_count or device_enum.vf_count: self._log_event( category=EventCategory.PLATFORM, description=f"Counted {device_enum.cpu_count} CPUs, {device_enum.gpu_count} GPUs, {device_enum.vf_count} VFs", @@ -124,6 +121,7 @@ def collect_data(self, args=None) -> tuple[TaskResult, Optional[DeviceEnumeratio ) self.result.message = f"Device Enumeration: {device_enum.model_dump(exclude_none=True)}" self.result.status = ExecutionStatus.OK + return self.result, device_enum else: self.result.message = "Device Enumeration info not found" self.result.status = ExecutionStatus.EXECUTION_FAILURE @@ -132,5 +130,4 @@ def collect_data(self, args=None) -> tuple[TaskResult, Optional[DeviceEnumeratio description=self.result.message, priority=EventPriority.CRITICAL, ) - - return self.result, device_enum + return self.result, None