[Dashboard] Memory View Group by Stack Trace and UI Overhaul (#10227)

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

@@ -355,7 +355,7 @@ export type MemoryTableSummary = {
   total_object_size: number;
   total_pinned_in_memory: number;
   total_used_by_pending_task: number;
-} | null;
+};
 
 export type MemoryTableEntry = {
   node_ip_address: string;
@@ -384,12 +384,12 @@ export type MemoryTableResponse = {
 // This doesn't return anything.
 export type StopMemoryTableResponse = {};
 
-export const getMemoryTable = (shouldObtainMemoryTable: boolean) => {
-  if (shouldObtainMemoryTable) {
-    return get<MemoryTableResponse>("/api/memory_table", {});
-  } else {
-    return null;
-  }
+export type MemoryGroupByKey = "node" | "stack_trace" | "";
+
+export const getMemoryTable = async (groupByKey: MemoryGroupByKey) => {
+  return get<MemoryTableResponse>("/api/memory_table", {
+    group_by: groupByKey,
+  });
 };
 
 export const stopMemoryTableCollection = () =>

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

@@ -0,0 +1,19 @@
+import ExpandLessIcon from "@material-ui/icons/ExpandLess";
+import ExpandMoreIcon from "@material-ui/icons/ExpandMore";
+import React from "react";
+
+type MinimizerProps = {
+  onClick: React.MouseEventHandler;
+};
+
+type ExpanderProps = {
+  onClick: React.MouseEventHandler;
+};
+
+export const Minimizer: React.FC<MinimizerProps> = ({ onClick }) => (
+  <ExpandLessIcon onClick={onClick} />
+);
+
+export const Expander: React.FC<ExpanderProps> = ({ onClick }) => (
+  <ExpandMoreIcon onClick={onClick} />
+);

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

@@ -0,0 +1,39 @@
+import { Grid, makeStyles, Tooltip } from "@material-ui/core";
+import React from "react";
+
+type LabeledDatumProps = {
+  label: string;
+  datum: any;
+  tooltip?: string;
+};
+
+const useLabeledDatumStyles = makeStyles({
+  tooltipLabel: {
+    textDecorationLine: "underline",
+    textDecorationColor: "#a6c3e3",
+    textDecorationThickness: "1px",
+    textDecorationStyle: "dotted",
+    cursor: "help",
+  },
+});
+
+const LabeledDatum: React.FC<LabeledDatumProps> = ({
+  label,
+  datum,
+  tooltip,
+}) => {
+  const classes = useLabeledDatumStyles();
+  const innerHtml = (
+    <Grid container item xs={6}>
+      <Grid item xs={6}>
+        <span className={tooltip && classes.tooltipLabel}>{label}</span>
+      </Grid>
+      <Grid item xs={6}>
+        <span>{datum}</span>
+      </Grid>
+    </Grid>
+  );
+  return tooltip ? <Tooltip title={tooltip}>{innerHtml}</Tooltip> : innerHtml;
+};
+
+export default LabeledDatum;

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

@@ -1,3 +1,5 @@
+import { useEffect, useRef } from "react";
+
 export const getWeightedAverage = (
   input: {
     weight: number;
@@ -24,3 +26,24 @@ export const filterObj = (obj: Object, filterFn: any) =>
 
 export const mapObj = (obj: Object, filterFn: any) =>
   Object.fromEntries(Object.entries(obj).map(filterFn));
+
+export const useInterval = (callback: Function, delayMs: number) => {
+  const savedCallback = useRef<any>();
+  const intervalId = useRef<any>();
+  useEffect(() => {
+    savedCallback.current = callback;
+  }, [callback]);
+  useEffect(() => {
+    const tick = () => savedCallback?.current();
+    intervalId.current = setInterval(tick, delayMs);
+    savedCallback.current();
+    return () => {
+      if (intervalId.current) {
+        clearInterval(intervalId.current);
+      }
+    };
+  }, [callback, delayMs]);
+  return intervalId.current
+    ? () => clearInterval(intervalId.current)
+    : () => null;
+};

python/ray/dashboard/client/src/pages/dashboard/Dashboard.tsx

@@ -9,13 +9,7 @@ import {
 } from "@material-ui/core";
 import React from "react";
 import { connect } from "react-redux";
-import {
-  getMemoryTable,
-  getNodeInfo,
-  getRayletInfo,
-  getTuneAvailability,
-  stopMemoryTableCollection,
-} from "../../api";
+import { getNodeInfo, getRayletInfo, getTuneAvailability } from "../../api";
 import { StoreState } from "../../store";
 import LastUpdated from "./LastUpdated";
 import LogicalView from "./logical-view/LogicalView";
@@ -44,7 +38,6 @@ const styles = (theme: Theme) =>
 const mapStateToProps = (state: StoreState) => ({
   tab: state.dashboard.tab,
   tuneAvailability: state.dashboard.tuneAvailability,
-  shouldObtainMemoryTable: state.dashboard.shouldObtainMemoryTable,
 });
 
 const mapDispatchToProps = dashboardActions;
@@ -64,25 +57,15 @@ class Dashboard extends React.Component<
   ];
 
   refreshInfo = async () => {
-    const { shouldObtainMemoryTable } = this.props;
     try {
-      const [
-        nodeInfo,
-        rayletInfo,
-        memoryTable,
-        tuneAvailability,
-      ] = await Promise.all([
+      const [nodeInfo, rayletInfo, tuneAvailability] = await Promise.all([
         getNodeInfo(),
         getRayletInfo(),
-        getMemoryTable(shouldObtainMemoryTable),
         getTuneAvailability(),
       ]);
       this.props.setNodeAndRayletInfo({ nodeInfo, rayletInfo });
       this.props.setTuneAvailability(tuneAvailability);
       this.props.setError(null);
-      if (shouldObtainMemoryTable) {
-        this.props.setMemoryTable(memoryTable);
-      }
     } catch (error) {
       this.props.setError(error.toString());
     } finally {
@@ -98,15 +81,8 @@ class Dashboard extends React.Component<
     clearTimeout(this.timeoutId);
   }
 
-  handleTabChange = async (event: React.ChangeEvent<{}>, value: number) => {
+  handleTabChange = async (event: React.ChangeEvent<{}>, value: number) =>
     this.props.setTab(value);
-    if (this.tabs[value].label === "Memory") {
-      this.props.setShouldObtainMemoryTable(true);
-    } else {
-      this.props.setShouldObtainMemoryTable(false);
-      await stopMemoryTableCollection();
-    }
-  };
 
   render() {
     const { classes, tab, tuneAvailability } = this.props;

python/ray/dashboard/client/src/pages/dashboard/logical-view/ActorDetailsPane.tsx

@@ -4,45 +4,10 @@ import {
   Grid,
   makeStyles,
   Theme,
-  Tooltip,
 } from "@material-ui/core";
 import React from "react";
 import { ActorState, InvalidStateType } from "../../../api";
-
-type LabeledDatumProps = {
-  label: string;
-  datum: any;
-  tooltip?: string;
-};
-
-const useLabeledDatumStyles = makeStyles({
-  label: {
-    textDecorationLine: "underline",
-    textDecorationColor: "#a6c3e3",
-    textDecorationThickness: "1px",
-    textDecorationStyle: "dotted",
-    cursor: "help",
-  },
-});
-
-const LabeledDatum: React.FC<LabeledDatumProps> = ({
-  label,
-  datum,
-  tooltip,
-}) => {
-  const classes = useLabeledDatumStyles();
-  const innerHtml = (
-    <Grid container item xs={6}>
-      <Grid item xs={6}>
-        <span className={classes.label}>{label}</span>
-      </Grid>
-      <Grid item xs={6}>
-        <span>{datum}</span>
-      </Grid>
-    </Grid>
-  );
-  return tooltip ? <Tooltip title={tooltip}>{innerHtml}</Tooltip> : innerHtml;
-};
+import LabeledDatum from "../../../common/LabeledDatum";
 
 type ActorStateReprProps = {
   state: ActorState;

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

@@ -1,145 +1,68 @@
 import {
+  Box,
   Button,
-  Checkbox,
   createStyles,
-  FormControlLabel,
+  FormControl,
+  InputLabel,
   makeStyles,
-  Table,
-  TableBody,
+  MenuItem,
+  Select,
   Theme,
+  Typography,
 } from "@material-ui/core";
 import PauseIcon from "@material-ui/icons/Pause";
 import PlayArrowIcon from "@material-ui/icons/PlayArrow";
-import React, { useState } from "react";
+import SubdirectoryArrowRightIcon from "@material-ui/icons/SubdirectoryArrowRight";
+import React, { useCallback, useEffect, useRef, useState } from "react";
 import { useDispatch, useSelector } from "react-redux";
 import {
-  MemoryTableEntry,
-  MemoryTableGroups,
+  getMemoryTable,
+  MemoryGroupByKey,
+  MemoryTableResponse,
   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 ExpanderRow from "./ExpanderRow";
 import MemoryRowGroup from "./MemoryRowGroup";
-import { MemoryTableRow } from "./MemoryTableRow";
 
-const DEFAULT_ENTRIES_PER_GROUP = 10;
-const DEFAULT_UNGROUPED_ENTRIES = 25;
-
-type GroupedMemoryRowsProps = {
-  memoryTableGroups: MemoryTableGroups;
-  order: Order;
-  orderBy: keyof MemoryTableEntry | null;
+const groupTitle = (groupKey: string, groupBy: MemoryGroupByKey) => {
+  if (groupBy === "node") {
+    return <Typography variant="h6">{`Node ${groupKey}`}</Typography>;
+  }
+  if (groupBy === "stack_trace") {
+    return <PyStackTrace stackTrace={groupKey} />;
+  }
+  if (groupBy === "") {
+    return <Typography variant="h6">All entries</Typography>;
+  }
+  return <Typography variant="h6">Unknown Group</Typography>;
 };
 
-const GroupedMemoryRows: React.FC<GroupedMemoryRowsProps> = ({
-  memoryTableGroups,
-  order,
-  orderBy,
-}) => {
-  const comparator = orderBy && getComparator(order, orderBy);
-  return (
-    <React.Fragment>
-      {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}
-            initialVisibleEntries={DEFAULT_ENTRIES_PER_GROUP}
-          />
-        );
-      })}
-    </React.Fragment>
-  );
+const PyStackTrace: React.FC<{ stackTrace: string }> = ({ stackTrace }) => {
+  const stackFrames = stackTrace.split(" | ");
+  const renderedFrames = stackFrames.map((frame, i) => (
+    <Typography
+      variant={i === 0 ? "h6" : "subtitle2"}
+      style={{ marginLeft: `${i}em` }}
+    >
+      {i !== 0 && <SubdirectoryArrowRightIcon />}
+      {frame}
+    </Typography>
+  ));
+  return <Box>{renderedFrames}</Box>;
 };
 
-type UngroupedMemoryRowsProps = {
-  memoryTableGroups: MemoryTableGroups;
-  order: Order;
-  orderBy: memoryColumnId | null;
-};
-
-const UngroupedMemoryRows: React.FC<UngroupedMemoryRowsProps> = ({
-  memoryTableGroups,
-  order,
-  orderBy,
-}) => {
-  const [visibleEntries, setVisibleEntries] = useState(
-    DEFAULT_UNGROUPED_ENTRIES,
-  );
-  const onExpand = () => setVisibleEntries(visibleEntries + 10);
-  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 (
-    <React.Fragment>
-      {" "}
-      {sortedEntries.slice(0, visibleEntries).map((memoryTableEntry, index) => (
-        <MemoryTableRow
-          memoryTableEntry={memoryTableEntry}
-          key={index.toString()}
-        />
-      ))}
-      <ExpanderRow onExpand={onExpand} />
-    </React.Fragment>
-  );
-};
-
-type memoryColumnId =
-  | "node_ip_address"
-  | "pid"
-  | "type"
-  | "object_ref"
-  | "object_size"
-  | "reference_type"
-  | "call_site";
-
-const memoryHeaderInfo: HeaderInfo<memoryColumnId>[] = [
-  { id: "node_ip_address", label: "IP Address", numeric: true, sortable: true },
-  { id: "pid", label: "pid", numeric: true, sortable: true },
-  { id: "type", label: "Type", numeric: false, sortable: true },
-  { id: "object_ref", label: "Object Ref", numeric: false, sortable: true },
-  {
-    id: "object_size",
-    label: "Object Size (B)",
-    numeric: true,
-    sortable: true,
-  },
-  {
-    id: "reference_type",
-    label: "Reference Type",
-    numeric: false,
-    sortable: true,
-  },
-  { id: "call_site", label: "Call Site", numeric: false, sortable: true },
-];
+const MEMORY_POLLING_INTERVAL_MS = 4000;
 
 const useMemoryInfoStyles = makeStyles((theme: Theme) =>
   createStyles({
-    table: {
-      marginTop: theme.spacing(1),
-    },
-    cell: {
+    pauseButton: {
+      margin: theme.spacing(1),
       padding: theme.spacing(1),
-      textAlign: "center",
+      float: "right",
+    },
+    select: {
+      minWidth: "7em",
     },
   }),
 );
@@ -150,84 +73,105 @@ const memoryInfoSelector = (state: StoreState) => ({
   shouldObtainMemoryTable: state.dashboard.shouldObtainMemoryTable,
 });
 
-const MemoryInfo: React.FC<{}> = () => {
-  const { memoryTable, shouldObtainMemoryTable } = useSelector(
-    memoryInfoSelector,
-  );
-  const dispatch = useDispatch();
-  const toggleMemoryCollection = async () => {
-    dispatch(
-      dashboardActions.setShouldObtainMemoryTable(!shouldObtainMemoryTable),
-    );
-    if (shouldObtainMemoryTable) {
-      await stopMemoryTableCollection();
-    }
+const fetchMemoryTable = (
+  groupByKey: MemoryGroupByKey,
+  setResults: (mtr: MemoryTableResponse) => void,
+) => {
+  return async () => {
+    const resp = await getMemoryTable(groupByKey);
+    setResults(resp);
   };
+};
+
+const MemoryInfo: React.FC<{}> = () => {
+  const { memoryTable } = useSelector(memoryInfoSelector);
+  const dispatch = useDispatch();
+
+  const [paused, setPaused] = useState(false);
+  const pauseButtonIcon = paused ? <PlayArrowIcon /> : <PauseIcon />;
 
-  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<memoryColumnId | null>(null);
+  const [groupBy, setGroupBy] = useState<MemoryGroupByKey>("node");
+
+  // Set up polling memory data
+  const fetchData = useCallback(
+    fetchMemoryTable(groupBy, (resp) =>
+      dispatch(dashboardActions.setMemoryTable(resp)),
+    ),
+    [groupBy],
+  );
+  const intervalId = useRef<any>(null);
+  useEffect(() => {
+    if (!intervalId.current && !paused) {
+      fetchData();
+      intervalId.current = setInterval(fetchData, MEMORY_POLLING_INTERVAL_MS);
+    }
+    const cleanup = () => {
+      if (intervalId.current) {
+        clearInterval(intervalId.current);
+        intervalId.current = null;
+      }
+    };
+    return cleanup;
+  }, [paused, fetchData]);
+
+  if (!memoryTable) {
+    return (
+      <Typography variant="h5" align="center">
+        Loading memory information
+      </Typography>
+    );
+  }
+
+  const children = Object.entries(memoryTable.group)
+    .sort(([key1], [key2]) => (key1 < key2 ? -1 : 1))
+    .map(([groupKey, memoryGroup]) => (
+      <MemoryRowGroup
+        key={groupKey}
+        groupKey={groupKey}
+        groupTitle={groupTitle(groupKey, groupBy)}
+        entries={memoryGroup.entries}
+        summary={memoryGroup.summary}
+        initialExpanded={false}
+        initialVisibleEntries={10}
+      />
+    ));
   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={(_, property) => {
-                if (property === orderBy) {
-                  toggleOrder();
-                } else {
-                  setOrderBy(property);
-                  setOrder("asc");
-                }
-              }}
-              headerInfo={memoryHeaderInfo}
-              firstColumnEmpty={false}
-            />
-            <TableBody>
-              {isGrouped ? (
-                <GroupedMemoryRows
-                  memoryTableGroups={memoryTable.group}
-                  order={order}
-                  orderBy={orderBy}
-                />
-              ) : (
-                <UngroupedMemoryRows
-                  memoryTableGroups={memoryTable.group}
-                  order={order}
-                  orderBy={orderBy}
-                />
-              )}
-            </TableBody>
-          </Table>
-        </React.Fragment>
-      ) : (
-        <div>No Memory Table Information Provided</div>
-      )}
-    </React.Fragment>
+    <Box>
+      <FormControl>
+        <InputLabel shrink id="group-by-label">
+          Group by
+        </InputLabel>
+        <Select
+          labelId="group-by-label"
+          value={groupBy}
+          className={classes.select}
+          onChange={(e: any) => setGroupBy(e.target.value)}
+          color="primary"
+          displayEmpty
+        >
+          <MenuItem value="">
+            <em>None</em>
+          </MenuItem>
+          <MenuItem value={"node"}>Node IP Address</MenuItem>
+          <MenuItem value={"stack_trace"}>Stack Trace</MenuItem>
+        </Select>
+      </FormControl>
+      <Button
+        color="primary"
+        className={classes.pauseButton}
+        onClick={() => {
+          if (!paused) {
+            stopMemoryTableCollection();
+          }
+          setPaused(!paused);
+        }}
+      >
+        {pauseButtonIcon}
+        {paused ? "Resume Collection" : "Pause Collection"}
+      </Button>
+      {children}
+    </Box>
   );
 };
 

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

@@ -1,35 +1,37 @@
-import { createStyles, makeStyles, TableRow, Theme } from "@material-ui/core";
-import AddIcon from "@material-ui/icons/Add";
-import RemoveIcon from "@material-ui/icons/Remove";
-import React, { useState } from "react";
-import { MemoryTableEntry, MemoryTableSummary } from "../../../api";
 import {
-  ExpandableStyledTableCell,
-  StyledTableCell,
-} from "../../../common/TableCell";
-import ExpanderRow from "./ExpanderRow";
+  Box,
+  createStyles,
+  makeStyles,
+  Paper,
+  styled,
+  Theme,
+} from "@material-ui/core";
+import React, { ReactChild, useState } from "react";
+import { MemoryTableEntry, MemoryTableSummary } from "../../../api";
+import { Expander, Minimizer } from "../../../common/ExpandControls";
 import MemorySummary from "./MemorySummary";
-import { MemoryTableRow } from "./MemoryTableRow";
+import MemoryTable from "./MemoryTable";
+
+const CenteredBox = styled(Box)({
+  textAlign: "center",
+});
 
 const useMemoryRowGroupStyles = makeStyles((theme: Theme) =>
   createStyles({
-    expandCollapseCell: {
-      cursor: "pointer",
-    },
-    expandCollapseIcon: {
-      color: theme.palette.text.secondary,
-      fontSize: "1.5em",
-      verticalAlign: "middle",
-    },
-    extraInfo: {
-      fontFamily: "SFMono-Regular,Consolas,Liberation Mono,Menlo,monospace",
-      whiteSpace: "pre",
+    container: {
+      marginTop: theme.spacing(2),
+      marginBottom: theme.spacing(2),
+      paddingTop: theme.spacing(2),
+      paddingBottom: theme.spacing(1),
+      paddingLeft: theme.spacing(2),
+      paddingRight: theme.spacing(2),
     },
   }),
 );
 
 type MemoryRowGroupProps = {
   groupKey: string;
+  groupTitle: ReactChild;
   summary: MemoryTableSummary;
   entries: MemoryTableEntry[];
   initialExpanded: boolean;
@@ -38,6 +40,7 @@ type MemoryRowGroupProps = {
 
 const MemoryRowGroup: React.FC<MemoryRowGroupProps> = ({
   groupKey,
+  groupTitle,
   entries,
   summary,
   initialExpanded,
@@ -45,53 +48,32 @@ const MemoryRowGroup: React.FC<MemoryRowGroupProps> = ({
 }) => {
   const classes = useMemoryRowGroupStyles();
   const [expanded, setExpanded] = useState(initialExpanded);
-  const [visibleEntries, setVisibleEntries] = useState(initialVisibleEntries);
+  const [numVisibleEntries, setNumVisibleEntries] = useState(
+    initialVisibleEntries,
+  );
   const toggleExpanded = () => setExpanded(!expanded);
-
-  const features = [
-    "node_ip_address",
-    "pid",
-    "type",
-    "object_ref",
-    "object_size",
-    "reference_type",
-    "call_site",
-  ];
-
+  const showMoreEntries = () => setNumVisibleEntries(numVisibleEntries + 10);
+  const visibleEntries = entries.slice(0, numVisibleEntries);
   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 && (
+    <Paper key={groupKey} className={classes.container}>
+      {groupTitle}
+      <MemorySummary initialExpanded={false} memoryTableSummary={summary} />
+      {expanded ? (
         <React.Fragment>
-          <MemorySummary initialExpanded={false} memoryTableSummary={summary} />
-          {entries.slice(0, visibleEntries).map((memoryTableEntry, index) => {
-            return (
-              <MemoryTableRow
-                memoryTableEntry={memoryTableEntry}
-                key={`${index}`}
-              />
-            );
-          })}
-          <ExpanderRow
-            onExpand={() => setVisibleEntries(visibleEntries + 10)}
-          />
+          <MemoryTable tableEntries={visibleEntries} />
+          <CenteredBox>
+            {entries.length > numVisibleEntries && (
+              <Expander onClick={showMoreEntries} />
+            )}
+            <Minimizer onClick={toggleExpanded} />
+          </CenteredBox>
         </React.Fragment>
+      ) : (
+        <CenteredBox>
+          <Expander onClick={toggleExpanded} />
+        </CenteredBox>
       )}
-    </React.Fragment>
+    </Paper>
   );
 };
 

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

@@ -1,91 +1,59 @@
-import {
-  createStyles,
-  TableCell,
-  TableRow,
-  Theme,
-  withStyles,
-  WithStyles,
-} from "@material-ui/core";
+import { createStyles, Grid, makeStyles, Theme } from "@material-ui/core";
 import React from "react";
 import { MemoryTableSummary } from "../../../api";
+import { formatByteAmount } from "../../../common/formatUtils";
+import LabeledDatum from "../../../common/LabeledDatum";
 
-const styles = (theme: Theme) =>
+const useMemorySummaryStyles = makeStyles((theme: Theme) =>
   createStyles({
-    cell: {
-      padding: theme.spacing(1),
-      textAlign: "center",
-      "&:last-child": {
-        paddingRight: theme.spacing(1),
-      },
-    },
-    expandCollapseCell: {
-      cursor: "pointer",
-    },
     expandCollapseIcon: {
       color: theme.palette.text.secondary,
       fontSize: "1.5em",
       verticalAlign: "middle",
     },
-    extraInfo: {
-      fontFamily: "SFMono-Regular,Consolas,Liberation Mono,Menlo,monospace",
-      whiteSpace: "pre",
+    container: {
+      padding: theme.spacing(1),
+      margin: theme.spacing(1),
     },
-  });
+  }),
+);
 
-type Props = {
+type MemorySummaryProps = {
   memoryTableSummary: MemoryTableSummary;
   initialExpanded: boolean;
 };
 
-type State = {
-  expanded: boolean;
+const MemorySummary: React.FC<MemorySummaryProps> = ({
+  memoryTableSummary,
+}) => {
+  const classes = useMemorySummaryStyles();
+  const memoryData = [
+    [
+      "Total Local Reference Count",
+      `${memoryTableSummary.total_local_ref_count}`,
+    ],
+    ["Pinned in Memory Count", `${memoryTableSummary.total_pinned_in_memory}`],
+    [
+      "Total Used by Pending Tasks Count",
+      `${memoryTableSummary.total_used_by_pending_task}`,
+    ],
+    [
+      "Total Captured in Objects Count",
+      `${memoryTableSummary.total_captured_in_objects}`,
+    ],
+    [
+      "Total Memory Used by Objects",
+      `${formatByteAmount(memoryTableSummary.total_object_size, "mebibyte")}`,
+    ],
+    ["Total Actor Handle Count", `${memoryTableSummary.total_actor_handles}`],
+  ];
+
+  return (
+    <Grid container className={classes.container}>
+      {memoryData.map(([label, value]) => (
+        <LabeledDatum key={label} label={label} datum={value} />
+      ))}
+    </Grid>
+  );
 };
-
-class MemorySummary extends React.Component<
-  Props & WithStyles<typeof styles>,
-  State
-> {
-  state: State = {
-    expanded: this.props.initialExpanded,
-  };
-
-  toggleExpand = () => {
-    this.setState((state) => ({
-      expanded: !state.expanded,
-    }));
-  };
-
-  render() {
-    const { classes, memoryTableSummary } = this.props;
-
-    const memorySummaries =
-      memoryTableSummary !== null
-        ? [
-            "", // Padding
-            `Total Local Reference Count: ${memoryTableSummary.total_local_ref_count}`,
-            `Total Pinned In Memory Count: ${memoryTableSummary.total_pinned_in_memory}`,
-            `Total Used By Pending Tasks Count: ${memoryTableSummary.total_used_by_pending_task}`,
-            `Total Caputed In Objects Count: ${memoryTableSummary.total_captured_in_objects}`,
-            `Total Object Size: ${memoryTableSummary.total_object_size} B`,
-            `Total Actor Handle Count: ${memoryTableSummary.total_actor_handles}`,
-            "", // Padding
-          ]
-        : ["No Summary Provided"];
-
-    return (
-      memoryTableSummary !== null && (
-        <React.Fragment>
-          <TableRow hover>
-            {memorySummaries.map((summary, index) => (
-              <TableCell key={index} className={classes.cell}>
-                {summary}
-              </TableCell>
-            ))}
-          </TableRow>
-        </React.Fragment>
-      )
-    );
-  }
-}
-
-export default withStyles(styles)(MemorySummary);
+export default MemorySummary;

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

@@ -0,0 +1,105 @@
+import {
+  createStyles,
+  makeStyles,
+  Paper,
+  Table,
+  TableBody,
+  Theme,
+} from "@material-ui/core";
+import React from "react";
+import { MemoryTableEntry } from "../../../api";
+import SortableTableHead, {
+  HeaderInfo,
+} from "../../../common/SortableTableHead";
+import { getComparator, Order, stableSort } from "../../../common/tableUtils";
+
+import { MemoryTableRow } from "./MemoryTableRow";
+
+const useMemoryTableStyles = makeStyles((theme: Theme) =>
+  createStyles({
+    container: {
+      margin: theme.spacing(1),
+      padding: theme.spacing(1),
+    },
+    cell: {
+      padding: theme.spacing(1),
+      textAlign: "center",
+    },
+  }),
+);
+
+type memoryColumnId =
+  | "node_ip_address"
+  | "pid"
+  | "type"
+  | "object_ref"
+  | "object_size"
+  | "reference_type"
+  | "call_site";
+
+const memoryHeaderInfo: HeaderInfo<memoryColumnId>[] = [
+  {
+    id: "node_ip_address",
+    label: "IP Address",
+    numeric: false,
+    sortable: true,
+  },
+  { id: "pid", label: "PID", numeric: false, sortable: true },
+  { id: "type", label: "Type", numeric: false, sortable: true },
+  { id: "object_ref", label: "Object Ref", numeric: false, sortable: true },
+  {
+    id: "object_size",
+    label: "Object Size",
+    numeric: false,
+    sortable: true,
+  },
+  {
+    id: "reference_type",
+    label: "Reference Type",
+    numeric: false,
+    sortable: true,
+  },
+  { id: "call_site", label: "Call Site", numeric: false, sortable: true },
+];
+
+type MemoryTableProps = {
+  tableEntries: MemoryTableEntry[];
+};
+
+const MemoryTable: React.FC<MemoryTableProps> = ({ tableEntries }) => {
+  const toggleOrder = () => setOrder(order === "asc" ? "desc" : "asc");
+  const classes = useMemoryTableStyles();
+  const [order, setOrder] = React.useState<Order>("asc");
+  const [orderBy, setOrderBy] = React.useState<memoryColumnId | null>(null);
+  const comparator = orderBy && getComparator(order, orderBy);
+  const sortedTableEntries = comparator
+    ? stableSort(tableEntries, comparator)
+    : tableEntries;
+  const tableRows = sortedTableEntries.map((tableEntry) => (
+    <MemoryTableRow memoryTableEntry={tableEntry} key={tableEntry.object_ref} />
+  ));
+  // Todo(max) add in sorting code
+  return (
+    <Paper className={classes.container} elevation={2}>
+      <Table>
+        <SortableTableHead
+          orderBy={orderBy}
+          order={order}
+          onRequestSort={(_, property) => {
+            if (property === orderBy) {
+              toggleOrder();
+            } else {
+              setOrderBy(property);
+              setOrder("asc");
+            }
+          }}
+          headerInfo={memoryHeaderInfo}
+          firstColumnEmpty={false}
+        />
+        <TableBody>{tableRows}</TableBody>
+      </Table>
+    </Paper>
+  );
+};
+
+export default MemoryTable;

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

@@ -1,6 +1,7 @@
 import { TableRow } from "@material-ui/core";
 import React from "react";
 import { MemoryTableEntry } from "../../../api";
+import { formatByteAmount } from "../../../common/formatUtils";
 import { StyledTableCell } from "../../../common/TableCell";
 
 type Props = {
@@ -13,7 +14,7 @@ export const MemoryTableRow = (props: Props) => {
   const object_size =
     memoryTableEntry.object_size === -1
       ? "?"
-      : `${memoryTableEntry.object_size}  B`;
+      : formatByteAmount(memoryTableEntry.object_size, "mebibyte");
   const memoryTableEntryValues = [
     memoryTableEntry.node_ip_address,
     memoryTableEntry.pid,

python/ray/dashboard/dashboard.py

@@ -31,7 +31,8 @@ from ray.core.generated import core_worker_pb2
 from ray.core.generated import core_worker_pb2_grpc
 from ray.dashboard.interface import BaseDashboardController
 from ray.dashboard.interface import BaseDashboardRouteHandler
-from ray.dashboard.memory import construct_memory_table, MemoryTable
+from ray.dashboard.memory import construct_memory_table, MemoryTable, \
+     GroupByType, SortingType
 from ray.dashboard.metrics_exporter.client import Exporter
 from ray.dashboard.metrics_exporter.client import MetricsExportClient
 from ray.dashboard.node_stats import NodeStats
@@ -175,7 +176,9 @@ class DashboardController(BaseDashboardController):
     def get_raylet_info(self):
         return self._construct_raylet_info()
 
-    def get_memory_table_info(self) -> MemoryTable:
+    def get_memory_table_info(self,
+                              group_by=GroupByType.NODE_ADDRESS,
+                              sort_by=SortingType.OBJECT_SIZE) -> MemoryTable:
         # Collecting memory info adds big overhead to the cluster.
         # This must be collected only when it is necessary.
         self.raylet_stats.include_memory_info = True
@@ -184,7 +187,8 @@ class DashboardController(BaseDashboardController):
             data["nodeId"]: data.get("workersStats")
             for data in D.values()
         }
-        self.memory_table = construct_memory_table(workers_info_by_node)
+        self.memory_table = construct_memory_table(
+            workers_info_by_node, group_by=group_by, sort_by=sort_by)
         return self.memory_table
 
     def stop_collecting_memory_table_info(self):
@@ -280,7 +284,19 @@ class DashboardRouteHandler(BaseDashboardRouteHandler):
         return await json_response(self.is_dev, result=result)
 
     async def memory_table_info(self, req) -> aiohttp.web.Response:
-        memory_table = self.dashboard_controller.get_memory_table_info()
+        group_by = req.query.get("group_by")
+        sort_by = req.query.get("sort_by")
+        kwargs = {}
+        try:
+            if group_by:
+                kwargs["group_by"] = GroupByType(group_by)
+            if sort_by:
+                kwargs["sort_by"] = SortingType(sort_by)
+        except ValueError as e:
+            return aiohttp.web.HTTPBadRequest(reason=str(e))
+
+        memory_table = self.dashboard_controller.get_memory_table_info(
+            **kwargs)
         return await json_response(self.is_dev, result=memory_table.__dict__())
 
     async def stop_collecting_memory_table_info(self,

python/ray/dashboard/memory.py

@@ -41,7 +41,8 @@ class SortingType(Enum):
 
 
 class GroupByType(Enum):
-    NODE_ADDRESS = 2
+    NODE_ADDRESS = "node"
+    STACK_TRACE = "stack_trace"
 
 
 class ReferenceType:
@@ -94,6 +95,8 @@ class MemoryTableEntry:
     def group_key(self, group_by_type: GroupByType) -> str:
         if group_by_type == GroupByType.NODE_ADDRESS:
             return self.node_address
+        elif group_by_type == GroupByType.STACK_TRACE:
+            return self.call_site
         else:
             raise ValueError(
                 "group by type {} is invalid.".format(group_by_type))
@@ -272,7 +275,9 @@ class MemoryTable:
         return self.__repr__()
 
 
-def construct_memory_table(workers_info_by_node: dict) -> MemoryTable:
+def construct_memory_table(workers_info_by_node: dict,
+                           group_by: GroupByType = GroupByType.NODE_ADDRESS,
+                           sort_by=SortingType.OBJECT_SIZE) -> MemoryTable:
     memory_table_entries = []
     for node_id, worker_infos in workers_info_by_node.items():
         for worker_info in worker_infos:
@@ -290,5 +295,6 @@ def construct_memory_table(workers_info_by_node: dict) -> MemoryTable:
                     pid=pid)
                 if memory_table_entry.is_valid():
                     memory_table_entries.append(memory_table_entry)
-    memory_table = MemoryTable(memory_table_entries)
+    memory_table = MemoryTable(
+        memory_table_entries, group_by_type=group_by, sort_by_type=sort_by)
     return memory_table