Benchmarks: Micro benchmark - Add numa support for nvbandwidth (#742)

**Description**
Add numa support for nvbandwidth.

superbench/benchmarks/micro_benchmarks/ib_loopback_performance.py

@@ -88,13 +88,6 @@ class IBLoopbackBenchmark(MicroBenchmarkWithInvoke):
             default=['write'],
             help='The ib command used to run, e.g., {}.'.format(' '.join(list(self.__support_ib_commands.keys()))),
         )
-        self._parser.add_argument(
-            '--numa',
-            type=int,
-            default=0,
-            required=False,
-            help='The index of numa node.',
-        )
         self._parser.add_argument(
             '--gid_index',
             type=int,
@@ -103,7 +96,7 @@ class IBLoopbackBenchmark(MicroBenchmarkWithInvoke):
             help='Test uses GID with GID index taken from command.',
         )
 
-    def __get_arguments_from_env(self):
+    def _get_arguments_from_env(self):
         """Read environment variables from runner used for parallel and fill in ib_index and numa_node_index.
 
         Get 'PROC_RANK'(rank of current process) 'IB_DEVICES' 'NUMA_NODES' environment variables
@@ -128,7 +121,7 @@ class IBLoopbackBenchmark(MicroBenchmarkWithInvoke):
         Return:
             True if _preprocess() succeed.
         """
-        if not super()._preprocess() or not self.__get_arguments_from_env():
+        if not super()._preprocess() or not self._get_arguments_from_env():
             return False
 
         # Format the arguments

superbench/benchmarks/micro_benchmarks/micro_base.py

@@ -116,6 +116,29 @@ class MicroBenchmarkWithInvoke(MicroBenchmark):
             default=False,
             help='Tolerant failure for sub microbenchmark.',
         )
+        self._parser.add_argument(
+            '--numa',
+            type=int,
+            required=False,
+            help='The index of numa node.',
+        )
+
+    def _get_arguments_from_env(self):
+        """Read environment variables from runner used for parallel and fill in numa_node_index.
+
+        Get 'PROC_RANK'(rank of current process) 'NUMA_NODES' environment variables
+        Get numa_node_index according to 'NUMA_NODES'['PROC_RANK']
+        Note: The config from env variables will overwrite the configs defined in the command line
+        """
+        try:
+            if os.getenv('PROC_RANK'):
+                rank = int(os.getenv('PROC_RANK'))
+                if os.getenv('NUMA_NODES'):
+                    self._args.numa = int(os.getenv('NUMA_NODES').split(',')[rank])
+            return True
+        except BaseException:
+            logger.error('The proc_rank is out of index of devices - benchmark: {}.'.format(self._name))
+            return False
 
     def _set_binary_path(self):
         """Search the binary from self._args.bin_dir or from system environment path and set the binary directory.

superbench/benchmarks/micro_benchmarks/nvbandwidth.py

@@ -66,7 +66,11 @@ class NvBandwidthBenchmark(MicroBenchmarkWithInvoke):
         self._parser.add_argument(
             '--disable_affinity',
             action='store_true',
-            help='Disable automatic CPU affinity control. Default is False.',
+            help=(
+                'Disable automatic CPU affinity control. Default is False. '
+                'If user would like to bind the process to specific NUMA node, '
+                'please use --disable_affinity along with --numa argument.'
+            ),
         )
 
         self._parser.add_argument(
@@ -92,7 +96,7 @@ class NvBandwidthBenchmark(MicroBenchmarkWithInvoke):
         if not super()._preprocess():
             return False
 
-        if not self._set_binary_path():
+        if not self._set_binary_path() or not self._get_arguments_from_env():
             return False
 
         # Construct the command for nvbandwidth
@@ -111,6 +115,8 @@ class NvBandwidthBenchmark(MicroBenchmarkWithInvoke):
 
         if self._args.disable_affinity:
             command += ' --disableAffinity'
+            if self._args.numa is not None:
+                command = f'numactl --cpunodebind={self._args.numa} --membind={self._args.numa} ' + command
 
         if self._args.use_mean:
             command += ' --useMean'

tests/benchmarks/micro_benchmarks/test_nvbandwidth.py

@@ -3,6 +3,8 @@
 
 """Tests for nvbandwidth benchmark."""
 
+import os
+
 import unittest
 
 from tests.helper import decorator
@@ -52,6 +54,14 @@ class TestNvBandwidthBenchmark(BenchmarkTestCase, unittest.TestCase):
         assert ('--useMean' in benchmark._commands[0])
         assert ('--testSamples 100' in benchmark._commands[0])
 
+        # Test preprocess with numa
+        os.environ['NUMA_NODES'] = '0'
+        os.environ['PROC_RANK'] = '0'
+        benchmark = benchmark_class(benchmark_name, parameters=parameters)
+        assert benchmark._preprocess()
+        assert benchmark.return_code == ReturnCode.SUCCESS
+        assert ('numactl --cpunodebind=0 --membind=0' in benchmark._commands[0])
+
     @decorator.load_data('tests/data/nvbandwidth_results.log')
     def test_nvbandwidth_result_parsing_real_output(self, results):
         """Test NV Bandwidth benchmark result parsing."""