Skip to content

Commit 1dd4bdc

Browse files
alexandraBaraalexandraBara
committed
removed duplicate cmd call
1 parent 10a8fde commit 1dd4bdc

File tree

2 files changed

+84
-59
lines changed

2 files changed

+84
-59
lines changed

‎nodescraper/plugins/inband/device_enumeration/device_enumeration_collector.py‎

Lines changed: 80 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@
2626
from typing import Optional
2727

2828
from nodescraper.base import InBandDataCollector
29-
from nodescraper.connection.inband.inband import CommandArtifact, TextFileArtifact
29+
from nodescraper.connection.inband.inband import TextFileArtifact
3030
from nodescraper.enums import EventCategory, EventPriority, ExecutionStatus, OSFamily
3131
from nodescraper.models import TaskResult
3232

@@ -38,7 +38,6 @@ class DeviceEnumerationCollector(InBandDataCollector[DeviceEnumerationDataModel,
3838

3939
DATA_MODEL = DeviceEnumerationDataModel
4040

41-
CMD_CPU_COUNT_LINUX = "lscpu | grep Socket | awk '{ print $2 }'"
4241
CMD_GPU_COUNT_LINUX = "lspci -d {vendorid_ep}: | grep -i 'VGA\\|Display\\|3D' | wc -l"
4342
CMD_VF_COUNT_LINUX = "lspci -d {vendorid_ep}: | grep -i 'Virtual Function' | wc -l"
4443
CMD_LSCPU_LINUX = "/usr/bin/lscpu"
@@ -52,33 +51,14 @@ class DeviceEnumerationCollector(InBandDataCollector[DeviceEnumerationDataModel,
5251
'powershell -Command "(Get-VMHostPartitionableGpu | Measure-Object).Count"'
5352
)
5453

55-
def _warning(
56-
self,
57-
description: str,
58-
command: CommandArtifact,
59-
category: EventCategory = EventCategory.PLATFORM,
60-
):
61-
self._log_event(
62-
category=category,
63-
description=description,
64-
data={
65-
"command": command.command,
66-
"stdout": command.stdout,
67-
"stderr": command.stderr,
68-
"exit_code": command.exit_code,
69-
},
70-
priority=EventPriority.WARNING,
71-
)
72-
7354
def collect_data(self, args=None) -> tuple[TaskResult, Optional[DeviceEnumerationDataModel]]:
7455
"""
7556
Read CPU and GPU count
7657
On Linux, use lscpu and lspci
7758
On Windows, use WMI and hyper-v cmdlets
7859
"""
7960
if self.system_info.os_family == OSFamily.LINUX:
80-
# Count CPU sockets
81-
cpu_count_res = self._run_sut_cmd(self.CMD_CPU_COUNT_LINUX)
61+
lscpu_res = self._run_sut_cmd(self.CMD_LSCPU_LINUX, log_artifact=False)
8262

8363
# Count all AMD GPUs
8464
vendor_id = format(self.system_info.vendorid_ep, "x")
@@ -89,9 +69,6 @@ def collect_data(self, args=None) -> tuple[TaskResult, Optional[DeviceEnumeratio
8969
# Count AMD Virtual Functions
9070
vf_count_res = self._run_sut_cmd(self.CMD_VF_COUNT_LINUX.format(vendorid_ep=vendor_id))
9171

92-
# Collect lscpu output
93-
lscpu_res = self._run_sut_cmd(self.CMD_LSCPU_LINUX, log_artifact=False)
94-
9572
# Collect lshw output
9673
lshw_res = self._run_sut_cmd(self.CMD_LSHW_LINUX, sudo=True, log_artifact=False)
9774
else:
@@ -101,36 +78,87 @@ def collect_data(self, args=None) -> tuple[TaskResult, Optional[DeviceEnumeratio
10178

10279
device_enum = DeviceEnumerationDataModel()
10380

104-
if cpu_count_res.exit_code == 0:
105-
device_enum.cpu_count = int(cpu_count_res.stdout)
81+
if self.system_info.os_family == OSFamily.LINUX:
82+
if lscpu_res.exit_code == 0 and lscpu_res.stdout:
83+
# Extract socket count from lscpu output
84+
for line in lscpu_res.stdout.splitlines():
85+
if line.startswith("Socket(s):"):
86+
try:
87+
device_enum.cpu_count = int(line.split(":")[1].strip())
88+
break
89+
except (ValueError, IndexError):
90+
self._log_event(
91+
category=EventCategory.PLATFORM,
92+
description="Cannot parse CPU count from lscpu output",
93+
data={
94+
"command": lscpu_res.command,
95+
"exit_code": lscpu_res.exit_code,
96+
"stderr": lscpu_res.stderr,
97+
},
98+
priority=EventPriority.WARNING,
99+
)
100+
device_enum.lscpu_output = lscpu_res.stdout
101+
self._log_event(
102+
category=EventCategory.PLATFORM,
103+
description="Collected lscpu output",
104+
priority=EventPriority.INFO,
105+
)
106+
else:
107+
self._log_event(
108+
category=EventCategory.PLATFORM,
109+
description="Cannot collect lscpu output",
110+
data={
111+
"command": lscpu_res.command,
112+
"exit_code": lscpu_res.exit_code,
113+
"stderr": lscpu_res.stderr,
114+
},
115+
priority=EventPriority.WARNING,
116+
)
106117
else:
107-
self._warning(description="Cannot determine CPU count", command=cpu_count_res)
118+
if cpu_count_res.exit_code == 0:
119+
device_enum.cpu_count = int(cpu_count_res.stdout)
120+
else:
121+
self._log_event(
122+
category=EventCategory.PLATFORM,
123+
description="Cannot determine CPU count",
124+
data={
125+
"command": cpu_count_res.command,
126+
"exit_code": cpu_count_res.exit_code,
127+
"stderr": cpu_count_res.stderr,
128+
},
129+
priority=EventPriority.WARNING,
130+
)
108131

109132
if gpu_count_res.exit_code == 0:
110133
device_enum.gpu_count = int(gpu_count_res.stdout)
111134
else:
112-
self._warning(description="Cannot determine GPU count", command=gpu_count_res)
135+
self._log_event(
136+
category=EventCategory.PLATFORM,
137+
description="Cannot determine GPU count",
138+
data={
139+
"command": gpu_count_res.command,
140+
"exit_code": gpu_count_res.exit_code,
141+
"stderr": gpu_count_res.stderr,
142+
},
143+
priority=EventPriority.WARNING,
144+
)
113145

114146
if vf_count_res.exit_code == 0:
115147
device_enum.vf_count = int(vf_count_res.stdout)
116148
else:
117-
self._warning(
118-
description="Cannot determine VF count",
119-
command=vf_count_res,
149+
self._log_event(
120150
category=EventCategory.SW_DRIVER,
151+
description="Cannot determine VF count",
152+
data={
153+
"command": vf_count_res.command,
154+
"exit_code": vf_count_res.exit_code,
155+
"stderr": vf_count_res.stderr,
156+
},
157+
priority=EventPriority.WARNING,
121158
)
122159

160+
# Collect lshw output on Linux
123161
if self.system_info.os_family == OSFamily.LINUX:
124-
if lscpu_res.exit_code == 0 and lscpu_res.stdout:
125-
device_enum.lscpu_output = lscpu_res.stdout
126-
self._log_event(
127-
category=EventCategory.PLATFORM,
128-
description="Collected lscpu output",
129-
priority=EventPriority.INFO,
130-
)
131-
else:
132-
self._warning(description="Cannot collect lscpu output", command=lscpu_res)
133-
134162
if lshw_res.exit_code == 0 and lshw_res.stdout:
135163
device_enum.lshw_output = lshw_res.stdout
136164
self.result.artifacts.append(
@@ -142,7 +170,16 @@ def collect_data(self, args=None) -> tuple[TaskResult, Optional[DeviceEnumeratio
142170
priority=EventPriority.INFO,
143171
)
144172
else:
145-
self._warning(description="Cannot collect lshw output", command=lshw_res)
173+
self._log_event(
174+
category=EventCategory.PLATFORM,
175+
description="Cannot collect lshw output",
176+
data={
177+
"command": lshw_res.command,
178+
"exit_code": lshw_res.exit_code,
179+
"stderr": lshw_res.stderr,
180+
},
181+
priority=EventPriority.WARNING,
182+
)
146183

147184
if device_enum.cpu_count or device_enum.gpu_count or device_enum.vf_count:
148185
log_data = device_enum.model_dump(

‎test/unit/plugin/test_device_enumeration_collector.py‎

Lines changed: 4 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -58,9 +58,9 @@ def test_collect_linux(system_info, device_enumeration_collector):
5858
side_effect=[
5959
MagicMock(
6060
exit_code=0,
61-
stdout="2",
61+
stdout=lscpu_output,
6262
stderr="",
63-
command="lscpu | grep Socket | awk '{ print $2 }'",
63+
command="/usr/bin/lscpu",
6464
),
6565
MagicMock(
6666
exit_code=0,
@@ -74,12 +74,6 @@ def test_collect_linux(system_info, device_enumeration_collector):
7474
stderr="",
7575
command="lspci -d 1002: | grep -i 'Virtual Function' | wc -l",
7676
),
77-
MagicMock(
78-
exit_code=0,
79-
stdout=lscpu_output,
80-
stderr="",
81-
command="/usr/bin/lscpu",
82-
),
8377
MagicMock(
8478
exit_code=0,
8579
stdout=lshw_output,
@@ -140,9 +134,9 @@ def test_collect_error(system_info, device_enumeration_collector):
140134
side_effect=[
141135
MagicMock(
142136
exit_code=1,
143-
stdout="some output",
137+
stdout="",
144138
stderr="command failed",
145-
command="lscpu | grep Socket | awk '{ print $2 }'",
139+
command="/usr/bin/lscpu",
146140
),
147141
MagicMock(
148142
exit_code=1,
@@ -156,12 +150,6 @@ def test_collect_error(system_info, device_enumeration_collector):
156150
stderr="command failed",
157151
command="lspci -d 1002: | grep -i 'Virtual Function' | wc -l",
158152
),
159-
MagicMock(
160-
exit_code=1,
161-
stdout="",
162-
stderr="command failed",
163-
command="/usr/bin/lscpu",
164-
),
165153
MagicMock(
166154
exit_code=1,
167155
stdout="",

0 commit comments

Comments
 (0)