Fix prefill for S3 connections

opendatahub-io · Feb 5, 2025 · cca6d03 · cca6d03
1 parent 09b7d60
commit cca6d03
Show file tree

Hide file tree

Showing 8 changed files with 68 additions and 1 deletion.
diff --git a/frontend/src/concepts/connectionTypes/fields/ConnectionTypeFormFields.tsx b/frontend/src/concepts/connectionTypes/fields/ConnectionTypeFormFields.tsx
@@ -64,6 +64,7 @@ const ConnectionTypeFormFields: React.FC<Props> = ({
     return unmatched;
   }, [connectionValues, fields]);
 
+  console.log('connectionValues', connectionValues);
   const renderDataFields = (dataFields: ConnectionTypeDataField[]) =>
     dataFields.map((field, i) => {
       const id = `field-${field.envVar}`;

diff --git a/frontend/src/concepts/connectionTypes/utils.ts b/frontend/src/concepts/connectionTypes/utils.ts
@@ -14,7 +14,7 @@ import {
   ConnectionTypeValueType,
 } from '~/concepts/connectionTypes/types';
 import { enumIterator } from '~/utilities/utils';
-import { AWSDataEntry } from '~/pages/projects/types';
+import { AWSDataEntry, EnvVariableDataEntry } from '~/pages/projects/types';
 import { AwsKeys } from '~/pages/projects/dataConnections/const';
 
 export const isConnectionTypeDataFieldType = (
@@ -224,6 +224,24 @@ export const getDefaultValues = (
   return defaults;
 };
 
+export const getMRConnectionValues = (
+  connectionValues: EnvVariableDataEntry[] | string,
+): { [key: string]: ConnectionTypeValueType } => {
+  const defaults: {
+    [key: string]: ConnectionTypeValueType;
+  } = {};
+  if (typeof connectionValues !== 'string') {
+    connectionValues.map(
+      (connectionValue) => (defaults[connectionValue.key] = connectionValue.value),
+    );
+    return defaults;
+  }
+  if (typeof connectionValues === 'string') {
+    defaults.URI = connectionValues;
+  }
+  return defaults;
+};
+
 export const withRequiredFields = (
   connectionType?: ConnectionTypeConfigMapObj,
   envVars?: string[],

diff --git a/frontend/src/concepts/modelRegistry/types.ts b/frontend/src/concepts/modelRegistry/types.ts
@@ -1,4 +1,5 @@
 import { K8sAPIOptions } from '~/k8sTypes';
+import { ModelLocationType } from '~/pages/modelRegistry/screens/RegisterModel/useRegisterModelData';
 
 export enum ModelState {
   LIVE = 'LIVE',
@@ -100,6 +101,7 @@ export type ModelArtifact = ModelRegistryBase & {
   storagePath?: string;
   modelFormatVersion?: string;
   serviceAccountName?: string;
+  modelLocationType?: ModelLocationType;
   artifactType: string;
 };
 

diff --git a/frontend/src/pages/modelRegistry/screens/RegisterModel/utils.ts b/frontend/src/pages/modelRegistry/screens/RegisterModel/utils.ts
@@ -99,6 +99,7 @@ export const registerVersion = async (
       author,
       modelFormatName: formData.sourceModelFormat,
       modelFormatVersion: formData.sourceModelFormatVersion,
+      modelLocationType: formData.modelLocationType,
       // TODO fill in the name of the data connection we used to prefill if we used one
       // TODO this should be done as part of https://issues.redhat.com/browse/RHOAIENG-9914
       // storageKey: 'TODO',

diff --git a/...odelRegistry/screens/RegisteredModels/usePrefillDeployModalConnectionFromModelRegistry.ts b/...odelRegistry/screens/RegisteredModels/usePrefillDeployModalConnectionFromModelRegistry.ts
@@ -61,6 +61,7 @@ const usePrefillDeployModalConnectionFromModelRegistry = (
           setCreateData('storage', {
             awsData: prefilledAWSData,
             dataConnection: '',
+            connection: registeredModelDeployInfo.modelLocationType,
             path: storageFields.path,
             type: InferenceServiceStorageType.NEW_STORAGE,
             alert: {

diff --git a/frontend/src/pages/modelRegistry/screens/RegisteredModels/useRegisteredModelDeployInfo.ts b/frontend/src/pages/modelRegistry/screens/RegisteredModels/useRegisteredModelDeployInfo.ts
@@ -2,11 +2,13 @@ import React from 'react';
 import useModelArtifactsByVersionId from '~/concepts/modelRegistry/apiHooks/useModelArtifactsByVersionId';
 import useRegisteredModelById from '~/concepts/modelRegistry/apiHooks/useRegisteredModelById';
 import { ModelVersion } from '~/concepts/modelRegistry/types';
+import { uriToObjectStorageFields } from '~/concepts/modelRegistry/utils';
 
 export type RegisteredModelDeployInfo = {
   modelName: string;
   modelFormat?: string;
   modelArtifactUri?: string;
+  modelLocationType?: string;
   modelArtifactStorageKey?: string;
   modelVersionId?: string;
   registeredModelId?: string;
@@ -40,13 +42,22 @@ const useRegisteredModelDeployInfo = (
       };
     }
     const modelArtifact = modelArtifactList.items[0];
+    const storageFields = uriToObjectStorageFields(modelArtifact.uri || '');
+    let modelLocationType;
+    if (!storageFields) {
+      modelLocationType = 'uri-v1';
+    }
+    if (storageFields) {
+      modelLocationType = 's3';
+    }
     return {
       registeredModelDeployInfo: {
         modelName,
         modelFormat: modelArtifact.modelFormatName
           ? `${modelArtifact.modelFormatName} - ${modelArtifact.modelFormatVersion ?? ''}`
           : undefined,
         modelArtifactUri: modelArtifact.uri,
+        modelLocationType,
         modelArtifactStorageKey: modelArtifact.storageKey,
         modelVersionId: modelVersion.id,
         registeredModelId: modelVersion.registeredModelId,

diff --git a/frontend/src/pages/modelServing/screens/projects/InferenceServiceModal/ConnectionSection.tsx b/frontend/src/pages/modelServing/screens/projects/InferenceServiceModal/ConnectionSection.tsx
@@ -32,6 +32,7 @@ import {
   getConnectionTypeDisplayName,
   getConnectionTypeRef,
   getDefaultValues,
+  getMRConnectionValues,
   isConnectionTypeDataField,
   isUriConnection,
   isUriConnectionType,
@@ -201,12 +202,42 @@ const NewConnectionField: React.FC<NewConnectionFieldProps> = ({
       ? withRequiredFields(connectionTypes[0], S3ConnectionTypeKeys)
       : undefined,
   );
+
   const { data: nameDescData, onDataChange: setNameDescData } = useK8sNameDescriptionFieldData();
 
   const [connectionValues, setConnectionValues] = React.useState<{
     [key: string]: ConnectionTypeValueType;
   }>(enabledConnectionTypes.length === 1 ? getDefaultValues(enabledConnectionTypes[0]) : {});
 
+  React.useEffect(() => {
+    const locationType = data.storage.connection;
+    if (locationType) {
+      if (locationType === 's3') {
+        setSelectedConnectionType(
+          withRequiredFields(
+            connectionTypes.find((t) => getResourceNameFromK8sResource(t) === locationType),
+            S3ConnectionTypeKeys,
+          ),
+        );
+        // setConnectionValues(getMRConnectionValues(data.storage.awsData));
+      }
+      if (data.storage.uri) {
+        setSelectedConnectionType(
+          withRequiredFields(
+            connectionTypes.find(
+              (t) => getResourceNameFromK8sResource(t) === data.storage.connection,
+            ),
+            ['URI'],
+          ),
+        );
+        setConnectionValues(getMRConnectionValues(data.storage.uri));
+      }
+    }
+  }, [data.storage.connection, connectionTypes, data.storage.uri, data.storage.awsData]);
+
+  console.log('data.storage.awsData', data.storage.awsData);
+  console.log('connectionValues', connectionValues);
+
   const [validations, setValidations] = React.useState<{
     [key: string]: boolean;
   }>({});

diff --git a/frontend/src/pages/modelServing/screens/types.ts b/frontend/src/pages/modelServing/screens/types.ts
@@ -1,6 +1,7 @@
 import { AlertVariant } from '@patternfly/react-core';
 import { Connection } from '~/concepts/connectionTypes/types';
 import { SecretKind, ServingContainer, ServingRuntimeKind } from '~/k8sTypes';
+import { ModelLocationType } from '~/pages/modelRegistry/screens/RegisterModel/useRegisterModelData';
 import { DataConnection, EnvVariableDataEntry } from '~/pages/projects/types';
 import { ContainerResources } from '~/types';
 
@@ -92,6 +93,7 @@ export type InferenceServiceStorage = {
   type: InferenceServiceStorageType;
   path: string;
   dataConnection: string;
+  connection?: ModelLocationType | string;
   uri?: string;
   awsData: EnvVariableDataEntry[];
   alert?: {