From 42fe2f9ae9c4d65cf74b66b12c0ea2b1c4da7e48 Mon Sep 17 00:00:00 2001 From: Michael McKinsey Date: Thu, 29 Oct 2026 15:50:15 -0700 Subject: [PATCH] Add unit test for update_inclusive_columns --- .../tests/test_update_inclusive_columns.py | 313 ++++++++++++++++++ 1 file changed, 313 insertions(+) create mode 100644 thicket/tests/test_update_inclusive_columns.py diff --git a/thicket/tests/test_update_inclusive_columns.py b/thicket/tests/test_update_inclusive_columns.py new file mode 100644 index 00000000..c201929b --- /dev/null +++ b/thicket/tests/test_update_inclusive_columns.py @@ -0,0 +1,313 @@ +# Copyright 2022 Lawrence Livermore National Security, LLC and other +# Thicket Project Developers. See the top-level LICENSE file for details. +# +# SPDX-License-Identifier: MIT + +"""Tests of update_inclusive_columns, which Thicket inherits from Hatchet's +GraphFrame, on the MultiIndex dataframes of a Thicket. + +The rows of a Thicket's performance data are indexed by "node", usually +"profile", and any other index levels of the original GraphFrames, such as +"rank". After concat_thickets(axis="columns"), the columns are a MultiIndex too. +Inclusive values must be computed separately for each combination of the index +levels other than "node". +""" + +import numpy as np +import pandas as pd +from hatchet import GraphFrame +from hatchet.frame import Frame +from hatchet.graph import Graph +from hatchet.node import Node + +from thicket import Thicket +from thicket.query import Query + +# A -> B -> C, A -> B -> D, A -> E +TREE = [("A", "B"), ("B", "C"), ("B", "D"), ("A", "E")] +TREE_TIMES = {"A": 1.0, "B": 2.0, "C": 3.0, "D": 4.0, "E": 5.0} + + +def _make_graphframe(edges, times): + """Build a GraphFrame by hand, so that its graph can be any DAG. + + Arguments: + edges (list): (parent, child) pairs of node names + times (dict): maps a node name, or a (name, rank) tuple, to the exclusive + "time" of that row + + Returns: + (GraphFrame): GraphFrame indexed by "node", or by "node" and "rank" + """ + keys = [key if isinstance(key, tuple) else (key,) for key in times] + names = sorted({key[0] for key in keys}) + nodes = {name: Node(Frame({"name": name, "type": "function"})) for name in names} + for parent, child in edges: + nodes[parent].add_child(nodes[child]) + nodes[child].add_parent(nodes[parent]) + graph = Graph([node for node in nodes.values() if not node.parents]) + # Nodes hash by their ID, so set the IDs before indexing by node + graph.enumerate_traverse() + + index_names = ["node", "rank"][: len(keys[0])] + rows = [ + (nodes[key[0]],) + key[1:] + (key[0], time) + for key, time in zip(keys, times.values()) + ] + dataframe = pd.DataFrame(rows, columns=index_names + ["name", "time"]) + dataframe = dataframe.set_index(index_names).sort_index() + return GraphFrame(graph, dataframe, exc_metrics=["time"], inc_metrics=[]) + + +def _concat_profiles(graphframes, **kwargs): + """Make each GraphFrame a profile of one Thicket. + + The metadata column "run" is the position of the profile's GraphFrame in + graphframes. Keyword arguments are passed to Thicket.concat_thickets. + """ + thickets = [] + for run, gf in enumerate(graphframes): + tk = Thicket.thicketize_graphframe(gf, "run {}".format(run)) + # thicketize_graphframe leaves the metadata table empty, which squash() + # rejects because its profiles don't match the performance data's + tk.metadata = pd.DataFrame( + {"run": [run]}, index=pd.Index(tk.profile, name=tk.profile_idx_name) + ) + thickets.append(tk) + return Thicket.concat_thickets(thickets, disable_tqdm=True, **kwargs) + + +def _scaled(values, factor): + return {key: factor * value for key, value in values.items()} + + +def _assert_column(tk, column, *runs): + """Check a column of the performance data of a Thicket from _concat_profiles. + + Arguments: + tk (Thicket): Thicket to check + column (str): column of tk.dataframe + runs (dict): for each run, a dict that maps a node name, or a (name, rank) + tuple, to the expected value. The rows of the column must be exactly + the rows given. + """ + expected = {} + for run, values in enumerate(runs): + for key, value in values.items(): + key = key if isinstance(key, tuple) else (key,) + expected[(key[0], run) + key[1:]] = value + + run_of_profile = tk.metadata["run"] + index = [ + (node.frame["name"], run_of_profile[profile]) + tuple(rest) + for node, profile, *rest in tk.dataframe.index + ] + actual = pd.Series( + tk.dataframe[column].to_numpy(), index=pd.MultiIndex.from_tuples(index) + ) + pd.testing.assert_series_equal( + actual.sort_index(), pd.Series(expected).sort_index(), check_names=False + ) + + +def _reference_inclusive(tk, exc_metric): + """Compute the inclusive values of an exclusive metric independently. + + The inclusive value of a row is the sum of exc_metric over the row's node and + all of its descendants, each counted once, over the rows that have the same + values as the row in all other index levels (e.g., the same profile). NaNs + are skipped, but a sum of only NaNs is NaN. + """ + exc = tk.dataframe[exc_metric] + nodes = exc.index.get_level_values("node") + other_levels = [level for level in exc.index.names if level != "node"] + parts = [] + for node in tk.graph.traverse(): + subgraph = exc[nodes.isin(list(node.traverse()))] + sums = subgraph.groupby(level=other_levels).sum(min_count=1) + own_rows = exc[nodes == node].index + parts.append( + pd.Series( + sums.reindex(own_rows.droplevel("node")).to_numpy(), index=own_rows + ) + ) + return pd.concat(parts).reindex(exc.index) + + +def _check_against_reference(tk): + """Run update_inclusive_columns on tk and compare each inclusive column with + _reference_inclusive.""" + exc_metrics = list(tk.exc_metrics) + assert exc_metrics + + tk.update_inclusive_columns() + + for exc_metric in exc_metrics: + if isinstance(exc_metric, tuple): # concat_thickets(axis="columns") + inc_metric = exc_metric[:-1] + (exc_metric[-1] + " (inc)",) + else: + inc_metric = exc_metric + " (inc)" + assert inc_metric in tk.inc_metrics + pd.testing.assert_series_equal( + tk.dataframe[inc_metric], + _reference_inclusive(tk, exc_metric), + check_names=False, + rtol=1e-12, + ) + + +def test_update_inclusive_columns_per_profile(): + """Each profile's inclusive values only sum that profile's exclusive values.""" + tk = _concat_profiles( + [ + _make_graphframe(TREE, TREE_TIMES), + _make_graphframe(TREE, _scaled(TREE_TIMES, 10)), + ] + ) + assert tk.dataframe.index.names == ["node", "profile"] + exclusive = tk.dataframe["time"].copy() + + tk.update_inclusive_columns() + + inclusive = {"A": 15.0, "B": 9.0, "C": 3.0, "D": 4.0, "E": 5.0} + _assert_column(tk, "time (inc)", inclusive, _scaled(inclusive, 10)) + assert tk.exc_metrics == ["time"] + assert tk.inc_metrics == ["time (inc)"] + # The exclusive column and the rows are unchanged + pd.testing.assert_series_equal(tk.dataframe["time"], exclusive) + + # Inclusive values are recomputed from the exclusive values, so calling the + # function again does not change them + tk.update_inclusive_columns() + _assert_column(tk, "time (inc)", inclusive, _scaled(inclusive, 10)) + assert tk.inc_metrics == ["time (inc)"] + + +def test_update_inclusive_columns_union(fill_perfdata): + """For profiles with different call trees, a node's inclusive value only sums + the descendants present in that profile. Rows that fill_perfdata filled with + NaN stay NaN, and rows that are missing are not added.""" + run0 = _make_graphframe( + [("A", "B"), ("B", "C"), ("A", "E")], + {"A": 1.0, "B": 2.0, "C": 3.0, "E": 5.0}, + ) + run1 = _make_graphframe( + [("A", "B"), ("B", "D"), ("A", "E")], + {"A": 10.0, "B": 20.0, "D": 40.0, "E": 50.0}, + ) + tk = _concat_profiles([run0, run1], fill_perfdata=fill_perfdata) + + tk.update_inclusive_columns() + + inclusive0 = {"A": 11.0, "B": 5.0, "C": 3.0, "E": 5.0} + inclusive1 = {"A": 120.0, "B": 60.0, "D": 40.0, "E": 50.0} + if fill_perfdata: + inclusive0["D"] = np.nan + inclusive1["C"] = np.nan + _assert_column(tk, "time (inc)", inclusive0, inclusive1) + + +def test_update_inclusive_columns_graph(): + """For a graph that is not a tree, each descendant of a node is counted once, + including descendants of a node that has several parents.""" + # A -> B -> D -> F, A -> C -> D, and a second root R -> S + edges = [("A", "B"), ("A", "C"), ("B", "D"), ("C", "D"), ("D", "F"), ("R", "S")] + times = {"A": 1.0, "B": 2.0, "C": 3.0, "D": 4.0, "F": 5.0, "R": 6.0, "S": 7.0} + tk = _concat_profiles( + [_make_graphframe(edges, times), _make_graphframe(edges, _scaled(times, 10))] + ) + assert len(tk.graph.roots) == 2 + assert len(tk.get_node("D").parents) == 2 + + tk.update_inclusive_columns() + + inclusive = { + "A": 15.0, + "B": 11.0, + "C": 12.0, + "D": 9.0, + "F": 5.0, + "R": 13.0, + "S": 7.0, + } + _assert_column(tk, "time (inc)", inclusive, _scaled(inclusive, 10)) + + +def test_update_inclusive_columns_extra_index_level(): + """With an index level after "profile", such as "rank" for per-rank data, + inclusive values are computed for each profile and rank.""" + edges = [("A", "B"), ("B", "C")] + times = { + ("A", 0): 1.0, + ("B", 0): 2.0, + ("C", 0): 3.0, + ("A", 1): 4.0, + ("B", 1): 5.0, + ("C", 1): 6.0, + } + tk = _concat_profiles( + [_make_graphframe(edges, times), _make_graphframe(edges, _scaled(times, 10))] + ) + assert tk.dataframe.index.names == ["node", "profile", "rank"] + + tk.update_inclusive_columns() + + inclusive = { + ("A", 0): 6.0, + ("B", 0): 5.0, + ("C", 0): 3.0, + ("A", 1): 15.0, + ("B", 1): 11.0, + ("C", 1): 6.0, + } + _assert_column(tk, "time (inc)", inclusive, _scaled(inclusive, 10)) + + +def test_squash_updates_inclusive_columns(): + """Squash (here, after a query) recomputes inclusive values for each profile + from the new graph.""" + tk = _concat_profiles( + [ + _make_graphframe(TREE, TREE_TIMES), + _make_graphframe(TREE, _scaled(TREE_TIMES, 10)), + ] + ) + tk.update_inclusive_columns() + + # Remove B, so that C and D become children of A + query = Query().match( + ".", lambda row: row["name"].apply(lambda name: name != "B").all() + ) + squashed = tk.query(query) + + inclusive = {"A": 13.0, "C": 3.0, "D": 4.0, "E": 5.0} + _assert_column(squashed, "time (inc)", inclusive, _scaled(inclusive, 10)) + + +def test_update_inclusive_columns_caliper( + rajaperf_unique_tunings, intersection, fill_perfdata +): + """Three profiles whose call trees differ (two CUDA tunings and one + sequential variant).""" + tk = Thicket.from_caliperreader( + rajaperf_unique_tunings, + intersection=intersection, + fill_perfdata=fill_perfdata, + disable_tqdm=True, + ) + _check_against_reference(tk) + + +def test_update_inclusive_columns_column_multiindex(thicket_axis_columns): + """After concat_thickets(axis="columns"), metrics are (header, metric) tuples + and each header's inclusive values are computed separately.""" + _, _, tk = thicket_axis_columns + assert isinstance(tk.dataframe.columns, pd.MultiIndex) + _check_against_reference(tk) + + +def test_update_inclusive_columns_timeseries(example_timeseries): + """A timeseries Thicket is indexed by ("node", "iteration").""" + tk = Thicket.from_timeseries(example_timeseries) + assert tk.dataframe.index.names == ["node", "iteration"] + _check_against_reference(tk)