Skip CpuSchedulingMode when running on ARM as SMT is unsupported on ARM.

PiperOrigin-RevId: 739983300
Change-Id: I805665609f53853228f54c8d6f29e61dedf41666

Author: liyuying0000

fleetbench/parallel/cpu.py

@@ -17,12 +17,14 @@
 import enum
 import math
 import os
+import subprocess
 from typing import Tuple
 
 from absl import logging
 import psutil
 
 
+# TODO: Consider using `lscpu -p` to get CPU information.
 class CpuInfo:
   """Class for reading and parsing CPU information from the cpuinfo file.
 
@@ -316,3 +318,15 @@ def Utilization(cpus: list[int]) -> Tuple[float, dict[int, float], int]:
       usable_utilization,
       busy_cpus,
   )
+
+
+def GetCPUArch() -> str:
+  """Returns the CPU architecture."""
+  try:
+    cpu_arch = subprocess.run(
+        ["uname", "-m"], capture_output=True, text=True, check=True
+    )
+    return cpu_arch.stdout.strip()
+  except subprocess.CalledProcessError:
+    logging.warning("Could not determine machine architecture. Using x86.")
+    return "x86_64"

fleetbench/parallel/cpu_test.py

@@ -22,7 +22,7 @@
 from fleetbench.parallel import cpu
 
 
-class CpuTest(absltest.TestCase):
+class CpuTest(parameterized.TestCase):
 
   def test_Available(self):
     self.assertGreaterEqual(len(cpu.Available()), 1)
@@ -38,6 +38,23 @@ def test_Utilization(self, mock_cpu_percent):
     with self.assertRaises(ValueError):
       cpu.Utilization([])
 
+  @parameterized.named_parameters(
+      ("x86", "x86_64"),
+      ("arm", "aarch64"),
+  )
+  @mock.patch("subprocess.run")
+  def test_get_cpu_arch_success(self, cpu_arch, mock_run):
+    """Test when subprocess.run succeeds."""
+    mock_process = mock.Mock()
+    mock_process.stdout = cpu_arch + "\n"
+    mock_process.returncode = 0
+    mock_run.return_value = mock_process
+
+    self.assertEqual(cpu.GetCPUArch(), cpu_arch)
+    mock_run.assert_called_once_with(
+        ["uname", "-m"], capture_output=True, text=True, check=True
+    )
+
 
 class SchedlingModeTest(parameterized.TestCase):
 

fleetbench/parallel/parallel_bench.py

@@ -147,12 +147,21 @@ def main(argv: Sequence[str]) -> None:
   shutil.rmtree(_TEMP_ROOT.value, ignore_errors=True)
   os.makedirs(_TEMP_ROOT.value, exist_ok=True)
 
-  scheduling_mode = cpu.CpuSchedulingMode(
-      num_cpus=_NUM_CPUS.value,
-      utilization=_UTILIZATION.value,
-      hyperthreading_mode=_HYPERTHREADING_MODE.value,
-  )
-  cpus, target_utilization = scheduling_mode.SelectCPURangeAndSetUtilization()
+  cpu_arch = cpu.GetCPUArch()
+  if cpu_arch == "aarch64":
+    logging.info("Running on ARM. SMT is unsupport.")
+    cpus = list(cpu.Available())[: _NUM_CPUS.value]
+    target_utilization = _UTILIZATION.value
+  elif cpu_arch == "x86_64":
+    logging.info("Running on x86. Adjusting CPU scheduling...")
+    scheduling_mode = cpu.CpuSchedulingMode(
+        num_cpus=_NUM_CPUS.value,
+        utilization=_UTILIZATION.value,
+        hyperthreading_mode=_HYPERTHREADING_MODE.value,
+    )
+    cpus, target_utilization = scheduling_mode.SelectCPURangeAndSetUtilization()
+  else:
+    raise ValueError(f"Unsupported CPU architecture: {cpu_arch}")
 
   bench = parallel_bench_lib.ParallelBench(
       cpus=cpus,