[ML] Functional tests - support tiny trained models

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

@@ -66,6 +66,7 @@ export type PutTrainedModelConfig = {
     model_aliases?: string[];
   } & Record<string, unknown>;
   tags?: string[];
+  model_type?: TrainedModelType;
   inference_config?: Record<string, unknown>;
   input: { field_names: string[] };
 } & XOR<

x-pack/test/functional/apps/ml/short_tests/model_management/model_list.ts

@@ -6,14 +6,33 @@
  */
 
 import { FtrProviderContext } from '../../../../ftr_provider_context';
+import { TrainedModelName } from '../../../../services/ml/api';
 
 export default function ({ getService }: FtrProviderContext) {
   const ml = getService('ml');
 
+  const tinyTrainedModels = [
+    'fill_mask',
+    'ner',
+    'pass_through',
+    'text_classification',
+    'text_embedding',
+    'zero_shot',
+  ].map((type) => ({
+    id: `tiny_${type}`,
+    name: `pt_tiny_${type}` as TrainedModelName,
+    desription: `Tiny/Dummy PyTorch model (${type})`,
+    modelTypes: ['pytorch', type],
+  }));
+
   describe('trained models', function () {
     before(async () => {
-      await ml.trainedModels.createTestTrainedModels('classification', 15, true);
-      await ml.trainedModels.createTestTrainedModels('regression', 15);
+      for (const model of tinyTrainedModels) {
+        await ml.api.importTrainedModel(model.id, model.name);
+      }
+
+      await ml.api.createTestTrainedModels('classification', 15, true);
+      await ml.api.createTestTrainedModels('regression', 15);
     });
 
     after(async () => {
@@ -56,7 +75,7 @@ export default function ({ getService }: FtrProviderContext) {
           'should display the stats bar with the total number of models'
         );
         // +1 because of the built-in model
-        await ml.trainedModels.assertStats(31);
+        await ml.trainedModels.assertStats(37);
 
         await ml.testExecution.logTestStep('should display the table');
         await ml.trainedModels.assertTableExists();
@@ -81,6 +100,16 @@ export default function ({ getService }: FtrProviderContext) {
         await ml.trainedModelsTable.assertPipelinesTabContent(false);
       });
 
+      for (const model of tinyTrainedModels) {
+        it(`renders expanded row content correctly for imported tiny model ${model.id} without pipelines`, async () => {
+          await ml.trainedModelsTable.ensureRowIsExpanded(model.id);
+          await ml.trainedModelsTable.assertDetailsTabContent();
+          await ml.trainedModelsTable.assertInferenceConfigTabContent();
+          await ml.trainedModelsTable.assertStatsTabContent();
+          await ml.trainedModelsTable.assertPipelinesTabContent(false);
+        });
+      }
+
       it('displays the built-in model and no actions are enabled', async () => {
         await ml.testExecution.logTestStep('should display the model in the table');
         await ml.trainedModelsTable.filterWithSearchString(builtInModelData.modelId, 1);

x-pack/test/functional/services/ml/api.ts

@@ -9,7 +9,6 @@ import type * as estypes from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
 import expect from '@kbn/expect';
 import { ProvidedType } from '@kbn/test';
 import fs from 'fs';
-import path from 'path';
 import { Calendar } from '@kbn/ml-plugin/server/models/calendar';
 import { Annotation } from '@kbn/ml-plugin/common/types/annotations';
 import { DataFrameAnalyticsConfig } from '@kbn/ml-plugin/public/application/data_frame_analytics/common';
@@ -29,6 +28,16 @@ import { FtrProviderContext } from '../../ftr_provider_context';
 export type MlApi = ProvidedType<typeof MachineLearningAPIProvider>;
 
 type ModelType = 'regression' | 'classification';
+export type TrainedModelName =
+  | 'pt_tiny_fill_mask'
+  | 'pt_tiny_ner'
+  | 'pt_tiny_pass_through'
+  | 'pt_tiny_text_classification'
+  | 'pt_tiny_text_embedding'
+  | 'pt_tiny_zero_shot';
+export interface TrainedModelVocabulary {
+  vocabulary: string[];
+}
 
 export function MachineLearningAPIProvider({ getService }: FtrProviderContext) {
   const es = getService('es');
@@ -1135,6 +1144,38 @@ export function MachineLearningAPIProvider({ getService }: FtrProviderContext) {
       return model;
     },
 
+    async createTrainedModelVocabularyES(modelId: string, body: TrainedModelVocabulary) {
+      log.debug(`Creating vocabulary for trained model "${modelId}"`);
+      const { body: responseBody, status } = await esSupertest
+        .put(`/_ml/trained_models/${modelId}/vocabulary`)
+        .send(body);
+      this.assertResponseStatusCode(200, status, responseBody);
+
+      log.debug('> Trained model vocabulary created');
+    },
+
+    /**
+     * For the purpose of the functional tests where we only deal with very
+     * small models, we assume that the model definition can be uploaded as
+     * one part.
+     */
+    async uploadTrainedModelDefinitionES(modelId: string, modelDefinitionPath: string) {
+      log.debug(`Uploading definition for trained model "${modelId}"`);
+
+      const body = {
+        total_definition_length: fs.statSync(modelDefinitionPath).size,
+        definition: fs.readFileSync(modelDefinitionPath).toString('base64'),
+        total_parts: 1,
+      };
+
+      const { body: responseBody, status } = await esSupertest
+        .put(`/_ml/trained_models/${modelId}/definition/0`)
+        .send(body);
+      this.assertResponseStatusCode(200, status, responseBody);
+
+      log.debug('> Trained model definition uploaded');
+    },
+
     async deleteTrainedModelES(modelId: string) {
       log.debug(`Creating trained model with id "${modelId}"`);
       const { body: model, status } = await esSupertest.delete(`/_ml/trained_models/${modelId}`);
@@ -1149,24 +1190,9 @@ export function MachineLearningAPIProvider({ getService }: FtrProviderContext) {
       count: number = 10,
       withIngestPipelines = false
     ) {
-      const compressedDefinition = this.getCompressedModelDefinition(modelType);
+      const modelIds = new Array(count).fill(null).map((_v, i) => `dfa_${modelType}_model_n_${i}`);
 
-      const modelIds = new Array(count).fill(null).map((v, i) => `dfa_${modelType}_model_n_${i}`);
-
-      const models = modelIds.map((id) => {
-        return {
-          model_id: id,
-          body: {
-            compressed_definition: compressedDefinition,
-            inference_config: {
-              [modelType]: {},
-            },
-            input: {
-              field_names: ['common_field'],
-            },
-          } as PutTrainedModelConfig,
-        };
-      });
+      const models = modelIds.map((id) => this.createTestTrainedModelConfig(id, modelType));
 
       for (const model of models) {
         await this.createTrainedModel(model.model_id, model.body);
@@ -1178,7 +1204,7 @@ export function MachineLearningAPIProvider({ getService }: FtrProviderContext) {
       return modelIds;
     },
 
-    async createTestTrainedModelConfig(modelId: string, modelType: ModelType) {
+    createTestTrainedModelConfig(modelId: string, modelType: ModelType) {
       const compressedDefinition = this.getCompressedModelDefinition(modelType);
 
       return {
@@ -1201,16 +1227,44 @@ export function MachineLearningAPIProvider({ getService }: FtrProviderContext) {
      */
     getCompressedModelDefinition(modelType: ModelType) {
       return fs.readFileSync(
-        path.resolve(
-          __dirname,
-          'resources',
-          'trained_model_definitions',
-          `minimum_valid_config_${modelType}.json.gz.b64`
+        require.resolve(
+          `./resources/trained_model_definitions/minimum_valid_config_${modelType}.json.gz.b64`
         ),
         'utf-8'
       );
     },
 
+    getTrainedModelConfig(modelName: TrainedModelName) {
+      const configFileContent = fs.readFileSync(
+        require.resolve(`./resources/trained_model_definitions/${modelName}/config.json`),
+        'utf-8'
+      );
+      return JSON.parse(configFileContent) as PutTrainedModelConfig;
+    },
+
+    getTrainedModelVocabulary(modelName: TrainedModelName) {
+      const vocabularyFileContent = fs.readFileSync(
+        require.resolve(`./resources/trained_model_definitions/${modelName}/vocabulary.json`),
+        'utf-8'
+      );
+      return JSON.parse(vocabularyFileContent) as TrainedModelVocabulary;
+    },
+
+    getTrainedModelDefinitionPath(modelName: TrainedModelName) {
+      return require.resolve(
+        `./resources/trained_model_definitions/${modelName}/traced_pytorch_model.pt`
+      );
+    },
+
+    async importTrainedModel(modelId: string, modelName: TrainedModelName) {
+      await this.createTrainedModel(modelId, this.getTrainedModelConfig(modelName));
+      await this.createTrainedModelVocabularyES(modelId, this.getTrainedModelVocabulary(modelName));
+      await this.uploadTrainedModelDefinitionES(
+        modelId,
+        this.getTrainedModelDefinitionPath(modelName)
+      );
+    },
+
     async createModelAlias(modelId: string, modelAlias: string) {
       log.debug(`Creating alias for model "${modelId}"`);
       const { body, status } = await esSupertest.put(

x-pack/test/functional/services/ml/index.ts

@@ -124,7 +124,7 @@ export function MachineLearningProvider(context: FtrProviderContext) {
   const testResources = MachineLearningTestResourcesProvider(context, api);
   const alerting = MachineLearningAlertingProvider(context, api, commonUI);
   const swimLane = SwimLaneProvider(context);
-  const trainedModels = TrainedModelsProvider(context, api, commonUI);
+  const trainedModels = TrainedModelsProvider(context, commonUI);
   const trainedModelsTable = TrainedModelsTableProvider(context, commonUI);
   const mlNodesPanel = MlNodesPanelProvider(context);
 

x-pack/test/functional/services/ml/resources/trained_model_definitions/pt_tiny_fill_mask/config.json

@@ -0,0 +1,7 @@
+{
+  "description": "Tiny/Dummy PyTorch model (fill_mask)",
+  "model_type": "pytorch",
+  "inference_config": {
+    "fill_mask": {}
+  }
+}

x-pack/test/functional/services/ml/resources/trained_model_definitions/pt_tiny_fill_mask/vocabulary.json

@@ -0,0 +1,16 @@
+{
+  "vocabulary": [
+    "[PAD]",
+    "[UNK]",
+    "[CLS]",
+    "[SEP]",
+    "[MASK]",
+    "Hello",
+    "world",
+    "car",
+    "bike",
+    "bee",
+    "bird",
+    "and"
+ ]
+}

x-pack/test/functional/services/ml/resources/trained_model_definitions/pt_tiny_ner/config.json

@@ -0,0 +1,7 @@
+{
+  "description": "Tiny/Dummy PyTorch model (ner)",
+  "model_type": "pytorch",
+  "inference_config": {
+    "ner": {}
+  }
+}

x-pack/test/functional/services/ml/resources/trained_model_definitions/pt_tiny_ner/vocabulary.json

@@ -0,0 +1,21 @@
+{
+  "vocabulary": [
+    "[PAD]",
+    "[UNK]",
+    "[CLS]",
+    "[SEP]",
+    "[MASK]",
+    "Hello",
+    "world",
+    "car",
+    "bike",
+    "bee",
+    "bird",
+    "and",
+    "my",
+    "name",
+    "is",
+    "I'm",
+    "Spartacus"
+ ]
+}

x-pack/test/functional/services/ml/resources/trained_model_definitions/pt_tiny_pass_through/config.json

@@ -0,0 +1,13 @@
+{
+  "description": "Tiny/Dummy PyTorch model (pass_through)",
+  "model_type": "pytorch",
+  "inference_config": {
+    "pass_through": {
+      "tokenization": {
+        "bert": {
+          "with_special_tokens": false
+        }
+      }
+    }
+  }
+}

x-pack/test/functional/services/ml/resources/trained_model_definitions/pt_tiny_pass_through/vocabulary.json

@@ -0,0 +1,5 @@
+{
+  "vocabulary": [
+    "[UNK]", "[PAD]", "there", "is", "no", "spoon"
+  ]
+}

x-pack/test/functional/services/ml/resources/trained_model_definitions/pt_tiny_text_classification/config.json

@@ -0,0 +1,9 @@
+{
+  "description": "Tiny/Dummy PyTorch model (text_classification)",
+  "model_type": "pytorch",
+  "inference_config": {
+    "text_classification": {
+      "classification_labels": ["POSITIVE", "NEGATIVE"]
+    }
+  }
+}

x-pack/test/functional/services/ml/resources/trained_model_definitions/pt_tiny_text_classification/vocabulary.json

@@ -0,0 +1,16 @@
+{
+  "vocabulary": [
+    "[PAD]",
+    "[UNK]",
+    "[CLS]",
+    "[SEP]",
+    "[MASK]",
+    "Hello",
+    "world",
+    "car",
+    "bike",
+    "bee",
+    "bird",
+    "and"
+ ]
+}

x-pack/test/functional/services/ml/resources/trained_model_definitions/pt_tiny_text_embedding/config.json

@@ -0,0 +1,7 @@
+{
+  "description": "Tiny/Dummy PyTorch model (text_embedding)",
+  "model_type": "pytorch",
+  "inference_config": {
+    "text_embedding": {}
+  }
+}

x-pack/test/functional/services/ml/resources/trained_model_definitions/pt_tiny_text_embedding/vocabulary.json

@@ -0,0 +1,16 @@
+{
+  "vocabulary": [
+    "[PAD]",
+    "[UNK]",
+    "[CLS]",
+    "[SEP]",
+    "[MASK]",
+    "Hello",
+    "world",
+    "car",
+    "bike",
+    "bee",
+    "bird",
+    "and"
+ ]
+}

x-pack/test/functional/services/ml/resources/trained_model_definitions/pt_tiny_zero_shot/config.json

@@ -0,0 +1,9 @@
+{
+  "description": "Tiny/Dummy PyTorch model (zero_shot)",
+  "model_type": "pytorch",
+  "inference_config": {
+    "zero_shot_classification": {
+      "classification_labels": ["entailment", "neutral", "contradiction"]
+    }
+  }
+}