.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "examples/list_gpus.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_examples_list_gpus.py>`
        to download the full example code.

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_examples_list_gpus.py:


Simple script for listing GPU devices
=====================================

.. GENERATED FROM PYTHON SOURCE LINES 7-74

.. code-block:: Python


    import cuda.bindings.runtime as cudart

    from pyhwloc import from_this_system
    from pyhwloc.hwobject import GetTypeDepth, ObjType, OsDevice
    from pyhwloc.topology import TypeFilter


    def list_with_nvml() -> None:
        import pynvml

        from pyhwloc.nvml import get_cpu_affinity, get_device

        try:
            pynvml.nvmlInit()

            with from_this_system().set_io_types_filter(TypeFilter.KEEP_ALL) as topo:
                status, cnt = cudart.cudaGetDeviceCount()
                assert status == cudart.cudaError_t.cudaSuccess
                for i in range(cnt):
                    # Create a NVML device
                    dev = get_device(topo, i)
                    # Get the hwloc native device type
                    osdev = dev.get_osdev()
                    assert osdev is not None
                    print(osdev, ":", osdev.format_attr())
                    print("cpuset:", dev.get_affinity().to_sched_set())
                    print("cpuset from nvml:", get_cpu_affinity(i).to_sched_set())
        finally:
            pynvml.nvmlShutdown()


    def list_with_cudart() -> None:
        from pyhwloc.cuda_runtime import get_device

        with from_this_system().set_io_types_filter(TypeFilter.KEEP_ALL) as topo:
            status, cnt = cudart.cudaGetDeviceCount()
            assert status == cudart.cudaError_t.cudaSuccess
            for i in range(cnt):
                # Create a CUDA runtime device
                dev = get_device(topo, i)
                # Get the hwloc native device type
                osdev = dev.get_osdev()
                assert osdev is not None
                print(osdev, ":", osdev.format_attr())
                print("cpuset:", dev.get_affinity().to_sched_set())


    def list_with_osdev() -> None:
        # GPU is categorized as IO device.
        # Hwloc lists devices from both NVML and CUDA.
        with from_this_system().set_io_types_filter(TypeFilter.KEEP_ALL) as topo:
            # Look for OS devices (which include GPUs)
            for obj in topo.iter_objs_by_type(ObjType.OS_DEVICE):
                # Check if it's a GPU device
                if isinstance(obj, OsDevice) and obj.is_gpu():
                    assert obj.depth == GetTypeDepth.OS_DEVICE
                    print(obj, ":", obj.format_attr())


    if __name__ == "__main__":
        print("List by hwloc OS device\n")
        list_with_osdev()
        print("\nList by CUDA\n")
        list_with_cudart()
        print("\nList by NVML\n")
        list_with_nvml()


.. _sphx_glr_download_examples_list_gpus.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: list_gpus.ipynb <list_gpus.ipynb>`

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: list_gpus.py <list_gpus.py>`

    .. container:: sphx-glr-download sphx-glr-download-zip

      :download:`Download zipped: list_gpus.zip <list_gpus.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_