diff --git a/react/package.json b/react/package.json index 9605d10be2..ae6d392e82 100644 --- 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 index 02faf2a342..28f782ead0 100644 --- a/react/pnpm-lock.yaml +++ b/react/pnpm-lock.yaml @@ -96,8 +96,8 @@ importers: specifier: ^4.17.21 version: 4.17.21 lucide-react: - specifier: ^0.383.0 - version: 0.383.0(react@18.3.1) + specifier: ^0.438.0 + version: 0.438.0(react@18.3.1) markdown-to-jsx: specifier: ^7.4.7 version: 7.4.7(react@18.3.1) @@ -219,6 +219,9 @@ importers: '@types/uuid': specifier: ^9.0.8 version: 9.0.8 + ajv: + specifier: ^8.17.1 + version: 8.17.1 babel-plugin-named-exports-order: specifier: ^0.0.2 version: 0.0.2 @@ -5617,10 +5620,10 @@ packages: lru-cache@5.1.1: resolution: {integrity: sha512-KpNARQA3Iwv+jTA0utUVVbrh+Jlrr1Fv0e56GGzAFOXN7dk/FviaDW8LHmK52DlcH4WP2n6gI8vN1aesBFgo9w==} - lucide-react@0.383.0: - resolution: {integrity: sha512-13xlG0CQCJtzjSQYwwJ3WRqMHtRj3EXmLlorrARt7y+IHnxUCp3XyFNL1DfaGySWxHObDvnu1u1dV+0VMKHUSg==} + lucide-react@0.438.0: + resolution: {integrity: sha512-uq6yCB+IzVfgIPMK8ibkecXSWTTSOMs9UjUgZigfrDCVqgdwkpIgYg1fSYnf0XXF2AoSyCJZhoZXQwzoai7VGw==} peerDependencies: - react: ^16.5.1 || ^17.0.0 || ^18.0.0 + react: ^16.5.1 || ^17.0.0 || ^18.0.0 || ^19.0.0-rc lz-string@1.5.0: resolution: {integrity: sha512-h5bgJWpxJNswbU7qCrV0tIKQCaS3blPDrqKWx+QxzuzL1zGUzij9XCWLrSLsJPu5t+eWA/ycetzYAO5IOMcWAQ==} @@ -15641,7 +15644,7 @@ snapshots: dependencies: yallist: 3.1.1 - lucide-react@0.383.0(react@18.3.1): + lucide-react@0.438.0(react@18.3.1): dependencies: react: 18.3.1 diff --git a/react/src/components/ResourceAllocationFormItems.tsx b/react/src/components/ResourceAllocationFormItems.tsx index 5d52127d94..7d06023e48 100644 --- 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 index 9670535791..56d04c6010 100644 --- 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 = { - [key in string]: V; + [key in KnownAcceleratorResourceSlotName]: V; +} & { + [key in BaseResourceSlotName]: V; }; const ResourceNumber: React.FC = ({ type, @@ -112,7 +104,7 @@ const MWCIconWrap: React.FC<{ size?: number; children: string }> = ({ }; interface AccTypeIconProps extends Omit, 'src'> { - type: string; + type: ResourceSlotName | string; showIcon?: boolean; showUnit?: boolean; showTooltip?: boolean; @@ -126,33 +118,27 @@ export const ResourceTypeIcon: React.FC = ({ showTooltip = true, ...props }) => { - const { t } = useTranslation(); - - const resourceTypeIconSrcMap: ResourceTypeInfo< - [ReactElement | string, string] - > = { - cpu: [ - developer_board, - t('session.core'), - ], - mem: [memory, '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': [view_module, 'TPU'], - 'ipu.device': [view_module, '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 = { + cpu: developer_board, + mem: memory, + 'cuda.device': '/resources/icons/file_type_cuda.svg', + 'cuda.shares': '/resources/icons/file_type_cuda.svg', + 'rocm.device': '/resources/icons/rocm.svg', + 'tpu.device': view_module, + 'ipu.device': view_module, + '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 + ] ?? ; + const content = - typeof resourceTypeIconSrcMap[type]?.[0] === 'string' ? ( + typeof targetIcon === 'string' ? ( = ({ ...(props.style || {}), }} // @ts-ignore - src={resourceTypeIconSrcMap[type]?.[0] || ''} + src={resourceTypeIconSrcMap[type] || ''} alt={type} /> ) : ( - - {resourceTypeIconSrcMap[type]?.[0] || type} - + {targetIcon || type} ); return showTooltip ? ( - // {content} ) : ( {content} diff --git a/react/src/components/ResourcePresetSelect.tsx b/react/src/components/ResourcePresetSelect.tsx index fc12bc1ad3..e293a4a182 100644 --- 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 = ({ // @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 = ({ > {_.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 index f19f7c141d..6aa9446216 100644 --- 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 new file mode 100644 index 0000000000..842b942329 --- /dev/null +++ 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 index 50d9dc42d6..63eb9bf381 100644 --- 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: () => { diff --git a/react/src/hooks/useResourceLimitAndRemaining.test.ts b/react/src/hooks/useResourceLimitAndRemaining.test.ts new file mode 100644 index 0000000000..7fb3a5e0ff --- /dev/null +++ b/react/src/hooks/useResourceLimitAndRemaining.test.ts @@ -0,0 +1,49 @@ +import { isMatchingMaxPerContainer } from './useResourceLimitAndRemaining'; +import exp from 'constants'; + +describe('getConfigName', () => { + test('should match unknown devices', () => { + expect( + isMatchingMaxPerContainer('maxCUDADevicesPerContainer', 'cuda.device'), + ).toBe(true); + expect( + isMatchingMaxPerContainer('maxCUDASharesPerContainer', 'cuda.shares'), + ).toBe(true); + expect( + isMatchingMaxPerContainer('maxROCMDevicesPerContainer', 'rocm.device'), + ).toBe(true); + expect( + isMatchingMaxPerContainer('maxTPUDevicesPerContainer', 'tpu.device'), + ).toBe(true); + expect( + isMatchingMaxPerContainer('maxIPUDevicesPerContainer', 'ipu.device'), + ).toBe(true); + expect( + isMatchingMaxPerContainer('maxATOMDevicesPerContainer', 'atom.device'), + ).toBe(true); + expect( + isMatchingMaxPerContainer( + 'maxATOMPLUSDevicesPerContainer', + 'atom-plus.device', + ), + ).toBe(true); + expect( + isMatchingMaxPerContainer( + 'maxGaudi2DevicesPerContainer', + 'gaudi2.device', + ), + ).toBe(true); + expect( + isMatchingMaxPerContainer( + 'maxWarboyDevicesPerContainer', + 'warboy.device', + ), + ).toBe(true); + expect( + isMatchingMaxPerContainer( + 'maxHyperaccelLPUDevicesPerContainer', + 'hyperaccel-lpu.device', + ), + ).toBe(true); + }); +}); diff --git a/react/src/hooks/useResourceLimitAndRemaining.tsx b/react/src/hooks/useResourceLimitAndRemaining.tsx index 8416102579..4da0a17674 100644 --- a/react/src/hooks/useResourceLimitAndRemaining.tsx +++ b/react/src/hooks/useResourceLimitAndRemaining.tsx @@ -5,7 +5,24 @@ import { addNumberWithUnits, iSizeToSize } from '../helper'; import { useResourceSlots } from '../hooks/backendai'; import { useSuspenseTanQuery } from './reactQueryAlias'; import _ from 'lodash'; +import { useMemo } from 'react'; +const maxPerContainerRegex = /^max([A-Za-z0-9]+)PerContainer$/; + +export const isMatchingMaxPerContainer = (configName: string, key: string) => { + const match = configName.match(maxPerContainerRegex); + if (match) { + const configLowerCase = match[1].toLowerCase(); + const keyLowerCase = key.replaceAll(/[.-]/g, '').toLowerCase(); + // Because some accelerator names are not the same as the config name, we need to check if the config name is a substring of the accelerator name + // cuda.shares => maxCUDASharesPerContainer + // cuda.device => maxCUDADevicesPerContainer (Not maxCUDADevicePerContainer) + return ( + configLowerCase === keyLowerCase || configLowerCase === keyLowerCase + 's' + ); + } + return false; +}; export interface MergedResourceLimits { accelerators: { [key: string]: @@ -195,6 +212,14 @@ export const useResourceLimitAndRemaining = ({ }, ), }; + const perContainerConfigs = useMemo( + () => + _.omitBy(baiClient._config, (value, key) => { + return !maxPerContainerRegex.test(key); + }), + [baiClient._config], + ); + const resourceLimits: MergedResourceLimits = { cpu: resourceSlots?.cpu === undefined @@ -265,19 +290,11 @@ export const useResourceLimitAndRemaining = ({ accelerators: _.reduce( acceleratorSlots, (result, value, key) => { - const configName = - { - 'cuda.device': 'maxCUDADevicesPerContainer', - 'cuda.shares': 'maxCUDASharesPerContainer', - 'rocm.device': 'maxROCMDevicesPerContainer', - 'tpu.device': 'maxTPUDevicesPerContainer', - 'ipu.device': 'maxIPUDevicesPerContainer', - 'atom.device': 'maxATOMDevicesPerContainer', - 'atom-plus.device': 'maxATOMPlusDevicesPerContainer', - 'gaudi2.device': 'maxGaudi2DevicesPerContainer', - 'warboy.device': 'maxWarboyDevicesPerContainer', - 'hyperaccel-lpu.device': 'maxHyperaccelLPUDevicesPerContainer', // FIXME: add maxLPUDevicesPerContainer to config - }[key] || 'cuda.device'; // FIXME: temporally `cuda.device` config, when undefined + const perContainerLimit = + _.find(perContainerConfigs, (configValue, configName) => { + return isMatchingMaxPerContainer(configName, key); + }) ?? baiClient._config['cuda.device']; // FIXME: temporally `cuda.device` config, when undefined + result[key] = { min: parseInt( _.filter( @@ -288,7 +305,7 @@ export const useResourceLimitAndRemaining = ({ )?.[0]?.min || '0', ), max: _.min([ - baiClient._config[configName] || 8, + perContainerLimit || 8, // scaling group all cpu (using + remaining), string type resourceGroupResourceSize.accelerators[key], ]), diff --git a/resources/device_metadata.json b/resources/device_metadata.json index 6650dbf3bb..2c41bef2d6 100644 --- a/resources/device_metadata.json +++ b/resources/device_metadata.json @@ -1,4 +1,5 @@ { + "$schema": "./device_metadata.schema.json", "deviceInfo": { "cpu": { "slot_name": "cpu", diff --git a/resources/device_metadata.schema.json b/resources/device_metadata.schema.json new file mode 100644 index 0000000000..4b35549ada --- /dev/null +++ b/resources/device_metadata.schema.json @@ -0,0 +1,32 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "type": "object", + "properties": { + "deviceInfo": { + "type": "object", + "patternProperties": { + "^[a-zA-Z0-9.-]+$": { + "type": "object", + "properties": { + "slot_name": { "type": "string" }, + "description": { "type": "string" }, + "human_readable_name": { "type": "string" }, + "display_unit": { "type": "string" }, + "number_format": { + "type": "object", + "properties": { + "binary": { "type": "boolean" }, + "round_length": { "type": "integer" } + }, + "required": ["binary", "round_length"] + }, + "display_icon": { "type": "string" } + }, + "required": ["slot_name", "description", "human_readable_name", "display_unit", "number_format", "display_icon"] + } + }, + "additionalProperties": false + } + }, + "required": ["deviceInfo"] +} \ No newline at end of file