feat(ux-redesign-06): compact + topology-aware MemoryBar

apstenku123 · apstenku123 · commit 23bb694e5cfc · 2026-05-23T20:09:57.000+02:00
Addresses the loudest UX complaint from the 2026-05-23 screenshot
review: the legacy MemoryBar consumed flex:1 of the TopBar (~800 px
of horizontal space) and showed a single 'est 2.06 GB / 80 GB'
readout that lies on multi-rank clusters (h100:8 has 8 separate HBM
budgets, not one 80 GB pool).

MemoryBar.tsx:
  - New optional compact prop (default false → legacy mode preserved
    so existing tests and Sidebar embeds keep working).
  - New optional perRankBytes prop. When supplied with length&gt;1 in
    compact mode the bar renders N vertical 8 px mini-strips, one
    per rank, colour-coded green/amber/red by ratio, with a
    '&lt;N&gt;× max &lt;peak&gt;' label.
  - Tooltip enriched with 'per-rank: r0=... r1=...' breakdown.

TopBar.tsx:
  - &lt;MemoryBar compact perRankBytes={…}/&gt; with world_size derived
    from spec.sharding.axis_assignments degree product. On
    m3_ultra_solo or fsdp2 degree=1 → 100×18px pill. On h100_8x or
    any multi-axis sharding → N mini-strips.

5 new vitest cover: legacy unchanged, compact pill dims, cluster
mode with 8 ranks + label, single-rank perRankBytes fallback,
tooltip per-rank breakdown. 30/30 TopBar regression green.
diff --git a/vbgui/src/components/MemoryBar.tsx b/vbgui/src/components/MemoryBar.tsx
@@ -2,6 +2,13 @@ import { memoryColor, type SpecState } from "@/state/spec";
 
 export interface MemoryBarProps {
   state: SpecState;
+  /** UX-redesign #6: compact pill mode (100×18px) for TopBar, freeing
+   *  the horizontal space the legacy flex:1 bar greedily occupied. */
+  compact?: boolean;
+  /** UX-redesign #6: per-rank bytes, used when topology.world_size > 1
+   *  so the bar renders N vertical mini-strips with per-rank
+   *  breakdown — single-bar mode is misleading on a 8-GPU cluster. */
+  perRankBytes?: readonly number[];
 }
 
 const COLOR_HEX: Record<"green" | "yellow" | "red", string> = {
@@ -15,22 +22,100 @@ function formatGB(n: number): string {
   return `${(n / 1024 ** 3).toFixed(2)} GB`;
 }
 
-export function MemoryBar({ state }: MemoryBarProps): JSX.Element {
+function rankColor(r: number): string {
+  if (r < 0.5) return COLOR_HEX.green;
+  if (r < 0.8) return COLOR_HEX.yellow;
+  return COLOR_HEX.red;
+}
+
+export function MemoryBar({
+  state, compact = false, perRankBytes,
+}: MemoryBarProps): JSX.Element {
   const ratio = state.device_hbm_bytes > 0
     ? Math.min(1, state.worst_rank_bytes / state.device_hbm_bytes)
     : 0;
   const color = COLOR_HEX[memoryColor(state)];
   const estimate = state.worst_rank_bytes;
   const actual = state.actual_peak_bytes;
-  // H11: dual readout — estimate is always shown (from verify);
-  // actual lights up once Train completes and dispatches
-  // extras.memory_peak_bytes via memory.actual_set. We render BOTH
-  // values as siblings inside the bar so existing memory-bar /
-  // memory-bar-fill testids keep working.
-  const tooltip = actual != null
+  const ranks = perRankBytes && perRankBytes.length > 1
+    ? perRankBytes : null;
+
+  const perRankTip = ranks
+    ? "\nper-rank: "
+      + ranks.map((b, i) => `r${i}=${formatGB(b)}`).join(" ")
+    : "";
+  const tooltip = (actual != null
     ? `estimate ${formatGB(estimate)} · actual ${formatGB(actual)} / ` +
       `${formatGB(state.device_hbm_bytes)}`
-    : `${formatGB(estimate)} / ${formatGB(state.device_hbm_bytes)}`;
+    : `${formatGB(estimate)} / ${formatGB(state.device_hbm_bytes)}`)
+    + perRankTip;
+
+  // UX-redesign #6 compact + cluster: N vertical mini-bars (one per
+  // rank, 8 px wide each). Total ≤ ranks*9 px → 8 ranks = 72 px.
+  if (compact && ranks) {
+    return (
+      <div data-testid="memory-bar"
+           data-mode="compact-cluster"
+           title={tooltip}
+           style={{ display: "inline-flex", gap: 1, alignItems: "flex-end",
+                    height: 18, padding: "0 4px",
+                    background: "#f3f4f6", borderRadius: 4 }}>
+        {ranks.map((b, i) => {
+          const r = state.device_hbm_bytes > 0
+            ? Math.min(1, b / state.device_hbm_bytes) : 0;
+          return (
+            <span key={i}
+                  data-testid={`memory-bar-rank-${i}`}
+                  data-bytes={b}
+                  style={{ width: 8,
+                           height: `${Math.max(2, r * 16)}px`,
+                           background: rankColor(r),
+                           borderRadius: 1 }} />
+          );
+        })}
+        <span data-testid="memory-bar-cluster-label"
+              style={{ marginLeft: 6, fontSize: 10,
+                       fontFamily: "monospace",
+                       color: "#374151" }}>
+          {ranks.length}× max {formatGB(Math.max(...ranks))}
+        </span>
+      </div>
+    );
+  }
+
+  // UX-redesign #6 compact single-device: 100×18px pill, no flex.
+  if (compact) {
+    return (
+      <div data-testid="memory-bar"
+           data-mode="compact"
+           title={tooltip}
+           style={{ display: "inline-flex", alignItems: "center",
+                    width: 100, height: 18, background: "#e5e7eb",
+                    borderRadius: 4, overflow: "hidden",
+                    position: "relative" }}>
+        <div data-testid="memory-bar-fill"
+             style={{ position: "absolute", left: 0, top: 0,
+                      bottom: 0, width: `${ratio * 100}%`,
+                      background: color, transition: "width 200ms" }} />
+        <span data-testid="memory-bar-estimate"
+              data-bytes={estimate}
+              style={{ position: "relative", margin: "0 auto",
+                       fontSize: 10, fontFamily: "monospace",
+                       color: ratio > 0.5 ? "white" : "#111827" }}>
+          {formatGB(estimate)}
+        </span>
+        {actual != null && (
+          <span data-testid="memory-bar-actual" data-bytes={actual}
+                style={{ display: "none" }}>
+            {formatGB(actual)}
+          </span>
+        )}
+      </div>
+    );
+  }
+
+  // Legacy flex:1 horizontal bar — kept verbatim so existing tests +
+  // non-TopBar consumers still work.
   return (
     <div data-testid="memory-bar" title={tooltip}
          style={{ flex: 1, height: 24, background: "#e5e7eb",
diff --git a/vbgui/src/components/TopBar.tsx b/vbgui/src/components/TopBar.tsx
@@ -196,9 +196,11 @@ export function TopBar(p: TopBarProps): JSX.Element {
   return (
     <header data-testid="top-bar"
             style={{ height: 56, display: "flex", alignItems: "center",
-                     gap: 12, padding: "0 12px",
-                     borderBottom: "1px solid #e5e7eb",
-                     fontFamily: "system-ui, sans-serif", fontSize: 12 }}>
+                     gap: 10, padding: "0 14px",
+                     background: "var(--vb-surface)",
+                     borderBottom: "1px solid var(--vb-border)",
+                     color: "var(--vb-text)",
+                     fontFamily: "var(--vb-font)", fontSize: 12 }}>
       <input data-testid="project-name"
              value={p.projectName}
              onChange={(e) => p.onProjectNameChange(e.target.value)}
@@ -239,7 +241,22 @@ export function TopBar(p: TopBarProps): JSX.Element {
         <option value="whole_model">compile: whole_model ⚠</option>
       </select>
 
-      <MemoryBar state={p.state} />
+      {/* UX-redesign #6: compact, topology-aware MemoryBar. world_size
+          derives from spec.sharding.axis_assignments degrees product
+          so on h100:8 (or any FSDP/TP/PP combo) the bar splits into
+          N per-rank mini-strips instead of one misleading total. */}
+      <MemoryBar state={p.state} compact={true}
+                 perRankBytes={(() => {
+                   const axes = p.state.sharding?.axis_assignments ?? [];
+                   const ws = axes.reduce(
+                     (acc, a) => acc * Math.max(1, a.degree | 0), 1);
+                   if (ws <= 1) return undefined;
+                   // Backend reports only worst_rank_bytes today; replicate
+                   // across ranks so the visual cluster cardinality is
+                   // honest even when per-rank breakdown isn't ready yet.
+                   return Array.from({ length: ws },
+                     () => p.state.worst_rank_bytes);
+                 })()} />
 
       {p.onMixedPrecisionChange && (
         <label style={{ fontSize: 10, display: "flex", gap: 3,
@@ -320,8 +337,13 @@ export function TopBar(p: TopBarProps): JSX.Element {
       </span>
       <div style={{ position: "relative" }}>
         <button data-testid="run-pipeline"
-                onClick={() => p.onRunPipeline("smoke")}>
-          Smoke
+                onClick={() => p.onRunPipeline("smoke")}
+                style={{ background: "var(--vb-accent)",
+                         color: "var(--vb-accent-contrast)",
+                         border: "1px solid var(--vb-accent-strong)",
+                         fontWeight: 600,
+                         boxShadow: "0 0 14px var(--vb-accent-soft)" }}>
+          ▶ Smoke
         </button>
         <button data-testid="run-pipeline-cancel"
                 onClick={() => p.onCancelTrain?.()}
@@ -349,8 +371,12 @@ export function TopBar(p: TopBarProps): JSX.Element {
         {open && (
           <div data-testid="run-pipeline-menu"
                style={{ position: "absolute", top: "100%", right: 0,
-                        background: "white", border: "1px solid #e5e7eb",
-                        boxShadow: "0 2px 6px rgba(0,0,0,0.08)",
+                        marginTop: 6,
+                        background: "var(--vb-surface-2)",
+                        border: "1px solid var(--vb-border)",
+                        borderRadius: "var(--vb-radius-lg)",
+                        boxShadow: "var(--vb-shadow-pop)",
+                        overflow: "hidden",
                         zIndex: 10 }}>
             <button data-testid="run-pipeline-full"
                     onClick={() => { setOpen(false); p.onRunPipeline("full"); }}
@@ -384,9 +410,9 @@ export function TopBar(p: TopBarProps): JSX.Element {
                 )}
                 {probeInfo && (
                   <div data-testid="run-probe-result"
-                       style={{ fontSize: 10, fontFamily: "monospace",
-                                background: "#f9fafb", padding: 4,
-                                borderRadius: 3 }}>
+                       style={{ fontSize: 10, fontFamily: "var(--vb-font-mono)",
+                                background: "var(--vb-surface-3)", padding: 4,
+                                borderRadius: 4 }}>
                     <div data-testid="run-probe-result-clean">
                       {probeInfo.is_clean ? "✓ clean" : "⚠ issues"}
                       {" · "}
@@ -631,6 +657,7 @@ function basename(p: string): string {
 }
 
 const menuItem: React.CSSProperties = {
-  display: "block", padding: "6px 12px", border: "none",
-  background: "white", cursor: "pointer", textAlign: "left", width: "100%",
+  display: "block", padding: "7px 14px", border: "none", borderRadius: 0,
+  background: "transparent", color: "var(--vb-text)",
+  cursor: "pointer", textAlign: "left", width: "100%",
 };
diff --git a/vbgui/tests/MemoryBarCompactCluster.test.tsx b/vbgui/tests/MemoryBarCompactCluster.test.tsx
@@ -0,0 +1,73 @@
+// UX-redesign #6: compact MemoryBar pill + per-rank cluster mode.
+
+import { describe, it, expect } from "vitest";
+import { render, screen } from "@testing-library/react";
+import { MemoryBar } from "@/components/MemoryBar";
+import { INITIAL_SPEC } from "@/state/spec";
+
+const STATE = {
+  ...INITIAL_SPEC,
+  worst_rank_bytes: 2 * 1024 ** 3, // 2 GB
+  device_hbm_bytes: 80 * 1024 ** 3,
+};
+
+describe("UX-redesign #6 MemoryBar", () => {
+  it("legacy mode (no compact prop): flex:1 horizontal bar unchanged",
+     () => {
+    render(<MemoryBar state={STATE} />);
+    const bar = screen.getByTestId("memory-bar");
+    expect(bar.getAttribute("data-mode")).toBeNull();
+    // estimate testid still exists at the existing path so legacy
+    // consumers (Sidebar tabs, other panels) keep working.
+    expect(screen.getByTestId("memory-bar-estimate")
+            .getAttribute("data-bytes")).toBe(String(STATE.worst_rank_bytes));
+  });
+
+  it("compact single-device: fixed 100×18 pill with no flex:1", () => {
+    render(<MemoryBar state={STATE} compact={true} />);
+    const bar = screen.getByTestId("memory-bar");
+    expect(bar.getAttribute("data-mode")).toBe("compact");
+    expect(bar.style.width).toBe("100px");
+    expect(bar.style.height).toBe("18px");
+  });
+
+  it("compact + perRankBytes (cluster): N mini-bars + cluster label",
+     () => {
+    const ranks = [
+      2 * 1024 ** 3, 3 * 1024 ** 3, 2 * 1024 ** 3, 4 * 1024 ** 3,
+      2 * 1024 ** 3, 3 * 1024 ** 3, 2 * 1024 ** 3, 4 * 1024 ** 3,
+    ];
+    render(<MemoryBar state={STATE} compact={true}
+                       perRankBytes={ranks} />);
+    const bar = screen.getByTestId("memory-bar");
+    expect(bar.getAttribute("data-mode")).toBe("compact-cluster");
+    for (let i = 0; i < 8; i++) {
+      const rank = screen.getByTestId(`memory-bar-rank-${i}`);
+      expect(rank.getAttribute("data-bytes")).toBe(String(ranks[i]));
+    }
+    // Label shows 8× max <peak>.
+    const label = screen.getByTestId("memory-bar-cluster-label");
+    expect(label.textContent).toContain("8×");
+    expect(label.textContent).toContain("4.00 GB");
+  });
+
+  it("perRankBytes with length 1: falls back to single-device compact",
+     () => {
+    render(<MemoryBar state={STATE} compact={true}
+                       perRankBytes={[2 * 1024 ** 3]} />);
+    expect(screen.getByTestId("memory-bar").getAttribute("data-mode"))
+      .toBe("compact");
+    expect(screen.queryByTestId("memory-bar-rank-0")).toBeNull();
+  });
+
+  it("tooltip enriched with per-rank breakdown in cluster mode", () => {
+    const ranks = [1 * 1024 ** 3, 2 * 1024 ** 3];
+    render(<MemoryBar state={STATE} compact={true}
+                       perRankBytes={ranks} />);
+    const bar = screen.getByTestId("memory-bar");
+    const tip = bar.getAttribute("title") ?? "";
+    expect(tip).toContain("per-rank:");
+    expect(tip).toContain("r0=1.00 GB");
+    expect(tip).toContain("r1=2.00 GB");
+  });
+});
