Sortable/Groupable Memory Dashboard (#9014)

python/ray/dashboard/client/src/api.ts

@@ -362,13 +362,17 @@ export type MemoryTableEntry = {
   call_site: string;
 };
 
+export type MemoryTableGroups = {
+  [groupKey: string]: MemoryTableGroup;
+};
+
+export type MemoryTableGroup = {
+  entries: MemoryTableEntry[];
+  summary: MemoryTableSummary;
+}
+
 export type MemoryTableResponse = {
-  group: {
-    [groupKey: string]: {
-      entries: MemoryTableEntry[];
-      summary: MemoryTableSummary;
-    };
-  };
+  group: MemoryTableGroups;
   summary: MemoryTableSummary;
 };
 

python/ray/dashboard/client/src/common/SortableTableHead.tsx

@@ -0,0 +1,78 @@
+import {
+  createStyles,
+  makeStyles,
+  TableHead,
+  TableRow,
+  TableSortLabel,
+  Theme,
+} from "@material-ui/core";
+import React from "react";
+import { StyledTableCell } from "./TableCell";
+import { Order } from "./tableUtils";
+
+const useSortableTableHeadStyles = makeStyles((theme: Theme) =>
+  createStyles({
+    visuallyHidden: {
+      border: 0,
+      clip: "rect(0 0 0 0)",
+      height: 1,
+      margin: -1,
+      overflow: "hidden",
+      padding: 0,
+      position: "absolute",
+      top: 20,
+      width: 1,
+    },
+  }),
+);
+
+export type HeaderInfo<T> = {
+  id: keyof T;
+  label: string;
+  numeric: boolean;
+};
+
+type SortableTableHeadProps<T> = {
+  onRequestSort: (event: React.MouseEvent<unknown>, property: keyof T) => void;
+  order: Order;
+  orderBy: string | null;
+  headerInfo: HeaderInfo<T>[];
+};
+
+const SortableTableHead = <T,>(props: SortableTableHeadProps<T>) => {
+  const { order, orderBy, onRequestSort, headerInfo } = props;
+  const classes = useSortableTableHeadStyles();
+  const createSortHandler = (property: keyof T) => (
+    event: React.MouseEvent<unknown>,
+  ) => {
+    onRequestSort(event, property);
+  };
+  return (
+    <TableHead>
+      <TableRow>
+        {headerInfo.map((headerInfo) => (
+          <StyledTableCell
+            key={headerInfo.label}
+            align={headerInfo.numeric ? "right" : "left"}
+            sortDirection={orderBy === headerInfo.id ? order : false}
+          >
+            <TableSortLabel
+              active={orderBy === headerInfo.id}
+              direction={orderBy === headerInfo.id ? order : "asc"}
+              onClick={createSortHandler(headerInfo.id)}
+            >
+              {headerInfo.label}
+              {orderBy === headerInfo.id ? (
+                <span className={classes.visuallyHidden}>
+                  {order === "desc" ? "sorted descending" : "sorted ascending"}
+                </span>
+              ) : null}
+            </TableSortLabel>
+          </StyledTableCell>
+        ))}
+      </TableRow>
+    </TableHead>
+  );
+};
+
+export default SortableTableHead;

python/ray/dashboard/client/src/common/TableCell.tsx

@@ -0,0 +1,13 @@
+import { TableCell } from "@material-ui/core";
+import { styled } from "@material-ui/core/styles";
+
+export const StyledTableCell = styled(TableCell)(({ theme }) => ({
+  padding: theme.spacing(1),
+  textAlign: "center",
+}));
+
+export const ExpandableStyledTableCell = styled(TableCell)(({ theme }) => ({
+  padding: theme.spacing(1),
+  textAlign: "center",
+  cursor: "pointer",
+}));

python/ray/dashboard/client/src/common/tableUtils.ts

@@ -0,0 +1,38 @@
+export const descendingComparator = <T>(a: T, b: T, orderBy: keyof T) => {
+  if (b[orderBy] < a[orderBy]) {
+    return -1;
+  }
+  if (b[orderBy] > a[orderBy]) {
+    return 1;
+  }
+  return 0;
+};
+
+export type Order = "asc" | "desc";
+
+export const getComparator = <Key extends keyof any>(
+  order: Order,
+  orderBy: Key,
+): ((
+  a: { [key in Key]: number | string },
+  b: { [key in Key]: number | string },
+) => number) => {
+  return order === "desc"
+    ? (a, b) => descendingComparator(a, b, orderBy)
+    : (a, b) => -descendingComparator(a, b, orderBy);
+};
+
+export const stableSort = <T>(
+  array: T[],
+  comparator: (a: T, b: T) => number,
+) => {
+  const stabilizedThis = array.map((el, index) => [el, index] as [T, number]);
+  stabilizedThis.sort((a, b) => {
+    const order = comparator(a[0], b[0]);
+    if (order !== 0) {
+      return order;
+    }
+    return a[1] - b[1];
+  });
+  return stabilizedThis.map((el) => el[0]);
+};

python/ray/dashboard/client/src/pages/dashboard/memory/Memory.tsx

@@ -1,25 +1,81 @@
 import {
   Button,
+  Checkbox,
   createStyles,
+  FormControlLabel,
+  makeStyles,
   Table,
   TableBody,
-  TableCell,
-  TableHead,
-  TableRow,
   Theme,
-  WithStyles,
-  withStyles,
 } from "@material-ui/core";
 import PauseIcon from "@material-ui/icons/Pause";
 import PlayArrowIcon from "@material-ui/icons/PlayArrow";
-import React from "react";
-import { connect } from "react-redux";
-import { stopMemoryTableCollection } from "../../../api";
+import React, { useState } from "react";
+import { useDispatch, useSelector } from "react-redux";
+import {
+  MemoryTableEntry,
+  MemoryTableGroups,
+  stopMemoryTableCollection,
+} from "../../../api";
+import SortableTableHead, {
+  HeaderInfo,
+} from "../../../common/SortableTableHead";
+import { getComparator, Order, stableSort } from "../../../common/tableUtils";
 import { StoreState } from "../../../store";
-import { dashboardActions } from "../state";
 import MemoryRowGroup from "./MemoryRowGroup";
+import { MemoryTableRow } from "./MemoryTableRow";
 
-const styles = (theme: Theme) =>
+const makeGroupedEntries = (memoryTableGroups: MemoryTableGroups, order: Order, orderBy: keyof MemoryTableEntry | null) => {
+  const comparator = orderBy && getComparator(order, orderBy);
+  return Object.entries(memoryTableGroups).map(([groupKey, group]) => {
+    const sortedEntries = comparator
+      ? stableSort(group.entries, comparator)
+      : group.entries;
+    
+    return <MemoryRowGroup
+      groupKey={groupKey}
+      summary={group.summary}
+      entries={sortedEntries}
+      initialExpanded={true}
+    />
+  });
+};
+
+const makeUngroupedEntries = (
+  memoryTableGroups: MemoryTableGroups,
+  order: Order,
+  orderBy: keyof MemoryTableEntry | null,
+) => {
+  const allEntries = Object.values(memoryTableGroups).reduce(
+    (allEntries: Array<MemoryTableEntry>, memoryTableGroup) => {
+      const groupEntries = memoryTableGroup.entries;
+      return allEntries.concat(groupEntries);
+    },
+    [],
+  );
+  const sortedEntries =
+    orderBy === null
+      ? allEntries
+      : stableSort(allEntries, getComparator(order, orderBy));
+  return sortedEntries.map((memoryTableEntry, index) => (
+    <MemoryTableRow
+      memoryTableEntry={memoryTableEntry}
+      key={`mem-row-${index}`}
+    />
+  ));
+};
+
+const memoryHeaderInfo: HeaderInfo<MemoryTableEntry>[] = [
+  { id: "node_ip_address", label: "IP Address", numeric: true },
+  { id: "pid", label: "pid", numeric: true },
+  { id: "type", label: "Type", numeric: false },
+  { id: "object_id", label: "Object ID", numeric: false },
+  { id: "object_size", label: "Object Size (B)", numeric: true },
+  { id: "reference_type", label: "Reference Type", numeric: false },
+  { id: "call_site", label: "Call Site", numeric: false },
+];
+
+const useMemoryInfoStyles = makeStyles((theme: Theme) =>
   createStyles({
     table: {
       marginTop: theme.spacing(1),
@@ -28,96 +84,83 @@ const styles = (theme: Theme) =>
       padding: theme.spacing(1),
       textAlign: "center",
     },
-  });
+  }),
+);
 
-const mapStateToProps = (state: StoreState) => ({
+const memoryInfoSelector = (state: StoreState) => ({
   tab: state.dashboard.tab,
   memoryTable: state.dashboard.memoryTable,
   shouldObtainMemoryTable: state.dashboard.shouldObtainMemoryTable,
 });
 
-const mapDispatchToProps = dashboardActions;
-
-type State = {
-  // If memory table is captured, it should stop renewing memory table.
-  pauseMemoryTable: boolean;
-};
-
-class MemoryInfo extends React.Component<
-  WithStyles<typeof styles> &
-    ReturnType<typeof mapStateToProps> &
-    typeof mapDispatchToProps,
-  State
-> {
-  handlePauseMemoryTable = async () => {
-    const { shouldObtainMemoryTable } = this.props;
-    this.props.setShouldObtainMemoryTable(!shouldObtainMemoryTable);
+const MemoryInfo: React.FC<{}> = () => {
+  const { memoryTable, shouldObtainMemoryTable } = useSelector(
+    memoryInfoSelector,
+  );
+  const { setShouldObtainMemoryTable } = useDispatch();
+  const toggleMemoryCollection = async () => {
+    setShouldObtainMemoryTable(!shouldObtainMemoryTable);
     if (shouldObtainMemoryTable) {
       await stopMemoryTableCollection();
     }
   };
 
-  renderIcon = () => {
-    if (this.props.shouldObtainMemoryTable) {
-      return <PauseIcon />;
-    } else {
-      return <PlayArrowIcon />;
-    }
-  };
+  const pauseButtonIcon = shouldObtainMemoryTable ? (
+    <PauseIcon />
+  ) : (
+    <PlayArrowIcon />
+  );
+  const classes = useMemoryInfoStyles();
+  const [isGrouped, setIsGrouped] = useState(true);
+  const [order, setOrder] = React.useState<Order>("asc");
+  const toggleOrder = () => setOrder(order === "asc" ? "desc" : "asc");
+  const [orderBy, setOrderBy] = React.useState<keyof MemoryTableEntry | null>(
+    null,
+  );
+  return (
+    <React.Fragment>
+      {memoryTable !== null ? (
+        <React.Fragment>
+          <Button color="primary" onClick={toggleMemoryCollection}>
+            {pauseButtonIcon}
+            {shouldObtainMemoryTable ? "Pause Collection" : "Resume Collection"}
+          </Button>
+          <FormControlLabel
+            control={
+              <Checkbox
+                checked={isGrouped}
+                onChange={() => setIsGrouped(!isGrouped)}
+                color="primary"
+              />
+            }
+            label="Group by host"
+          />
+          <Table className={classes.table}>
+            <SortableTableHead
+              orderBy={orderBy || ""}
+              order={order}
+              onRequestSort={(event, property) => {
+                if (property === orderBy) {
+                  toggleOrder();
+                } else {
+                  setOrderBy(property);
+                  setOrder("asc");
+                }
+              }}
+              headerInfo={memoryHeaderInfo}
+            />
+            <TableBody>
+              {isGrouped
+                ? makeGroupedEntries(memoryTable.group, order, orderBy)
+                : makeUngroupedEntries(memoryTable.group, order, orderBy)}
+            </TableBody>
+          </Table>
+        </React.Fragment>
+      ) : (
+        <div>No Memory Table Information Provided</div>
+      )}
+    </React.Fragment>
+  );
+};
 
-  render() {
-    const { classes, memoryTable } = this.props;
-    const memoryTableHeaders = [
-      "", // Padding
-      "IP Address",
-      "Pid",
-      "Type",
-      "Object ID",
-      "Object Size",
-      "Reference Type",
-      "Call Site",
-    ];
-    return (
-      <React.Fragment>
-        {memoryTable !== null ? (
-          <React.Fragment>
-            <Button color="primary" onClick={this.handlePauseMemoryTable}>
-              {this.renderIcon()}
-              {this.props.shouldObtainMemoryTable
-                ? "Pause Collection"
-                : "Resume Collection"}
-            </Button>
-            <Table className={classes.table}>
-              <TableHead>
-                <TableRow>
-                  {memoryTableHeaders.map((header, index) => (
-                    <TableCell key={index} className={classes.cell}>
-                      {header}
-                    </TableCell>
-                  ))}
-                </TableRow>
-              </TableHead>
-              <TableBody>
-                {Object.keys(memoryTable.group).map((group_key, index) => (
-                  <MemoryRowGroup
-                    key={index}
-                    groupKey={group_key}
-                    memoryTableGroups={memoryTable.group}
-                    initialExpanded={true}
-                  />
-                ))}
-              </TableBody>
-            </Table>
-          </React.Fragment>
-        ) : (
-          <div>No Memory Table Information Provided</div>
-        )}
-      </React.Fragment>
-    );
-  }
-}
-
-export default connect(
-  mapStateToProps,
-  mapDispatchToProps,
-)(withStyles(styles)(MemoryInfo));
+export default MemoryInfo;

python/ray/dashboard/client/src/pages/dashboard/memory/MemoryRowGroup.tsx

@@ -1,28 +1,25 @@
 import {
   createStyles,
-  TableCell,
   TableRow,
   Theme,
-  withStyles,
-  WithStyles,
+  makeStyles
 } from "@material-ui/core";
 import AddIcon from "@material-ui/icons/Add";
 import RemoveIcon from "@material-ui/icons/Remove";
-import classNames from "classnames";
-import React from "react";
+import React, { useState } from "react";
 import {
   MemoryTableEntry,
-  MemoryTableResponse,
   MemoryTableSummary,
 } from "../../../api";
+import {
+  ExpandableStyledTableCell,
+  StyledTableCell,
+} from "../../../common/TableCell";
 import MemorySummary from "./MemorySummary";
+import { MemoryTableRow } from "./MemoryTableRow";
 
-const styles = (theme: Theme) =>
+const useMemoryRowGroupStyles = makeStyles((theme: Theme) =>
   createStyles({
-    cell: {
-      padding: theme.spacing(1),
-      textAlign: "center",
-    },
     expandCollapseCell: {
       cursor: "pointer",
     },
@@ -35,106 +32,62 @@ const styles = (theme: Theme) =>
       fontFamily: "SFMono-Regular,Consolas,Liberation Mono,Menlo,monospace",
       whiteSpace: "pre",
     },
-  });
+  }));
 
-type Props = {
+type MemoryRowGroupProps = {
   groupKey: string;
-  memoryTableGroups: MemoryTableResponse["group"];
+  summary: MemoryTableSummary;
+  entries: MemoryTableEntry[];
   initialExpanded: boolean;
 };
 
-type State = {
-  expanded: boolean;
+const MemoryRowGroup: React.FC<MemoryRowGroupProps> = ({ groupKey, entries, summary, initialExpanded}) => {
+  const classes = useMemoryRowGroupStyles();
+  const [expanded, setExpanded] = useState(initialExpanded);
+  const toggleExpanded = () => setExpanded(!expanded);
+
+  const features = [
+    "node_ip_address",
+    "pid",
+    "type",
+    "object_id",
+    "object_size",
+    "reference_type",
+    "call_site",
+  ];
+
+  return (
+    <React.Fragment>
+      <TableRow hover>
+        <ExpandableStyledTableCell onClick={toggleExpanded}>
+          {!expanded ? (
+            <AddIcon className={classes.expandCollapseIcon} />
+          ) : (
+            <RemoveIcon className={classes.expandCollapseIcon} />
+          )}
+        </ExpandableStyledTableCell>
+        {features.map((feature, index) => (
+          <StyledTableCell key={index}>
+            {// TODO(sang): For now, it is always grouped by node_ip_address.
+            feature === "node_ip_address" ? groupKey : ""}
+          </StyledTableCell>
+        ))}
+      </TableRow>
+      {expanded && (
+        <React.Fragment>
+          <MemorySummary initialExpanded={false} memoryTableSummary={summary} />
+          {entries.map((memoryTableEntry, index) => {
+            return (
+              <MemoryTableRow
+                memoryTableEntry={memoryTableEntry}
+                key={`${index}`}
+              />
+            );
+          })}
+        </React.Fragment>
+      )}
+    </React.Fragment>
+  );
 };
 
-class MemoryRowGroup extends React.Component<
-  Props & WithStyles<typeof styles>,
-  State
-> {
-  state: State = {
-    expanded: this.props.initialExpanded,
-  };
-
-  toggleExpand = () => {
-    this.setState((state) => ({
-      expanded: !state.expanded,
-    }));
-  };
-
-  render() {
-    const { classes, groupKey, memoryTableGroups } = this.props;
-    const { expanded } = this.state;
-
-    const features = [
-      "node_ip_address",
-      "pid",
-      "type",
-      "object_id",
-      "object_size",
-      "reference_type",
-      "call_site",
-    ];
-
-    const memoryTableGroup = memoryTableGroups[groupKey];
-    const entries: Array<MemoryTableEntry> = memoryTableGroup["entries"];
-    const summary: MemoryTableSummary = memoryTableGroup["summary"];
-
-    return (
-      <React.Fragment>
-        <TableRow hover>
-          <TableCell
-            className={classNames(classes.cell, classes.expandCollapseCell)}
-            onClick={this.toggleExpand}
-          >
-            {!expanded ? (
-              <AddIcon className={classes.expandCollapseIcon} />
-            ) : (
-              <RemoveIcon className={classes.expandCollapseIcon} />
-            )}
-          </TableCell>
-          {features.map((feature, index) => (
-            <TableCell className={classes.cell} key={index}>
-              {// TODO(sang): For now, it is always grouped by node_ip_address.
-              feature === "node_ip_address" ? groupKey : ""}
-            </TableCell>
-          ))}
-        </TableRow>
-        {expanded && (
-          <React.Fragment>
-            <MemorySummary
-              initialExpanded={false}
-              memoryTableSummary={summary}
-            />
-            {entries.map((memoryTableEntry, index) => {
-              const object_size =
-                memoryTableEntry.object_size === -1
-                  ? "?"
-                  : `${memoryTableEntry.object_size}  B`;
-              const memoryTableEntryValues = [
-                "", // Padding
-                memoryTableEntry.node_ip_address,
-                memoryTableEntry.pid,
-                memoryTableEntry.type,
-                memoryTableEntry.object_id,
-                object_size,
-                memoryTableEntry.reference_type,
-                memoryTableEntry.call_site,
-              ];
-              return (
-                <TableRow hover key={index}>
-                  {memoryTableEntryValues.map((value, index) => (
-                    <TableCell key={index} className={classes.cell}>
-                      {value}
-                    </TableCell>
-                  ))}
-                </TableRow>
-              );
-            })}
-          </React.Fragment>
-        )}
-      </React.Fragment>
-    );
-  }
-}
-
-export default withStyles(styles)(MemoryRowGroup);
+export default MemoryRowGroup;

python/ray/dashboard/client/src/pages/dashboard/memory/MemoryTableRow.tsx

@@ -0,0 +1,33 @@
+import { TableRow } from "@material-ui/core";
+import React from "react";
+import { MemoryTableEntry } from "../../../api";
+import { StyledTableCell } from "../../../common/TableCell";
+
+type Props = {
+  memoryTableEntry: MemoryTableEntry;
+  key: string;
+};
+
+export const MemoryTableRow = (props: Props) => {
+  const { memoryTableEntry, key } = props;
+  const object_size =
+    memoryTableEntry.object_size === -1
+      ? "?"
+      : `${memoryTableEntry.object_size}  B`;
+  const memoryTableEntryValues = [
+    memoryTableEntry.node_ip_address,
+    memoryTableEntry.pid,
+    memoryTableEntry.type,
+    memoryTableEntry.object_id,
+    object_size,
+    memoryTableEntry.reference_type,
+    memoryTableEntry.call_site,
+  ];
+  return (
+    <TableRow hover key={key}>
+      {memoryTableEntryValues.map((value, index) => (
+        <StyledTableCell key={`${key}-${index}`}>{value}</StyledTableCell>
+      ))}
+    </TableRow>
+  );
+};

python/ray/dashboard/client/src/pages/dashboard/node-info/NodeInfo.tsx

@@ -15,6 +15,7 @@ import { connect } from "react-redux";
 import { RayletInfoResponse } from "../../../api";
 import { sum } from "../../../common/util";
 import { StoreState } from "../../../store";
+import {} from "../../../common/tableUtils";
 import Errors from "./dialogs/errors/Errors";
 import Logs from "./dialogs/logs/Logs";
 import NodeRowGroup from "./NodeRowGroup";