def test_cpu_affinity_and_cuda_visible_devices():
affinity = dict()
for i in range(get_n_gpus()):
# The negative here would be `device = 0` as required for CUDA runtime
# calls.
device = nvml_device_index(0, cuda_visible_devices(i))
affinity[device] = get_cpu_affinity(device)
for i in range(get_n_gpus()):
assert get_cpu_affinity(i) == affinity[i]
def new_worker_spec(self):
""" Returns the info for a next worker
This method is related with `cluster.scale`
used to change the number of worker in the cluster
"""
worker_count = self.worker_count
gpu_count = self.gpu_count
spec = copy.deepcopy(self.new_spec)
spec['options'].update({
'env': {'CUDA_VISIBLE_DEVICES': str(gpu_count)},
'plugins': {CPUAffinity(get_cpu_affinity(gpu_count))}
})
return {worker_count: spec}
def test_cpu_affinity():
for i in range(get_n_gpus()):
affinity = get_cpu_affinity(i)
os.sched_setaffinity(0, affinity)
assert os.sched_getaffinity(0) == set(affinity)
