plot: add class to parse Hwloc topology

plot/hwloc.py

+import xml.etree.ElementTree as ET
+
+
+class HwlocTopology:
+    def __init__(self, xml_filename, levels_to_remove=[]):
+        tree = ET.parse(xml_filename)
+        xml_machine = tree.getroot().find("object")
+        assert(xml_machine.attrib["type"] == "Machine")
+
+        self.machine = self.__prune_machine(self.__handle_object_dict(xml_machine))
+
+        for l in levels_to_remove:
+            self.machine = self.__remove_machine_element(self.machine, l)
+
+        assert(self.machine['children'][0]['type'] == "Package")
+        self.nb_sockets = len(self.machine['children'])
+
+        assert(self.machine['children'][0]['children'][0]['type'] == "NUMANode")
+        self.nb_numa_nodes_per_socket = len(self.machine['children'][0]['children'])
+
+        assert(self.machine['children'][0]['children'][0]['children'][0]['type'] == "Core")
+        self.nb_cores_per_numa_node = len(self.machine['children'][0]['children'][0]['children'])
+
+
+    @property
+    def nb_numa_nodes_total(self):
+        return self.nb_sockets * self.nb_numa_nodes_per_socket
+
+
+    @property
+    def nb_cores_total(self):
+        return self.nb_sockets * self.nb_numa_nodes_per_socket * self.nb_cores_per_numa_node
+
+
+    @classmethod
+    def __handle_object_dict(cls, obj, keep_all=False):
+        """
+        Convert the XML topology of a machine into more convenient Python objects.
+        """
+        t = obj.attrib["type"]
+        properties = obj.attrib
+
+        d = {
+            "type": t,
+            "properties": properties,
+            "children": []
+        }
+        for o in obj.findall("object"):
+            if keep_all or o.attrib["type"] in ["PU", "Core", "L1iCache", "L1Cache", "L2Cache", "NUMANode", "Group", "L3Cache", "Package"]:
+                d["children"].append(cls.__handle_object_dict(o, keep_all))
+        return d
+
+
+    @classmethod
+    def __prune_machine(cls, obj):
+        """
+        Improve the Python object representing topology:
+        - do not consider hyperthreads as a level, but as core's property
+        - remove {L2,L1,L1i}Cache levels (as in remove_machine_element())
+        - consider NUMA nodes as a level (instead of a Group or a Package
+          containing a NUMA node and cores on the same level)
+        """
+        if obj["type"] == "Core":
+            return {
+                "type": "Core",
+                "properties": {"nb_threads": len(obj["children"])},
+                "children": []
+            }
+        elif obj["type"] in ["L2Cache", "L1Cache", "L1iCache"]:
+            assert(len(obj["children"]) == 1)
+            return cls.__prune_machine(obj["children"][0])
+        elif obj["type"] == "Group":
+            assert(obj["children"][0]["type"] == "NUMANode")
+            return {
+                "type": "NUMANode",
+                "properties": obj["children"][0]["properties"],
+                "children": [cls.__prune_machine(c) for c in obj["children"][1:]]
+            }
+        elif obj["type"] == "Package" and obj["children"][0]["type"] == "NUMANode":
+            return {
+                "type": obj["type"],
+                "properties": obj["properties"],
+                "children": [{
+                    "type": "NUMANode",
+                    "properties": obj["children"][0]["properties"],
+                    "children": [cls.__prune_machine(c) for c in obj["children"][1:]]
+                }]
+            }
+        else:
+            return {
+                "type": obj["type"],
+                "properties": obj["properties"],
+                "children": [cls.__prune_machine(c) for c in obj["children"]]
+            }
+
+
+    @classmethod
+    def __remove_machine_element(cls, obj, element_type):
+        """
+        Remove the hierarchy level of type 'element_type' in the topology:
+        all children of the element to be removed become children of element's parent.
+
+        Useful if you want to ignore L3Cache, for instance.
+        """
+        new_children = []
+        for c in obj["children"]:
+            if c["type"] == element_type:
+                new_children += c["children"][:]
+
+        if len(new_children) > 0:
+            obj["children"] = new_children
+            return obj
+        else:
+            return {
+                "type": obj["type"],
+                "properties": obj["properties"],
+                "children": [cls.__remove_machine_element(c, element_type) for c in obj["children"]]
+            }
+
+
+    @classmethod
+    def __print_machine(cls, obj, depth=0):
+        print(" " * depth, obj["type"])
+        for c in obj["children"]:
+            cls.__print_machine(c, depth+1)
+
+    def print(self):
+        self.__print_machine(self.machine)

plot/test_hwloc.py

+import unittest
+from hwloc import *
+
+class TestHwlocTopology(unittest.TestCase):
+    def test_henri_subnuma(self):
+        topo = HwlocTopology("test_data/henri-subnuma.xml", ["L3Cache"])
+        self.assertEqual(topo.nb_sockets, 2)
+        self.assertEqual(topo.nb_numa_nodes_total, 4)
+        self.assertEqual(topo.nb_cores_total, 36)
+
+    def test_henri(self):
+        topo = HwlocTopology("test_data/henri.xml", ["L3Cache"])
+        self.assertEqual(topo.nb_sockets, 2)
+        self.assertEqual(topo.nb_numa_nodes_total, 2)
+        self.assertEqual(topo.nb_cores_total, 36)
+
+
+if __name__ == '__main__':
+    unittest.main()