Spaces:
Running
Running
| import os | |
| import gradio as gr | |
| import numpy as np | |
| import pandas as pd | |
| import datetime | |
| import plotly.express as px | |
| import plotly.graph_objects as go | |
| import datasets | |
| ##### GPU PLOT ##### | |
| def split_multi_users(dfs): | |
| df = dfs.copy() | |
| df["usernames"] = df["username"].apply(lambda x: x.split(", ")) | |
| df["count"] = 1 | |
| new_df = [] | |
| for row in df.to_dict(orient="records"): | |
| gpu_users_num = len(row["usernames"]) | |
| for username in row["usernames"]: | |
| new_row = row.copy() | |
| new_row["count"] = 1 / gpu_users_num | |
| new_row["username"] = username | |
| new_df.append(new_row) | |
| df = pd.DataFrame(new_df) | |
| return df | |
| def plot_gpus(): | |
| ### Load Data | |
| dfs = datasets.load_dataset("pluslab/PLUS_Lab_GPUs_Data", data_files="gpus.csv", download_mode='force_redownload')["train"].to_pandas() | |
| dfs = dfs.drop(columns=["Unnamed: 0"]) | |
| dfs = dfs.fillna("FREE") | |
| dfs_plot = split_multi_users(dfs) | |
| fig = px.bar( | |
| dfs_plot, x="count", y="server", color="username", | |
| title=f"Last Updated {min(dfs['timestamp'])}", | |
| color_discrete_map={ | |
| "FREE": "black", | |
| }, | |
| text=dfs_plot['username'].astype(str) + "<br>" + dfs_plot['device'].astype(str), | |
| ) | |
| fig.update_layout( | |
| yaxis={'categoryorder': 'array', 'categoryarray': dfs_plot["server"].unique()[::-1]}, | |
| barcornerradius=5, | |
| ) | |
| fig.update_traces(textposition='inside', insidetextanchor='middle') | |
| # print(dfs_plot) | |
| return fig, dfs | |
| ##### DISK PLOT ##### | |
| def _pick_col(df, candidates): | |
| norm = {c.strip().lower(): c for c in df.columns} | |
| for cand in candidates: | |
| cand = cand.strip().lower() | |
| if cand in norm: | |
| return norm[cand] | |
| return None | |
| def _kblocks_to_tib(kblocks): # shown as "TB" per your convention | |
| return kblocks / (1024**3) | |
| def _kblocks_to_gib(kblocks): # shown as "GB" | |
| return kblocks / (1024**2) | |
| def plot_disks(alert_threshold_pct=99.0): | |
| df = datasets.load_dataset( | |
| "pluslab/PLUS_Lab_GPUs_Data", | |
| data_files="disks.csv", | |
| download_mode="force_redownload", | |
| )["train"].to_pandas() | |
| if "Unnamed: 0" in df.columns: | |
| df = df.drop(columns=["Unnamed: 0"]) | |
| df = pd.concat([df[df['Mounted'] != '/data2'], df[df['Mounted'] == '/data2'].drop_duplicates(subset=['Mounted'])]) # Keep one of /data2 | |
| server_col = _pick_col(df, ["server"]) | |
| fs_col = _pick_col(df, ["filesystem"]) | |
| blocks_col = _pick_col(df, ["1k-blocks", "1k blocks", "blocks"]) | |
| used_col = _pick_col(df, ["used"]) | |
| avail_col = _pick_col(df, ["available", "avail"]) | |
| mount_col = _pick_col(df, ["mounted", "mounted on", "mount", "mountpoint"]) | |
| required = [server_col, fs_col, blocks_col, used_col, avail_col] | |
| if any(c is None for c in required): | |
| raise ValueError(f"Missing required columns. Found: {list(df.columns)}") | |
| for c in [blocks_col, used_col, avail_col]: | |
| df[c] = pd.to_numeric(df[c], errors="coerce") | |
| # Label | |
| if mount_col is not None: | |
| df["Label"] = df[server_col].astype(str) + " • " + df[mount_col].astype(str) | |
| else: | |
| df["Label"] = df[server_col].astype(str) + " • " + df[fs_col].astype(str) | |
| # Totals & pct | |
| df["Total_kb"] = df[used_col] + df[avail_col] | |
| df["Used_pct"] = (df[used_col] / df["Total_kb"]) * 100.0 | |
| df["Used_pct"] = df["Used_pct"].clip(0, 100) | |
| df["Avail_pct"] = (100.0 - df["Used_pct"]).clip(0, 100) | |
| # Sizes | |
| df["Used_TB"] = _kblocks_to_tib(df[used_col]) | |
| df["Total_TB"] = _kblocks_to_tib(df["Total_kb"]) | |
| df["Avail_GB"] = _kblocks_to_gib(df[avail_col]) # <-- GB for hovers | |
| # Alerts | |
| df["ALERT"] = df["Used_pct"] > alert_threshold_pct | |
| # Sort | |
| # df = df.sort_values("Total_kb", ascending=False).reset_index(drop=True) | |
| df = df.sort_values("Mounted", ascending=False).reset_index(drop=True) | |
| y_labels = np.where(df["ALERT"].to_numpy(), "⚠ " + df["Label"], df["Label"]) | |
| used_text = [f"{u:.2f} TB ({p:.0f}%)" for u, p in zip(df["Used_TB"], df["Used_pct"])] | |
| total_annot = [f"{t:.2f} TB" for t in df["Total_TB"]] | |
| avail_gb_0 = [f"{g:.0f} GB" for g in df["Avail_GB"]] | |
| # Colors | |
| COLOR_TOTAL = "#CBD5E1" | |
| COLOR_USED = "#2563EB" | |
| COLOR_FREE = "#94A3B8" | |
| COLOR_ALERT = "#F59E0B" | |
| used_colors = np.where(df["ALERT"].to_numpy(), COLOR_ALERT, COLOR_USED) | |
| fig = go.Figure() | |
| # Gray background hover: Available in GB (0dp) | |
| fig.add_trace( | |
| go.Bar( | |
| y=y_labels, | |
| x=[100] * len(df), | |
| base=0, | |
| orientation="h", | |
| marker=dict(color=COLOR_TOTAL), | |
| opacity=0.40, | |
| hovertemplate="<b>%{y}</b><br>Available: %{customdata}<br><extra></extra>", | |
| customdata=avail_gb_0, | |
| showlegend=False, | |
| ) | |
| ) | |
| # Used hover: Available in GB (0dp) too | |
| fig.add_trace( | |
| go.Bar( | |
| y=y_labels, | |
| x=df["Used_pct"], | |
| base=0, | |
| name=f"Used (>{alert_threshold_pct:.0f}% highlighted)", | |
| orientation="h", | |
| marker=dict(color=used_colors), | |
| text=used_text, | |
| textposition="inside", | |
| insidetextanchor="middle", | |
| hovertemplate=( | |
| "<b>%{y}</b><br>" | |
| "Used: %{customdata[0]} (%{customdata[3]:.2f}%)<br>" | |
| "Available: %{customdata[1]}<br>" | |
| "Total: %{customdata[2]}<br>" | |
| "<extra></extra>" | |
| ), | |
| customdata=np.stack( | |
| [ | |
| df["Used_TB"].map(lambda v: f"{v:.2f} TB").to_numpy(), | |
| df["Avail_GB"].map(lambda v: f"{v:.0f} GB").to_numpy(), # <-- changed | |
| df["Total_TB"].map(lambda v: f"{v:.2f} TB").to_numpy(), | |
| df["Used_pct"].to_numpy(), | |
| ], | |
| axis=1, | |
| ), | |
| ) | |
| ) | |
| # Available hover: Available in GB (0dp) | |
| fig.add_trace( | |
| go.Bar( | |
| y=y_labels, | |
| x=df["Avail_pct"], | |
| base=df["Used_pct"], | |
| name="Available", | |
| orientation="h", | |
| marker=dict(color=COLOR_FREE), | |
| hovertemplate=( | |
| "<b>%{y}</b><br>" | |
| "Available: %{customdata[0]}<br>" | |
| "Used: %{customdata[1]}<br>" | |
| "Total: %{customdata[2]}<br>" | |
| "<extra></extra>" | |
| ), | |
| customdata=np.stack( | |
| [ | |
| df["Avail_GB"].map(lambda v: f"{v:.0f} GB").to_numpy(), # <-- changed | |
| df["Used_TB"].map(lambda v: f"{v:.2f} TB").to_numpy(), | |
| df["Total_TB"].map(lambda v: f"{v:.2f} TB").to_numpy(), | |
| ], | |
| axis=1, | |
| ), | |
| ) | |
| ) | |
| # Total annotation (TB, 2dp) | |
| for y, ttxt, is_alert in zip(y_labels, total_annot, df["ALERT"].to_numpy()): | |
| fig.add_annotation( | |
| x=100, | |
| y=y, | |
| text=ttxt, | |
| showarrow=False, | |
| xanchor="left", | |
| yanchor="middle", | |
| xshift=6, | |
| font=dict(color=("#B45309" if is_alert else "#334155")), | |
| ) | |
| fig.update_layout( | |
| barmode="overlay", | |
| template="plotly_white", | |
| title=f"Disk usage (alerts: Used > {alert_threshold_pct:.0f}%)", | |
| xaxis=dict(range=[0, 100], ticksuffix="%", title="Percent of total"), | |
| yaxis_title="", | |
| height=max(420, 28 * len(df)), | |
| margin=dict(l=280, r=120, t=60, b=40), | |
| legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1), | |
| barcornerradius=4, | |
| ) | |
| fig.update_yaxes(autorange="reversed") | |
| 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 | |
| demo = gr.Interface( | |
| fn=plot_figs, | |
| inputs = [ | |
| ], | |
| outputs = [ | |
| gr.Plot(label="GPU Status", elem_classes="plotcss"), | |
| gr.Plot(label="Disk Status", elem_classes="plotcss"), | |
| gr.Dataframe(label="GPU Status Details"), | |
| gr.Dataframe(label="Disk Status Details"), | |
| ], | |
| live=True, | |
| flagging_options=[], | |
| css=".plotcss {max-width: 820px !important;}", | |
| delete_cache=(1, 1) | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(debug=False) | |