feat: Add Resource Allocation stats to trial detail page (#9533)

determined-ai · Jun 28, 2024 · 875e326 · 875e326
1 parent 1e9dc42
commit 875e326
Show file tree

Hide file tree

Showing 15 changed files with 515 additions and 322 deletions.
diff --git a/harness/determined/common/api/bindings.py b/harness/determined/common/api/bindings.py
diff --git a/master/internal/api_tasks.go b/master/internal/api_tasks.go
@@ -329,19 +329,31 @@ func (a *apiServer) GetTaskAcceleratorData(
 		return nil, err
 	}
 
-	var res []model.AcceleratorData
-	err := db.Bun().NewSelect().
-		ColumnExpr("alloc_acc.*").
+	res := []struct {
+		model.AcceleratorData
+		ResourcePool string
+	}{}
+
+	err := db.Bun().NewSelect().ColumnExpr("alloc_acc.*").
+		Column("resource_pool").
 		TableExpr("allocation_accelerators alloc_acc").
 		Join("LEFT JOIN allocations alloc ON alloc_acc.allocation_id = alloc.allocation_id").
+		Order("start_time DESC").
 		Where("alloc.task_id = ?", req.TaskId).Scan(ctx, &res)
 	if err != nil {
 		return nil, fmt.Errorf("querying allocation accelerators: %w", err)
 	}
 
 	var accelerationData []*apiv1.AcceleratorData
 	for _, r := range res {
-		accelerationData = append(accelerationData, r.Proto())
+		accelerationData = append(accelerationData, &apiv1.AcceleratorData{
+			ContainerId:      r.ContainerID,
+			AllocationId:     string(r.AllocationID),
+			NodeName:         r.NodeName,
+			AcceleratorType:  r.AcceleratorType,
+			AcceleratorUuids: r.AcceleratorUuids,
+			ResourcePool:     r.ResourcePool,
+		})
 	}
 	return &apiv1.GetTaskAcceleratorDataResponse{AcceleratorData: accelerationData}, nil
 }

diff --git a/master/internal/api_tasks_intg_test.go b/master/internal/api_tasks_intg_test.go
@@ -423,4 +423,6 @@ func TestGetAllocationAcceleratorData(t *testing.T) {
 	require.Len(t, resp.AcceleratorData, 1, "incorrect number of allocation accelerator data returned")
 	require.Equal(t, resp.AcceleratorData[0].AllocationId,
 		aID1.String(), "failed to get the correct allocation's accelerator data")
+	require.Equal(t, resp.AcceleratorData[0].ResourcePool,
+		a1.ResourcePool, "failed to get the correct allocation's resource pool data")
 }
diff --git a/proto/pkg/apiv1/trial.pb.go b/proto/pkg/apiv1/trial.pb.go
diff --git a/proto/src/determined/api/v1/trial.proto b/proto/src/determined/api/v1/trial.proto
@@ -595,6 +595,8 @@ message AcceleratorData {
   string accelerator_type = 5;
   // An array of UUIDs of the accelerators associated with the allocation.
   repeated string accelerator_uuids = 6;
+  // The name of the resource pool.
+  string resource_pool = 7;
 }
 
 // Arguments to an all gather.

diff --git a/webui/react/src/pages/ResourcePool/Topology.module.scss b/webui/react/src/pages/ResourcePool/Topology.module.scss
@@ -10,52 +10,51 @@
   display: flex;
   flex-flow: row wrap;
   justify-content: flex-start;
+}
+.node {
+  align-items: flex-start;
+  color: var(--theme-status-active-strong);
+  display: flex;
+  flex-direction: column;
+  margin-bottom: 20px;
+  margin-right: 30px;
 
-  .node {
-    align-items: flex-start;
-    color: var(--theme-status-active-strong);
+  .nodeName {
+    display: inline-block;
+    font-size: var(--single-line-size-l);
+    margin-bottom: 6px;
+    overflow: hidden;
+    text-overflow: ellipsis;
+    text-transform: uppercase;
+    vertical-align: bottom;
+    white-space: nowrap;
+  }
+  .nodeCluster {
+    align-items: center;
+    border: 2px solid var(--theme-status-active-strong);
+    box-sizing: border-box;
     display: flex;
-    flex-direction: column;
-    margin-bottom: 20px;
-    margin-right: 30px;
+    justify-content: space-around;
+    padding: 6px;
 
-    .nodeName {
-      display: inline-block;
-      font-size: var(--single-line-size-l);
-      margin-bottom: 6px;
-      overflow: hidden;
-      text-overflow: ellipsis;
-      text-transform: uppercase;
-      vertical-align: bottom;
-      white-space: nowrap;
-    }
-    .nodeCluster {
-      align-items: center;
-      border: 2px solid var(--theme-status-active-strong);
-      box-sizing: border-box;
-      display: flex;
-      justify-content: space-around;
-      padding: 6px;
-
-      .nodeSlot {
-        background-color: var(--theme-status-inactive-weak);
-        height: 20px;
-        margin-right: 6px;
-        width: 6px;
+    .nodeSlot {
+      background-color: var(--theme-status-inactive-weak);
+      height: 20px;
+      margin-right: 6px;
+      width: 6px;
 
-        &:last-of-type {
-          margin-right: 0;
-        }
-      }
-      .singleSlot {
-        width: 90px;
-      }
-      .coupleSlot {
-        width: 42px;
-      }
-      .active {
-        background-color: var(--theme-status-active-strong);
+      &:last-of-type {
+        margin-right: 0;
       }
     }
+    .singleSlot {
+      width: 90px;
+    }
+    .coupleSlot {
+      width: 42px;
+    }
+    .active {
+      background-color: var(--theme-status-active-strong);
+    }
   }
 }
diff --git a/webui/react/src/pages/ResourcePool/Topology.tsx b/webui/react/src/pages/ResourcePool/Topology.tsx
@@ -1,31 +1,30 @@
 import Tooltip from 'hew/Tooltip';
+import { range } from 'lodash';
 import React, { PropsWithChildren, useEffect, useMemo, useRef, useState } from 'react';
 
 import Section from 'components/Section';
-import { Agent, Resource, SlotsRecord } from 'types';
+import { Agent } from 'types';
 
 import css from './Topology.module.scss';
 
 interface NodeElementProps {
   name: string;
-  resources: Resource[];
-  slots?: SlotsRecord;
+  numOfSlots: number;
 }
 
 interface Props {
   nodes: Agent[];
 }
 
-const NodeElement: React.FC<PropsWithChildren<NodeElementProps>> = ({ name, slots, resources }) => {
+export const NodeElement: React.FC<PropsWithChildren<NodeElementProps>> = ({
+  name,
+  numOfSlots,
+}) => {
   const [containerWidth, setContainerWidth] = useState(0);
   const shouldTruncate = useMemo(() => name.length > 5, [name]);
   const slotsContainer = useRef<HTMLSpanElement>(null);
-  const slotsData = useMemo(
-    () => (slots !== undefined ? Object.values(slots) : resources),
-    [slots, resources],
-  );
-  const singleSlot = slotsData.length === 1;
-  const coupleSlot = slotsData.length === 2;
+  const singleSlot = numOfSlots === 1;
+  const coupleSlot = numOfSlots === 2;
   const styles = [css.nodeSlot];
 
   if (singleSlot) styles.push(css.singleSlot);
@@ -47,11 +46,8 @@ const NodeElement: React.FC<PropsWithChildren<NodeElementProps>> = ({ name, slot
         <span className={css.nodeName}>{name}</span>
       )}
       <span className={css.nodeCluster} ref={slotsContainer}>
-        {slotsData.map(({ container }, idx) => (
-          <span
-            className={`${styles.join(' ')} ${container ? css.active : ''}`}
-            key={`slot${idx}`}
-          />
+        {range(numOfSlots).map((idx) => (
+          <span className={`${styles.join(' ')} ${css.active}`} key={`slot${idx}`} />
         ))}
       </span>
     </div>
@@ -63,7 +59,13 @@ const Topology: React.FC<PropsWithChildren<Props>> = ({ nodes }) => {
     <Section title="Topology">
       <div className={`${css.mainContainer} ${css.nodesContainer}`}>
         {nodes.map(({ id, resources, slots }) => {
-          return <NodeElement key={id} name={id} resources={resources} slots={slots} />;
+          return (
+            <NodeElement
+              key={id}
+              name={id}
+              numOfSlots={(slots !== undefined ? Object.values(slots) : resources).length}
+            />
+          );
         })}
       </div>
     </Section>

diff --git a/webui/react/src/pages/TrialDetails/TrialInfoBox.module.scss b/webui/react/src/pages/TrialDetails/TrialInfoBox.module.scss
@@ -4,3 +4,19 @@
     flex-direction: column;
   }
 }
+.dataRow {
+  display: flex;
+  justify-content: space-between;
+  width: 100%;
+
+  .dashline {
+    border-bottom: 1px dashed var(--theme-ix-border-weak);
+    flex-grow: 1;
+    margin: 4px 8px;
+  }
+  .pools {
+    a {
+      margin-left: 4px;
+    }
+  }
+}
diff --git a/webui/react/src/pages/TrialDetails/TrialInfoBox.test.tsx b/webui/react/src/pages/TrialDetails/TrialInfoBox.test.tsx
@@ -4,6 +4,7 @@ import { ConfirmationProvider } from 'hew/useConfirm';
 
 import {} from 'stores/cluster';
 
+import { ThemeProvider } from 'components/ThemeProvider';
 import { ExperimentBase, TrialDetails } from 'types';
 
 import TrialInfoBox from './TrialInfoBox';
@@ -12,9 +13,11 @@ vi.useFakeTimers();
 const setup = (trial: TrialDetails, experiment: ExperimentBase) => {
   render(
     <UIProvider theme={DefaultTheme.Light}>
-      <ConfirmationProvider>
-        <TrialInfoBox experiment={experiment} trial={trial} />
-      </ConfirmationProvider>
+      <ThemeProvider>
+        <ConfirmationProvider>
+          <TrialInfoBox experiment={experiment} trial={trial} />
+        </ConfirmationProvider>
+      </ThemeProvider>
     </UIProvider>,
   );
 };