Merge pull request #42 from NVIDIA/bhagedorn/fix-unhashable-config-groupby

Fix unhashable (dict/list) config params crashing aggregate_data

Author: bastianhagedorn

nsight/transformation.py

@@ -49,20 +49,6 @@ def aggregate_data(
     # Note: When num_args=0, we need an empty list (not all columns via [-0:])
     func_fields = df.columns[-num_args:].tolist() if num_args > 0 else []
 
-    # Function to convert non-sortable columns to tuples or strings
-    def convert_non_sortable_columns(dframe: pd.DataFrame) -> pd.DataFrame:
-        for col in dframe.columns:
-            try:
-                # Try sorting the column to check if it's sortable.
-                sorted(dframe[col].dropna())
-            except (TypeError, ValueError):
-                # If sorting fails, convert the column to string
-                dframe[col] = dframe[col].astype(str)
-        return dframe
-
-    # Convert non-sortable columns before grouping
-    df = convert_non_sortable_columns(df)
-
     # Preserve original order by adding an index column
     df = df.reset_index(drop=True)
     df["_original_order"] = df.index
@@ -109,12 +95,41 @@ def convert_non_sortable_columns(dframe: pd.DataFrame) -> pd.DataFrame:
                 )(col),
             )
 
+    # Group keys must be hashable (pandas factorizes them) and sortable
+    # (groupby sorts them). Keys that are neither — e.g. dict- or list-valued
+    # config parameters — are grouped on a temporary stringified key, and the
+    # original objects are recovered with a "first" aggregation so the output
+    # preserves them (a dict stays a dict instead of becoming its repr).
+    def _is_groupable(series: pd.Series) -> bool:
+        non_null = series.dropna()
+        try:
+            set(non_null)
+            sorted(non_null)
+        except (TypeError, ValueError):
+            return False
+        return True
+
+    sortkey_cols: list[str] = []
+    group_keys: list[str] = []
+    func_field_keys: list[str] = []
+    for col in groupby_columns + func_fields:
+        if _is_groupable(df[col]):
+            key = col
+        else:
+            key = f"{col}__sortkey__"
+            df[key] = df[col].astype(str)
+            sortkey_cols.append(key)
+            named_aggs[col] = (col, "first")  # preserve the original objects
+        group_keys.append(key)
+        if col in func_fields:
+            func_field_keys.append(key)
+
     # Ensure Value column is numeric (explode can leave it as object dtype
     # even though the underlying values are always numeric from NCU metrics)
     df["Value"] = pd.to_numeric(df["Value"])
 
     # Apply aggregation with named aggregation
-    groupby_df = df.groupby(groupby_columns + func_fields, dropna=False)
+    groupby_df = df.groupby(group_keys, dropna=False)
     agg_df = groupby_df.agg(**named_aggs).reset_index()
 
     # Compute 95% confidence intervals
@@ -149,8 +164,9 @@ def convert_non_sortable_columns(dframe: pd.DataFrame) -> pd.DataFrame:
             normalize_against in agg_df["Annotation"].values
         ), f"Annotation '{normalize_against}' not found in data."
 
-        # Columns of normalization dataframe to merge on
-        merge_on = [*func_fields, "Metric"]
+        # Columns of normalization dataframe to merge on (use the stringified
+        # keys for any unhashable config columns so the merge can hash them)
+        merge_on = [*func_field_keys, "Metric"]
 
         # Create a DataFrame to hold the normalization values
         normalization_df = agg_df[agg_df["Annotation"] == normalize_against][
@@ -178,4 +194,8 @@ def convert_non_sortable_columns(dframe: pd.DataFrame) -> pd.DataFrame:
         )
     )
 
+    # Drop the temporary stringified group keys; the original columns remain.
+    if sortkey_cols:
+        agg_df = agg_df.drop(columns=sortkey_cols)
+
     return agg_df

tests/test_transformation.py

@@ -0,0 +1,82 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Unit tests for nsight.transformation.aggregate_data (no GPU required)."""
+
+from typing import Any
+
+import pandas as pd
+
+from nsight import transformation
+
+
+def _one_arg(config: Any) -> None:
+    """A function with a single config parameter (num_args == 1)."""
+    return None
+
+
+def _raw_df(config_values: list[Any], *, param_name: str = "config") -> pd.DataFrame:
+    """Synthetic raw frame; the config column is last, as aggregate_data expects."""
+    n = len(config_values)
+    data: dict[str, Any] = {
+        "Annotation": ["matmul"] * n,
+        "Metric": ["gpu__time_duration.sum"] * n,
+        "Value": [float(i + 1) for i in range(n)],
+        "Kernel": ["ampere_sgemm_128x128"] * n,
+        param_name: config_values,
+    }
+    return pd.DataFrame(data)
+
+
+def test_dict_config_single_row_does_not_crash() -> None:
+    # Regression: a single dict-config row used to crash groupby().agg() with
+    # "unhashable type: 'dict'". The output must keep the dict as a real dict.
+    dims = {"M": 512, "K": 512, "N": 512}
+    result = transformation.aggregate_data(_raw_df([dims]), _one_arg, None, False)
+    assert len(result) == 1
+    assert result["config"].iloc[0] == dims
+    assert isinstance(result["config"].iloc[0], dict)
+    assert result["config"].iloc[0]["M"] == 512  # usable as a dict, no parsing
+    assert result["NumRuns"].iloc[0] == 1
+    assert result["AvgValue"].iloc[0] == 1.0
+    # The temporary sort key must not leak into the output.
+    assert not any("sortkey" in c for c in result.columns)
+
+
+def test_dict_config_independent_of_run_count() -> None:
+    # Bumping runs only masked the bug (it added rows); every count must work.
+    dims = {"M": 256, "K": 256, "N": 256}
+    for runs in (1, 2, 5):
+        result = transformation.aggregate_data(
+            _raw_df([dims] * runs), _one_arg, None, False
+        )
+        assert len(result) == 1
+        assert result["NumRuns"].iloc[0] == runs
+        assert result["config"].iloc[0] == dims
+
+
+def test_distinct_dict_configs_stay_distinct() -> None:
+    # Different dicts form different groups, each preserved as a real dict.
+    d1 = {"M": 512, "K": 512, "N": 512}
+    d2 = {"M": 1024, "K": 1024, "N": 1024}
+    result = transformation.aggregate_data(_raw_df([d1, d2]), _one_arg, None, False)
+    assert len(result) == 2
+    configs = list(result["config"])
+    assert all(isinstance(c, dict) for c in configs)
+    assert d1 in configs and d2 in configs
+
+
+def test_list_config_single_row_does_not_crash() -> None:
+    # Lists are unhashable too, and must also be preserved as lists.
+    shape = [512, 512]
+    result = transformation.aggregate_data(_raw_df([shape]), _one_arg, None, False)
+    assert len(result) == 1
+    assert result["config"].iloc[0] == shape
+    assert isinstance(result["config"].iloc[0], list)
+
+
+def test_numeric_config_keeps_native_dtype() -> None:
+    # Sortable, hashable configs must not be needlessly coerced to strings.
+    result = transformation.aggregate_data(_raw_df([128]), _one_arg, None, False)
+    assert result["config"].iloc[0] == 128
+    assert pd.api.types.is_integer_dtype(result["config"])