Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix(FR-5): Set accelerator value to 0 when no accelerator available #3027

Merged
merged 1 commit into from
Jan 15, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 28 additions & 0 deletions react/src/components/FormItemControl.tsx
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
import { Form, FormItemProps } from 'antd';
import _ from 'lodash';
import React from 'react';

interface FormItemControlProps extends FormItemProps {
render?: (args: {
value: unknown;
setValue: (value: unknown) => void;
}) => React.ReactNode;
}
const FormItemControl: React.FC<FormItemControlProps> = (props) => {
const form = Form.useFormInstance();
return (
// Only set the label and required props for the UI to display
<Form.Item
label={props.label}
required={props.required || _.some(props.rules, 'required')}
>
{props.render?.({
value: form.getFieldValue(props.name),
setValue: (value) => form.setFieldValue(props.name, value),
})}
<Form.Item {...props} noStyle></Form.Item>
</Form.Item>
);
};

export default FormItemControl;
123 changes: 80 additions & 43 deletions react/src/components/ResourceAllocationFormItems.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ import AgentSelect from './AgentSelect';
import BAISelect from './BAISelect';
import DynamicUnitInputNumberWithSlider from './DynamicUnitInputNumberWithSlider';
import Flex from './Flex';
// import FormItemControl from './FormItemControl';
import {
Image,
ImageEnvironmentFormInput,
Expand Down Expand Up @@ -124,35 +125,65 @@ const ResourceAllocationFormItems: React.FC<
form,
preserve: true,
});
const currentEnvironmentManual = Form.useWatch(['environments', 'manual'], {
form,
preserve: true,
});

const [{ currentImageMinM, remaining, resourceLimits, checkPresetInfo }] =
useResourceLimitAndRemaining({
currentProjectName: currentProject.name,
currentResourceGroup: currentResourceGroupInForm || undefined, // global currentResourceGroup can be null
currentImage: currentImage,
});

const { mergedResourceSlots, resourceSlotsInRG: resourceSlots } =
useResourceSlotsDetails(currentResourceGroupInForm || undefined);
const { mergedResourceSlots, resourceSlotsInRG } = useResourceSlotsDetails(
currentResourceGroupInForm || undefined,
);

const acceleratorSlots = _.omitBy(resourceSlots, (value, key) => {
if (['cpu', 'mem', 'shmem'].includes(key)) return true;
// When undefined, it means that the resourceSlots are not loaded yet.
const acceleratorSlots = resourceSlotsInRG
? _.omitBy(resourceSlotsInRG, (value, key) => {
if (['cpu', 'mem', 'shmem'].includes(key)) return true;

if (
!resourceLimits.accelerators[key]?.max ||
resourceLimits.accelerators[key]?.max === 0
)
return true;
return false;
});
if (
!resourceLimits.accelerators[key]?.max ||
resourceLimits.accelerators[key]?.max === 0
)
return true;
return false;
})
: undefined;

// When undefined, it means that the image is not determined yet.
const currentImageAcceleratorLimits = useMemo(
() =>
_.filter(currentImage?.resource_limits, (limit) =>
limit ? !_.includes(['cpu', 'mem', 'shmem'], limit.key) : false,
),
[currentImage?.resource_limits],
currentImage
? _.filter(currentImage?.resource_limits, (limit) =>
limit ? !_.includes(['cpu', 'mem', 'shmem'], limit.key) : false,
)
: undefined,
[currentImage],
);

// Disable accelerator input when there is no accelerator slot or no accelerator required in the selected image
// TODO: use `supported_accelerators` information from the image instead of `currentImageAcceleratorLimits` (FR-55)
const isAcceleratorInputDisabled =
(!_.isUndefined(acceleratorSlots) && _.isEmpty(acceleratorSlots)) ||
(currentImageAcceleratorLimits &&
currentImageAcceleratorLimits.length === 0 &&
_.isEmpty(currentEnvironmentManual));

useEffect(() => {
if (isAcceleratorInputDisabled) {
form.setFieldsValue({
resource: {
accelerator: 0,
},
});
}
}, [isAcceleratorInputDisabled, form]);

const sessionSliderLimitAndRemaining = {
min: 1,
max: sessionLimitAndRemaining.max,
Expand All @@ -174,7 +205,7 @@ const ResourceAllocationFormItems: React.FC<
]);
// If the current accelerator type is not available,
// change accelerator type to the first supported accelerator
const nextAcceleratorType: string = acceleratorSlots[currentAcceleratorType]
const nextAcceleratorType = acceleratorSlots?.[currentAcceleratorType]
? currentAcceleratorType
: _.keys(acceleratorSlots)[0];

Expand Down Expand Up @@ -203,7 +234,10 @@ const ResourceAllocationFormItems: React.FC<
};

// NOTE: accelerator value setting is done inside the conditional statement
if (currentImageAcceleratorLimits.length > 0) {
if (
currentImageAcceleratorLimits &&
currentImageAcceleratorLimits.length > 0
) {
if (
_.find(
currentImageAcceleratorLimits,
Expand All @@ -229,7 +263,7 @@ const ResourceAllocationFormItems: React.FC<
// resourceSlots returns "all resources enable to allocate(including AI accelerator)"
// imageAcceleratorLimit returns "all resources that is supported in the selected image"
_.filter(currentImageAcceleratorLimits, (acceleratorInfo: any) =>
_.keys(resourceSlots).includes(acceleratorInfo?.key),
_.keys(resourceSlotsInRG).includes(acceleratorInfo?.key),
)[0]?.key;

if (nextImageSelectorType) {
Expand Down Expand Up @@ -274,7 +308,11 @@ const ResourceAllocationFormItems: React.FC<
});

// set to 0 when currentImage doesn't support any AI accelerator
if (currentImage && currentImageAcceleratorLimits.length === 0) {
if (
currentImage &&
currentImageAcceleratorLimits &&
currentImageAcceleratorLimits.length === 0
) {
form.setFieldValue(['resource', 'accelerator'], 0);
}

Expand All @@ -295,7 +333,7 @@ const ResourceAllocationFormItems: React.FC<
checkPresetInfo?.presets,
(preset) => preset.name === name,
);
const slots = _.pick(preset?.resource_slots, _.keys(resourceSlots));
const slots = _.pick(preset?.resource_slots, _.keys(resourceSlotsInRG));
const mem = convertBinarySizeUnit(
(slots?.mem || 0) + 'b',
'g',
Expand Down Expand Up @@ -372,7 +410,7 @@ const ResourceAllocationFormItems: React.FC<
// if resourceSlots is loaded, update the form based on the resourceSlots
if (
currentAllocationPreset === 'auto-select' &&
!_.isUndefined(resourceSlots)
!_.isUndefined(resourceSlotsInRG)
) {
if (
_.includes(
Expand Down Expand Up @@ -413,7 +451,7 @@ const ResourceAllocationFormItems: React.FC<
}, [
currentAllocationPreset,
allocatablePresetNames,
resourceSlots,
resourceSlotsInRG,
form,
enableResourcePresets,
// below are functions wrapped by useEventNotStable
Expand Down Expand Up @@ -490,7 +528,7 @@ const ResourceAllocationFormItems: React.FC<
return (
// getFieldValue('allocationPreset') === 'custom' && (
<>
{resourceSlots?.cpu && (
{resourceSlotsInRG?.cpu && (
<Form.Item
name={['resource', 'cpu']}
// initialValue={0}
Expand Down Expand Up @@ -592,7 +630,7 @@ const ResourceAllocationFormItems: React.FC<
/>
</Form.Item>
)}
{resourceSlots?.mem && (
{resourceSlotsInRG?.mem && (
<Form.Item
label={t('session.launcher.Memory')}
tooltip={{
Expand Down Expand Up @@ -932,12 +970,6 @@ const ResourceAllocationFormItems: React.FC<
'resource',
'acceleratorType',
]);
const isAcceleratorInputDisabled =
_.isEmpty(acceleratorSlots) ||
(currentImageAcceleratorLimits.length === 0 &&
_.isEmpty(
form.getFieldValue(['environments', 'manual']),
));
return (
<Form.Item
name={['resource', 'accelerator']}
Expand All @@ -952,7 +984,9 @@ const ResourceAllocationFormItems: React.FC<
}}
rules={[
{
required: currentImageAcceleratorLimits.length > 0,
required:
currentImageAcceleratorLimits &&
currentImageAcceleratorLimits.length > 0,
},
{
type: 'number',
Expand Down Expand Up @@ -1068,10 +1102,9 @@ const ResourceAllocationFormItems: React.FC<
formatter: (value = 0) => {
return `${value} ${mergedResourceSlots?.[currentAcceleratorType]?.display_unit || ''}`;
},
open:
currentImageAcceleratorLimits.length <= 0
? false
: undefined,
open: isAcceleratorInputDisabled
? false
: undefined,
},
}}
disabled={isAcceleratorInputDisabled}
Expand Down Expand Up @@ -1116,6 +1149,7 @@ const ResourceAllocationFormItems: React.FC<
mergedResourceSlots?.[name]
?.display_unit || 'UNIT',
disabled:
currentImageAcceleratorLimits &&
currentImageAcceleratorLimits.length >
0 &&
!_.find(
Expand Down Expand Up @@ -1490,14 +1524,17 @@ export const getAllocatablePresetNames = (
} else {
// When current image requires some accelerator,
// It's available if the preset has a required accelerator value that is larger than the current image's minimum value
return _.some(currentImageAcceleratorLimits, (limit) => {
return (
limit?.key &&
acceleratorResourceOfPreset[limit?.key] &&
_.toNumber(acceleratorResourceOfPreset[limit?.key]) >=
_.toNumber(limit?.min)
);
});
return (
currentImageAcceleratorLimits &&
_.some(currentImageAcceleratorLimits, (limit) => {
return (
limit?.key &&
acceleratorResourceOfPreset[limit?.key] &&
_.toNumber(acceleratorResourceOfPreset[limit?.key]) >=
_.toNumber(limit?.min)
);
})
);
}
}).map((preset) => preset.name);
return currentImageAcceleratorLimits.length === 0
Expand Down
7 changes: 5 additions & 2 deletions react/src/components/ResourceGroupSelect.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ import useControllableState from '../hooks/useControllableState';
import TextHighlighter from './TextHighlighter';
import { Select, SelectProps } from 'antd';
import _ from 'lodash';
import React, { useEffect, useTransition } from 'react';
import React, { useEffect, useState, useTransition } from 'react';

interface ResourceGroupSelectProps extends SelectProps {
projectName: string;
Expand Down Expand Up @@ -33,8 +33,11 @@ const ResourceGroupSelect: React.FC<ResourceGroupSelectProps> = ({
const [controllableValue, setControllableValueDoNotUseWithoutTransition] =
useControllableState(selectProps);
const [isPendingChangeTransition, startChangeTransition] = useTransition();

const [optimisticValue, setOptimisticValue] = useState();
const setControllableValueWithTransition = React.useCallback(
(v: typeof controllableValue, ...args: any[]) => {
setOptimisticValue(v);
startChangeTransition(() => {
setControllableValueDoNotUseWithoutTransition(v, ...args);
});
Expand Down Expand Up @@ -166,7 +169,7 @@ const ResourceGroupSelect: React.FC<ResourceGroupSelectProps> = ({
);
}}
{...selectProps}
value={controllableValue}
value={isPendingChangeTransition ? optimisticValue : controllableValue}
onChange={setControllableValueWithTransition}
/>
);
Expand Down
Loading