Update app.py

app.py

@@ -3,9 +3,11 @@ import gradio as gr
 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(", "))
@@ -21,9 +23,9 @@ def split_multi_users(dfs):
     df = pd.DataFrame(new_df)
     return df
 
-def plot_now():
+def plot_gpus():
     ### Load Data
-    dfs = datasets.load_dataset("pluslab/PLUS_Lab_GPUs_Data", download_mode='force_redownload')["train"].to_pandas()
+    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)
@@ -43,35 +45,194 @@ def plot_now():
     # print(dfs_plot)
     return fig, dfs
 
-def plot_history(sample=True, sampling_interval_minutes=180):
-    dfh = pd.read_pickle("hf://spaces/pluslab/PLUS_Lab_GPUs/history.pkl.gz", )
-    dfh = dfh.fillna("FREE")
-    dfh = split_multi_users(dfh)
-    dfh = dfh[["polling_timestamp", "username", "count"]]
-    dfh = dfh.groupby(["polling_timestamp", "username"]).sum()
-    dfh = dfh.reset_index()
-    dfh = dfh.sort_values(by=["polling_timestamp", "count"], ascending=False)
-    if sample:
-        unique_timestamps = dfh["polling_timestamp"].unique()
-        sampled_timestamps = [unique_timestamps[0]]
-        for i, t in enumerate(unique_timestamps[1:]):
-            diff = sampled_timestamps[-1] - t
-            if diff > datetime.timedelta(minutes=sampling_interval_minutes):
-                sampled_timestamps.append(t)
-        dfh = dfh[dfh["polling_timestamp"].isin(sampled_timestamps)]
-    fig = px.area(dfh, x="polling_timestamp", y="count", color='username', color_discrete_map={"FREE": "black",}, markers=True, line_shape='spline',)
-    return fig, dfh
+##### 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):
+    # KiB blocks -> TiB (so 104149210112 -> ~97.0)
+    return kblocks / (1024**3)
 
+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"])
+
+    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 (compute ourselves)
+    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 in TiB (shown as "TB")
+    df["Used_TB"]  = _kblocks_to_tib(df[used_col])
+    df["Avail_TB"] = _kblocks_to_tib(df[avail_col])
+    df["Total_TB"] = _kblocks_to_tib(df["Total_kb"])
+
+    # Alert rows
+    df["ALERT"] = df["Used_pct"] > alert_threshold_pct
+
+    # Sort by total desc
+    df = df.sort_values("Total_kb", ascending=False).reset_index(drop=True)
+
+    # Display text
+    used_text  = [f"{u:.1f} TB ({p:.0f}%)" for u, p in zip(df["Used_TB"], df["Used_pct"])]
+    total_text = [f"{t:.1f} TB" for t in df["Total_TB"]]
+    avail_text = [f"{a:.1f} TB" for a in df["Avail_TB"]]
+
+    # Pro palette + alert accent
+    COLOR_TOTAL = "#CBD5E1"  # slate-300
+    COLOR_USED  = "#2563EB"  # blue-600
+    COLOR_FREE  = "#94A3B8"  # slate-400
+    COLOR_ALERT = "#F59E0B"  # amber-500 (dashboard alert)
+    COLOR_OKTXT = "#0F172A"  # slate-900
+    COLOR_ALTXT = "#B45309"  # amber-700
+
+    # Used color per row (highlight alerts)
+    used_colors = np.where(df["ALERT"].to_numpy(), COLOR_ALERT, COLOR_USED)
+
+    # Add an icon to the y label for alerts
+    y_labels = np.where(df["ALERT"].to_numpy(), "⚠ " + df["Label"], df["Label"])
+
+    fig = go.Figure()
+
+    # Gray background (hover shows AVAILABLE)
+    fig.add_trace(
+        go.Bar(
+            y=y_labels,
+            x=[100] * len(df),
+            base=0,
+            name="(hover) Available",
+            orientation="h",
+            marker=dict(color=COLOR_TOTAL),
+            opacity=0.40,
+            hovertemplate="<b>%{y}</b><br>Available: %{customdata}<br><extra></extra>",
+            customdata=avail_text,
+            showlegend=False,
+        )
+    )
+
+    # Used (colored per-row; alert if >99%)
+    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"].to_numpy(),
+                    df["Avail_TB"].to_numpy(),
+                    df["Total_TB"].to_numpy(),
+                    df["Used_pct"].to_numpy(),
+                ],
+                axis=1,
+            ),
+        )
+    )
+
+    # Available
+    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_TB"].map(lambda v: f"{v:.2f} TB").to_numpy(),
+                    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 at far right (color it if alert)
+    for lbl, ttxt, is_alert in zip(y_labels, total_text, df["ALERT"].to_numpy()):
+        fig.add_annotation(
+            x=100,
+            y=lbl,
+            text=ttxt,
+            showarrow=False,
+            xanchor="left",
+            yanchor="middle",
+            xshift=6,
+            font=dict(color=(COLOR_ALTXT 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),
+    )
+    fig.update_yaxes(autorange="reversed")
+    return fig, df
+
+
+
+##### PLOT ALL #####
 def plot_figs():
-    fig_now, dfn = plot_now()
-    try:
-        fig_history, dfh = plot_history()
-    except Exception as e:
-        print(e)
-        fig_history = None
-        dfh = None
-    return fig_now, dfn, fig_history
+    fig_gpus, dfn = plot_gpus()
+    fig_disks, dfh = plot_disks()
+    return fig_gpus, dfn, fig_disks
 
 demo = gr.Interface(
     fn=plot_figs,
@@ -81,7 +242,7 @@ demo = gr.Interface(
     outputs = [
         gr.Plot(label="GPU Status", elem_classes="plotcss"),
         gr.Dataframe(label="GPU Status Details"),
-        gr.Plot(label="History", elem_classes="plotcss"),
+        gr.Plot(label="Disk Status", elem_classes="plotcss"),
     ],
     live=True,
     flagging_options=[],