Skip to content

Commit 4677880

Browse files
yomybabyyomybaby
committed
feat(FR-2867): add unified memory accelerator UX in Resource Allocation
1 parent 0177ac9 commit 4677880

3 files changed

Lines changed: 150 additions & 2 deletions

File tree

‎react/src/components/SessionFormItems/ResourceAllocationFormItems.test.ts‎

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,30 @@ import {
88
ResourcePreset,
99
} from '../../hooks/useResourceLimitAndRemaining';
1010
import { Image } from '../ImageEnvironmentSelectFormItems';
11-
import { getAllocatablePresetNames } from './ResourceAllocationFormItems';
11+
import {
12+
getAllocatablePresetNames,
13+
isUnifiedAcceleratorSlot,
14+
} from './ResourceAllocationFormItems';
15+
16+
describe('isUnifiedAcceleratorSlot', () => {
17+
it('returns true for slot names ending with .unified', () => {
18+
expect(isUnifiedAcceleratorSlot('cuda.unified')).toBe(true);
19+
expect(isUnifiedAcceleratorSlot('rocm.unified')).toBe(true);
20+
});
21+
22+
it('returns false for discrete accelerator slot names', () => {
23+
expect(isUnifiedAcceleratorSlot('cuda.shares')).toBe(false);
24+
expect(isUnifiedAcceleratorSlot('cuda.device')).toBe(false);
25+
expect(isUnifiedAcceleratorSlot('cuda.mem')).toBe(false);
26+
expect(isUnifiedAcceleratorSlot('rocm.device')).toBe(false);
27+
});
28+
29+
it('returns false for nullish or empty input', () => {
30+
expect(isUnifiedAcceleratorSlot(undefined)).toBe(false);
31+
expect(isUnifiedAcceleratorSlot(null)).toBe(false);
32+
expect(isUnifiedAcceleratorSlot('')).toBe(false);
33+
});
34+
});
1235

1336
describe('getAllocatablePresetNames', () => {
1437
const presets: Array<ResourcePreset> = [

‎react/src/components/SessionFormItems/ResourceAllocationFormItems.tsx‎

Lines changed: 125 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import {
77
addNumberWithUnits,
88
compareNumberWithUnits,
99
convertToBinaryUnit,
10+
InputSizeUnit,
1011
} from '../../helper';
1112
import { useSuspendedBackendaiClient } from '../../hooks';
1213
import { useResourceSlots } from '../../hooks/backendai';
@@ -62,6 +63,58 @@ export const isMinOversMaxValue = (min: number, max: number) => {
6263
return min >= max;
6364
};
6465

66+
/**
67+
* Returns true when the given accelerator slot name represents a unified
68+
* memory architecture, where the accelerator memory and the host memory
69+
* share a single physical pool. Identified by a `.unified` suffix on the
70+
* slot name (e.g. `cuda.unified`).
71+
*/
72+
export const isUnifiedAcceleratorSlot = (slotName?: string | null): boolean => {
73+
return !!slotName && _.endsWith(slotName, '.unified');
74+
};
75+
76+
/**
77+
* Maps a slot's `display_unit` (e.g. `GiB`, `MiB`, `TiB`) to the
78+
* corresponding `InputSizeUnit` accepted by `convertToBinaryUnit`. Falls
79+
* back to `'g'` (GiB) when the display unit is missing or not a recognized
80+
* binary size unit.
81+
*/
82+
export const displayUnitToInputSizeUnit = (
83+
displayUnit?: string | null,
84+
): InputSizeUnit => {
85+
const first = (_.first(displayUnit ?? '') ?? '').toLowerCase();
86+
const validUnits: ReadonlyArray<InputSizeUnit> = [
87+
'k',
88+
'm',
89+
'g',
90+
't',
91+
'p',
92+
'e',
93+
];
94+
const match = validUnits.find((u) => u === first);
95+
return match ?? 'g';
96+
};
97+
98+
/**
99+
* Returns the accelerator-field number derived from the host memory value
100+
* for a unified-memory slot. Reads the slot's `display_unit` from
101+
* `mergedResourceSlots` and converts `mem` to that unit. Falls back to
102+
* `'g'` (GiB) when the display unit is unavailable.
103+
*/
104+
export const getUnifiedAcceleratorValueFromMem = (
105+
mem: string | number | undefined,
106+
slotName: string | undefined,
107+
mergedResourceSlots:
108+
| Record<string, { display_unit?: string } | undefined>
109+
| undefined,
110+
): number => {
111+
const displayUnit = slotName
112+
? mergedResourceSlots?.[slotName]?.display_unit
113+
: undefined;
114+
const targetUnit = displayUnitToInputSizeUnit(displayUnit);
115+
return convertToBinaryUnit(mem || '0g', targetUnit)?.number ?? 0;
116+
};
117+
65118
export interface ResourceAllocationFormValue {
66119
resource: {
67120
cpu: number;
@@ -864,6 +917,32 @@ const ResourceAllocationFormItems: React.FC<
864917
) {
865918
runShmemAutomationRule(M_plus_S || '0g');
866919
}
920+
// When the active accelerator slot is a
921+
// unified-memory slot, the accelerator field
922+
// shares a single pool with host memory.
923+
// Derive the accelerator value from `mem`
924+
// at the source of change rather than via a
925+
// watcher effect, so submit-time form state
926+
// is always consistent.
927+
const activeAcceleratorType =
928+
form.getFieldValue([
929+
'resource',
930+
'acceleratorType',
931+
]);
932+
if (
933+
isUnifiedAcceleratorSlot(
934+
activeAcceleratorType,
935+
)
936+
) {
937+
form.setFieldValue(
938+
['resource', 'accelerator'],
939+
getUnifiedAcceleratorValueFromMem(
940+
M_plus_S,
941+
activeAcceleratorType,
942+
mergedResourceSlots,
943+
),
944+
);
945+
}
867946
}}
868947
/>
869948
</Form.Item>
@@ -921,6 +1000,10 @@ const ResourceAllocationFormItems: React.FC<
9211000
currentAcceleratorType as keyof typeof resourceSlots
9221001
] === 'unique';
9231002

1003+
const isUnifiedType = isUnifiedAcceleratorSlot(
1004+
currentAcceleratorType,
1005+
);
1006+
9241007
const isSingleCluster =
9251008
form.getFieldValue('cluster_size') < 2;
9261009
const hasQuantumSize = _.isNumber(
@@ -965,9 +1048,15 @@ const ResourceAllocationFormItems: React.FC<
9651048
/>
9661049
),
9671050
}}
1051+
extra={
1052+
isUnifiedType
1053+
? t('session.launcher.UnifiedAcceleratorMemoryNote')
1054+
: undefined
1055+
}
9681056
dependencies={[
9691057
['resource', 'acceleratorType'],
9701058
'cluster_size',
1059+
['resource', 'mem'],
9711060
]}
9721061
rules={[
9731062
{
@@ -1139,7 +1228,8 @@ const ResourceAllocationFormItems: React.FC<
11391228
},
11401229
}}
11411230
disabled={
1142-
supportedAcceleratorTypesInRGByImage?.length === 0
1231+
supportedAcceleratorTypesInRGByImage?.length ===
1232+
0 || isUnifiedType
11431233
}
11441234
min={0}
11451235
max={
@@ -1193,6 +1283,40 @@ const ResourceAllocationFormItems: React.FC<
11931283
};
11941284
},
11951285
)}
1286+
onChange={(nextType: string) => {
1287+
// Keep the accelerator field consistent
1288+
// at the moment the slot type changes,
1289+
// rather than relying on a watcher
1290+
// effect that runs after render.
1291+
if (isUnifiedAcceleratorSlot(nextType)) {
1292+
// Switching INTO a unified slot:
1293+
// mirror the current `mem` value
1294+
// converted to the slot's display
1295+
// unit.
1296+
form.setFieldValue(
1297+
['resource', 'accelerator'],
1298+
getUnifiedAcceleratorValueFromMem(
1299+
form.getFieldValue([
1300+
'resource',
1301+
'mem',
1302+
]),
1303+
nextType,
1304+
mergedResourceSlots,
1305+
),
1306+
);
1307+
} else {
1308+
// Switching OUT of a unified slot
1309+
// into a discrete one: reset to the
1310+
// discrete slot's min so the stale
1311+
// mirrored value does not bleed
1312+
// through.
1313+
form.setFieldValue(
1314+
['resource', 'accelerator'],
1315+
resourceLimits.accelerators[nextType]
1316+
?.min ?? 0,
1317+
);
1318+
}
1319+
}}
11961320
/>
11971321
</Form.Item>
11981322
) : undefined,

‎resources/i18n/en.json‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2882,6 +2882,7 @@
28822882
"TensorBoardPrepared": "TensorBoard app prepared.",
28832883
"TitleSession": "Session (Backend.AI)",
28842884
"TotalAllocation": "Total Allocation",
2885+
"UnifiedAcceleratorMemoryNote": "Accelerator memory is shared with host memory on this device and cannot be set separately.",
28852886
"UserResourceLimit": "User Resource Limit",
28862887
"UsingBootstrapScriptInfo": "Bootstrap script is included.",
28872888
"Version": "Version",

0 commit comments

Comments
 (0)