feat: Session usage monitor in react

react/src/components/SessionDetailContent.tsx

@@ -10,6 +10,7 @@ import SessionStatusTag from './ComputeSessionNodeItems/SessionStatusTag';
 import SessionTypeTag from './ComputeSessionNodeItems/SessionTypeTag';
 import Flex from './Flex';
 import ImageMetaIcon from './ImageMetaIcon';
+import SessionUsageMonitor from './SessionUsageMonitor';
 import { SessionDetailContentQuery } from './__generated__/SessionDetailContentQuery.graphql';
 import {
   Alert,
@@ -80,6 +81,7 @@ const SessionDetailContent: React.FC<{
             # fix: This fragment is not used in this component, but it is required by the SessionActionButtonsFragment.
             # It might be a bug in relay
             ...ContainerLogModalFragment
+            ...SessionUsageMonitorFragment
           }
           legacy_session: compute_session(id: $uuid) {
             image
@@ -178,11 +180,14 @@ const SessionDetailContent: React.FC<{
         <Descriptions.Item label={t('session.Agent')}>
           {session.agent_ids || '-'}
         </Descriptions.Item>
-        <Descriptions.Item label={t('session.Reservation')}>
+        <Descriptions.Item label={t('session.Reservation')} span={md ? 2 : 1}>
           <Flex gap={'xs'} wrap={'wrap'}>
             <SessionReservation sessionFrgmt={session} />
           </Flex>
         </Descriptions.Item>
+        <Descriptions.Item label={'Resource Usage'} span={md ? 2 : 1}>
+          <SessionUsageMonitor sessionFrgmt={session} col={md ? 2 : 1} />
+        </Descriptions.Item>
       </Descriptions>
     </Flex>
   ) : (

react/src/components/SessionUsageMonitor.tsx

@@ -0,0 +1,260 @@
+import { iSizeToSize, toFixedFloorWithoutTrailingZeros } from '../helper';
+import Flex from './Flex';
+import { ResourceTypeIcon } from './ResourceNumber';
+import { SessionUsageMonitorFragment$key } from './__generated__/SessionUsageMonitorFragment.graphql';
+import { Col, Progress, ProgressProps, Row, Typography, theme } from 'antd';
+import { createStyles } from 'antd-style';
+import graphql from 'babel-plugin-relay/macro';
+import _ from 'lodash';
+import { useMemo } from 'react';
+import { useFragment } from 'react-relay';
+
+const useStyles = createStyles(({ css, token }) => ({
+  progress: css`
+    .ant-progress-text {
+      font-size: ${token.fontSizeSM} !important;
+    }
+    .ant-progress-inner {
+    }
+    .ant-typography {
+      color: ${token.colorTextSecondary} !important;
+      font-size: ${token.fontSizeSM}px !important;
+    }
+  `,
+}));
+
+interface SessionUsageMonitorProps extends ProgressProps {
+  sessionFrgmt: SessionUsageMonitorFragment$key | null;
+  col?: 1 | 2 | 3 | 4;
+  size?: ProgressProps['size'];
+}
+
+// compute session node query returns accelerator type as `[device_name]_[util|mem]`
+// so we need to split via '_' and get the first part to map to the resource type in `resourceTypeIconSrcMap`
+
+type acceleratorType =
+  | 'cuda'
+  | 'rocm'
+  | 'tpu'
+  | 'ipu'
+  | 'atom'
+  | 'atom-plus'
+  | 'gaudi2'
+  | 'warboy'
+  | 'rngd'
+  | 'hyperaccel-lpu';
+
+const resourceTypeMap: {
+  [key in acceleratorType | 'cpu' | 'mem']: string;
+} = {
+  cpu: 'cpu',
+  mem: 'mem',
+  cuda: 'cuda.device',
+  rocm: 'rocm.device',
+  tpu: 'tpu.device',
+  ipu: 'ipu.device',
+  atom: 'atom.device',
+  'atom-plus': 'atom-plus.device',
+  gaudi2: 'gaudi2.device',
+  warboy: 'warboy.device',
+  rngd: 'rngd.device',
+  'hyperaccel-lpu': 'hyperaccel-lpu.device',
+};
+
+const gridColSpanMap = {
+  1: 24,
+  2: 12,
+  3: 8,
+  4: 6,
+};
+
+const SessionUsageMonitor: React.FC<SessionUsageMonitorProps> = ({
+  sessionFrgmt,
+  col = 1,
+  size = 'default',
+}) => {
+  const { token } = theme.useToken();
+  const { styles } = useStyles();
+
+  const kernel_nodes = useFragment(
+    graphql`
+      fragment SessionUsageMonitorFragment on ComputeSessionNode {
+        kernel_nodes {
+          edges {
+            node {
+              live_stat
+            }
+          }
+        }
+      }
+    `,
+    sessionFrgmt,
+  );
+
+  const liveStat = JSON.parse(
+    _.get(kernel_nodes, 'kernel_nodes.edges[0].node.live_stat') ?? '{}',
+  );
+  const sortedLiveStat = useMemo(
+    () =>
+      Object.keys(liveStat)
+        .sort((a, b) => {
+          const aUtil = a.includes('_util');
+          const bUtil = b.includes('_util');
+          const aMem = a.includes('_mem');
+          const bMem = b.includes('_mem');
+
+          if (aUtil && !bUtil) return -1;
+          if (!aUtil && bUtil) return 1;
+          if (aMem && !bMem) return -1;
+          if (!aMem && bMem) return 1;
+
+          return 0;
+        })
+        .reduce((acc: { [key: string]: any }, key) => {
+          acc[key] = liveStat[key];
+          return acc;
+        }, {}),
+    [liveStat],
+  );
+
+  return (
+    <>
+      <Row className={styles.progress} gutter={token.marginXL}>
+        {sortedLiveStat?.cpu_util ? (
+          <Col
+            span={_.get(gridColSpanMap, col)}
+            style={{ marginBottom: token.marginXS }}
+          >
+            <Flex direction="column" align="end">
+              <Flex
+                align="start"
+                gap={token.marginSM}
+                style={{ width: '100%' }}
+              >
+                <ResourceTypeIcon
+                  type={resourceTypeMap.cpu}
+                  showTooltip={true}
+                />
+                <Progress
+                  showInfo
+                  status="normal"
+                  percent={_.toNumber(
+                    toFixedFloorWithoutTrailingZeros(
+                      liveStat?.cpu_util?.pct || 0,
+                      1,
+                    ),
+                  )}
+                  strokeLinecap="butt"
+                  size={size === 'default' ? { height: token.sizeMD } : size}
+                  strokeColor={token.colorPrimary}
+                ></Progress>
+              </Flex>
+              {size === 'small' ? (
+                <Typography.Text>
+                  {`${liveStat?.cpu_util?.pct} %`}
+                </Typography.Text>
+              ) : null}
+            </Flex>
+          </Col>
+        ) : null}
+        {sortedLiveStat?.mem ? (
+          <Col
+            span={_.get(gridColSpanMap, col)}
+            style={{ marginBottom: token.marginXS }}
+          >
+            <Flex direction="column" align="end">
+              <Flex
+                align="start"
+                gap={token.marginSM}
+                style={{ width: '100%' }}
+              >
+                <ResourceTypeIcon
+                  type={resourceTypeMap.mem}
+                  showTooltip={true}
+                />
+                <Progress
+                  showInfo
+                  status="normal"
+                  percent={_.toNumber(
+                    toFixedFloorWithoutTrailingZeros(
+                      sortedLiveStat?.mem?.pct || 0,
+                      1,
+                    ),
+                  )}
+                  strokeLinecap="butt"
+                  size={{ height: 20 }}
+                  strokeColor={token.colorPrimary}
+                ></Progress>
+              </Flex>
+              <Typography.Text>
+                {`(mem) ${
+                  iSizeToSize(sortedLiveStat?.mem?.current, 'g')?.numberUnit
+                }iB
+                /
+                ${
+                  iSizeToSize(sortedLiveStat?.mem?.capacity, 'g')?.numberUnit
+                }iB`}
+              </Typography.Text>
+            </Flex>
+          </Col>
+        ) : null}
+        {/* accelerator type is in the format of `[device_name]_[util|mem]` */}
+        {_.map(sortedLiveStat, (value, key) => {
+          const device = key.split('_')?.[0] as keyof typeof resourceTypeMap;
+
+          return _.includes(
+            _.keysIn(_.omit(resourceTypeMap, 'cpu', 'mem')),
+            device,
+          ) ? (
+            <>
+              {/* only show accelerator's util and mem */}
+              {(_.includes(key, 'util') || _.includes(key, 'mem')) && (
+                <Col
+                  span={_.get(gridColSpanMap, col)}
+                  style={{ marginBottom: token.marginXS }}
+                >
+                  <Flex direction="column" align="end">
+                    <Flex
+                      align="start"
+                      gap={token.marginSM}
+                      style={{ width: '100%' }}
+                    >
+                      <ResourceTypeIcon
+                        type={resourceTypeMap[device]}
+                        showTooltip={true}
+                      />
+                      <Progress
+                        showInfo
+                        status="normal"
+                        percent={_.toNumber(
+                          toFixedFloorWithoutTrailingZeros(
+                            sortedLiveStat?.[device]?.pct || 0,
+                            1,
+                          ),
+                        )}
+                        strokeLinecap="butt"
+                        size={{ height: token.sizeMD }}
+                        strokeColor={token.colorPrimary}
+                      ></Progress>
+                    </Flex>
+                    <Typography.Text>
+                      {_.includes(key, 'mem') &&
+                        `${device} (mem) ${iSizeToSize(sortedLiveStat?.mem?.current, 'g')?.numberUnit}iB / ${iSizeToSize(sortedLiveStat?.mem?.capacity, 'g')?.numberUnit}iB`}
+                    </Typography.Text>
+                  </Flex>
+                </Col>
+              )}
+            </>
+          ) : null;
+        })}
+      </Row>
+      <Flex justify="end" style={{ width: '100%' }}>
+        <Typography.Text>
+          {`I/O Read: ${iSizeToSize(sortedLiveStat?.io_read?.current, 'g')?.numberUnit}B / Write: ${iSizeToSize(sortedLiveStat?.io_write?.current, 'g')?.numberUnit}B`}
+        </Typography.Text>
+      </Flex>
+    </>
+  );
+};
+
+export default SessionUsageMonitor;