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

.. only:: html

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

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

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

.. _sphx_glr_examples_gds_hops.py:


Count the number of hops between PCI storage and GPUs
=====================================================

This script counts the number of hops between NVIDIA GPUs and NVME drives and checks
whether the communication needs to go through the host socket.

Please see the `GDS document
<https://docs.nvidia.com/gpudirect-storage/configuration-guide/>`__ for a use case.

.. GENERATED FROM PYTHON SOURCE LINES 14-137

.. code-block:: Python


    from typing import TypeGuard, cast

    from pyhwloc import Topology
    from pyhwloc.hwobject import GetTypeDepth, Object, ObjType, OsDevice
    from pyhwloc.topology import TypeFilter


    def is_osdev_gpu(obj: Object) -> TypeGuard[OsDevice]:
        return isinstance(obj, OsDevice) and obj.is_gpu()


    def is_osdev_storage(obj: Object) -> TypeGuard[OsDevice]:
        return isinstance(obj, OsDevice) and obj.is_storage()


    def cnt_hops(gpu: Object, nvme: Object) -> tuple[int, bool]:
        """Count the number of hops, and check whether the host socket is required for
        communication.

        """
        print("GPU:", gpu, "<-> NVME:", nvme)
        gpu_ancestors = []
        while gpu.parent is not None:
            gpu_ancestors.append(gpu.parent)
            gpu = gpu.parent

        nvme_ancestors = []
        while nvme.parent is not None:
            nvme_ancestors.append(nvme.parent)
            nvme = nvme.parent

        # Find the common ancestor
        gpu_ancestors = list(reversed(gpu_ancestors))
        nvme_ancestors = list(reversed(nvme_ancestors))
        i, j = 0, 0
        assert gpu_ancestors[i] == nvme_ancestors[j]  # must have the same root
        while (
            i < len(gpu_ancestors)
            and j < len(nvme_ancestors)
            and gpu_ancestors[i] == nvme_ancestors[j]
        ):
            i += 1
            j += 1
        assert i >= 1 and j >= 1, (i, j)
        # It might be better to check whether this is a PCI bridge instead of checking
        # whether this is a host device.
        crossed = gpu_ancestors[i - 1].is_package() or gpu_ancestors[i - 1].is_machine()
        # Assuming GPU is an OS device attached to a PCI device. There are GPUs that are
        # connected through nvlink-C2C, which does not use PCI.
        assert gpu_ancestors[-1].is_pci_device() and nvme_ancestors[-1].is_pci_device()
        gpu_path = []
        for k in range(i - 1, len(gpu_ancestors)):
            gpu_path.append(str(gpu_ancestors[k]))
        print("- GPU path:", " -> ".join(gpu_path[::-1]))
        nvme_path = []

        for k in range(j - 1, len(nvme_ancestors)):
            nvme_path.append(str(nvme_ancestors[k]))
        print("- NVME path:", " -> ".join(nvme_path[::-1]))
        print()

        # +1 for the common ancestor. -1 for removing the PCI device. OS device is a
        # software device, which is "attached" to the PCI device. But they are the same
        # thing.
        n_hops = (len(gpu_ancestors) - i - 1) + (len(nvme_ancestors) - j - 1) + 1
        return n_hops, crossed


    def pprint(table: list[list]) -> None:
        """Print the result as a table."""
        for i in range(len(table)):
            if i == 0:
                print(" " * 7, end="")
            row = table[i]
            for j in range(len(row)):
                if j == 0:
                    print(f"{row[j]}", end="")
                else:
                    print(f"{row[j]:7},", end="")
            print()


    if __name__ == "__main__":
        with Topology.from_this_system().set_io_types_filter(TypeFilter.KEEP_ALL) as topo:
            # Look for GPUs
            gpus = []
            for obj in topo.iter_objs_by_type(ObjType.OS_DEVICE):
                # Check if it's a GPU device
                if is_osdev_gpu(obj):
                    assert obj.depth == GetTypeDepth.OS_DEVICE
                    assert obj.name is not None
                    if obj.name.lower().startswith("cuda"):
                        gpus.append(obj)

            # Look for NVME drives
            nvmes = []
            for obj in topo.iter_objs_by_type(ObjType.OS_DEVICE):
                if is_osdev_storage(obj):
                    assert obj.depth == GetTypeDepth.OS_DEVICE
                    assert obj.name is not None
                    if obj.name.lower().startswith("nvme"):
                        nvmes.append(obj)

            hops_rows: list[list[str | int]] = []
            cross_rows: list[list[str | bool]] = []
            hops_rows.append([cast(str, dev.name) for dev in nvmes])
            cross_rows.append([cast(str, dev.name) for dev in nvmes])
            for i, gpu in enumerate(gpus):
                assert gpu.name is not None
                hops_row: list[str | int] = [gpu.name]
                cross_row: list[str | bool] = [gpu.name]
                for j, nvme in enumerate(nvmes):
                    hops, cross = cnt_hops(gpu, nvme)
                    cross_row.append(cross)
                    hops_row.append(hops)
                hops_rows.append(hops_row)
                cross_rows.append(cross_row)

            print("Hops:")
            pprint(hops_rows)
            print("Cross CPU:")
            pprint(cross_rows)


.. _sphx_glr_download_examples_gds_hops.py:

.. only:: html

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

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

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

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

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

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

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


.. only:: html

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

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