Python pynvml.nvmlDeviceGetMemoryInfo方法代码示例

# 需要导入模块: import pynvml [as 别名]
# 或者: from pynvml import nvmlDeviceGetMemoryInfo [as 别名]
def getGPUUsage():
    try:
        pynvml.nvmlInit()
        count = pynvml.nvmlDeviceGetCount()
        if count == 0:
            return None

        result = {"driver": pynvml.nvmlSystemGetDriverVersion(),
                  "gpu_count": int(count)
                  }
        i = 0
        gpuData = []
        while i<count:
            handle = pynvml.nvmlDeviceGetHandleByIndex(i)
            mem = pynvml.nvmlDeviceGetMemoryInfo(handle)
            gpuData.append({"device_num": i, "name": pynvml.nvmlDeviceGetName(handle), "total": round(float(mem.total)/1000000000, 2), "used": round(float(mem.used)/1000000000, 2)})
            i = i+1

        result["devices"] = jsonpickle.encode(gpuData, unpicklable=False)
    except Exception as e:
        result = {"driver": "No GPU!", "gpu_count": 0, "devices": []}

    return result 

# 需要导入模块: import pynvml [as 别名]
# 或者: from pynvml import nvmlDeviceGetMemoryInfo [as 别名]
def gpu_info(self):
        # pip install nvidia-ml-py3
        if len(self.gpu_ids) >= 0 and torch.cuda.is_available():
            try:
                import pynvml
                pynvml.nvmlInit()
                self.config_dic['gpu_driver_version'] = pynvml.nvmlSystemGetDriverVersion()
                for gpu_id in self.gpu_ids:
                    handle = pynvml.nvmlDeviceGetHandleByIndex(gpu_id)
                    gpu_id_name = "gpu%s" % gpu_id
                    mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
                    gpu_utilize = pynvml.nvmlDeviceGetUtilizationRates(handle)
                    self.config_dic['%s_device_name' % gpu_id_name] = pynvml.nvmlDeviceGetName(handle)
                    self.config_dic['%s_mem_total' % gpu_id_name] = gpu_mem_total = round(mem_info.total / 1024 ** 3, 2)
                    self.config_dic['%s_mem_used' % gpu_id_name] = gpu_mem_used = round(mem_info.used / 1024 ** 3, 2)
                    # self.config_dic['%s_mem_free' % gpu_id_name] = gpu_mem_free = mem_info.free // 1024 ** 2
                    self.config_dic['%s_mem_percent' % gpu_id_name] = round((gpu_mem_used / gpu_mem_total) * 100, 1)
                    self._set_dict_smooth('%s_utilize_gpu' % gpu_id_name, gpu_utilize.gpu, 0.8)
                    # self.config_dic['%s_utilize_gpu' % gpu_id_name] = gpu_utilize.gpu
                    # self.config_dic['%s_utilize_memory' % gpu_id_name] = gpu_utilize.memory

                pynvml.nvmlShutdown()
            except Exception as e:
                print(e) 

# 需要导入模块: import pynvml [as 别名]
# 或者: from pynvml import nvmlDeviceGetMemoryInfo [as 别名]
def _get_vram(self):
        """ Obtain the total VRAM in Megabytes for each connected GPU.

        Returns
        -------
        list
             List of floats containing the total amount of VRAM in Megabytes for each connected GPU
             as corresponding to the values in :attr:`_handles
        """
        self._initialize()
        if self._device_count == 0:
            vram = list()
        elif self._is_plaidml:
            vram = self._plaid.vram
        elif IS_MACOS:
            vram = [pynvx.cudaGetMemTotal(handle, ignore=True) / (1024 * 1024)
                    for handle in self._handles]
        else:
            vram = [pynvml.nvmlDeviceGetMemoryInfo(handle).total /
                    (1024 * 1024)
                    for handle in self._handles]
        self._log("debug", "GPU VRAM: {}".format(vram))
        return vram 

# 需要导入模块: import pynvml [as 别名]
# 或者: from pynvml import nvmlDeviceGetMemoryInfo [as 别名]
def __query_mem(handle):
        """
        Query information on the memory of a GPU.

        Arguments:
            handle:
                NVML device handle.

        Returns:
            summaries (:obj:`dict`):
                Dictionary containing the memory values for ['mem_used', 'mem_free', 'mem_total'].
                All values are given in MiB as integers.
        """
        # Query information on the GPUs memory usage.
        info = nvml.nvmlDeviceGetMemoryInfo(handle)

        summaries = dict()
        bytes_mib = 1024.0 ** 2
        summaries['mem_used'] = int(info.used / bytes_mib)
        summaries['mem_free'] = int(info.free / bytes_mib)
        summaries['mem_total'] = int(info.total / bytes_mib)

        return summaries 

# 需要导入模块: import pynvml [as 别名]
# 或者: from pynvml import nvmlDeviceGetMemoryInfo [as 别名]
def get_appropriate_cuda(task_scale='s'):
    if task_scale not in {'s','m','l'}:
        logger.info('task scale wrong!')
        exit(2)
    import pynvml
    pynvml.nvmlInit()
    total_cuda_num = pynvml.nvmlDeviceGetCount()
    for i in range(total_cuda_num):
        logger.info(i)
        handle = pynvml.nvmlDeviceGetHandleByIndex(i)  # 这里的0是GPU id
        memInfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
        utilizationInfo = pynvml.nvmlDeviceGetUtilizationRates(handle)
        logger.info(i, 'mem:', memInfo.used / memInfo.total, 'util:',utilizationInfo.gpu)
        if memInfo.used / memInfo.total < 0.15 and utilizationInfo.gpu <0.2:
            logger.info(i,memInfo.used / memInfo.total)
            return 'cuda:'+str(i)

    if task_scale=='s':
        max_memory=2000
    elif task_scale=='m':
        max_memory=6000
    else:
        max_memory = 9000

    max_id = -1
    for i in range(total_cuda_num):
        handle = pynvml.nvmlDeviceGetHandleByIndex(0)  # 这里的0是GPU id
        memInfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
        utilizationInfo = pynvml.nvmlDeviceGetUtilizationRates(handle)
        if max_memory < memInfo.free:
            max_memory = memInfo.free
            max_id = i

    if id == -1:
        logger.info('no appropriate gpu, wait!')
        exit(2)

    return 'cuda:'+str(max_id)

        # if memInfo.used / memInfo.total < 0.5:
        #     return 

# 需要导入模块: import pynvml [as 别名]
# 或者: from pynvml import nvmlDeviceGetMemoryInfo [as 别名]
def mem_used_for(device_handle):
    """Get GPU device memory consumption in percent"""
    try:
        return pynvml.nvmlDeviceGetMemoryInfo(device_handle).used
    except pynvml.NVMLError:
        return None 

# 需要导入模块: import pynvml [as 别名]
# 或者: from pynvml import nvmlDeviceGetMemoryInfo [as 别名]
def mem_used_percent_for(device_handle):
    """Get GPU device memory consumption in percent"""
    try:
        memory_info = pynvml.nvmlDeviceGetMemoryInfo(device_handle)
        return memory_info.used * 100.0 / memory_info.total
    except pynvml.NVMLError:
        return None 

# 需要导入模块: import pynvml [as 别名]
# 或者: from pynvml import nvmlDeviceGetMemoryInfo [as 别名]
def _get_free_vram(self):
        """ Obtain the amount of VRAM that is available, in Megabytes, for each connected GPU.

        Returns
        -------
        list
             List of floats containing the amount of VRAM available, in Megabytes, for each
             connected GPU as corresponding to the values in :attr:`_handles

        Notes
        -----
        There is no useful way to get free VRAM on PlaidML. OpenCL loads and unloads VRAM as
        required, so this returns the total memory available per card for AMD cards, which us
        not particularly useful.

        """
        self._initialize()
        if self._is_plaidml:
            vram = self._plaid.vram
        elif IS_MACOS:
            vram = [pynvx.cudaGetMemFree(handle, ignore=True) / (1024 * 1024)
                    for handle in self._handles]
        else:
            vram = [pynvml.nvmlDeviceGetMemoryInfo(handle).free / (1024 * 1024)
                    for handle in self._handles]
        self._shutdown()
        self._log("debug", "GPU VRAM free: {}".format(vram))
        return vram 

# 需要导入模块: import pynvml [as 别名]
# 或者: from pynvml import nvmlDeviceGetMemoryInfo [as 别名]
def query_gpu(handle: int) -> Dict:
    memory = pynvml.nvmlDeviceGetMemoryInfo(handle)  # in Bytes
    utilization = pynvml.nvmlDeviceGetUtilizationRates(handle)

    return {
        "gpu_{}_memory_free".format(handle): int(memory.free),
        "gpu_{}_memory_used".format(handle): int(memory.used),
        "gpu_{}_utilization".format(handle): utilization.gpu,
    } 

# 需要导入模块: import pynvml [as 别名]
# 或者: from pynvml import nvmlDeviceGetMemoryInfo [as 别名]
def _crawl_in_system(self):
        '''
        nvidia-smi returns following: MEMORY, UTILIZATION, ECC, TEMPERATURE,
        POWER, CLOCK, COMPUTE, PIDS, PERFORMANCE, SUPPORTED_CLOCKS,
        PAGE_RETIREMENT, ACCOUNTING

        currently, following are requested based on dlaas requirements:
            utilization.gpu, utilization.memory,
            memory.total, memory.free, memory.used
        nvidia-smi --query-gpu=utilization.gpu,utilization.memory,\
            memory.total,memory.free,memory.used --format=csv,noheader,nounits
        '''

        if self._init_nvml() == -1:
            return

        self.inspect_arr = exec_dockerps()

        num_gpus = pynvml.nvmlDeviceGetCount()

        for gpuid in range(0, num_gpus):
            gpuhandle = pynvml.nvmlDeviceGetHandleByIndex(gpuid)
            temperature = pynvml.nvmlDeviceGetTemperature(
                gpuhandle, pynvml.NVML_TEMPERATURE_GPU)
            memory = pynvml.nvmlDeviceGetMemoryInfo(gpuhandle)
            mem_total = memory.total / 1024 / 1024
            mem_used = memory.used / 1024 / 1024
            mem_free = memory.free / 1024 / 1024
            power_draw = pynvml.nvmlDeviceGetPowerUsage(gpuhandle) / 1000
            power_limit = pynvml.nvmlDeviceGetEnforcedPowerLimit(
                gpuhandle) / 1000
            util = pynvml.nvmlDeviceGetUtilizationRates(gpuhandle)
            util_gpu = util.gpu
            util_mem = util.memory
            entry = {
                'utilization': {'gpu': util_gpu, 'memory': util_mem},
                'memory': {'total': mem_total, 'free': mem_free,
                           'used': mem_used},
                'temperature': temperature,
                'power': {'draw': power_draw, 'limit': power_limit}
            }
            key = self._get_feature_key(gpuhandle, gpuid)
            if gpuid == num_gpus - 1:
                self._shutdown_nvml()

            yield (key, entry, 'gpu')

        return