Update app.py

app.py

@@ -227,12 +227,114 @@ def plot_disks(alert_threshold_pct=99.0):
     return fig, df
 
 
+def plot_users_disks(topk_to_show=5):
+    df = datasets.load_dataset(
+        "pluslab/PLUS_Lab_GPUs_Data",
+        data_files="users_disks.csv",
+        download_mode="force_redownload",
+    )["train"].to_pandas()
+
+    df["server_path"] = df["server"] + "/" + df["path"]
+    df["server_path"] = df["server_path"].apply(
+        lambda x: x.replace(".cs.ucla.edu", "")
+    )
+    df = df[df["path"] != "home"]
+
+    # Ensure numeric
+    df["size_GB"] = pd.to_numeric(df["size_GB"], errors="coerce").fillna(0)
+
+    # Preserve original appearance order
+    original_order = df["server_path"].drop_duplicates()
+
+    heatmap_df = df.pivot_table(
+        index="server_path",
+        columns="username",
+        values="size_GB",
+        aggfunc="sum",
+        fill_value=0,
+    )
+
+    # Reindex to preserve original order first
+    heatmap_df = heatmap_df.reindex(original_order)
+
+    # ----------- Custom ordering logic -----------
+    local1_rows = [p for p in heatmap_df.index if "/local1" in p]
+    home_rows = [p for p in heatmap_df.index if "/home" in p]
+    other_rows = [
+        p for p in heatmap_df.index
+        if p not in local1_rows and p not in home_rows
+    ]
+
+    new_order =  home_rows + local1_rows + other_rows
+    heatmap_df = heatmap_df.reindex(new_order)
+    # ---------------------------------------------
+
+    # --- Top-K users per server_path (row), then union across rows ---
+    # Don't include any special rows (e.g., Total) if they exist yet
+    base_df = heatmap_df.copy()
+
+    # For each row, get top-K usernames by usage (ignore zeros)
+    top_users_per_row = base_df.apply(
+        lambda row: row[row > 0].nlargest(topk_to_show).index,
+        axis=1
+    )
+
+    # Union of all selected users across rows
+    selected_users = pd.Index(sorted(set().union(*top_users_per_row.tolist())))
+
+    # Keep only those users (columns) that exist
+    heatmap_df = heatmap_df.loc[:, heatmap_df.columns.intersection(selected_users)]
+    # ---------------------------------------------------------------
+
+    # Add Total row LAST
+    heatmap_df.loc["Total"] = heatmap_df.sum(axis=0)
+
+    # Sort users by total usage (descending)
+    heatmap_df = heatmap_df[
+        heatmap_df.loc["Total"].sort_values(ascending=False).index
+    ]
+
+    # -------- Reverse row order (keep Total last) --------
+    non_total_rows = [r for r in heatmap_df.index if r != "Total"]
+    reversed_rows = ["Total"] + list(reversed(non_total_rows))
+    heatmap_df = heatmap_df.reindex(reversed_rows)
+    # -----------------------------------------------------
+
+    # Text labels inside cells
+    text_values = heatmap_df.round(1).astype(str)
+
+    fig = go.Figure(
+        data=go.Heatmap(
+            z=heatmap_df.values,
+            x=heatmap_df.columns.tolist(),
+            y=heatmap_df.index.tolist(),
+            text=text_values.values,
+            texttemplate="%{text}",
+            textfont={"size": 12},
+            hovertemplate="path=%{y}<br>user=%{x}<br>size_GB=%{z:.2f}<extra></extra>",
+            colorscale="YlOrRd",
+        )
+    )
+
+    fig.update_layout(
+        title="Users Disk Usage Heatmap (GB)",
+        xaxis_title="username",
+        yaxis_title="server_path",
+    )
+
+    return fig, df
+
 
 ##### PLOT ALL #####
 def plot_figs():
-    fig_gpus, df_gpus = plot_gpus()
-    fig_disks, df_disks = plot_disks()
-    return fig_gpus, fig_disks, df_gpus, df_disks
+    fig_gpus, fig_disks, fig_users_disks, df_gpus, df_disks, df_users_disks = None, None, None, None, None, None
+    try:
+        fig_gpus, df_gpus = plot_gpus()
+        fig_disks, df_disks = plot_disks()
+        fig_users_disks, df_users_disks = plot_users_disks()
+    except Exception as e:
+        print(e)
+    return fig_gpus, fig_disks, fig_users_disks, df_gpus, df_disks, df_users_disks
 
 demo = gr.Interface(
     fn=plot_figs,
@@ -242,8 +344,10 @@ demo = gr.Interface(
     outputs = [
         gr.Plot(label="GPU Status", elem_classes="plotcss"),
         gr.Plot(label="Disk Status", elem_classes="plotcss"),
+        gr.Plot(label="Users Disk Status", elem_classes="plotcss"),
         gr.Dataframe(label="GPU Status Details"),
         gr.Dataframe(label="Disk Status Details"),
+        gr.Dataframe(label="Users Disk Status Details"),
     ],
     live=True,
     flagging_options=[],
@@ -252,4 +356,4 @@ demo = gr.Interface(
 )
 
 if __name__ == "__main__":
-    demo.launch(debug=False)
+    demo.launch(debug=False)