Source code for imaginaire.utils.gpu_affinity

# 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)