Refactor accelerator profile state hooks to no longer manage its own …

…selected split init and selected accelerator
Gkrumbach07 · Aug 16, 2024 · 74eae10 · 74eae10
1 parent a011b1b
commit 74eae10
Show file tree

Hide file tree

Showing 22 changed files with 382 additions and 265 deletions.
diff --git a/frontend/src/api/k8s/inferenceServices.ts b/frontend/src/api/k8s/inferenceServices.ts
@@ -14,6 +14,7 @@ import { translateDisplayNameForK8s } from '~/concepts/k8s/utils';
 import { applyK8sAPIOptions } from '~/api/apiMergeUtils';
 import { AcceleratorProfileState } from '~/utilities/useAcceleratorProfileState';
 import { ContainerResources } from '~/types';
+import { AcceleratorProfileSelectFieldState } from '~/pages/notebookController/screens/server/AcceleratorProfileSelectField';
 import { getModelServingProjects } from './projects';
 import { assemblePodSpecOptions } from './utils';
 
@@ -23,7 +24,8 @@ export const assembleInferenceService = (
   editName?: string,
   isModelMesh?: boolean,
   inferenceService?: InferenceServiceKind,
-  acceleratorState?: AcceleratorProfileState,
+  initialAcceleratorProfile?: AcceleratorProfileState,
+  selectedAcceleratorProfile?: AcceleratorProfileSelectFieldState,
 ): InferenceServiceKind => {
   const {
     storage,
@@ -143,7 +145,11 @@ export const assembleInferenceService = (
       },
     };
 
-    const { tolerations, resources } = assemblePodSpecOptions(resourceSettings, acceleratorState);
+    const { tolerations, resources } = assemblePodSpecOptions(
+      resourceSettings,
+      initialAcceleratorProfile,
+      selectedAcceleratorProfile,
+    );
 
     if (tolerations.length !== 0) {
       updateInferenceService.spec.predictor.tolerations = tolerations;
@@ -224,7 +230,8 @@ export const createInferenceService = (
   data: CreatingInferenceServiceObject,
   secretKey?: string,
   isModelMesh?: boolean,
-  acceleratorState?: AcceleratorProfileState,
+  initialAcceleratorProfile?: AcceleratorProfileState,
+  selectedAcceleratorProfile?: AcceleratorProfileSelectFieldState,
   dryRun = false,
 ): Promise<InferenceServiceKind> => {
   const inferenceService = assembleInferenceService(
@@ -233,7 +240,8 @@ export const createInferenceService = (
     undefined,
     isModelMesh,
     undefined,
-    acceleratorState,
+    initialAcceleratorProfile,
+    selectedAcceleratorProfile,
   );
   return k8sCreateResource<InferenceServiceKind>(
     applyK8sAPIOptions(
@@ -251,7 +259,8 @@ export const updateInferenceService = (
   existingData: InferenceServiceKind,
   secretKey?: string,
   isModelMesh?: boolean,
-  acceleratorState?: AcceleratorProfileState,
+  initialAcceleratorProfile?: AcceleratorProfileState,
+  selectedAcceleratorProfile?: AcceleratorProfileSelectFieldState,
   dryRun = false,
 ): Promise<InferenceServiceKind> => {
   const inferenceService = assembleInferenceService(
@@ -260,7 +269,8 @@ export const updateInferenceService = (
     existingData.metadata.name,
     isModelMesh,
     existingData,
-    acceleratorState,
+    initialAcceleratorProfile,
+    selectedAcceleratorProfile,
   );
 
   return k8sUpdateResource<InferenceServiceKind>(

diff --git a/frontend/src/api/k8s/notebooks.ts b/frontend/src/api/k8s/notebooks.ts
@@ -41,7 +41,8 @@ export const assembleNotebook = (
     description,
     notebookSize,
     envFrom,
-    acceleratorProfile,
+    initialAcceleratorProfile,
+    selectedAcceleratorProfile,
     image,
     volumes: formVolumes,
     volumeMounts: formVolumeMounts,
@@ -55,7 +56,8 @@ export const assembleNotebook = (
 
   const { affinity, tolerations, resources } = assemblePodSpecOptions(
     notebookSize.resources,
-    acceleratorProfile,
+    initialAcceleratorProfile,
+    selectedAcceleratorProfile,
     tolerationSettings,
     existingTolerations,
     undefined,
@@ -107,8 +109,7 @@ export const assembleNotebook = (
         'notebooks.opendatahub.io/last-image-selection': imageSelection,
         'notebooks.opendatahub.io/inject-oauth': 'true',
         'opendatahub.io/username': username,
-        'opendatahub.io/accelerator-name':
-          acceleratorProfile.acceleratorProfile?.metadata.name || '',
+        'opendatahub.io/accelerator-name': selectedAcceleratorProfile.profile?.metadata.name || '',
       },
       name: notebookId,
       namespace: projectName,

diff --git a/frontend/src/api/k8s/servingRuntimes.ts b/frontend/src/api/k8s/servingRuntimes.ts
@@ -19,6 +19,7 @@ import { getModelServingRuntimeName } from '~/pages/modelServing/utils';
 import { getDisplayNameFromK8sResource, translateDisplayNameForK8s } from '~/concepts/k8s/utils';
 import { applyK8sAPIOptions } from '~/api/apiMergeUtils';
 import { AcceleratorProfileState } from '~/utilities/useAcceleratorProfileState';
+import { AcceleratorProfileSelectFieldState } from '~/pages/notebookController/screens/server/AcceleratorProfileSelectField';
 import { getModelServingProjects } from './projects';
 import { assemblePodSpecOptions, getshmVolume, getshmVolumeMount } from './utils';
 
@@ -28,7 +29,8 @@ export const assembleServingRuntime = (
   servingRuntime: ServingRuntimeKind,
   isCustomServingRuntimesEnabled: boolean,
   isEditing?: boolean,
-  acceleratorProfileState?: AcceleratorProfileState,
+  initialAcceleratorProfile?: AcceleratorProfileState,
+  selectedAcceleratorProfile?: AcceleratorProfileSelectFieldState,
   isModelMesh?: boolean,
 ): ServingRuntimeKind => {
   const { name: displayName, numReplicas, modelSize, externalRoute, tokenAuth } = data;
@@ -71,7 +73,7 @@ export const assembleServingRuntime = (
         ...(isCustomServingRuntimesEnabled && {
           'opendatahub.io/template-display-name': getDisplayNameFromK8sResource(servingRuntime),
           'opendatahub.io/accelerator-name':
-            acceleratorProfileState?.acceleratorProfile?.metadata.name || '',
+            selectedAcceleratorProfile?.profile?.metadata.name || '',
         }),
       },
     };
@@ -80,8 +82,7 @@ export const assembleServingRuntime = (
       ...updatedServingRuntime.metadata,
       annotations: {
         ...annotations,
-        'opendatahub.io/accelerator-name':
-          acceleratorProfileState?.acceleratorProfile?.metadata.name || '',
+        'opendatahub.io/accelerator-name': selectedAcceleratorProfile?.profile?.metadata.name || '',
         ...(isCustomServingRuntimesEnabled && { 'openshift.io/display-name': displayName.trim() }),
       },
     };
@@ -107,7 +108,8 @@ export const assembleServingRuntime = (
 
   const { affinity, tolerations, resources } = assemblePodSpecOptions(
     resourceSettings,
-    acceleratorProfileState,
+    initialAcceleratorProfile,
+    selectedAcceleratorProfile,
     undefined,
     servingRuntime.spec.tolerations,
     undefined,
@@ -210,15 +212,17 @@ export const updateServingRuntime = (options: {
   existingData: ServingRuntimeKind;
   isCustomServingRuntimesEnabled: boolean;
   opts?: K8sAPIOptions;
-  acceleratorProfileState?: AcceleratorProfileState;
+  initialAcceleratorProfile?: AcceleratorProfileState;
+  selectedAcceleratorProfile?: AcceleratorProfileSelectFieldState;
   isModelMesh?: boolean;
 }): Promise<ServingRuntimeKind> => {
   const {
     data,
     existingData,
     isCustomServingRuntimesEnabled,
     opts,
-    acceleratorProfileState,
+    initialAcceleratorProfile,
+    selectedAcceleratorProfile,
     isModelMesh,
   } = options;
 
@@ -228,7 +232,8 @@ export const updateServingRuntime = (options: {
     existingData,
     isCustomServingRuntimesEnabled,
     true,
-    acceleratorProfileState,
+    initialAcceleratorProfile,
+    selectedAcceleratorProfile,
     isModelMesh,
   );
 
@@ -249,7 +254,8 @@ export const createServingRuntime = (options: {
   servingRuntime: ServingRuntimeKind;
   isCustomServingRuntimesEnabled: boolean;
   opts?: K8sAPIOptions;
-  acceleratorProfileState?: AcceleratorProfileState;
+  initialAcceleratorProfile?: AcceleratorProfileState;
+  selectedAcceleratorProfile?: AcceleratorProfileSelectFieldState;
   isModelMesh?: boolean;
 }): Promise<ServingRuntimeKind> => {
   const {
@@ -258,7 +264,8 @@ export const createServingRuntime = (options: {
     servingRuntime,
     isCustomServingRuntimesEnabled,
     opts,
-    acceleratorProfileState,
+    initialAcceleratorProfile,
+    selectedAcceleratorProfile,
     isModelMesh,
   } = options;
   const assembledServingRuntime = assembleServingRuntime(
@@ -267,7 +274,8 @@ export const createServingRuntime = (options: {
     servingRuntime,
     isCustomServingRuntimesEnabled,
     false,
-    acceleratorProfileState,
+    initialAcceleratorProfile,
+    selectedAcceleratorProfile,
     isModelMesh,
   );
 

diff --git a/frontend/src/api/k8s/utils.ts b/frontend/src/api/k8s/utils.ts
@@ -1,3 +1,4 @@
+import { AcceleratorProfileSelectFieldState } from '~/pages/notebookController/screens/server/AcceleratorProfileSelectField';
 import {
   PodAffinity,
   ContainerResources,
@@ -11,7 +12,8 @@ import { AcceleratorProfileState } from '~/utilities/useAcceleratorProfileState'
 
 export const assemblePodSpecOptions = (
   resourceSettings: ContainerResources,
-  acceleratorProfileState?: AcceleratorProfileState,
+  initialAcceleratorProfile?: AcceleratorProfileState,
+  selectedAcceleratorProfile?: AcceleratorProfileSelectFieldState,
   tolerationSettings?: TolerationSettings,
   existingTolerations?: Toleration[],
   affinitySettings?: PodAffinity,
@@ -27,39 +29,36 @@ export const assemblePodSpecOptions = (
     requests: { ...existingResources?.requests, ...resourceSettings.requests },
   };
 
-  if (
-    acceleratorProfileState?.additionalOptions?.useExisting &&
-    !acceleratorProfileState.useExisting
-  ) {
+  if (selectedAcceleratorProfile?.useExistingSettings) {
     resources = structuredClone(resourceSettings);
   }
 
   // Clear the last accelerator from the resources
-  if (acceleratorProfileState?.initialAcceleratorProfile) {
+  if (initialAcceleratorProfile?.acceleratorProfile) {
     if (resources.limits) {
-      delete resources.limits[acceleratorProfileState.initialAcceleratorProfile.spec.identifier];
+      delete resources.limits[initialAcceleratorProfile.acceleratorProfile.spec.identifier];
     }
     if (resources.requests) {
-      delete resources.requests[acceleratorProfileState.initialAcceleratorProfile.spec.identifier];
+      delete resources.requests[initialAcceleratorProfile.acceleratorProfile.spec.identifier];
     }
   }
 
   // Add back the new accelerator to the resources if count > 0
-  if (acceleratorProfileState?.acceleratorProfile && acceleratorProfileState.count > 0) {
+  if (selectedAcceleratorProfile?.profile && selectedAcceleratorProfile.count > 0) {
     if (resources.limits) {
-      resources.limits[acceleratorProfileState.acceleratorProfile.spec.identifier] =
-        acceleratorProfileState.count;
+      resources.limits[selectedAcceleratorProfile.profile.spec.identifier] =
+        selectedAcceleratorProfile.count;
     }
     if (resources.requests) {
-      resources.requests[acceleratorProfileState.acceleratorProfile.spec.identifier] =
-        acceleratorProfileState.count;
+      resources.requests[selectedAcceleratorProfile.profile.spec.identifier] =
+        selectedAcceleratorProfile.count;
     }
   }
 
   const tolerations = determineTolerations(
     tolerationSettings,
-    acceleratorProfileState,
-    existingTolerations,
+    initialAcceleratorProfile,
+    selectedAcceleratorProfile,
   );
   return { affinity, tolerations, resources };
 };

diff --git a/frontend/src/pages/modelServing/screens/projects/ModelMeshSection/ServingRuntimeDetails.tsx b/frontend/src/pages/modelServing/screens/projects/ModelMeshSection/ServingRuntimeDetails.tsx
@@ -20,7 +20,7 @@ type ServingRuntimeDetailsProps = {
 
 const ServingRuntimeDetails: React.FC<ServingRuntimeDetailsProps> = ({ obj, isvc }) => {
   const { dashboardConfig } = React.useContext(AppContext);
-  const [acceleratorProfile] = useServingAcceleratorProfile(obj, isvc);
+  const acceleratorProfile = useServingAcceleratorProfile(obj, isvc);
   const selectedAcceleratorProfile = acceleratorProfile.acceleratorProfile;
   const enabledAcceleratorProfiles = acceleratorProfile.acceleratorProfiles.filter(
     (ac) => ac.spec.enabled,
@@ -64,12 +64,12 @@ const ServingRuntimeDetails: React.FC<ServingRuntimeDetailsProps> = ({ obj, isvc
               }`
             : enabledAcceleratorProfiles.length === 0
             ? 'No accelerator enabled'
-            : acceleratorProfile.useExisting
+            : acceleratorProfile.unknownProfileDetected
             ? 'Unknown'
             : 'No accelerator selected'}
         </DescriptionListDescription>
       </DescriptionListGroup>
-      {!acceleratorProfile.useExisting && acceleratorProfile.acceleratorProfile && (
+      {!acceleratorProfile.unknownProfileDetected && acceleratorProfile.acceleratorProfile && (
         <DescriptionListGroup>
           <DescriptionListTerm>Number of accelerators</DescriptionListTerm>
           <DescriptionListDescription>{acceleratorProfile.count}</DescriptionListDescription>

diff --git a/...src/pages/modelServing/screens/projects/ServingRuntimeModal/ManageServingRuntimeModal.tsx b/...src/pages/modelServing/screens/projects/ServingRuntimeModal/ManageServingRuntimeModal.tsx
@@ -18,6 +18,8 @@ import DashboardModalFooter from '~/concepts/dashboard/DashboardModalFooter';
 import { NamespaceApplicationCase } from '~/pages/projects/types';
 import { ServingRuntimeEditInfo } from '~/pages/modelServing/screens/types';
 import { useAccessReview } from '~/api';
+import { AcceleratorProfileSelectFieldState } from '~/pages/notebookController/screens/server/AcceleratorProfileSelectField';
+import useGenericObjectState from '~/utilities/useGenericObjectState';
 import ServingRuntimeReplicaSection from './ServingRuntimeReplicaSection';
 import ServingRuntimeSizeSection from './ServingRuntimeSizeSection';
 import ServingRuntimeTemplateSection from './ServingRuntimeTemplateSection';
@@ -49,8 +51,13 @@ const ManageServingRuntimeModal: React.FC<ManageServingRuntimeModalProps> = ({
   editInfo,
 }) => {
   const [createData, setCreateData, resetData, sizes] = useCreateServingRuntimeObject(editInfo);
-  const [acceleratorProfileState, setAcceleratorProfileState, resetAcceleratorProfileData] =
-    useServingAcceleratorProfile(editInfo?.servingRuntime);
+  const initialAcceleratorProfile = useServingAcceleratorProfile(editInfo?.servingRuntime);
+  const [selectedAcceleratorProfile, setSelectedAcceleratorProfile] =
+    useGenericObjectState<AcceleratorProfileSelectFieldState>({
+      profile: undefined,
+      count: 0,
+      useExistingSettings: false,
+    });
   const [actionInProgress, setActionInProgress] = React.useState(false);
   const [error, setError] = React.useState<Error | undefined>();
 
@@ -78,7 +85,13 @@ const ManageServingRuntimeModal: React.FC<ManageServingRuntimeModalProps> = ({
     actionInProgress ||
     tokenErrors ||
     !inputValueValid ||
-    !isModelServerEditInfoChanged(createData, sizes, acceleratorProfileState, editInfo);
+    !isModelServerEditInfoChanged(
+      createData,
+      sizes,
+      initialAcceleratorProfile,
+      selectedAcceleratorProfile,
+      editInfo,
+    );
 
   const servingRuntimeSelected = React.useMemo(
     () =>
@@ -92,7 +105,6 @@ const ManageServingRuntimeModal: React.FC<ManageServingRuntimeModalProps> = ({
     setError(undefined);
     setActionInProgress(false);
     resetData();
-    resetAcceleratorProfileData();
   };
 
   const setErrorModal = (e: Error) => {
@@ -115,7 +127,8 @@ const ManageServingRuntimeModal: React.FC<ManageServingRuntimeModalProps> = ({
       namespace,
       editInfo,
       allowCreate,
-      acceleratorProfileState,
+      initialAcceleratorProfile,
+      selectedAcceleratorProfile,
       NamespaceApplicationCase.MODEL_MESH_PROMOTION,
       currentProject,
       undefined,
@@ -162,7 +175,7 @@ const ManageServingRuntimeModal: React.FC<ManageServingRuntimeModalProps> = ({
               setData={setCreateData}
               templates={servingRuntimeTemplates || []}
               isEditing={!!editInfo}
-              acceleratorProfileState={acceleratorProfileState}
+              selectedAcceleratorProfile={selectedAcceleratorProfile}
             />
           </StackItem>
           <StackItem>
@@ -179,8 +192,9 @@ const ManageServingRuntimeModal: React.FC<ManageServingRuntimeModalProps> = ({
               setData={setCreateData}
               sizes={sizes}
               servingRuntimeSelected={servingRuntimeSelected}
-              acceleratorProfileState={acceleratorProfileState}
-              setAcceleratorProfileState={setAcceleratorProfileState}
+              acceleratorProfileState={initialAcceleratorProfile}
+              selectedAcceleratorProfile={selectedAcceleratorProfile}
+              setSelectedAcceleratorProfile={setSelectedAcceleratorProfile}
               infoContent="Select a server size that will accommodate your largest model. See the product documentation for more information."
             />
           </StackItem>