fix: show/update supported AI accelerator based on selected environme…

…nt in session launcher
lablup · Sep 25, 2024 · 0a7646c · 0a7646c
1 parent 9821f06
commit 0a7646c
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 16 deletions.
diff --git a/react/src/components/ResourceAllocationFormItems.tsx b/react/src/components/ResourceAllocationFormItems.tsx
@@ -125,16 +125,7 @@ const ResourceAllocationFormItems: React.FC<
     currentResourceGroup || undefined,
   );
 
-  const acceleratorSlots = _.omitBy(resourceSlots, (value, key) => {
-    if (['cpu', 'mem', 'shmem'].includes(key)) return true;
-
-    if (
-      !resourceLimits.accelerators[key]?.max ||
-      resourceLimits.accelerators[key]?.max === 0
-    )
-      return true;
-    return false;
-  });
+  const acceleratorSlots = remaining.accelerators;
 
   const currentImageAcceleratorLimits = useMemo(
     () =>
@@ -348,7 +339,6 @@ const ResourceAllocationFormItems: React.FC<
       _.keys(acceleratorSlots),
       (value) => acceleratorObj[value] !== undefined,
     );
-
     let acceleratorSetting: {
       acceleratorType?: string;
       accelerator: number;

diff --git a/react/src/hooks/useResourceLimitAndRemaining.tsx b/react/src/hooks/useResourceLimitAndRemaining.tsx
@@ -108,7 +108,14 @@ export const useResourceLimitAndRemaining = ({
   const baiClient = useSuspendedBackendaiClient();
   const [resourceSlots] = useResourceSlots();
   const acceleratorSlots = _.omit(resourceSlots, ['cpu', 'mem', 'shmem']);
-
+  const currentImageSupportedAIAccelerator = _.filter(
+    _.map(currentImage?.resource_limits, 'key'),
+    (key) => !['cpu', 'mem', 'shmem'].includes(key),
+  );
+  const filteredAcceleratorSlots = _.omitBy(
+    acceleratorSlots,
+    (value, key) => !currentImageSupportedAIAccelerator.includes(key),
+  );
   const {
     data: checkPresetInfo,
     refetch,
@@ -130,7 +137,6 @@ export const useResourceLimitAndRemaining = ({
     staleTime: 1000,
     // suspense: !_.isEmpty(currentResourceGroup), //prevent flicking
   });
-
   const currentImageMinM =
     _.find(currentImage?.resource_limits, (i) => i?.key === 'mem')?.min || '0g';
 
@@ -178,7 +184,7 @@ export const useResourceLimitAndRemaining = ({
           )?.numberFixed + 'g'
         : undefined,
     accelerators: _.reduce(
-      acceleratorSlots,
+      filteredAcceleratorSlots,
       (result, value, key) => {
         result[key] =
           !_.isEmpty(
@@ -276,7 +282,7 @@ export const useResourceLimitAndRemaining = ({
               ]) + 'g',
           },
     accelerators: _.reduce(
-      acceleratorSlots,
+      filteredAcceleratorSlots,
       (result, value, key) => {
         const perContainerLimit =
           _.find(perContainerConfigs, (configValue, configName) => {
@@ -309,7 +315,7 @@ export const useResourceLimitAndRemaining = ({
   };
   const remaining: RemainingSlots = {
     accelerators: _.reduce(
-      acceleratorSlots,
+      filteredAcceleratorSlots,
       (result, value, key) => {
         result[key] =
           _.min([