Unverified Commit ad8e0143 authored by Yuting Jiang's avatar Yuting Jiang Committed by GitHub
Browse files

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

**Description**
Add numa support for nvbandwidth.
parent a7c4ed92
......@@ -88,13 +88,6 @@ def add_parser_arguments(self):
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 @@ def add_parser_arguments(self):
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 @@ def _preprocess(self):
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
......
......@@ -116,6 +116,29 @@ def add_parser_arguments(self):
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.
......
......@@ -66,7 +66,11 @@ def add_parser_arguments(self):
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 @@ def _preprocess(self):
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 @@ def _preprocess(self):
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'
......
......@@ -3,6 +3,8 @@
"""Tests for nvbandwidth benchmark."""
import os
import unittest
from tests.helper import decorator
......@@ -52,6 +54,14 @@ def test_nvbandwidth_preprocess(self):
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."""
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment