Fix nvme device info for linux machines

PiperOrigin-RevId: 833917856

Author: Arushi-07

perfkitbenchmarker/linux_virtual_machine.py

@@ -29,7 +29,7 @@
 import abc
 import collections
 import copy
-import json
+import io
 import logging
 import os
 import posixpath
@@ -41,7 +41,7 @@
 import uuid
 
 from absl import flags
-from packaging import version as packaging_version
+import pandas as pd
 from perfkitbenchmarker import background_tasks
 from perfkitbenchmarker import disk
 from perfkitbenchmarker import errors
@@ -2220,41 +2220,27 @@ def GetNVMEDeviceInfo(self):
     """Get the NVME disk device info, by querying the VM."""
     self.InstallPackages('nvme-cli')
     version_str, _ = self.RemoteCommand('sudo nvme --version')
-    version_num = version_str.split()[2]
-    # TODO(arushigaur): Version check can be removed and we can just parse
-    # the raw output.
-    if packaging_version.parse(version_num) >= packaging_version.parse(
-        '1.5'
-    ) and packaging_version.parse(version_num) < packaging_version.parse(
-        '2.11'
-    ):
-      stdout, _ = self.RemoteCommand('sudo nvme list --output-format json')
-      if not stdout:
-        return []
-      response = json.loads(stdout)
-      return response.get('Devices', [])
-    else:
-      # custom parsing for older OSes that do not ship nvme-cli ver 1.5+.
-      response = []
-      stdout, _ = self.RemoteCommand('sudo nvme list')
-      if 'No NVMe devices detected' in stdout:
-        return []
-      rows = stdout.splitlines()
-      delimiter_row = rows[1]  # row 0 is the column headers
-      delimiter_index = [0] + [
-          i for i in range(len(delimiter_row)) if delimiter_row[i] == ' '
-      ]
-      for row in rows[2:]:
-        device = {}
-        device_info = [
-            row[i:j]
-            for i, j in zip(delimiter_index, delimiter_index[1:] + [None])
-        ]
-        device['DevicePath'] = device_info[0].strip()
-        device['SerialNumber'] = device_info[1].strip()
-        device['ModelNumber'] = device_info[2].strip()
-        response.append(device)
-      return response
+    logging.info('nvme-cli version: %s', version_str.split()[2])
+    response = []
+    stdout, _ = self.RemoteCommand('sudo nvme list')
+    if 'No NVMe devices detected' in stdout:
+      return []
+    rows = stdout.splitlines()
+    header_row = rows[0]
+    delimiter_row = rows[1]
+    col_spans = [
+        (m.start(), m.end()) for m in re.finditer(r'-+', delimiter_row)
+    ]
+    data = '\n'.join([header_row] + rows[2:])
+    df_colspecs = pd.read_fwf(io.StringIO(data), colspecs=col_spans)
+    df_rows = df_colspecs.to_dict(orient='records')
+    for row in df_rows:
+      response.append({
+          'DevicePath': row.get('Node'),
+          'SerialNumber': row.get('SN'),
+          'ModelNumber': row.get('Model'),
+      })
+    return response
 
   def GenerateAndCaptureLogs(self) -> list[str]:
     """Generates and/or captures logs for this VM and returns the paths.

tests/linux_virtual_machine_test.py

@@ -854,6 +854,40 @@ def testDisableCstatesErrorCases(
       vm._DisableCstates()
     self.assertNotIn('disabled_cstates', vm.os_metadata)
 
+  @mock.patch.object(
+      pkb_common_test_case.TestLinuxVirtualMachine, 'RemoteCommand'
+  )
+  def test_get_nvme_device_info_with_devices(self, mock_remote_command):
+    vm = CreateTestLinuxVm()
+    nvme_list_output = """Node                  Generic               SN                   Model                                    Namespace  Usage                      Format           FW Rev
+--------------------- --------------------- -------------------- ---------------------------------------- ---------- -------------------------- ---------------- --------
+/dev/nvme0n1          /dev/ng0n1            vol06948076c9c135a44 Amazon Elastic Block Store               0x1          8.59  GB /   8.59  GB    512   B +  0 B   1.0
+/dev/nvme1n1          /dev/ng1n1            vol06d22d4fd149a2e68 Amazon Elastic Block Store               0x1         10.74  GB /  10.74  GB    512   B +  0 B   1.0
+"""
+    mock_remote_command.side_effect = [
+        ('nvme version 2.13', ''),  # nvme --version
+        (nvme_list_output, ''),  # nvme list
+    ]
+
+    expected_result = [
+        {
+            'DevicePath': '/dev/nvme0n1',
+            'SerialNumber': 'vol06948076c9c135a44',
+            'ModelNumber': 'Amazon Elastic Block Store',
+        },
+        {
+            'DevicePath': '/dev/nvme1n1',
+            'SerialNumber': 'vol06d22d4fd149a2e68',
+            'ModelNumber': 'Amazon Elastic Block Store',
+        },
+    ]
+
+    result = vm.GetNVMEDeviceInfo()
+
+    self.assertEqual(result, expected_result)
+    mock_remote_command.assert_any_call('sudo nvme --version')
+    mock_remote_command.assert_any_call('sudo nvme list')
+
 
 class RangeListUtilTest(parameterized.TestCase):