relay-throughput-testing/src/numa.py

import os
#
def parse_range_list(range_list_str):
	"""
	Take an input like '1-3,5,7-10' and return a list like [1,2,3,5,7,8,9,10].
	"""
	#
	range_strings = range_list_str.split(',')
	all_items = []
	for range_str in range_strings:
		if '-' in range_str:
			# in the form '12-34'
			range_ends = [int(x) for x in range_str.split('-')]
			assert(len(range_ends) == 2)
			all_items.extend(range(range_ends[0], range_ends[1]+1))
		else:
			# just a number
			all_items.append(int(range_str))
		#
	#
	return all_items
#
def generate_range_list(l):
	"""
	Take a list like [1,2,3,5,7,8,9,10] and return a string like '1-3,5,7-10'.
	"""
	#
	l = list(set(sorted(l)))
	ranges = []
	current_range_start = None
	#
	for index in range(len(l)):
		if current_range_start is None:
			current_range_start = l[index]
		else:
			if l[index] != l[index-1]+1:
				ranges.append((current_range_start, l[index-1]))
				current_range_start = l[index]
			#
		#
	#
	ranges.append((current_range_start, l[-1]))
	#
	return ','.join(['-'.join([str(y) for y in x]) if x[0] != x[1] else str(x[0]) for x in ranges])
#
def check_path_traversal(path, base_path):
	# this is not guarenteed to be secure
	if os.path.commonprefix([os.path.realpath(path), base_path]) != base_path:
		raise Exception('The path \'{}\' is not in the base path \'{}.\''.format(os.path.realpath(path), base_path))
	#
#
def get_online_numa_nodes():
	with open('/sys/devices/system/node/online', 'r') as f:
		online_nodes = f.read().strip()
		return parse_range_list(online_nodes)
	#
#
def get_processors_in_numa_node(node_id):
	path = '/sys/devices/system/node/node{}/cpulist'.format(node_id)
	check_path_traversal(path, '/sys/devices/system/node')
	with open(path, 'r') as f:
		return parse_range_list(f.read().strip())
	#
#
def get_thread_siblings(processor_id):
	path = '/sys/devices/system/cpu/cpu{}/topology/thread_siblings_list'.format(processor_id)
	check_path_traversal(path, '/sys/devices/system/cpu')
	with open(path, 'r') as f:
		return parse_range_list(f.read().strip())
	#
#
def get_numa_overview():
	numa_nodes = {}
	#
	for node_id in get_online_numa_nodes():
		numa_nodes[node_id] = {}
	#
	for node_id in numa_nodes:
		processors = get_processors_in_numa_node(node_id)
		#
		numa_nodes[node_id]['physical_cores'] = []
		for processor_id in processors:
			thread_siblings = sorted(get_thread_siblings(processor_id))
			#
			if thread_siblings not in numa_nodes[node_id]['physical_cores']:
				numa_nodes[node_id]['physical_cores'].append(thread_siblings)
			#
		#
	#
	return numa_nodes
#
if __name__ == '__main__':
	print(get_numa_overview())
#