refactor: centralize resource slot type definitions and improve devic…

…e metadata handling in `/react` (#2684) ### TL;DR Added support for dynamic accelerator types and improved resource allocation handling. ### What changed? - Introduced `ResourceSlotName` type to handle known and unknown accelerator types - Updated `ResourceTypeIcon` component to use dynamic resource slot details - Modified `ResourceAllocationFormItems` to use resource slot details for display units - Added JSON schema validation for device metadata - Implemented `isMatchingMaxPerContainer` function for flexible config matching - Updated resource limit calculations to support unknown accelerator types ### The steps to add a new accelerator type 1. Add the accelerator type to `knownAcceleratorResourceSlotNames` in `react/src/hooks/backendai.tsx`. 2. Add an icon for the new accelerator type to `resourceTypeIconSrcMap` in `react/src/components/ResourceNumber.tsx`. 3. Add new accelerator metadata to `resources/device_metadata.json`. After following only step 1, you got a TypeScript error related to step 2 and Jest test failures related to step 3. ### How to test? 1. Run the JSON schema validation test for device metadata 2. Test the resource allocation form with various accelerator types 3. Verify that unknown accelerator types are handled correctly in the UI 4. Check if the resource limits are calculated correctly for all accelerator types ### Why make this change? This change improves the flexibility and maintainability of the codebase by: 1. Supporting dynamic accelerator types without hardcoding 2. Ensuring consistency between device metadata and code 3. Improving resource limit calculations for various accelerator types 4. Enhancing the user interface to display correct units for different accelerators These improvements allow for easier addition of new accelerator types and better handling of resource allocation across different devices.
lablup · Sep 5, 2024 · d52c682 · d52c682
1 parent 95093ee
commit d52c682
Show file tree

Hide file tree

Showing 12 changed files with 218 additions and 101 deletions.
diff --git a/react/package.json b/react/package.json
@@ -32,7 +32,7 @@
     "jotai": "^2.9.2",
     "jotai-effect": "^1.0.0",
     "lodash": "^4.17.21",
-    "lucide-react": "^0.383.0",
+    "lucide-react": "^0.438.0",
     "markdown-to-jsx": "^7.4.7",
     "prettier": "^3.3.3",
     "prism-react-renderer": "^2.3.1",
@@ -109,6 +109,7 @@
     "@types/relay-runtime": "^14.1.24",
     "@types/relay-test-utils": "^14.1.4",
     "@types/uuid": "^9.0.8",
+    "ajv": "^8.17.1",
     "babel-plugin-named-exports-order": "^0.0.2",
     "babel-plugin-relay": "^16.2.0",
     "eslint": "^8.57.0",

diff --git a/react/pnpm-lock.yaml b/react/pnpm-lock.yaml
diff --git a/react/src/components/ResourceAllocationFormItems.tsx b/react/src/components/ResourceAllocationFormItems.tsx
@@ -24,7 +24,6 @@ import {
 } from './ImageEnvironmentSelectFormItems';
 import InputNumberWithSlider from './InputNumberWithSlider';
 import ResourceGroupSelectForCurrentProject from './ResourceGroupSelectForCurrentProject';
-import { ACCELERATOR_UNIT_MAP } from './ResourceNumber';
 import ResourcePresetSelect from './ResourcePresetSelect';
 import { CaretDownOutlined } from '@ant-design/icons';
 import {
@@ -928,7 +927,7 @@ const ResourceAllocationFormItems: React.FC<
                             },
                             tooltip: {
                               formatter: (value = 0) => {
-                                return `${value} ${ACCELERATOR_UNIT_MAP[currentAcceleratorType]}`;
+                                return `${value} ${resourceSlotsDetails?.[currentAcceleratorType]?.display_unit || ''}`;
                               },
                               open:
                                 currentImageAcceleratorLimits.length <= 0
@@ -992,7 +991,8 @@ const ResourceAllocationFormItems: React.FC<
                                       return {
                                         value: name,
                                         label:
-                                          ACCELERATOR_UNIT_MAP[name] || 'UNIT',
+                                          resourceSlotsDetails?.[name]
+                                            ?.display_unit || 'UNIT',
                                         disabled:
                                           currentImageAcceleratorLimits.length >
                                             0 &&

diff --git a/react/src/components/ResourceNumber.tsx b/react/src/components/ResourceNumber.tsx
@@ -1,32 +1,22 @@
 import { iSizeToSize } from '../helper';
-import { useResourceSlotsDetails } from '../hooks/backendai';
+import {
+  BaseResourceSlotName,
+  KnownAcceleratorResourceSlotName,
+  ResourceSlotName,
+  useResourceSlotsDetails,
+} from '../hooks/backendai';
 import { useCurrentResourceGroupValue } from '../hooks/useCurrentProject';
 import Flex from './Flex';
 import { Tooltip, Typography, theme } from 'antd';
 import _ from 'lodash';
+import { MicrochipIcon } from 'lucide-react';
 import React, { ReactElement } from 'react';
-import { useTranslation } from 'react-i18next';
-
-export const ACCELERATOR_UNIT_MAP: {
-  [key: string]: string;
-} = {
-  'cuda.device': 'GPU',
-  'cuda.shares': 'FGPU',
-  'rocm.device': 'GPU',
-  'tpu.device': 'TPU',
-  'ipu.device': 'IPU',
-  'atom.device': 'ATOM',
-  'atom-plus.device': 'ATOM+',
-  'gaudi2.device': 'Gaudi 2',
-  'warboy.device': 'Warboy',
-  'hyperaccel-lpu.device': 'Hyperaccel LPU',
-};
 
 export type ResourceOpts = {
   shmem?: number;
 };
 interface ResourceNumberProps {
-  type: string;
+  type: ResourceSlotName | string;
   extra?: ReactElement;
   opts?: ResourceOpts;
   value: string;
@@ -35,7 +25,9 @@ interface ResourceNumberProps {
 }
 
 type ResourceTypeInfo<V> = {
-  [key in string]: V;
+  [key in KnownAcceleratorResourceSlotName]: V;
+} & {
+  [key in BaseResourceSlotName]: V;
 };
 const ResourceNumber: React.FC<ResourceNumberProps> = ({
   type,
@@ -112,7 +104,7 @@ const MWCIconWrap: React.FC<{ size?: number; children: string }> = ({
 };
 interface AccTypeIconProps
   extends Omit<React.ImgHTMLAttributes<HTMLImageElement>, 'src'> {
-  type: string;
+  type: ResourceSlotName | string;
   showIcon?: boolean;
   showUnit?: boolean;
   showTooltip?: boolean;
@@ -126,33 +118,27 @@ export const ResourceTypeIcon: React.FC<AccTypeIconProps> = ({
   showTooltip = true,
   ...props
 }) => {
-  const { t } = useTranslation();
-
-  const resourceTypeIconSrcMap: ResourceTypeInfo<
-    [ReactElement | string, string]
-  > = {
-    cpu: [
-      <MWCIconWrap size={size}>developer_board</MWCIconWrap>,
-      t('session.core'),
-    ],
-    mem: [<MWCIconWrap size={size}>memory</MWCIconWrap>, 'GiB'],
-    'cuda.device': ['/resources/icons/file_type_cuda.svg', 'GPU'],
-    'cuda.shares': ['/resources/icons/file_type_cuda.svg', 'FGPU'],
-    'rocm.device': ['/resources/icons/rocm.svg', 'GPU'],
-    'tpu.device': [<MWCIconWrap size={size}>view_module</MWCIconWrap>, 'TPU'],
-    'ipu.device': [<MWCIconWrap size={size}>view_module</MWCIconWrap>, 'IPU'],
-    'atom.device': ['/resources/icons/rebel.svg', 'ATOM'],
-    'atom-plus.device': ['/resources/icons/rebel.svg', 'ATOM+'],
-    'gaudi2.device': ['/resources/icons/gaudi.svg', 'Gaudi 2'],
-    'warboy.device': ['/resources/icons/furiosa.svg', 'Warboy'],
-    'hyperaccel-lpu.device': [
-      '/resources/icons/npu_generic.svg',
-      'Hyperaccel LPU',
-    ],
+  const resourceTypeIconSrcMap: ResourceTypeInfo<ReactElement | string> = {
+    cpu: <MWCIconWrap size={size}>developer_board</MWCIconWrap>,
+    mem: <MWCIconWrap size={size}>memory</MWCIconWrap>,
+    'cuda.device': '/resources/icons/file_type_cuda.svg',
+    'cuda.shares': '/resources/icons/file_type_cuda.svg',
+    'rocm.device': '/resources/icons/rocm.svg',
+    'tpu.device': <MWCIconWrap size={size}>view_module</MWCIconWrap>,
+    'ipu.device': <MWCIconWrap size={size}>view_module</MWCIconWrap>,
+    'atom.device': '/resources/icons/rebel.svg',
+    'atom-plus.device': '/resources/icons/rebel.svg',
+    'gaudi2.device': '/resources/icons/gaudi.svg',
+    'warboy.device': '/resources/icons/furiosa.svg',
+    'hyperaccel-lpu.device': '/resources/icons/npu_generic.svg',
   };
 
+  const targetIcon = resourceTypeIconSrcMap[
+    type as KnownAcceleratorResourceSlotName
+  ] ?? <MicrochipIcon />;
+
   const content =
-    typeof resourceTypeIconSrcMap[type]?.[0] === 'string' ? (
+    typeof targetIcon === 'string' ? (
       <img
         {...props}
         style={{
@@ -161,17 +147,14 @@ export const ResourceTypeIcon: React.FC<AccTypeIconProps> = ({
           ...(props.style || {}),
         }}
         // @ts-ignore
-        src={resourceTypeIconSrcMap[type]?.[0] || ''}
+        src={resourceTypeIconSrcMap[type] || ''}
         alt={type}
       />
     ) : (
-      <Flex style={{ width: 16, height: 16 }}>
-        {resourceTypeIconSrcMap[type]?.[0] || type}
-      </Flex>
+      <Flex style={{ width: 16, height: 16 }}>{targetIcon || type}</Flex>
     );
 
   return showTooltip ? (
-    // <Tooltip title={showTooltip ? `${type} (${resourceTypeIconSrcMap[type][1]})` : undefined}>
     <Tooltip title={type}>{content}</Tooltip>
   ) : (
     <Flex style={{ pointerEvents: 'none' }}>{content}</Flex>

diff --git a/react/src/components/ResourcePresetSelect.tsx b/react/src/components/ResourcePresetSelect.tsx
@@ -1,6 +1,6 @@
 import { localeCompare } from '../helper';
 import { useUpdatableState } from '../hooks';
-import { useResourceSlots } from '../hooks/backendai';
+import { ResourceSlotName, useResourceSlots } from '../hooks/backendai';
 import useControllableState from '../hooks/useControllableState';
 import Flex from './Flex';
 import ResourceNumber from './ResourceNumber';
@@ -123,7 +123,7 @@ const ResourcePresetSelect: React.FC<ResourcePresetSelectProps> = ({
           // @ts-ignore
           options: _.map(resource_presets, (preset, index) => {
             const slotsInfo: {
-              [key in string]: string;
+              [key in ResourceSlotName]: string;
             } = JSON.parse(preset?.resource_slots);
             const disabled = allocatablePresetNames
               ? !allocatablePresetNames.includes(preset?.name || '')
@@ -145,7 +145,7 @@ const ResourcePresetSelect: React.FC<ResourcePresetSelectProps> = ({
                   >
                     {_.map(
                       _.omitBy(slotsInfo, (slot, key) =>
-                        _.isEmpty(resourceSlots[key]),
+                        _.isEmpty(resourceSlots[key as ResourceSlotName]),
                       ),
                       (slot, key) => {
                         return (

diff --git a/react/src/components/ServiceLauncherPageContent.tsx b/react/src/components/ServiceLauncherPageContent.tsx
@@ -9,6 +9,7 @@ import {
   useSuspendedBackendaiClient,
   useWebUINavigate,
 } from '../hooks';
+import { KnownAcceleratorResourceSlotName } from '../hooks/backendai';
 import { useSuspenseTanQuery, useTanMutation } from '../hooks/reactQueryAlias';
 import BAIModal, { DEFAULT_BAI_MODAL_Z_INDEX } from './BAIModal';
 import EnvVarFormList, { EnvVarFormListValue } from './EnvVarFormList';
@@ -56,20 +57,12 @@ interface ServiceCreateConfigResourceOptsType {
   shmem?: number | string;
 }
 
-interface ServiceCreateConfigResourceType {
+type ServiceCreateConfigResourceType = {
   cpu: number | string;
   mem: string;
-  'cuda.device'?: number | string;
-  'cuda.shares'?: number | string;
-  'rocm.device'?: number | string;
-  'tpu.device'?: number | string;
-  'ipu.device'?: number | string;
-  'atom.device'?: number | string;
-  'gaudi2.device'?: number | string;
-  'atom-plus.device'?: number | string;
-  'warboy.device'?: number | string;
-  'hyperaccel-lpu.device'?: number | string;
-}
+} & {
+  [key in KnownAcceleratorResourceSlotName]?: number | string;
+};
 export interface MountOptionType {
   mount_destination?: string;
   type?: string;

diff --git a/react/src/hooks/backendai.test.tsx b/react/src/hooks/backendai.test.tsx
@@ -0,0 +1,30 @@
+// import json
+import deviceMetadata from '../../../resources/device_metadata.json';
+import schema from '../../../resources/device_metadata.schema.json';
+import {
+  knownAcceleratorResourceSlotNames,
+  baseResourceSlotNames,
+} from './backendai';
+import Ajv from 'ajv';
+import _ from 'lodash';
+
+const ajv = new Ajv();
+
+describe('deviceMetadata JSON Schema Validation', () => {
+  it('should include all known accelerator names and base resource names', () => {
+    const strictSchema = _.cloneDeep(schema);
+
+    // @ts-ignore
+    strictSchema.properties.deviceInfo.required = [
+      ...knownAcceleratorResourceSlotNames,
+      ...baseResourceSlotNames,
+    ];
+
+    const validate = ajv.compile(strictSchema);
+    const valid = validate(deviceMetadata);
+    if (!valid) {
+      console.error(validate.errors);
+    }
+    expect(valid).toBe(true);
+  });
+});
diff --git a/react/src/hooks/backendai.tsx b/react/src/hooks/backendai.tsx
@@ -8,6 +8,26 @@ import {
 import _ from 'lodash';
 import { useEffect, useState } from 'react';
 
+export const baseResourceSlotNames = ['cpu', 'mem'] as const;
+export type BaseResourceSlotName = (typeof baseResourceSlotNames)[number];
+export const knownAcceleratorResourceSlotNames = [
+  'cuda.device',
+  'cuda.shares',
+  'rocm.device',
+  'tpu.device',
+  'ipu.device',
+  'atom.device',
+  'atom-plus.device',
+  'gaudi2.device',
+  'warboy.device',
+  'hyperaccel-lpu.device',
+] as const;
+export type KnownAcceleratorResourceSlotName =
+  (typeof knownAcceleratorResourceSlotNames)[number];
+
+export type ResourceSlotName =
+  | BaseResourceSlotName
+  | KnownAcceleratorResourceSlotName;
 export interface QuotaScope {
   id: string;
   quota_scope_id: string;
@@ -23,19 +43,7 @@ export const useResourceSlots = () => {
   const [key, checkUpdate] = useUpdatableState('first');
   const baiClient = useSuspendedBackendaiClient();
   const { data: resourceSlots } = useSuspenseTanQuery<{
-    cpu?: string;
-    mem?: string;
-    'cuda.shares'?: string;
-    'cuda.device'?: string;
-    'rocm.device'?: string;
-    'tpu.device'?: string;
-    'ipu.device'?: string;
-    'atom.device'?: string;
-    'atom-plus.device'?: string;
-    'gaudi2.device'?: string;
-    'warboy.device'?: string;
-    'hyperaccel-lpu.device'?: string;
-    [key: string]: string | undefined;
+    [key in ResourceSlotName]?: string;
   }>({
     queryKey: ['useResourceSlots', key],
     queryFn: () => {