Skip to content

Commit

Permalink
add pipelineReferences to ml processors results (#143701)
Browse files Browse the repository at this point in the history
Added a list of pipelines that reference the ml inference processors so
that we can determine if a processor is re-used on the pipelines page.
This will allow us to disable the delete action if the processor is used
more than 1 time.
  • Loading branch information
TattdCodeMonkey authored Oct 20, 2022
1 parent e01b6bd commit d08e119
Show file tree
Hide file tree
Showing 5 changed files with 144 additions and 45 deletions.
1 change: 1 addition & 0 deletions x-pack/plugins/enterprise_search/common/types/pipelines.ts
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ export interface InferencePipeline {
modelState: TrainedModelState;
modelStateReason?: string;
pipelineName: string;
pipelineReferences: string[];
types: string[];
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ export const DEFAULT_VALUES: InferencePipeline = {
modelId: 'sample-bert-ner-model',
modelState: TrainedModelState.Started,
pipelineName: 'Sample Processor',
pipelineReferences: [],
types: ['pytorch', 'ner'],
};

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ describe('TrainedModelHealth', () => {
modelId: 'sample-bert-ner-model',
modelState: TrainedModelState.NotDeployed,
pipelineName: 'Sample Processor',
pipelineReferences: [],
types: ['pytorch'],
};
it('renders model started', () => {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
* 2.0.
*/

import { errors } from '@elastic/elasticsearch';
import { ElasticsearchClient } from '@kbn/core/server';
import { MlTrainedModels } from '@kbn/ml-plugin/server';

Expand All @@ -13,7 +14,8 @@ import { InferencePipeline, TrainedModelState } from '../../../common/types/pipe
import {
fetchAndAddTrainedModelData,
getMlModelConfigsForModelIds,
fetchMlInferencePipelineProcessorNames,
getMlInferencePipelineProcessorNamesFromPipelines,
fetchMlInferencePipelines,
fetchMlInferencePipelineProcessors,
fetchPipelineProcessorInferenceData,
InferencePipelineData,
Expand Down Expand Up @@ -247,23 +249,35 @@ const trainedModelDataObject: Record<string, InferencePipeline> = {
modelId: 'trained-model-id-1',
modelState: TrainedModelState.NotDeployed,
pipelineName: 'ml-inference-pipeline-1',
pipelineReferences: ['my-index@ml-inference'],
types: ['lang_ident', 'ner'],
},
'trained-model-id-2': {
modelId: 'trained-model-id-2',
modelState: TrainedModelState.Started,
pipelineName: 'ml-inference-pipeline-2',
pipelineReferences: ['my-index@ml-inference'],
types: ['pytorch', 'ner'],
},
'ml-inference-pipeline-3': {
modelId: 'trained-model-id-1',
modelState: TrainedModelState.NotDeployed,
pipelineName: 'ml-inference-pipeline-3',
pipelineReferences: ['my-index@ml-inference'],
types: ['lang_ident', 'ner'],
},
};

describe('fetchMlInferencePipelineProcessorNames lib function', () => {
const notFoundResponse = { meta: { statusCode: 404 } };
const notFoundError = new errors.ResponseError({
body: notFoundResponse,
statusCode: 404,
headers: {},
meta: {} as any,
warnings: [],
});

describe('fetchMlInferencePipelines lib function', () => {
const mockClient = {
ingest: {
getPipeline: jest.fn(),
Expand All @@ -274,32 +288,58 @@ describe('fetchMlInferencePipelineProcessorNames lib function', () => {
jest.clearAllMocks();
});

it('should return pipeline processor names for the @ml-inference pipeline', async () => {
it('should return @ml-inference pipelines', async () => {
mockClient.ingest.getPipeline.mockImplementation(() => Promise.resolve(mockGetPipeline));

const expected = ['ml-inference-pipeline-1'];
const response = await fetchMlInferencePipelines(mockClient as unknown as ElasticsearchClient);

const response = await fetchMlInferencePipelineProcessorNames(
mockClient as unknown as ElasticsearchClient,
'my-index'
);
expect(mockClient.ingest.getPipeline).toHaveBeenCalledWith({ id: '*@ml-inference' });
expect(response).toEqual(mockGetPipeline);
});

expect(mockClient.ingest.getPipeline).toHaveBeenCalledWith({ id: 'my-index@ml-inference' });
expect(response).toEqual(expected);
it('should return an empty object when no @ml-inference pipelines found', async () => {
mockClient.ingest.getPipeline.mockImplementation(() => Promise.resolve({}));

const response = await fetchMlInferencePipelines(mockClient as unknown as ElasticsearchClient);

expect(response).toEqual({});
});

it('should return an empty array for a missing @ml-inference pipeline', async () => {
mockClient.ingest.getPipeline.mockImplementation(() => Promise.resolve(mockGetPipeline));
it('should return an empty object when getPipeline throws an error ', async () => {
mockClient.ingest.getPipeline.mockImplementation(() => Promise.reject(notFoundError));

const response = await fetchMlInferencePipelineProcessorNames(
mockClient as unknown as ElasticsearchClient,
'my-index-without-ml-inference-pipeline'
const response = await fetchMlInferencePipelines(mockClient as unknown as ElasticsearchClient);

expect(response).toEqual({});
});
});

describe('getMlInferencePipelineProcessorNamesFromPipelines', () => {
it('should return pipeline processor names for the @ml-inference pipeline', () => {
const expected = ['ml-inference-pipeline-1'];
const processorNames = getMlInferencePipelineProcessorNamesFromPipelines(
'my-index',
mockGetPipeline
);
expect(processorNames).toEqual(expected);
});
it('should return an empty array for a missing @ml-inference pipeline', () => {
const processorNames = getMlInferencePipelineProcessorNamesFromPipelines(
'my-index-without-ml-inference-pipeline',
mockGetPipeline
);

expect(mockClient.ingest.getPipeline).toHaveBeenCalledWith({
id: 'my-index-without-ml-inference-pipeline@ml-inference',
});
expect(response).toEqual([]);
expect(processorNames).toEqual([]);
});
it('should return an empty array for a pipeline missing processors', () => {
const processorNames = getMlInferencePipelineProcessorNamesFromPipelines(
'my-index-without-ml-inference-pipeline',
{
'my-index-without-ml-inference-pipeline': {},
}
);

expect(processorNames).toEqual([]);
});
});

Expand All @@ -322,21 +362,27 @@ describe('fetchPipelineProcessorInferenceData lib function', () => {
modelId: 'trained-model-id-1',
modelState: TrainedModelState.NotDeployed,
pipelineName: 'ml-inference-pipeline-1',
pipelineReferences: ['my-index@ml-inference', 'other-index@ml-inference'],
trainedModelName: 'trained-model-id-1',
types: [],
},
{
modelId: 'trained-model-id-2',
modelState: TrainedModelState.NotDeployed,
pipelineName: 'ml-inference-pipeline-2',
pipelineReferences: ['my-index@ml-inference'],
trainedModelName: 'trained-model-id-2',
types: [],
},
];

const response = await fetchPipelineProcessorInferenceData(
mockClient as unknown as ElasticsearchClient,
['ml-inference-pipeline-1', 'ml-inference-pipeline-2', 'non-ml-inference-pipeline']
['ml-inference-pipeline-1', 'ml-inference-pipeline-2', 'non-ml-inference-pipeline'],
{
'ml-inference-pipeline-1': ['my-index@ml-inference', 'other-index@ml-inference'],
'ml-inference-pipeline-2': ['my-index@ml-inference'],
}
);

expect(mockClient.ingest.getPipeline).toHaveBeenCalledWith({
Expand Down Expand Up @@ -377,13 +423,15 @@ describe('getMlModelConfigsForModelIds lib function', () => {
modelId: 'trained-model-id-1',
modelState: TrainedModelState.Started,
pipelineName: '',
pipelineReferences: [],
trainedModelName: 'trained-model-id-1',
types: ['pytorch', 'ner'],
},
'trained-model-id-2': {
modelId: 'trained-model-id-2',
modelState: TrainedModelState.Started,
pipelineName: '',
pipelineReferences: [],
trainedModelName: 'trained-model-id-2',
types: ['pytorch', 'ner'],
},
Expand Down Expand Up @@ -413,20 +461,23 @@ describe('getMlModelConfigsForModelIds lib function', () => {
modelId: 'trained-model-id-1',
modelState: TrainedModelState.Started,
pipelineName: '',
pipelineReferences: [],
trainedModelName: 'trained-model-id-1',
types: ['pytorch', 'ner'],
},
'trained-model-id-2': {
modelId: 'trained-model-id-2',
modelState: TrainedModelState.Started,
pipelineName: '',
pipelineReferences: [],
trainedModelName: 'trained-model-id-2',
types: ['pytorch', 'ner'],
},
'trained-model-id-3-in-other-space': {
modelId: undefined, // Redacted
modelState: TrainedModelState.Started,
pipelineName: '',
pipelineReferences: [],
trainedModelName: 'trained-model-id-3-in-other-space',
types: ['pytorch', 'ner'],
},
Expand Down Expand Up @@ -483,27 +534,31 @@ describe('fetchAndAddTrainedModelData lib function', () => {
modelId: 'trained-model-id-1',
modelState: TrainedModelState.NotDeployed,
pipelineName: 'ml-inference-pipeline-1',
pipelineReferences: [],
trainedModelName: 'trained-model-id-1',
types: [],
},
{
modelId: 'trained-model-id-2',
modelState: TrainedModelState.NotDeployed,
pipelineName: 'ml-inference-pipeline-2',
pipelineReferences: [],
trainedModelName: 'trained-model-id-2',
types: [],
},
{
modelId: 'trained-model-id-3',
modelState: TrainedModelState.NotDeployed,
pipelineName: 'ml-inference-pipeline-3',
pipelineReferences: [],
trainedModelName: 'trained-model-id-3',
types: [],
},
{
modelId: 'trained-model-id-4',
modelState: TrainedModelState.NotDeployed,
pipelineName: 'ml-inference-pipeline-4',
pipelineReferences: [],
trainedModelName: 'trained-model-id-4',
types: [],
},
Expand All @@ -514,13 +569,15 @@ describe('fetchAndAddTrainedModelData lib function', () => {
modelId: 'trained-model-id-1',
modelState: TrainedModelState.NotDeployed,
pipelineName: 'ml-inference-pipeline-1',
pipelineReferences: [],
trainedModelName: 'trained-model-id-1',
types: ['lang_ident', 'ner'],
},
{
modelId: 'trained-model-id-2',
modelState: TrainedModelState.Started,
pipelineName: 'ml-inference-pipeline-2',
pipelineReferences: [],
trainedModelName: 'trained-model-id-2',
types: ['pytorch', 'ner'],
},
Expand All @@ -529,13 +586,15 @@ describe('fetchAndAddTrainedModelData lib function', () => {
modelState: TrainedModelState.Failed,
modelStateReason: 'something is wrong, boom',
pipelineName: 'ml-inference-pipeline-3',
pipelineReferences: [],
trainedModelName: 'trained-model-id-3',
types: ['pytorch', 'text_classification'],
},
{
modelId: 'trained-model-id-4',
modelState: TrainedModelState.Starting,
pipelineName: 'ml-inference-pipeline-4',
pipelineReferences: [],
trainedModelName: 'trained-model-id-4',
types: ['pytorch', 'fill_mask'],
},
Expand Down Expand Up @@ -599,7 +658,7 @@ describe('fetchMlInferencePipelineProcessors lib function', () => {
);

expect(mockClient.ingest.getPipeline).toHaveBeenCalledWith({
id: 'index-with-no-ml-inference-pipeline@ml-inference',
id: '*@ml-inference',
});
expect(mockClient.ingest.getPipeline).toHaveBeenCalledTimes(1);
expect(mockClient.ml.getTrainedModels).toHaveBeenCalledTimes(0);
Expand All @@ -626,7 +685,7 @@ describe('fetchMlInferencePipelineProcessors lib function', () => {
);

expect(mockClient.ingest.getPipeline).toHaveBeenCalledWith({
id: 'my-index@ml-inference',
id: '*@ml-inference',
});
expect(mockClient.ingest.getPipeline).toHaveBeenCalledWith({
id: 'ml-inference-pipeline-1',
Expand Down Expand Up @@ -663,7 +722,7 @@ describe('fetchMlInferencePipelineProcessors lib function', () => {
);

expect(mockClient.ingest.getPipeline).toHaveBeenCalledWith({
id: 'my-index@ml-inference',
id: '*@ml-inference',
});
expect(mockClient.ingest.getPipeline).toHaveBeenCalledWith({
id: 'ml-inference-pipeline-1',
Expand Down Expand Up @@ -707,7 +766,7 @@ describe('fetchMlInferencePipelineProcessors lib function', () => {
);

expect(mockClient.ingest.getPipeline).toHaveBeenCalledWith({
id: 'my-index@ml-inference',
id: '*@ml-inference',
});
expect(mockClient.ingest.getPipeline).toHaveBeenCalledWith({
id: 'ml-inference-pipeline-1,ml-inference-pipeline-3',
Expand Down
Loading

0 comments on commit d08e119

Please sign in to comment.