def thread_target():
try:
logger.debug(f"thread '{name}' starting")
if _prctl_available:
prctl.set_name(name)
if _numa_available and numa.available():
numa.set_localalloc()
logger.debug(
f"Set NUMA local allocation policy on thread {name}")
worker(index)
logger.debug(f"thread '{name}' finishing")
except Exception:
logger.critical("Exception occured in thread; exiting")
logger.critical(traceback.format_exc())
# Communicate back the main thread that something bad has happened.
# This seems to be the only reliable way to do it.
_thread.interrupt_main()
# Now we still need to make sure that the main thread doesn't block
# on the queue.get/join (as it won't be interrupted). This is an attempt
# to make sure that it unblocks. May not be fool-proof though.
#
# TODO This doesn't really work. We can still block on pushing things
# onto the queue. We'll probably have to do something ourselves using
# timeouts and stuff to see if an error has occured.
if consumer:
while True:
try:
work_queue.task_done()
except ValueError:
break
else:
work_queue.put(None)
def check_numa():
try:
import numa
except:
return
if not numa.available():
return
if numa.get_max_node() > 0 and len(numa.get_run_on_node_mask()) > 1:
print("Warning: NUMA settings may be suboptimal!", file=sys.stderr)
def coremap():
try:
import numa
except ImportError:
print('This script requires the libnuma python bindings')
raise RuntimeError("Numa not available")
if not numa.available():
raise RuntimeError("Numa not available")
node_to_core = {int(i): deque([int(k) for k in numa.node_to_cpus(i)]) for i in range(numa.get_max_node() + 1)}
total_core = max(itertools.chain(*node_to_core.values())) + 1
return node_to_core, total_core
def load_numa():
""" Load information about core numbers and numa patterns """
if not numa.available():
raise Exception('Numa detection not available')
max_node = numa.get_max_node()
nodes = {}
for i in range(max_node + 1):
nodes[i] = list(numa.node_to_cpus(i))
return nodes
def _configure_numa(self):
self._numa_available = \
numa.available() and which('numactl') is not None
if not self._numa_available:
return
num_numa_nodes = numa.get_max_node() + 1
self._numa_cpu_map = {}
num_gpus = len(self._gpu_ids)
# Calculate how many CPUs to allocate for each GPU. Ensure this number
# is a power of 2.
num_cpus = 0
for i in range(num_numa_nodes):
num_cpus += len(numa.node_to_cpus(i))
num_cpus_per_gpu = min(MAX_CPUS_PER_GPU, max(num_cpus // num_gpus, 1))
num_cpus_per_gpu = pow(2, round(math.log(num_cpus_per_gpu, 2)))
# Find blocks of contiguous CPUs.
contiguous_blocks = []
for i in range(num_numa_nodes):
cpus = sorted(numa.node_to_cpus(i))
contiguous_block = [cpus[0]]
for j in range(1, len(cpus)):
if (cpus[j] - cpus[j - 1] == 1
and len(contiguous_block) < num_cpus_per_gpu):
contiguous_block.append(cpus[j])
else:
contiguous_blocks.append(
(contiguous_block, len(contiguous_block)))
contiguous_block = [cpus[j]]
if len(contiguous_block) > 0:
contiguous_blocks.append(
(contiguous_block, len(contiguous_block)))
contiguous_blocks.sort(key=lambda x: x[-1], reverse=True)
# Assign CPUs to GPUs.
block_idx = 0
for i in range(num_gpus):
self._numa_cpu_map[i] = []
while len(self._numa_cpu_map[i]) < num_cpus_per_gpu:
self._numa_cpu_map[i] += contiguous_blocks[block_idx][0]
block_idx = (block_idx + 1) % len(contiguous_blocks)
self._logger.info('GPU {gpu} assigned CPUs {cpus}'.format(
gpu=i, cpus=str(self._numa_cpu_map[i])))
def _lazy_cpu_and_mem_set_init(self):
# Implicitly assume lock is already held
if len(self._numa_nodes) != 0:
# Init already happened
return
if (self._available_cpu_ids is None or self._cpus_per_job is None
or self._use_memset_of_nearest_node is None):
raise Exception('Cannot do init. One or more params were None')
import numa
if not numa.available():
raise Exception('NUMA not available')
numa_nodes = list(range(0, numa.get_max_node() + 1))
cpu_count = 0
for numa_node in numa_nodes:
cpus = numa.node_to_cpus(numa_node)
for cpu_id in cpus:
if cpu_id in self._available_cpu_ids:
try:
self._numa_nodes[numa_node].add(cpu_id)
except KeyError:
self._numa_nodes[numa_node] = set()
self._numa_nodes[numa_node].add(cpu_id)
try:
self._numa_node_pool[numa_node].add(cpu_id)
except KeyError:
self._numa_node_pool[numa_node] = set()
self._numa_node_pool[numa_node].add(cpu_id)
_logger.info(
'Putting CPU {} in NUMA node {} in resource pool'.
format(cpu_id, numa_node))
cpu_count += 1
else:
_logger.info(
'CPU {} in NUMA node {} is NOT IN resource pool'.
format(cpu_id, numa_node))
if cpu_count == 0:
raise Exception('Found no available CPUs')
if cpu_count != len(self._available_cpu_ids):
raise Exception(
'Mismatch between provided available CPU ids and what was found on system'
)
assert len(self._numa_node_pool) == len(self._numa_nodes)
def test_available(self):
self.failUnlessEqual(True, numa.available())
def test_available(self):
self.assertEqual(True, numa.available())
def test_available(self):
self.failUnlessEqual(True, numa.available())