[ML] Trained Models: Optimize trained models Kibana API

x-pack/plugins/ml/common/types/trained_models.ts

@@ -5,14 +5,23 @@
  * 2.0.
  */
 import type * as estypes from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
-import type { TrainedModelType } from '@kbn/ml-trained-models-utils';
+import type { MlInferenceConfigCreateContainer } from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
+import type {
+  InferenceAPIConfigResponse,
+  ModelDefinitionResponse,
+  ModelState,
+  TrainedModelType,
+} from '@kbn/ml-trained-models-utils';
+import {
+  BUILT_IN_MODEL_TAG,
+  ELASTIC_MODEL_TAG,
+  TRAINED_MODEL_TYPE,
+} from '@kbn/ml-trained-models-utils';
 import type {
   DataFrameAnalyticsConfig,
   FeatureImportanceBaseline,
   TotalFeatureImportance,
 } from '@kbn/ml-data-frame-analytics-utils';
-import type { IndexName, IndicesIndexState } from '@elastic/elasticsearch/lib/api/types';
-import type { InferenceAPIConfigResponse } from '@kbn/ml-trained-models-utils';
 import type { XOR } from './common';
 import type { MlSavedObjectType } from './saved_objects';
 
@@ -95,33 +104,12 @@ export type PutTrainedModelConfig = {
 >; // compressed_definition and definition are mutually exclusive
 
 export type TrainedModelConfigResponse = estypes.MlTrainedModelConfig & {
-  /**
-   * Associated pipelines. Extends response from the ES endpoint.
-   */
-  pipelines?: Record<string, PipelineDefinition> | null;
-  origin_job_exists?: boolean;
-
-  metadata?: {
-    analytics_config: DataFrameAnalyticsConfig;
+  metadata?: estypes.MlTrainedModelConfig['metadata'] & {
+    analytics_config?: DataFrameAnalyticsConfig;
     input: unknown;
     total_feature_importance?: TotalFeatureImportance[];
     feature_importance_baseline?: FeatureImportanceBaseline;
-    model_aliases?: string[];
   } & Record<string, unknown>;
-  model_id: string;
-  model_type: TrainedModelType;
-  tags: string[];
-  version: string;
-  inference_config?: Record<string, any>;
-  indices?: Array<Record<IndexName, IndicesIndexState | null>>;
-  /**
-   * Whether the model has inference services
-   */
-  hasInferenceServices?: boolean;
-  /**
-   * Inference services associated with the model
-   */
-  inference_apis?: InferenceAPIConfigResponse[];
 };
 
 export interface PipelineDefinition {
@@ -309,3 +297,125 @@ export interface ModelDownloadState {
   total_parts: number;
   downloaded_parts: number;
 }
+
+export type Stats = Omit<TrainedModelStat, 'model_id' | 'deployment_stats'>;
+
+/**
+ * Additional properties for all items in the Trained models table
+ * */
+interface BaseModelItem {
+  type?: string[];
+  tags: string[];
+  /**
+   * Whether the model has inference services
+   */
+  hasInferenceServices?: boolean;
+  /**
+   * Inference services associated with the model
+   */
+  inference_apis?: InferenceAPIConfigResponse[];
+  /**
+   * Associated pipelines. Extends response from the ES endpoint.
+   */
+  pipelines?: Record<string, PipelineDefinition>;
+  /**
+   * Indices with associated pipelines that have inference processors utilizing the model deployments.
+   */
+  indices?: string[];
+}
+
+/** Common properties for existing NLP models and NLP model download configs */
+interface BaseNLPModelItem extends BaseModelItem {
+  disclaimer?: string;
+  recommended?: boolean;
+  supported?: boolean;
+  state: ModelState | undefined;
+  downloadState?: ModelDownloadState;
+}
+
+/** Model available for download */
+export type ModelDownloadItem = BaseNLPModelItem &
+  Omit<ModelDefinitionResponse, 'version' | 'config'> & {
+    putModelConfig?: object;
+    softwareLicense?: string;
+  };
+/** Trained NLP model, i.e. pytorch model returned by the trained_models API */
+export type NLPModelItem = BaseNLPModelItem &
+  TrainedModelItem & {
+    stats: Stats & { deployment_stats: TrainedModelDeploymentStatsResponse[] };
+    /**
+     * Description of the current model state
+     */
+    stateDescription?: string;
+    /**
+     * Deployment ids extracted from the deployment stats
+     */
+    deployment_ids: string[];
+  };
+
+export function isBaseNLPModelItem(item: unknown): item is BaseNLPModelItem {
+  return (
+    typeof item === 'object' &&
+    item !== null &&
+    'type' in item &&
+    Array.isArray(item.type) &&
+    item.type.includes(TRAINED_MODEL_TYPE.PYTORCH)
+  );
+}
+
+export function isNLPModelItem(item: unknown): item is NLPModelItem {
+  return isExistingModel(item) && item.model_type === TRAINED_MODEL_TYPE.PYTORCH;
+}
+
+export const isElasticModel = (item: TrainedModelConfigResponse) =>
+  item.tags.includes(ELASTIC_MODEL_TAG);
+
+export type ExistingModelBase = TrainedModelConfigResponse & BaseModelItem;
+
+/** Any model returned by the trained_models API, e.g. lang_ident, elser, dfa model */
+export type TrainedModelItem = ExistingModelBase & { stats: Stats };
+
+/** Trained DFA model */
+export type DFAModelItem = Omit<TrainedModelItem, 'inference_config'> & {
+  origin_job_exists?: boolean;
+  inference_config?: Pick<MlInferenceConfigCreateContainer, 'classification' | 'regression'>;
+  metadata?: estypes.MlTrainedModelConfig['metadata'] & {
+    analytics_config: DataFrameAnalyticsConfig;
+    input: unknown;
+    total_feature_importance?: TotalFeatureImportance[];
+    feature_importance_baseline?: FeatureImportanceBaseline;
+  } & Record<string, unknown>;
+};
+
+export type TrainedModelWithPipelines = TrainedModelItem & {
+  pipelines: Record<string, PipelineDefinition>;
+};
+
+export function isExistingModel(item: unknown): item is TrainedModelItem {
+  return (
+    typeof item === 'object' &&
+    item !== null &&
+    'model_type' in item &&
+    'create_time' in item &&
+    !!item.create_time
+  );
+}
+
+export function isDFAModelItem(item: unknown): item is DFAModelItem {
+  return isExistingModel(item) && item.model_type === TRAINED_MODEL_TYPE.TREE_ENSEMBLE;
+}
+
+export function isModelDownloadItem(item: TrainedModelUIItem): item is ModelDownloadItem {
+  return 'putModelConfig' in item && !!item.type?.includes(TRAINED_MODEL_TYPE.PYTORCH);
+}
+
+export const isBuiltInModel = (item: TrainedModelConfigResponse | TrainedModelUIItem) =>
+  item.tags.includes(BUILT_IN_MODEL_TAG);
+/**
+ * This type represents a union of different model entities:
+ * - Any existing trained model returned by the API, e.g., lang_ident_model_1, DFA models, etc.
+ * - Hosted model configurations available for download, e.g., ELSER or E5
+ * - NLP models already downloaded into Elasticsearch
+ * - DFA models
+ */
+export type TrainedModelUIItem = TrainedModelItem | ModelDownloadItem | NLPModelItem | DFAModelItem;

x-pack/plugins/ml/public/application/components/ml_inference/add_inference_pipeline_flyout.tsx

@@ -20,7 +20,7 @@ import {
 import { i18n } from '@kbn/i18n';
 import { extractErrorProperties } from '@kbn/ml-error-utils';
 
-import type { ModelItem } from '../../model_management/models_list';
+import type { DFAModelItem } from '../../../../common/types/trained_models';
 import type { AddInferencePipelineSteps } from './types';
 import { ADD_INFERENCE_PIPELINE_STEPS } from './constants';
 import { AddInferencePipelineFooter } from '../shared';
@@ -39,7 +39,7 @@ import { useFetchPipelines } from './hooks/use_fetch_pipelines';
 
 export interface AddInferencePipelineFlyoutProps {
   onClose: () => void;
-  model: ModelItem;
+  model: DFAModelItem;
 }
 
 export const AddInferencePipelineFlyout: FC<AddInferencePipelineFlyoutProps> = ({

x-pack/plugins/ml/public/application/components/ml_inference/components/processor_configuration.tsx

@@ -25,7 +25,7 @@ import {
 import { i18n } from '@kbn/i18n';
 import { FormattedMessage } from '@kbn/i18n-react';
 import { CodeEditor } from '@kbn/code-editor';
-import type { ModelItem } from '../../../model_management/models_list';
+import type { DFAModelItem } from '../../../../../common/types/trained_models';
 import {
   EDIT_MESSAGE,
   CANCEL_EDIT_MESSAGE,
@@ -56,9 +56,9 @@ interface Props {
   condition?: string;
   fieldMap: MlInferenceState['fieldMap'];
   handleAdvancedConfigUpdate: (configUpdate: Partial<MlInferenceState>) => void;
-  inferenceConfig: ModelItem['inference_config'];
-  modelInferenceConfig: ModelItem['inference_config'];
-  modelInputFields: ModelItem['input'];
+  inferenceConfig: DFAModelItem['inference_config'];
+  modelInferenceConfig: DFAModelItem['inference_config'];
+  modelInputFields: DFAModelItem['input'];
   modelType?: InferenceModelTypes;
   setHasUnsavedChanges: React.Dispatch<React.SetStateAction<boolean>>;
   tag?: string;

x-pack/plugins/ml/public/application/components/ml_inference/state.ts

@@ -6,10 +6,10 @@
  */
 
 import { getAnalysisType } from '@kbn/ml-data-frame-analytics-utils';
+import type { DFAModelItem } from '../../../../common/types/trained_models';
 import type { MlInferenceState } from './types';
-import type { ModelItem } from '../../model_management/models_list';
 
-export const getModelType = (model: ModelItem): string | undefined => {
+export const getModelType = (model: DFAModelItem): string | undefined => {
   const analysisConfig = model.metadata?.analytics_config?.analysis;
   return analysisConfig !== undefined ? getAnalysisType(analysisConfig) : undefined;
 };
@@ -54,13 +54,17 @@ export const getDefaultOnFailureConfiguration = (): MlInferenceState['onFailure'
   },
 ];
 
-export const getInitialState = (model: ModelItem): MlInferenceState => {
+export const getInitialState = (model: DFAModelItem): MlInferenceState => {
   const modelType = getModelType(model);
   let targetField;
 
   if (modelType !== undefined) {
     targetField = model.inference_config
-      ? `ml.inference.${model.inference_config[modelType].results_field}`
+      ? `ml.inference.${
+          model.inference_config[
+            modelType as keyof Exclude<DFAModelItem['inference_config'], undefined>
+          ]!.results_field
+        }`
       : undefined;
   }
 

x-pack/plugins/ml/public/application/data_frame_analytics/pages/components/analytics_selector/analytics_id_selector.tsx

@@ -154,6 +154,7 @@ export function AnalyticsIdSelector({
   async function fetchAnalyticsModels() {
     setIsLoading(true);
     try {
+      // FIXME should if fetch all trained models?
       const response = await trainedModelsApiService.getTrainedModels();
       setTrainedModels(response);
     } catch (e) {

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

@@ -30,12 +30,12 @@ import { FormattedMessage } from '@kbn/i18n-react';
 import React, { type FC, useMemo, useState } from 'react';
 import { groupBy } from 'lodash';
 import { ElandPythonClient } from '@kbn/inference_integration_flyout';
+import type { ModelDownloadItem } from '../../../common/types/trained_models';
 import { usePermissionCheck } from '../capabilities/check_capabilities';
 import { useMlKibana } from '../contexts/kibana';
-import type { ModelItem } from './models_list';
 
 export interface AddModelFlyoutProps {
-  modelDownloads: ModelItem[];
+  modelDownloads: ModelDownloadItem[];
   onClose: () => void;
   onSubmit: (modelId: string) => void;
 }
@@ -138,7 +138,7 @@ export const AddModelFlyout: FC<AddModelFlyoutProps> = ({ onClose, onSubmit, mod
 };
 
 interface ClickToDownloadTabContentProps {
-  modelDownloads: ModelItem[];
+  modelDownloads: ModelDownloadItem[];
   onModelDownload: (modelId: string) => void;
 }
 

x-pack/plugins/ml/public/application/model_management/create_pipeline_for_model/create_pipeline_for_model_flyout.tsx

@@ -21,7 +21,7 @@ import { i18n } from '@kbn/i18n';
 import { extractErrorProperties } from '@kbn/ml-error-utils';
 import type { SupportedPytorchTasksType } from '@kbn/ml-trained-models-utils';
 
-import type { ModelItem } from '../models_list';
+import type { TrainedModelItem } from '../../../../common/types/trained_models';
 import type { AddInferencePipelineSteps } from '../../components/ml_inference/types';
 import { ADD_INFERENCE_PIPELINE_STEPS } from '../../components/ml_inference/constants';
 import { AddInferencePipelineFooter } from '../../components/shared';
@@ -40,7 +40,7 @@ import { useTestTrainedModelsContext } from '../test_models/test_trained_models_
 
 export interface CreatePipelineForModelFlyoutProps {
   onClose: (refreshList?: boolean) => void;
-  model: ModelItem;
+  model: TrainedModelItem;
 }
 
 export const CreatePipelineForModelFlyout: FC<CreatePipelineForModelFlyoutProps> = ({

x-pack/plugins/ml/public/application/model_management/create_pipeline_for_model/state.ts

@@ -7,8 +7,8 @@
 
 import type * as estypes from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
 import type { IngestInferenceProcessor } from '@elastic/elasticsearch/lib/api/types';
+import type { TrainedModelItem } from '../../../../common/types/trained_models';
 import { getDefaultOnFailureConfiguration } from '../../components/ml_inference/state';
-import type { ModelItem } from '../models_list';
 
 export interface InferecePipelineCreationState {
   creatingPipeline: boolean;
@@ -26,7 +26,7 @@ export interface InferecePipelineCreationState {
 }
 
 export const getInitialState = (
-  model: ModelItem,
+  model: TrainedModelItem,
   initialPipelineConfig: estypes.IngestPipeline | undefined
 ): InferecePipelineCreationState => ({
   creatingPipeline: false,

x-pack/plugins/ml/public/application/model_management/create_pipeline_for_model/test_trained_model.tsx

@@ -12,13 +12,13 @@ import { i18n } from '@kbn/i18n';
 import { FormattedMessage } from '@kbn/i18n-react';
 import type * as estypes from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
 
-import type { ModelItem } from '../models_list';
+import type { TrainedModelItem } from '../../../../common/types/trained_models';
 import { TestTrainedModelContent } from '../test_models/test_trained_model_content';
 import { useMlKibana } from '../../contexts/kibana';
 import { type InferecePipelineCreationState } from './state';
 
 interface ContentProps {
-  model: ModelItem;
+  model: TrainedModelItem;
   handlePipelineConfigUpdate: (configUpdate: Partial<InferecePipelineCreationState>) => void;
   externalPipelineConfig?: estypes.IngestPipeline;
 }

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

@@ -22,14 +22,15 @@ import {
   EuiSpacer,
 } from '@elastic/eui';
 import { isPopulatedObject } from '@kbn/ml-is-populated-object';
+import type { TrainedModelItem, TrainedModelUIItem } from '../../../common/types/trained_models';
+import { isExistingModel } from '../../../common/types/trained_models';
 import { type WithRequired } from '../../../common/types/common';
 import { useTrainedModelsApiService } from '../services/ml_api_service/trained_models';
 import { useToastNotificationService } from '../services/toast_notification_service';
 import { DeleteSpaceAwareItemCheckModal } from '../components/delete_space_aware_item_check_modal';
-import { type ModelItem } from './models_list';
 
 interface DeleteModelsModalProps {
-  models: ModelItem[];
+  models: TrainedModelUIItem[];
   onClose: (refreshList?: boolean) => void;
 }
 
@@ -42,11 +43,14 @@ export const DeleteModelsModal: FC<DeleteModelsModalProps> = ({ models, onClose
 
   const modelIds = models.map((m) => m.model_id);
 
-  const modelsWithPipelines = models.filter((m) => isPopulatedObject(m.pipelines)) as Array<
-    WithRequired<ModelItem, 'pipelines'>
-  >;
+  const modelsWithPipelines = models.filter(
+    (m): m is WithRequired<TrainedModelItem, 'pipelines'> =>
+      isExistingModel(m) && isPopulatedObject(m.pipelines)
+  );
 
-  const modelsWithInferenceAPIs = models.filter((m) => m.hasInferenceServices);
+  const modelsWithInferenceAPIs = models.filter(
+    (m): m is TrainedModelItem => isExistingModel(m) && !!m.hasInferenceServices
+  );
 
   const inferenceAPIsIDs: string[] = modelsWithInferenceAPIs.flatMap((model) => {
     return (model.inference_apis ?? []).map((inference) => inference.inference_id);