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

doc/source/ray-dashboard.rst

@@ -3,10 +3,12 @@ Ray Dashboard
 Ray's built-in dashboard provides metrics, charts, and other features that help
 Ray users to understand Ray clusters and libraries.
 
-Through the dashboard, you can
-
+The dashboard lets you:
 - View cluster metrics.
-- Visualize the actor relationships and statistics.
+- See errors and exceptions at a glance.
+- View logs across many machines in a single pane.
+- Understand Ray memory utilization and debug memory errors.
+- See per-actor resource usage, executed tasks, logs, and more.
 - Kill actors and profile your Ray jobs.
 - See Tune jobs and trial information.
 - Detect cluster anomalies and debug them.
@@ -64,7 +66,7 @@ Memory View
 The memory view shows you:
 
 - The state of Ray objects, including their size, reference type, and call site.
-- A summary of reference types and object sizes in use.
+- The aggregate amount of memory being used by various groups, such as line of code, or the node.
 
 .. image:: https://raw.githubusercontent.com/ray-project/images/master/docs/dashboard/Memory-view-basic.png
     :align: center
@@ -141,6 +143,9 @@ is bigger than the total gpus available in this cluster (2 GPUs).
 Debugging ObjectStoreFullError and Memory Leak
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 You can view information for Ray objects in the memory tab. It is useful to debug memory leaks, especially `ObjectStoreFullError`.
+
+One common cause of these memory errors is that there are objects which never go out of scope. In order to find these, you can go to the Memory View, then select to "Group By Stack Trace." This groups memory entries by their stack traces up to three frames deep. If you see a group which is growing without bound, you might want to examine that line of code to see if you intend to keep that reference around.
+
 Note that this is the same information as displayed in the `ray memory command <https://docs.ray.io/en/latest/memory-management.html#debugging-using-ray-memory>`_. For details about the information contained in the table, please see the `ray memory` documentation.
 
 Inspect Memory Usage
@@ -318,7 +323,7 @@ Memory
 
 **IP Address**: Node IP Address where a Ray object is pinned.
 
-**Pid**: ID of a process where a Ray object is being used.
+**PID**: ID of a process where a Ray object is being used.
 
 **Type**: Type of a process. It is either a driver or worker.
 
@@ -328,7 +333,7 @@ Memory
 
 **Reference Type**: Reference types of Ray objects. Checkout the `ray memory command <https://docs.ray.io/en/latest/memory-management.html#debugging-using-ray-memory>`_ to learn each reference type.
 
-**Call Site**: Call site where this Ray object is referenced.
+**Call Site**: Call site where this Ray object is referenced, up to three stack frames deep.
 
 Ray Config
 ~~~~~~~~~~~~

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