From ccb92507b986c7e9d5123f467139acf1f928999d Mon Sep 17 00:00:00 2001
From: Yu Zhao <zyu0321@gmail.com>
Date: Mon, 26 Aug 2024 19:22:47 -0400
Subject: [PATCH] Modify test case based on new change from upstream

---
 .../k8s/__tests__/inferenceServices.spec.ts   |  6 +++++
 .../NIMServiceModal/DeployNIMServiceModal.tsx | 27 ++++++++++++++-----
 2 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/frontend/src/api/k8s/__tests__/inferenceServices.spec.ts b/frontend/src/api/k8s/__tests__/inferenceServices.spec.ts
index 0dda0c16a9..27f38b7e2d 100644
--- a/frontend/src/api/k8s/__tests__/inferenceServices.spec.ts
+++ b/frontend/src/api/k8s/__tests__/inferenceServices.spec.ts
@@ -186,6 +186,7 @@ describe('assembleInferenceService', () => {
       undefined,
       false,
       undefined,
+      undefined,
       acceleratorProfileState,
       selectedAcceleratorProfile,
     );
@@ -218,6 +219,7 @@ describe('assembleInferenceService', () => {
       undefined,
       true,
       undefined,
+      undefined,
       acceleratorProfileState,
       selectedAcceleratorProfile,
     );
@@ -251,6 +253,7 @@ describe('assembleInferenceService', () => {
       undefined,
       false,
       undefined,
+      undefined,
       acceleratorProfileState,
       selectedAcceleratorProfile,
     );
@@ -279,6 +282,7 @@ describe('assembleInferenceService', () => {
       undefined,
       true,
       undefined,
+      undefined,
       acceleratorProfileState,
       selectedAcceleratorProfile,
     );
@@ -321,6 +325,7 @@ describe('assembleInferenceService', () => {
       undefined,
       false,
       undefined,
+      undefined,
       acceleratorProfileState,
       selectedAcceleratorProfile,
     );
@@ -373,6 +378,7 @@ describe('assembleInferenceService', () => {
       undefined,
       true,
       undefined,
+      undefined,
       acceleratorProfileState,
       selectedAcceleratorProfile,
     );
diff --git a/frontend/src/pages/modelServing/screens/projects/NIMServiceModal/DeployNIMServiceModal.tsx b/frontend/src/pages/modelServing/screens/projects/NIMServiceModal/DeployNIMServiceModal.tsx
index 9f91ffc759..267fa0a429 100644
--- a/frontend/src/pages/modelServing/screens/projects/NIMServiceModal/DeployNIMServiceModal.tsx
+++ b/frontend/src/pages/modelServing/screens/projects/NIMServiceModal/DeployNIMServiceModal.tsx
@@ -38,6 +38,8 @@ import { getDisplayNameFromK8sResource, translateDisplayNameForK8s } from '~/con
 import { useAccessReview } from '~/api';
 import { SupportedArea, useIsAreaAvailable } from '~/concepts/areas';
 import KServeAutoscalerReplicaSection from '~/pages/modelServing/screens/projects/kServeModal/KServeAutoscalerReplicaSection';
+import useGenericObjectState from '~/utilities/useGenericObjectState';
+import { AcceleratorProfileSelectFieldState } from '~/pages/notebookController/screens/server/AcceleratorProfileSelectField';
 
 const NIM_SECRET_NAME = 'nvidia-nim-secrets';
 const NIM_NGC_SECRET_NAME = 'ngc-secret';
@@ -92,11 +94,19 @@ const DeployNIMServiceModal: React.FC<DeployNIMServiceModalProps> = ({
   const isInferenceServiceNameWithinLimit =
     translateDisplayNameForK8s(createDataInferenceService.name).length <= 253;
 
-  const [acceleratorProfileState, setAcceleratorProfileState, resetAcceleratorProfileData] =
-    useServingAcceleratorProfile(
-      editInfo?.servingRuntimeEditInfo?.servingRuntime,
-      editInfo?.inferenceServiceEditInfo,
-    );
+  const acceleratorProfileState = useServingAcceleratorProfile(
+    editInfo?.servingRuntimeEditInfo?.servingRuntime,
+    editInfo?.inferenceServiceEditInfo,
+  );
+  const [
+    selectedAcceleratorProfile,
+    setSelectedAcceleratorProfile,
+    resetSelectedAcceleratorProfile,
+  ] = useGenericObjectState<AcceleratorProfileSelectFieldState>({
+    profile: undefined,
+    count: 0,
+    useExistingSettings: false,
+  });
   const customServingRuntimesEnabled = useCustomServingRuntimesEnabled();
   const [allowCreate] = useAccessReview({
     ...accessReviewResource,
@@ -142,7 +152,7 @@ const DeployNIMServiceModal: React.FC<DeployNIMServiceModalProps> = ({
     setActionInProgress(false);
     resetDataServingRuntime();
     resetDataInferenceService();
-    resetAcceleratorProfileData();
+    resetSelectedAcceleratorProfile();
     setAlertVisible(true);
   };
 
@@ -172,6 +182,7 @@ const DeployNIMServiceModal: React.FC<DeployNIMServiceModalProps> = ({
       editInfo?.servingRuntimeEditInfo,
       false,
       acceleratorProfileState,
+      selectedAcceleratorProfile,
       NamespaceApplicationCase.KSERVE_PROMOTION,
       projectContext?.currentProject,
       servingRuntimeName,
@@ -184,6 +195,7 @@ const DeployNIMServiceModal: React.FC<DeployNIMServiceModalProps> = ({
       servingRuntimeName,
       false,
       acceleratorProfileState,
+      selectedAcceleratorProfile,
       allowCreate,
       editInfo?.secrets,
       false,
@@ -296,7 +308,8 @@ const DeployNIMServiceModal: React.FC<DeployNIMServiceModalProps> = ({
               sizes={sizes}
               servingRuntimeSelected={servingRuntimeSelected}
               acceleratorProfileState={acceleratorProfileState}
-              setAcceleratorProfileState={setAcceleratorProfileState}
+              selectedAcceleratorProfile={selectedAcceleratorProfile}
+              setSelectedAcceleratorProfile={setSelectedAcceleratorProfile}
               infoContent="Select a server size that will accommodate your largest model. See the product documentation for more information."
             />
           </StackItem>