feat(FR-2867): add unified memory accelerator UX in Resource Allocation

Author: yomybaby

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

@@ -8,7 +8,83 @@ import {
   ResourcePreset,
 } from '../../hooks/useResourceLimitAndRemaining';
 import { Image } from '../ImageEnvironmentSelectFormItems';
-import { getAllocatablePresetNames } from './ResourceAllocationFormItems';
+import {
+  displayUnitToInputSizeUnit,
+  getAllocatablePresetNames,
+  getUnifiedAcceleratorValueFromMem,
+  isUnifiedAcceleratorSlot,
+} from './ResourceAllocationFormItems';
+
+describe('displayUnitToInputSizeUnit', () => {
+  it('maps known binary display units to InputSizeUnit', () => {
+    expect(displayUnitToInputSizeUnit('GiB')).toBe('g');
+    expect(displayUnitToInputSizeUnit('MiB')).toBe('m');
+    expect(displayUnitToInputSizeUnit('TiB')).toBe('t');
+    expect(displayUnitToInputSizeUnit('KiB')).toBe('k');
+    expect(displayUnitToInputSizeUnit('PiB')).toBe('p');
+    expect(displayUnitToInputSizeUnit('EiB')).toBe('e');
+  });
+
+  it('falls back to "g" for unknown or empty values', () => {
+    expect(displayUnitToInputSizeUnit('foo')).toBe('g');
+    expect(displayUnitToInputSizeUnit('')).toBe('g');
+    expect(displayUnitToInputSizeUnit(null)).toBe('g');
+    expect(displayUnitToInputSizeUnit(undefined)).toBe('g');
+  });
+});
+
+describe('getUnifiedAcceleratorValueFromMem', () => {
+  const slotsGiB = { 'cuda.unified': { display_unit: 'GiB' } };
+  const slotsMiB = { 'cuda.unified': { display_unit: 'MiB' } };
+  const slotsNoMeta = { 'cuda.unified': {} };
+
+  it('returns mem converted to the slot display unit (GiB)', () => {
+    expect(
+      getUnifiedAcceleratorValueFromMem('8g', 'cuda.unified', slotsGiB),
+    ).toBe(8);
+  });
+
+  it('returns mem converted to the slot display unit (MiB)', () => {
+    expect(
+      getUnifiedAcceleratorValueFromMem('8g', 'cuda.unified', slotsMiB),
+    ).toBe(8192);
+  });
+
+  it('falls back to GiB when slot metadata has no display_unit', () => {
+    expect(
+      getUnifiedAcceleratorValueFromMem('8g', 'cuda.unified', slotsNoMeta),
+    ).toBe(8);
+  });
+
+  it('returns 0 for undefined or zero mem', () => {
+    expect(
+      getUnifiedAcceleratorValueFromMem(undefined, 'cuda.unified', slotsGiB),
+    ).toBe(0);
+    expect(getUnifiedAcceleratorValueFromMem(0, 'cuda.unified', slotsGiB)).toBe(
+      0,
+    );
+  });
+});
+
+describe('isUnifiedAcceleratorSlot', () => {
+  it('returns true for slot names ending with .unified', () => {
+    expect(isUnifiedAcceleratorSlot('cuda.unified')).toBe(true);
+    expect(isUnifiedAcceleratorSlot('rocm.unified')).toBe(true);
+  });
+
+  it('returns false for discrete accelerator slot names', () => {
+    expect(isUnifiedAcceleratorSlot('cuda.shares')).toBe(false);
+    expect(isUnifiedAcceleratorSlot('cuda.device')).toBe(false);
+    expect(isUnifiedAcceleratorSlot('cuda.mem')).toBe(false);
+    expect(isUnifiedAcceleratorSlot('rocm.device')).toBe(false);
+  });
+
+  it('returns false for nullish or empty input', () => {
+    expect(isUnifiedAcceleratorSlot(undefined)).toBe(false);
+    expect(isUnifiedAcceleratorSlot(null)).toBe(false);
+    expect(isUnifiedAcceleratorSlot('')).toBe(false);
+  });
+});
 
 describe('getAllocatablePresetNames', () => {
   const presets: Array<ResourcePreset> = [

react/src/components/SessionFormItems/ResourceAllocationFormItems.tsx

@@ -7,6 +7,7 @@ import {
   addNumberWithUnits,
   compareNumberWithUnits,
   convertToBinaryUnit,
+  InputSizeUnit,
 } from '../../helper';
 import { useSuspendedBackendaiClient } from '../../hooks';
 import { useResourceSlots } from '../../hooks/backendai';
@@ -62,6 +63,58 @@ export const isMinOversMaxValue = (min: number, max: number) => {
   return min >= max;
 };
 
+/**
+ * Returns true when the given accelerator slot name represents a unified
+ * memory architecture, where the accelerator memory and the host memory
+ * share a single physical pool. Identified by a `.unified` suffix on the
+ * slot name (e.g. `cuda.unified`).
+ */
+export const isUnifiedAcceleratorSlot = (slotName?: string | null): boolean => {
+  return !!slotName && _.endsWith(slotName, '.unified');
+};
+
+/**
+ * Maps a slot's `display_unit` (e.g. `GiB`, `MiB`, `TiB`) to the
+ * corresponding `InputSizeUnit` accepted by `convertToBinaryUnit`. Falls
+ * back to `'g'` (GiB) when the display unit is missing or not a recognized
+ * binary size unit.
+ */
+export const displayUnitToInputSizeUnit = (
+  displayUnit?: string | null,
+): InputSizeUnit => {
+  const first = (_.first(displayUnit ?? '') ?? '').toLowerCase();
+  const validUnits: ReadonlyArray<InputSizeUnit> = [
+    'k',
+    'm',
+    'g',
+    't',
+    'p',
+    'e',
+  ];
+  const match = validUnits.find((u) => u === first);
+  return match ?? 'g';
+};
+
+/**
+ * Returns the accelerator-field number derived from the host memory value
+ * for a unified-memory slot. Reads the slot's `display_unit` from
+ * `mergedResourceSlots` and converts `mem` to that unit. Falls back to
+ * `'g'` (GiB) when the display unit is unavailable.
+ */
+export const getUnifiedAcceleratorValueFromMem = (
+  mem: string | number | undefined,
+  slotName: string | undefined,
+  mergedResourceSlots:
+    | Record<string, { display_unit?: string } | undefined>
+    | undefined,
+): number => {
+  const displayUnit = slotName
+    ? mergedResourceSlots?.[slotName]?.display_unit
+    : undefined;
+  const targetUnit = displayUnitToInputSizeUnit(displayUnit);
+  return convertToBinaryUnit(mem || '0g', targetUnit)?.number ?? 0;
+};
+
 export interface ResourceAllocationFormValue {
   resource: {
     cpu: number;
@@ -267,6 +320,28 @@ const ResourceAllocationFormItems: React.FC<
     }
   };
 
+  // Centralized helper that keeps the `resource.accelerator` field in sync
+  // with `resource.mem` whenever the active accelerator slot is a unified
+  // one. Call this from every code path that mutates `resource.mem` or
+  // `resource.acceleratorType` so submit-time form state is consistent.
+  const syncUnifiedAcceleratorIfNeeded = useEventNotStable(() => {
+    const activeAcceleratorType = form.getFieldValue([
+      'resource',
+      'acceleratorType',
+    ]);
+    if (!isUnifiedAcceleratorSlot(activeAcceleratorType)) {
+      return;
+    }
+    form.setFieldValue(
+      ['resource', 'accelerator'],
+      getUnifiedAcceleratorValueFromMem(
+        form.getFieldValue(['resource', 'mem']),
+        activeAcceleratorType,
+        mergedResourceSlots,
+      ),
+    );
+  });
+
   const ensureValidAcceleratorType = useEventNotStable(() => {
     const currentAcceleratorType = form.getFieldValue([
       'resource',
@@ -283,6 +358,9 @@ const ResourceAllocationFormItems: React.FC<
         acceleratorType: nextAcceleratorType || currentAcceleratorType,
       },
     });
+    // The accelerator type may have changed; mirror `mem` into the
+    // accelerator field if the resolved type is a unified slot.
+    syncUnifiedAcceleratorIfNeeded();
   });
 
   const updateResourceFieldsBasedOnImage = useEventNotStable(
@@ -388,6 +466,9 @@ const ResourceAllocationFormItems: React.FC<
       if (form.getFieldValue('enabledAutomaticShmem')) {
         runShmemAutomationRule(form.getFieldValue(['resource', 'mem']) || '0g');
       }
+      // mem and/or acceleratorType may have changed above; keep the
+      // accelerator field consistent for unified slots.
+      syncUnifiedAcceleratorIfNeeded();
       form
         .validateFields(['resource'], {
           recursive: true,
@@ -445,6 +526,9 @@ const ResourceAllocationFormItems: React.FC<
       if (!hasPresetShmem) {
         runShmemAutomationRule(mem || '0g');
       }
+      // mem and/or acceleratorType may have changed above; keep the
+      // accelerator field consistent for unified slots.
+      syncUnifiedAcceleratorIfNeeded();
 
       form
         .validateFields(['resource'], {
@@ -864,6 +948,13 @@ const ResourceAllocationFormItems: React.FC<
                                 ) {
                                   runShmemAutomationRule(M_plus_S || '0g');
                                 }
+                                // When the active accelerator slot is a
+                                // unified-memory slot, the accelerator field
+                                // shares a single pool with host memory.
+                                // Derive the accelerator value from `mem` at
+                                // the source of change so submit-time form
+                                // state is always consistent.
+                                syncUnifiedAcceleratorIfNeeded();
                               }}
                             />
                           </Form.Item>
@@ -921,6 +1012,10 @@ const ResourceAllocationFormItems: React.FC<
                         currentAcceleratorType as keyof typeof resourceSlots
                       ] === 'unique';
 
+                    const isUnifiedType = isUnifiedAcceleratorSlot(
+                      currentAcceleratorType,
+                    );
+
                     const isSingleCluster =
                       form.getFieldValue('cluster_size') < 2;
                     const hasQuantumSize = _.isNumber(
@@ -965,9 +1060,15 @@ const ResourceAllocationFormItems: React.FC<
                             />
                           ),
                         }}
+                        extra={
+                          isUnifiedType
+                            ? t('session.launcher.UnifiedAcceleratorMemoryNote')
+                            : undefined
+                        }
                         dependencies={[
                           ['resource', 'acceleratorType'],
                           'cluster_size',
+                          ['resource', 'mem'],
                         ]}
                         rules={[
                           {
@@ -1139,7 +1240,8 @@ const ResourceAllocationFormItems: React.FC<
                             },
                           }}
                           disabled={
-                            supportedAcceleratorTypesInRGByImage?.length === 0
+                            supportedAcceleratorTypesInRGByImage?.length ===
+                              0 || isUnifiedType
                           }
                           min={0}
                           max={
@@ -1193,6 +1295,58 @@ const ResourceAllocationFormItems: React.FC<
                                         };
                                       },
                                     )}
+                                    onChange={(nextType: string) => {
+                                      // Changing the slot type mutates the
+                                      // active allocation; the previously
+                                      // selected preset no longer matches.
+                                      form.setFieldValue(
+                                        'allocationPreset',
+                                        'custom',
+                                      );
+                                      // Keep the accelerator field consistent
+                                      // at the moment the slot type changes,
+                                      // rather than relying on a watcher
+                                      // effect that runs after render.
+                                      if (isUnifiedAcceleratorSlot(nextType)) {
+                                        // Switching INTO a unified slot:
+                                        // mirror the current `mem` value
+                                        // converted to the slot's display
+                                        // unit. Write acceleratorType first
+                                        // so the shared helper sees the new
+                                        // active slot when it reads from the
+                                        // form.
+                                        form.setFieldValue(
+                                          ['resource', 'acceleratorType'],
+                                          nextType,
+                                        );
+                                        syncUnifiedAcceleratorIfNeeded();
+                                      } else if (
+                                        isUnifiedAcceleratorSlot(
+                                          currentAcceleratorType,
+                                        )
+                                      ) {
+                                        // Switching OUT of a unified slot
+                                        // into a discrete one: reset to the
+                                        // discrete slot's min so the stale
+                                        // mirrored GiB value from unified
+                                        // mode does not bleed through as a
+                                        // device count.
+                                        //
+                                        // Discrete-to-discrete switches are
+                                        // intentionally NOT reset here:
+                                        // ensureValidAcceleratorType clamps
+                                        // the existing value if it falls
+                                        // outside the new slot's range, so
+                                        // the user's current allocation is
+                                        // preserved across discrete slot
+                                        // changes.
+                                        form.setFieldValue(
+                                          ['resource', 'accelerator'],
+                                          resourceLimits.accelerators[nextType]
+                                            ?.min ?? 0,
+                                        );
+                                      }
+                                    }}
                                   />
                                 </Form.Item>
                               ) : undefined,

resources/i18n/en.json

@@ -2882,6 +2882,7 @@
       "TensorBoardPrepared": "TensorBoard app prepared.",
       "TitleSession": "Session (Backend.AI)",
       "TotalAllocation": "Total Allocation",
+      "UnifiedAcceleratorMemoryNote": "Accelerator memory is shared with host memory on this device and cannot be set separately.",
       "UserResourceLimit": "User Resource Limit",
       "UsingBootstrapScriptInfo": "Bootstrap script is included.",
       "Version": "Version",