# Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# This work is made available under the Nvidia Source Code License-NC.
# To view a copy of this license, check out LICENSE.md
import math
import os
import pynvml
pynvml.nvmlInit()
[docs]def systemGetDriverVersion():
r"""Get Driver Version"""
return pynvml.nvmlSystemGetDriverVersion()
[docs]def deviceGetCount():
r"""Get number of devices"""
return pynvml.nvmlDeviceGetCount()
[docs]class device(object):
r"""Device used for nvml."""
_nvml_affinity_elements = math.ceil(os.cpu_count() / 64)
def __init__(self, device_idx):
super().__init__()
self.handle = pynvml.nvmlDeviceGetHandleByIndex(device_idx)
[docs] def getName(self):
r"""Get obect name"""
return pynvml.nvmlDeviceGetName(self.handle)
[docs] def getCpuAffinity(self):
r"""Get CPU affinity"""
affinity_string = ''
for j in pynvml.nvmlDeviceGetCpuAffinity(
self.handle, device._nvml_affinity_elements):
# assume nvml returns list of 64 bit ints
affinity_string = '{:064b}'.format(j) + affinity_string
affinity_list = [int(x) for x in affinity_string]
affinity_list.reverse() # so core 0 is in 0th element of list
return [i for i, e in enumerate(affinity_list) if e != 0]
[docs]def set_affinity(gpu_id=None):
r"""Set GPU affinity
Args:
gpu_id (int): Which gpu device.
"""
if gpu_id is None:
gpu_id = int(os.getenv('LOCAL_RANK', 0))
dev = device(gpu_id)
os.sched_setaffinity(0, dev.getCpuAffinity())
# list of ints
# representing the logical cores this process is now affinitied with
return os.sched_getaffinity(0)