refactor: centralize resource slot type definitions and improve device metadata handling in /react

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",

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 &&

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>

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 (

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;

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);
+  });
+});

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: () => {