From 40b6062cd9a2cbf055560f3fb17bf4a33a296960 Mon Sep 17 00:00:00 2001
From: yomybaby <621215+yomybaby@users.noreply.github.com>
Date: Mon, 9 Sep 2024 08:37:53 +0000
Subject: [PATCH] fix: merged resource limit is affected by remaining in the
 resource group. (#2696)

### TL;DR

Fixed and refactored resource limit and remaining calculations in useResourceLimitAndRemaining hook.

This PR resolves the issue where the merged resource limit in useResourceLimitAndRemaining is affected by the remaining in the resource group.

Now, the maximum value is determined as the minimum value of the following:

- pre-container configuration (e.g., maxCPUCoresPerContainer)
- keypair resource limit
- group resource limit

The useResourceLimitAndRemaining should consider the domain resource limit and the total slot size of the resource group. This will be handled by another PR.

### What changed?

- Updated ResourceLimits type to use ResourceSlotName as key
- Commented out resourceGroupResourceSize calculations for cpu and mem
- Added keypair and group limits for accelerators
- Updated type assertions in remaining resource calculations

### How to test?

1. Test resource allocation functionality in the UI
2. Verify that resource limits are correctly applied for different resource types (CPU, memory, accelerators)
3. Check that remaining resources are calculated accurately
4. Ensure that the changes don't introduce any regressions in resource management
---
 .../hooks/useResourceLimitAndRemaining.tsx    | 43 ++++++++-----------
 1 file changed, 19 insertions(+), 24 deletions(-)

diff --git a/react/src/hooks/useResourceLimitAndRemaining.tsx b/react/src/hooks/useResourceLimitAndRemaining.tsx
index 4da0a1767..8266692b8 100644
--- a/react/src/hooks/useResourceLimitAndRemaining.tsx
+++ b/react/src/hooks/useResourceLimitAndRemaining.tsx
@@ -2,7 +2,7 @@ import { useSuspendedBackendaiClient } from '.';
 import { Image } from '../components/ImageEnvironmentSelectFormItems';
 import { AUTOMATIC_DEFAULT_SHMEM } from '../components/ResourceAllocationFormItems';
 import { addNumberWithUnits, iSizeToSize } from '../helper';
-import { useResourceSlots } from '../hooks/backendai';
+import { ResourceSlotName, useResourceSlots } from '../hooks/backendai';
 import { useSuspenseTanQuery } from './reactQueryAlias';
 import _ from 'lodash';
 import { useMemo } from 'react';
@@ -47,9 +47,7 @@ export interface MergedResourceLimits {
 }
 
 type ResourceLimits = {
-  cpu: string | 'Infinity' | 'NaN';
-  mem: string | 'Infinity' | 'NaN';
-  'cuda.device': string | 'Infinity' | 'NaN';
+  [key in ResourceSlotName]?: string | 'Infinity' | 'NaN';
 };
 type ResourceUsing = ResourceLimits;
 type ResourceRemaining = ResourceLimits;
@@ -237,7 +235,7 @@ export const useResourceLimitAndRemaining = ({
                 : baiClient._config.maxCPUCoresPerContainer,
               limitParser(checkPresetInfo?.keypair_limits.cpu),
               limitParser(checkPresetInfo?.group_limits.cpu),
-              resourceGroupResourceSize?.cpu,
+              // resourceGroupResourceSize?.cpu,
             ]),
           },
     mem:
@@ -273,20 +271,10 @@ export const useResourceLimitAndRemaining = ({
                     'g',
                   )?.number,
                 // scaling group all mem (using + remaining), string type
-                resourceGroupResourceSize?.mem &&
-                  iSizeToSize(resourceGroupResourceSize?.mem + '', 'g')?.number,
+                // resourceGroupResourceSize?.mem &&
+                //   iSizeToSize(resourceGroupResourceSize?.mem + '', 'g')?.number,
               ]) + 'g',
           },
-    // shmem:
-    //   resourceSlots?.mem === undefined
-    //     ? undefined
-    //     : {
-    //         min: _.max([
-    //           _.find(currentImage?.resource_limits, (i) => i?.key === 'shmem')
-    //             ?.min,
-    //           '64m',
-    //         ]),
-    //       },
     accelerators: _.reduce(
       acceleratorSlots,
       (result, value, key) => {
@@ -306,8 +294,12 @@ export const useResourceLimitAndRemaining = ({
           ),
           max: _.min([
             perContainerLimit || 8,
+            limitParser(
+              checkPresetInfo?.keypair_limits[key as ResourceSlotName],
+            ),
+            limitParser(checkPresetInfo?.group_limits[key as ResourceSlotName]),
             // scaling group all cpu (using + remaining), string type
-            resourceGroupResourceSize.accelerators[key],
+            // resourceGroupResourceSize.accelerators[key],
           ]),
         };
         return result;
@@ -321,12 +313,15 @@ export const useResourceLimitAndRemaining = ({
       (result, value, key) => {
         result[key] =
           _.min([
-            // @ts-ignore
-            _.toNumber(checkPresetInfo?.keypair_remaining[key]),
-            // @ts-ignore
-            _.toNumber(checkPresetInfo?.group_remaining[key]),
-            // @ts-ignore
-            _.toNumber(checkPresetInfo?.scaling_group_remaining[key]),
+            _.toNumber(
+              checkPresetInfo?.keypair_remaining[key as ResourceSlotName],
+            ),
+            _.toNumber(
+              checkPresetInfo?.group_remaining[key as ResourceSlotName],
+            ),
+            _.toNumber(
+              checkPresetInfo?.scaling_group_remaining[key as ResourceSlotName],
+            ),
           ]) ?? Number.MAX_SAFE_INTEGER;
         return result;
       },