Class NVMLContext

tensorpack/utils/nvml.py:149–206

Creates a context to query information Example: with NVMLContext() as ctx: num_gpus = ctx.num_devices() for device in ctx.devices(): print(device.memory()) print(device.utilization())

class NVMLContext(object):
    """Creates a context to query information

    Example:

        with NVMLContext() as ctx:
            num_gpus = ctx.num_devices()
            for device in ctx.devices():
                print(device.memory())
                print(device.utilization())

    """
    def __enter__(self):
        """Create a new context """
        _NVML.load()
        _check_return(_NVML.get_function("nvmlInit_v2")())
        return self

    def __exit__(self, type, value, tb):
        """Destroy current context"""
        _check_return(_NVML.get_function("nvmlShutdown")())

    def num_devices(self):
        """Get number of devices """
        c_count = c_uint()
        _check_return(_NVML.get_function(
            "nvmlDeviceGetCount_v2")(byref(c_count)))
        return c_count.value

    def devices(self):
        """
        Returns:
            [NvidiaDevice]: a list of devices
        """
        return [self.device(i) for i in range(self.num_devices())]

    def device(self, idx):
        """Get a specific GPU device

        Args:
            idx: index of device

        Returns:
            NvidiaDevice: single GPU device
        """
        num_dev = self.num_devices()
        assert idx < num_dev, "Cannot obtain device {}: NVML only found {} devices.".format(idx, num_dev)

        class GpuDevice(Structure):
            pass

        c_nvmlDevice_t = POINTER(GpuDevice)

        c_index = c_uint(idx)
        device = c_nvmlDevice_t()
        _check_return(_NVML.get_function(
            "nvmlDeviceGetHandleByIndex_v2")(c_index, byref(device)))
        return NvidiaDevice(device)