Bug - Fix bug of detecting if gpu_index is none #275

**Description** Fix bug of detecting if gpu_index is none.
microsoft · Dec 24, 2021 · ce1481d · ce1481d
1 parent 3d8721a
commit ce1481d
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 2 deletions.
diff --git a/superbench/benchmarks/micro_benchmarks/ib_validation_performance.py b/superbench/benchmarks/micro_benchmarks/ib_validation_performance.py
@@ -249,7 +249,7 @@ def __prepare_general_ib_command_params(self):
             msg_size = '-s ' + str(self._args.msg_size)
         # Add GPUDirect for ib command
         gpu_enable = ''
-        if self._args.gpu_index:
+        if self._args.gpu_index is not None:
             gpu = GPU()
             if gpu.vendor == 'nvidia':
                 gpu_enable = ' --use_cuda={gpu_index}'.format(gpu_index=str(self._args.gpu_index))

diff --git a/tests/benchmarks/micro_benchmarks/test_ib_traffic_performance.py b/tests/benchmarks/micro_benchmarks/test_ib_traffic_performance.py
@@ -117,8 +117,9 @@ def read_config(filename):
 
         Path(test_config_file).unlink()
 
+    @mock.patch('superbench.common.devices.GPU.vendor', new_callable=mock.PropertyMock)
     @mock.patch('superbench.common.utils.network.get_ib_devices')
-    def test_ib_traffic_performance(self, mock_ib_devices):
+    def test_ib_traffic_performance(self, mock_ib_devices, mock_gpu):
         """Test ib-traffic benchmark."""
         # Test without ib devices
         # Check registry.
@@ -168,6 +169,22 @@ def test_ib_traffic_performance(self, mock_ib_devices):
         command = benchmark._bin_name + benchmark._commands[0].split(benchmark._bin_name)[1]
         assert (command == expect_command)
 
+        parameters = '--ib_index 0 --iters 2000 --pattern one-to-one --hostfile hostfile --gpu_index 0'
+        mock_gpu.return_value = 'nvidia'
+        benchmark = benchmark_class(benchmark_name, parameters=parameters)
+        ret = benchmark._preprocess()
+        expect_command = 'ib_validation --hostfile hostfile --cmd_prefix "ib_write_bw -F ' + \
+            '--iters=2000 -d mlx5_0 -a --use_cuda=0" --input_config ' + os.getcwd() + '/config.txt'
+        command = benchmark._bin_name + benchmark._commands[0].split(benchmark._bin_name)[1]
+        assert (command == expect_command)
+        mock_gpu.return_value = 'amd'
+        benchmark = benchmark_class(benchmark_name, parameters=parameters)
+        ret = benchmark._preprocess()
+        expect_command = 'ib_validation --hostfile hostfile --cmd_prefix "ib_write_bw -F ' + \
+            '--iters=2000 -d mlx5_0 -a --use_rocm=0" --input_config ' + os.getcwd() + '/config.txt'
+        command = benchmark._bin_name + benchmark._commands[0].split(benchmark._bin_name)[1]
+        assert (command == expect_command)
+
         # Custom config
         config = ['0,1', '1,0;0,1', '0,1;1,0', '1,0;0,1']
         with open('test_config.txt', 'w') as f: