[Dashboard] Ray memory frontend (#8563)

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

@@ -283,3 +283,47 @@ export const setTuneExperiment = (experiment: string) =>
 
 export const enableTuneTensorBoard = () =>
   post<{}>("/api/enable_tune_tensorboard", {});
+
+export type MemoryTableSummary = {
+  total_actor_handles: number;
+  total_captured_in_objects: number;
+  total_local_ref_count: number;
+  // The measurement is B.
+  total_object_size: number;
+  total_pinned_in_memory: number;
+  total_used_by_pending_task: number;
+} | null;
+
+export type MemoryTableEntry = {
+  node_ip_address: string;
+  pid: number;
+  type: string;
+  object_id: string;
+  object_size: number;
+  reference_type: string;
+  call_site: string;
+};
+
+export type MemoryTableResponse = {
+  group: {
+    [groupKey: string]: {
+      entries: MemoryTableEntry[];
+      summary: MemoryTableSummary;
+    };
+  };
+  summary: MemoryTableSummary;
+};
+
+// 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 const stopMemoryTableCollection = () =>
+  get<StopMemoryTableResponse>("/api/stop_memory_table", {});

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

@@ -9,7 +9,13 @@ import {
 } from "@material-ui/core";
 import React from "react";
 import { connect } from "react-redux";
-import { getNodeInfo, getRayletInfo, getTuneAvailability } from "../../api";
+import {
+  getNodeInfo,
+  getRayletInfo,
+  getTuneAvailability,
+  getMemoryTable,
+  stopMemoryTableCollection,
+} from "../../api";
 import { StoreState } from "../../store";
 import LastUpdated from "./LastUpdated";
 import LogicalView from "./logical-view/LogicalView";
@@ -17,6 +23,7 @@ import NodeInfo from "./node-info/NodeInfo";
 import RayConfig from "./ray-config/RayConfig";
 import { dashboardActions } from "./state";
 import Tune from "./tune/Tune";
+import MemoryInfo from "./memory/Memory";
 
 const styles = (theme: Theme) =>
   createStyles({
@@ -37,6 +44,7 @@ const styles = (theme: Theme) =>
 const mapStateToProps = (state: StoreState) => ({
   tab: state.dashboard.tab,
   tuneAvailability: state.dashboard.tuneAvailability,
+  shouldObtainMemoryTable: state.dashboard.shouldObtainMemoryTable,
 });
 
 const mapDispatchToProps = dashboardActions;
@@ -47,48 +55,66 @@ class Dashboard extends React.Component<
     typeof mapDispatchToProps
 > {
   timeoutId = 0;
+  tabs = [
+    { label: "Machine view", component: NodeInfo },
+    { label: "Logical view", component: LogicalView },
+    { label: "Memory", component: MemoryInfo },
+    { label: "Ray config", component: RayConfig },
+    { label: "Tune", component: Tune },
+  ];
 
-  refreshNodeAndRayletInfo = async () => {
+  refreshInfo = async () => {
+    let { shouldObtainMemoryTable } = this.props;
     try {
-      const [nodeInfo, rayletInfo, tuneAvailability] = await Promise.all([
+      const [
+        nodeInfo,
+        rayletInfo,
+        memoryTable,
+        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 {
-      this.timeoutId = window.setTimeout(this.refreshNodeAndRayletInfo, 1000);
+      this.timeoutId = window.setTimeout(this.refreshInfo, 1000);
     }
   };
 
   async componentDidMount() {
-    await this.refreshNodeAndRayletInfo();
+    await this.refreshInfo();
   }
 
   componentWillUnmount() {
     clearTimeout(this.timeoutId);
   }
 
-  handleTabChange = (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;
-    const tabs = [
-      { label: "Machine view", component: NodeInfo },
-      { label: "Logical view", component: LogicalView },
-      { label: "Ray config", component: RayConfig },
-      { label: "Tune", component: Tune },
-    ];
+    let tabs = this.tabs.slice();
 
     // if Tune information is not available, remove Tune tab from the dashboard
     if (tuneAvailability === null || !tuneAvailability.available) {
-      tabs.splice(3);
+      tabs.splice(4);
     }
 
     const SelectedComponent = tabs[tab].component;

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

@@ -0,0 +1,123 @@
+import {
+  createStyles,
+  Table,
+  TableBody,
+  TableCell,
+  TableHead,
+  TableRow,
+  Theme,
+  withStyles,
+  WithStyles,
+  Button,
+} from "@material-ui/core";
+import PauseIcon from "@material-ui/icons/Pause";
+import PlayArrowIcon from "@material-ui/icons/PlayArrow";
+import React from "react";
+import { StoreState } from "../../../store";
+import { connect } from "react-redux";
+import MemoryRowGroup from "./MemoryRowGroup";
+import { dashboardActions } from "../state";
+import { stopMemoryTableCollection } from "../../../api";
+
+const styles = (theme: Theme) =>
+  createStyles({
+    table: {
+      marginTop: theme.spacing(1),
+    },
+    cell: {
+      padding: theme.spacing(1),
+      textAlign: "center",
+    },
+  });
+
+const mapStateToProps = (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);
+    if (shouldObtainMemoryTable) {
+      await stopMemoryTableCollection();
+    }
+  };
+
+  renderIcon = () => {
+    if (this.props.shouldObtainMemoryTable) {
+      return <PauseIcon />;
+    } else {
+      return <PlayArrowIcon />;
+    }
+  };
+
+  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));

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

@@ -0,0 +1,142 @@
+import {
+  createStyles,
+  TableCell,
+  TableRow,
+  Theme,
+  withStyles,
+  WithStyles,
+} 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 {
+  MemoryTableResponse,
+  MemoryTableEntry,
+  MemoryTableSummary,
+} from "../../../api";
+import MemorySummary from "./MemorySummary";
+
+const styles = (theme: Theme) =>
+  createStyles({
+    cell: {
+      padding: theme.spacing(1),
+      textAlign: "center",
+    },
+    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",
+    },
+  });
+
+type Props = {
+  groupKey: string;
+  memoryTableGroups: MemoryTableResponse["group"];
+  initialExpanded: boolean;
+};
+
+type State = {
+  expanded: boolean;
+};
+
+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);

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

@@ -0,0 +1,91 @@
+import {
+  createStyles,
+  TableCell,
+  TableRow,
+  Theme,
+  withStyles,
+  WithStyles,
+} from "@material-ui/core";
+import React from "react";
+import { MemoryTableSummary } from "../../../api";
+
+const styles = (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",
+    },
+  });
+
+type Props = {
+  memoryTableSummary: MemoryTableSummary;
+  initialExpanded: boolean;
+};
+
+type State = {
+  expanded: boolean;
+};
+
+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);

python/ray/dashboard/client/src/pages/dashboard/state.ts

@@ -5,6 +5,7 @@ import {
   RayletInfoResponse,
   TuneAvailabilityResponse,
   TuneJobResponse,
+  MemoryTableResponse,
 } from "../../api";
 
 const name = "dashboard";
@@ -18,6 +19,8 @@ type State = {
   tuneAvailability: TuneAvailabilityResponse | null;
   lastUpdatedAt: number | null;
   error: string | null;
+  memoryTable: MemoryTableResponse | null;
+  shouldObtainMemoryTable: boolean;
 };
 
 const initialState: State = {
@@ -29,6 +32,8 @@ const initialState: State = {
   tuneAvailability: null,
   lastUpdatedAt: null,
   error: null,
+  memoryTable: null,
+  shouldObtainMemoryTable: false,
 };
 
 const slice = createSlice({
@@ -66,6 +71,15 @@ const slice = createSlice({
     setError: (state, action: PayloadAction<string | null>) => {
       state.error = action.payload;
     },
+    setMemoryTable: (
+      state,
+      action: PayloadAction<MemoryTableResponse | null>,
+    ) => {
+      state.memoryTable = action.payload;
+    },
+    setShouldObtainMemoryTable: (state, action: PayloadAction<boolean>) => {
+      state.shouldObtainMemoryTable = action.payload;
+    },
   },
 });
 

python/ray/dashboard/dashboard.py

@@ -241,8 +241,6 @@ class DashboardController(BaseDashboardController):
         return self.memory_table
 
     def stop_collecting_memory_table_info(self):
-        # Reset memory table.
-        self.memory_table = MemoryTable([])
         self.raylet_stats.include_memory_info = False
 
     def tune_info(self):

python/ray/dashboard/memory.py

@@ -290,6 +290,5 @@ 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)
     return memory_table