Spaces:
Runtime error
Runtime error
File size: 1,736 Bytes
f670afc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 |
# 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()
def systemGetDriverVersion():
r"""Get Driver Version"""
return pynvml.nvmlSystemGetDriverVersion()
def deviceGetCount():
r"""Get number of devices"""
return pynvml.nvmlDeviceGetCount()
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)
def getName(self):
r"""Get obect name"""
return pynvml.nvmlDeviceGetName(self.handle)
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]
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)
|