From 105ad02148f3231addca0128147472740eef865a Mon Sep 17 00:00:00 2001 From: Aditya Ranjan Date: Wed, 23 Oct 2024 23:50:38 -0400 Subject: [PATCH] bugfix for load_imbalance (#133) Co-authored-by: adityaranjan Co-authored-by: Alexander Movsesyan <92064328+movsesyanae@users.noreply.github.com> --- pipit/trace.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/pipit/trace.py b/pipit/trace.py index fcc9f75d..75946931 100644 --- a/pipit/trace.py +++ b/pipit/trace.py @@ -585,10 +585,14 @@ def load_imbalance(self, metric="time.exc", num_processes=1): for function in functions: curr_series = flat_profile.loc[function] - top_n = curr_series.sort_values(ascending=False).iloc[0:num_display] + top_n = curr_series.sort_values(by=metric, ascending=False).iloc[ + 0:num_display + ] - imbalance_dict[mean_metric].append(curr_series.mean()) - imbalance_dict[imb_metric].append(top_n.values[0] / curr_series.mean()) + imbalance_dict[mean_metric].append(curr_series.mean().values[0]) + imbalance_dict[imb_metric].append( + (top_n.values[0] / curr_series.mean()).values[0] + ) imbalance_dict[imb_ranks].append(list(top_n.index)) imbalance_df = pd.DataFrame(imbalance_dict)