[PATCH v6 2/4] usertools/cpu_layout: print out NUMA nodes

[Anatoly Burakov]

In traditional NUMA case, NUMA nodes and physical sockets were used
interchangeably, but there are cases where there can be multiple NUMA
nodes per socket, as well as all CPU's being assigned NUMA node 0 even in
cases of multiple sockets. Use sysfs to print out NUMA information.

Signed-off-by: Anatoly Burakov <anatoly.burakov at intel.com>
---

Notes:
    v5 -> v6:
    - Track NUMA changes per socket to avoid issues with missing cores
    
    v2 -> v3:
    - Sort imports alphabetically

 usertools/cpu_layout.py | 35 ++++++++++++++++++++++++++++++-----
 1 file changed, 30 insertions(+), 5 deletions(-)

diff --git a/usertools/cpu_layout.py b/usertools/cpu_layout.py
index e133fb8ad3..5972cfecdb 100755
--- a/usertools/cpu_layout.py
+++ b/usertools/cpu_layout.py
@@ -5,6 +5,7 @@
 
 """Display CPU topology information."""
 
+import glob
 import typing as T
 
 
@@ -29,12 +30,21 @@ def read_sysfs(path: str) -> str:
         return fd.read().strip()
 
 
+def read_numa_node(base: str) -> int:
+    """Read the NUMA node of a CPU."""
+    node_glob = f"{base}/node*"
+    node_dirs = glob.glob(node_glob)
+    if not node_dirs:
+        return 0  # default to node 0
+    return int(node_dirs[0].split("node")[1])
+
+
 def print_row(row: T.Tuple[str, ...], col_widths: T.List[int]) -> None:
     """Print a row of a table with the given column widths."""
     first, *rest = row
     w_first, *w_rest = col_widths
     first_end = " " * 4
-    rest_end = " " * 10
+    rest_end = " " * 4
 
     print(first.ljust(w_first), end=first_end)
     for cell, width in zip(rest, w_rest):
@@ -56,6 +66,7 @@ def main() -> None:
     sockets_s: T.Set[int] = set()
     cores_s: T.Set[int] = set()
     core_map: T.Dict[T.Tuple[int, int], T.List[int]] = {}
+    numa_map: T.Dict[int, int] = {}
     base_path = "/sys/devices/system/cpu"
 
     cpus = range_expand(read_sysfs(f"{base_path}/online"))
@@ -64,12 +75,14 @@ def main() -> None:
         lcore_base = f"{base_path}/cpu{cpu}"
         core = int(read_sysfs(f"{lcore_base}/topology/core_id"))
         socket = int(read_sysfs(f"{lcore_base}/topology/physical_package_id"))
+        node = read_numa_node(lcore_base)
 
         cores_s.add(core)
         sockets_s.add(socket)
         key = (socket, core)
         core_map.setdefault(key, [])
         core_map[key].append(cpu)
+        numa_map[cpu] = node
 
     cores = sorted(cores_s)
     sockets = sorted(sockets_s)
@@ -78,24 +91,36 @@ def main() -> None:
 
     print("cores = ", cores)
     print("sockets = ", sockets)
+    print("numa = ", sorted(set(numa_map.values())))
     print()
 
-    # Core, [Socket, Socket, ...]
-    heading_strs = "", *[f"Socket {s}" for s in sockets]
+    # Core, [NUMA, Socket, NUMA, Socket, ...]
+    heading_strs = "", *[v for s in sockets for v in ("", f"Socket {s}")]
     sep_strs = tuple("-" * len(hstr) for hstr in heading_strs)
     rows: T.List[T.Tuple[str, ...]] = []
 
+    # track NUMA changes per socket
+    prev_numa: T.Dict[int, T.Optional[int]] = {socket: None for socket in sockets}
     for c in cores:
         # Core,
         row: T.Tuple[str, ...] = (f"Core {c}",)
 
-        # [lcores, lcores, ...]
+        # [NUMA, lcores, NUMA, lcores, ...]
         for s in sockets:
             try:
                 lcores = core_map[(s, c)]
+
+                numa = numa_map[lcores[0]]
+                numa_changed = prev_numa[s] != numa
+                prev_numa[s] = numa
+
+                if numa_changed:
+                    row += (f"NUMA {numa}",)
+                else:
+                    row += ("",)
                 row += (str(lcores),)
             except KeyError:
-                row += ("",)
+                row += ("", "")
         rows += [row]
 
     # find max widths for each column, including header and rows
-- 
2.43.5