Skip to content

Commit

Permalink
Refactor: Remove all hard-coded resource slot names in /react
Browse files Browse the repository at this point in the history
  • Loading branch information
yomybaby committed Sep 4, 2024
1 parent 08487ce commit 3086120
Show file tree
Hide file tree
Showing 12 changed files with 209 additions and 92 deletions.
1 change: 1 addition & 0 deletions react/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
3 changes: 3 additions & 0 deletions react/pnpm-lock.yaml

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

6 changes: 3 additions & 3 deletions react/src/components/ResourceAllocationFormItems.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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 &&
Expand Down
76 changes: 29 additions & 47 deletions react/src/components/ResourceNumber.tsx
Original file line number Diff line number Diff line change
@@ -1,32 +1,21 @@
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 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;
Expand All @@ -35,7 +24,9 @@ interface ResourceNumberProps {
}

type ResourceTypeInfo<V> = {
[key in string]: V;
[key in KnownAcceleratorResourceSlotName]: V;
} & {
[key in BaseResourceSlotName]: V;
};
const ResourceNumber: React.FC<ResourceNumberProps> = ({
type,
Expand Down Expand Up @@ -112,7 +103,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;
Expand All @@ -126,33 +117,24 @@ 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 content =
typeof resourceTypeIconSrcMap[type]?.[0] === 'string' ? (
typeof resourceTypeIconSrcMap[type as KnownAcceleratorResourceSlotName] ===
'string' ? (
<img
{...props}
style={{
Expand All @@ -161,17 +143,17 @@ 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}
{resourceTypeIconSrcMap[type as KnownAcceleratorResourceSlotName] ||
type}
</Flex>
);

return showTooltip ? (
// <Tooltip title={showTooltip ? `${type} (${resourceTypeIconSrcMap[type][1]})` : undefined}>
<Tooltip title={type}>{content}</Tooltip>
) : (
<Flex style={{ pointerEvents: 'none' }}>{content}</Flex>
Expand Down
6 changes: 3 additions & 3 deletions react/src/components/ResourcePresetSelect.tsx
Original file line number Diff line number Diff line change
@@ -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';
Expand Down Expand Up @@ -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 || '')
Expand All @@ -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 (
Expand Down
17 changes: 5 additions & 12 deletions react/src/components/ServiceLauncherPageContent.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -57,20 +58,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;
Expand Down
30 changes: 30 additions & 0 deletions react/src/hooks/backendai.test.tsx
Original file line number Diff line number Diff line change
@@ -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', () => {

Check failure on line 14 in react/src/hooks/backendai.test.tsx

View workflow job for this annotation

GitHub Actions / Tests annotations (🧪 jest-coverage-report-action)

deviceMetadata JSON Schema Validation > should include all known accelerator names and base resource names

Error: expect(received).toBe(expected) // Object.is equality Expected: true Received: false at Object.<anonymous> (/home/runner/work/backend.ai-webui/backend.ai-webui/react/src/hooks/backendai.test.tsx:28:19) at Promise.then.completed (/home/runner/work/backend.ai-webui/backend.ai-webui/react/node_modules/.pnpm/jest-circus@27.5.1/node_modules/jest-circus/build/utils.js:391:28) at new Promise (<anonymous>) at callAsyncCircusFn (/home/runner/work/backend.ai-webui/backend.ai-webui/react/node_modules/.pnpm/jest-circus@27.5.1/node_modules/jest-circus/build/utils.js:316:10) at _callCircusTest (/home/runner/work/backend.ai-webui/backend.ai-webui/react/node_modules/.pnpm/jest-circus@27.5.1/node_modules/jest-circus/build/run.js:218:40) at processTicksAndRejections (node:internal/process/task_queues:95:5) at _runTest (/home/runner/work/backend.ai-webui/backend.ai-webui/react/node_modules/.pnpm/jest-circus@27.5.1/node_modules/jest-circus/build/run.js:155:3) at _runTestsForDescribeBlock (/home/runner/work/backend.ai-webui/backend.ai-webui/react/node_modules/.pnpm/jest-circus@27.5.1/node_modules/jest-circus/build/run.js:66:9) at _runTestsForDescribeBlock (/home/runner/work/backend.ai-webui/backend.ai-webui/react/node_modules/.pnpm/jest-circus@27.5.1/node_modules/jest-circus/build/run.js:60:9) at run (/home/runner/work/backend.ai-webui/backend.ai-webui/react/node_modules/.pnpm/jest-circus@27.5.1/node_modules/jest-circus/build/run.js:25:3) at runAndTransformResultsToJestFormat (/home/runner/work/backend.ai-webui/backend.ai-webui/react/node_modules/.pnpm/jest-circus@27.5.1/node_modules/jest-circus/build/legacy-code-todo-rewrite/jestAdapterInit.js:170:21) at jestAdapter (/home/runner/work/backend.ai-webui/backend.ai-webui/react/node_modules/.pnpm/jest-circus@27.5.1/node_modules/jest-circus/build/legacy-code-todo-rewrite/jestAdapter.js:82:19) at runTestInternal (/home/runner/work/backend.ai-webui/backend.ai-webui/react/node_modules/.pnpm/jest-runner@27.5.1/node_modules/jest-runner/build/runTest.js:389:16) at runTest (/home/runner/work/backend.ai-webui/backend.ai-webui/react/node_modules/.pnpm/jest-runner@27.5.1/node_modules/jest-runner/build/runTest.js:475:34) at Object.worker (/home/runner/work/backend.ai-webui/backend.ai-webui/react/node_modules/.pnpm/jest-runner@27.5.1/node_modules/jest-runner/build/testWorker.js:133:12)
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);
});
});
35 changes: 22 additions & 13 deletions react/src/hooks/backendai.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,27 @@ 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',
'new-unknown-device',
] as const;
export type KnownAcceleratorResourceSlotName =
(typeof knownAcceleratorResourceSlotNames)[number];

export type ResourceSlotName =
| BaseResourceSlotName
| KnownAcceleratorResourceSlotName;
export interface QuotaScope {
id: string;
quota_scope_id: string;
Expand All @@ -23,19 +44,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: () => {
Expand Down
49 changes: 49 additions & 0 deletions react/src/hooks/useResourceLimitAndRemaining.test.ts
Original file line number Diff line number Diff line change
@@ -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);
});
});
Loading

0 comments on commit 3086120

Please sign in to comment.