[ML] Adds ELSER config to the Trained Models UI

x-pack/packages/ml/trained_models_utils/src/constants/trained_models.ts

@@ -39,4 +39,34 @@ export const BUILT_IN_MODEL_TYPE = i18n.translate(
   { defaultMessage: 'built-in' }
 );
 
+export const CURATED_MODEL_TYPE = i18n.translate(
+  'xpack.ml.trainedModels.modelsList.curatedModelLabel',
+  { defaultMessage: 'curated' }
+);
+
 export const BUILT_IN_MODEL_TAG = 'prepackaged';
+
+export const CURATED_MODEL_TAG = 'curated';
+
+export const CURATED_MODEL_DEFINITIONS = {
+  '.elser_model_1_SNAPSHOT': {
+    config: {
+      input: {
+        field_names: ['text_field'],
+      },
+    },
+    description: i18n.translate('xpack.ml.trainedModels.modelsList.elserDescription', {
+      defaultMessage: 'Elastic Learned Sparse EncodeR',
+    }),
+  },
+} as const;
+
+export const MODEL_STATE = {
+  ...DEPLOYMENT_STATE,
+  DOWNLOADING: i18n.translate('xpack.ml.trainedModels.modelsList.downloadingStateLabel', {
+    defaultMessage: 'downloading',
+  }),
+  DOWNLOADED: i18n.translate('xpack.ml.trainedModels.modelsList.downloadedStateLabel', {
+    defaultMessage: 'downloaded',
+  }),
+} as const;

x-pack/plugins/ml/public/application/model_management/model_actions.tsx

@@ -15,6 +15,10 @@ import {
   DEPLOYMENT_STATE,
   TRAINED_MODEL_TYPE,
 } from '@kbn/ml-trained-models-utils';
+import {
+  CURATED_MODEL_TAG,
+  MODEL_STATE,
+} from '@kbn/ml-trained-models-utils/src/constants/trained_models';
 import { useTrainedModelsApiService } from '../services/ml_api_service/trained_models';
 import { getUserConfirmationProvider } from './force_stop_dialog';
 import { useToastNotificationService } from '../services/toast_notification_service';
@@ -154,7 +158,7 @@ export function useModelActions({
         type: 'icon',
         isPrimary: true,
         enabled: (item) => {
-          return canStartStopTrainedModels && !isLoading;
+          return canStartStopTrainedModels && !isLoading && item.state !== MODEL_STATE.DOWNLOADING;
         },
         available: (item) => item.model_type === TRAINED_MODEL_TYPE.PYTORCH,
         onClick: async (item) => {
@@ -317,6 +321,47 @@ export function useModelActions({
           }
         },
       },
+      {
+        name: i18n.translate('xpack.ml.inference.modelsList.downloadModelActionLabel', {
+          defaultMessage: 'Download model',
+        }),
+        description: i18n.translate('xpack.ml.inference.modelsList.downloadModelActionLabel', {
+          defaultMessage: 'Download model',
+        }),
+        'data-test-subj': 'mlModelsTableRowDownloadModelAction',
+        icon: 'download',
+        type: 'icon',
+        isPrimary: true,
+        available: (item) => item.tags.includes(CURATED_MODEL_TAG),
+        enabled: (item) => !item.state && !isLoading,
+        onClick: async (item) => {
+          try {
+            onLoading(true);
+            await trainedModelsApiService.putConfig(item.model_id, item.putConfig!);
+            displaySuccessToast(
+              i18n.translate('xpack.ml.trainedModels.modelsList.downloadSuccess', {
+                defaultMessage: 'Download of the model "{modelId}" has been started successfully.',
+                values: {
+                  modelId: item.model_id,
+                },
+              })
+            );
+            // Need to fetch model state updates
+            await fetchModels();
+          } catch (e) {
+            displayErrorToast(
+              e,
+              i18n.translate('xpack.ml.trainedModels.modelsList.downloadFailed', {
+                defaultMessage: 'Failed to download "{modelId}"',
+                values: {
+                  modelId: item.model_id,
+                },
+              })
+            );
+            onLoading(false);
+          }
+        },
+      },
       {
         name: (model) => {
           const enabled = !isPopulatedObject(model.pipelines);
@@ -350,7 +395,7 @@ export function useModelActions({
         onClick: (model) => {
           onModelsDeleteRequest([model.model_id]);
         },
-        available: (item) => canDeleteTrainedModels && !isBuiltInModel(item),
+        available: (item) => canDeleteTrainedModels && !isBuiltInModel(item) && !item.putConfig,
         enabled: (item) => {
           // TODO check for permissions to delete ingest pipelines.
           // ATM undefined means pipelines fetch failed server-side.

x-pack/plugins/ml/public/application/model_management/models_list.tsx

@@ -33,6 +33,12 @@ import {
   DEPLOYMENT_STATE,
 } from '@kbn/ml-trained-models-utils';
 import { isDefined } from '@kbn/ml-is-defined';
+import {
+  CURATED_MODEL_DEFINITIONS,
+  CURATED_MODEL_TAG,
+  CURATED_MODEL_TYPE,
+  MODEL_STATE,
+} from '@kbn/ml-trained-models-utils/src/constants/trained_models';
 import { useModelActions } from './model_actions';
 import { ModelsTableToConfigMapping } from '.';
 import { ModelsBarStats, StatsBar } from '../components/stats_bar';
@@ -62,6 +68,8 @@ export type ModelItem = TrainedModelConfigResponse & {
   stats?: Stats & { deployment_stats: TrainedModelDeploymentStatsResponse[] };
   pipelines?: ModelPipelines['pipelines'] | null;
   deployment_ids: string[];
+  putConfig?: object;
+  state: string;
 };
 
 export type ModelItemFull = Required<ModelItem>;
@@ -134,6 +142,35 @@ export const ModelsList: FC<Props> = ({
     []
   );
 
+  const isCuratedModel = useCallback(
+    (item: ModelItem) => item.tags.includes(CURATED_MODEL_TAG),
+    []
+  );
+
+  /**
+   * Checks if the model download complete.
+   */
+  const isDownloadComplete = useCallback(
+    async (modelId: string): Promise<boolean> => {
+      try {
+        const response = await trainedModelsApiService.getTrainedModels(modelId, {
+          include: 'definition_status',
+        });
+        // @ts-ignore
+        return !!response[0]?.fully_defined;
+      } catch (error) {
+        displayErrorToast(
+          error,
+          i18n.translate('xpack.ml.trainedModels.modelsList.downloadStatusCheckErrorMessage', {
+            defaultMessage: 'Failed to check download status',
+          })
+        );
+      }
+      return false;
+    },
+    [trainedModelsApiService, displayErrorToast]
+  );
+
   /**
    * Fetches trained models.
    */
@@ -158,6 +195,7 @@ export const ModelsList: FC<Props> = ({
                   model.model_type,
                   ...Object.keys(model.inference_config),
                   ...(isBuiltInModel(model as ModelItem) ? [BUILT_IN_MODEL_TYPE] : []),
+                  ...(isCuratedModel(model as ModelItem) ? [CURATED_MODEL_TYPE] : []),
                 ],
               }
             : {}),
@@ -239,7 +277,26 @@ export const ModelsList: FC<Props> = ({
           model.deployment_ids = modelStats
             .map((v) => v.deployment_stats?.deployment_id)
             .filter(isDefined);
+          model.state = model.stats.deployment_stats?.some(
+            (v) => v.state === DEPLOYMENT_STATE.STARTED
+          )
+            ? DEPLOYMENT_STATE.STARTED
+            : '';
         });
+
+        const curatedModels = models.filter((model) =>
+          CURATED_MODEL_DEFINITIONS.hasOwnProperty(model.model_id)
+        );
+        if (curatedModels.length > 0) {
+          for (const model of curatedModels) {
+            if (model.state === MODEL_STATE.STARTED) {
+              // no need to check for the download status if the model has been deployed
+              continue;
+            }
+            const isDownloaded = await isDownloadComplete(model.model_id);
+            model.state = isDownloaded ? MODEL_STATE.DOWNLOADED : MODEL_STATE.DOWNLOADING;
+          }
+        }
       }
 
       return true;
@@ -310,21 +367,26 @@ export const ModelsList: FC<Props> = ({
       align: 'left',
       width: '40px',
       isExpander: true,
-      render: (item: ModelItem) => (
-        <EuiButtonIcon
-          onClick={toggleDetails.bind(null, item)}
-          aria-label={
-            itemIdToExpandedRowMap[item.model_id]
-              ? i18n.translate('xpack.ml.trainedModels.modelsList.collapseRow', {
-                  defaultMessage: 'Collapse',
-                })
-              : i18n.translate('xpack.ml.trainedModels.modelsList.expandRow', {
-                  defaultMessage: 'Expand',
-                })
-          }
-          iconType={itemIdToExpandedRowMap[item.model_id] ? 'arrowDown' : 'arrowRight'}
-        />
-      ),
+      render: (item: ModelItem) => {
+        if (!item.stats) {
+          return null;
+        }
+        return (
+          <EuiButtonIcon
+            onClick={toggleDetails.bind(null, item)}
+            aria-label={
+              itemIdToExpandedRowMap[item.model_id]
+                ? i18n.translate('xpack.ml.trainedModels.modelsList.collapseRow', {
+                    defaultMessage: 'Collapse',
+                  })
+                : i18n.translate('xpack.ml.trainedModels.modelsList.expandRow', {
+                    defaultMessage: 'Expand',
+                  })
+            }
+            iconType={itemIdToExpandedRowMap[item.model_id] ? 'arrowDown' : 'arrowRight'}
+          />
+        );
+      },
       'data-test-subj': 'mlModelsTableRowDetailsToggle',
     },
     {
@@ -368,17 +430,13 @@ export const ModelsList: FC<Props> = ({
       'data-test-subj': 'mlModelsTableColumnType',
     },
     {
+      field: 'state',
       name: i18n.translate('xpack.ml.trainedModels.modelsList.stateHeader', {
         defaultMessage: 'State',
       }),
       align: 'left',
       truncateText: false,
-      render: (model: ModelItem) => {
-        const state = model.stats?.deployment_stats?.some(
-          (v) => v.state === DEPLOYMENT_STATE.STARTED
-        )
-          ? DEPLOYMENT_STATE.STARTED
-          : '';
+      render: (state: string) => {
         return state ? <EuiBadge color="hollow">{state}</EuiBadge> : null;
       },
       'data-test-subj': 'mlModelsTableColumnDeploymentState',
@@ -510,6 +568,22 @@ export const ModelsList: FC<Props> = ({
       : {}),
   };
 
+  const resultItems = useMemo<ModelItem[]>(() => {
+    const idSet = new Set(items.map((i) => i.model_id));
+    const notDownloaded: ModelItem[] = Object.entries(CURATED_MODEL_DEFINITIONS)
+      .filter(([modelId]) => !idSet.has(modelId))
+      .map(([modelId, modelDefinition]) => {
+        return {
+          model_id: modelId,
+          type: [CURATED_MODEL_TYPE],
+          tags: [CURATED_MODEL_TAG],
+          putConfig: modelDefinition.config,
+          description: modelDefinition.description,
+        } as ModelItem;
+      });
+    return [...items, ...notDownloaded];
+  }, [items]);
+
   return (
     <>
       <SavedObjectsWarning onCloseFlyout={fetchModelsData} forceRefresh={isLoading} />
@@ -531,7 +605,7 @@ export const ModelsList: FC<Props> = ({
           isExpandable={true}
           itemIdToExpandedRowMap={itemIdToExpandedRowMap}
           isSelectable={false}
-          items={items}
+          items={resultItems}
           itemId={ModelsTableToConfigMapping.id}
           loading={isLoading}
           search={search}

x-pack/plugins/ml/public/application/services/ml_api_service/trained_models.ts

@@ -196,6 +196,14 @@ export function trainedModelsApiProvider(httpService: HttpService) {
         query: { type, node, showClosedJobs },
       });
     },
+
+    putConfig(modelId: string, config: object) {
+      return httpService.http<estypes.MlPutTrainedModelResponse>({
+        path: `${apiBasePath}/trained_models/${modelId}`,
+        method: 'PUT',
+        body: JSON.stringify(config),
+      });
+    },
   };
 }