Skip to content

Commit

Permalink
Search indices pipeline deploy ELSER v2 (#167157)
Browse files Browse the repository at this point in the history
  • Loading branch information
@afoucret
afoucret authored Sep 26, 2023
1 parent ec20d61 commit ea8b7c0
Show file tree
Hide file tree
Showing 17 changed files with 91 additions and 49 deletions.
1 change: 0 additions & 1 deletion x-pack/plugins/enterprise_search/common/ml_inference_pipeline/index.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,6 @@ import {
export const TEXT_EXPANSION_TYPE = SUPPORTED_PYTORCH_TASKS.TEXT_EXPANSION;
export const TEXT_EXPANSION_FRIENDLY_TYPE = 'ELSER';
export const ML_INFERENCE_PREFIX = 'ml.inference.';
export const ELSER_MODEL_ID = '.elser_model_1';

export interface MlInferencePipelineParams {
description?: string;
Expand Down
2 changes: 2 additions & 0 deletions ...ck/plugins/enterprise_search/public/applications/__mocks__/kea_logic/kibana_logic.mock.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ import { uiSettingsServiceMock } from '@kbn/core-ui-settings-browser-mocks';
import { dataPluginMock } from '@kbn/data-plugin/public/mocks';

import { LensPublicStart } from '@kbn/lens-plugin/public';
import { mlPluginMock } from '@kbn/ml-plugin/public/mocks';
import { securityMock } from '@kbn/security-plugin/public/mocks';
import { sharePluginMock } from '@kbn/share-plugin/public/mocks';

Expand Down Expand Up @@ -60,6 +61,7 @@ export const mockKibanaValues = {
setChromeIsVisible: jest.fn(),
setDocTitle: jest.fn(),
share: sharePluginMock.createStartContract(),
ml: mlPluginMock.createStartContract(),
uiSettings: uiSettingsServiceMock.createStartContract(),
};

Expand Down
15 changes: 8 additions & 7 deletions ...rise_search_content/api/ml_models/text_expansion/create_text_expansion_model_api_logic.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,22 +4,23 @@
* 2.0; you may not use this file except in compliance with the Elastic License
* 2.0.
*/
import { ELSER_MODEL_ID } from '../../../../../../common/ml_inference_pipeline';
import { Actions, createApiLogic } from '../../../../shared/api_logic/create_api_logic';
import { HttpLogic } from '../../../../shared/http';

export type CreateTextExpansionModelArgs = undefined;
export interface CreateTextExpansionModelArgs {
modelId: string;
}

export interface CreateTextExpansionModelResponse {
deploymentState: string;
modelId: string;
}

export const createTextExpansionModel = async (): Promise<CreateTextExpansionModelResponse> => {
const route = `/internal/enterprise_search/ml/models/${ELSER_MODEL_ID}`;
return await HttpLogic.values.http.post<CreateTextExpansionModelResponse>(route, {
body: undefined,
});
export const createTextExpansionModel = async ({
modelId,
}: CreateTextExpansionModelArgs): Promise<CreateTextExpansionModelResponse> => {
const route = `/internal/enterprise_search/ml/models/${modelId}`;
return await HttpLogic.values.http.post<CreateTextExpansionModelResponse>(route);
};

export const CreateTextExpansionModelApiLogic = createApiLogic(
Expand Down
9 changes: 5 additions & 4 deletions ...prise_search_content/api/ml_models/text_expansion/fetch_text_expansion_model_api_logic.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,11 +4,12 @@
* 2.0; you may not use this file except in compliance with the Elastic License
* 2.0.
*/
import { ELSER_MODEL_ID } from '../../../../../../common/ml_inference_pipeline';
import { Actions, createApiLogic } from '../../../../shared/api_logic/create_api_logic';
import { HttpLogic } from '../../../../shared/http';

export type FetchTextExpansionModelArgs = undefined;
export interface FetchTextExpansionModelArgs {
modelId: string;
}

export interface FetchTextExpansionModelResponse {
deploymentState: string;
Expand All @@ -18,9 +19,9 @@ export interface FetchTextExpansionModelResponse {
threadsPerAllocation: number;
}

export const fetchTextExpansionModelStatus = async () => {
export const fetchTextExpansionModelStatus = async ({ modelId }: FetchTextExpansionModelArgs) => {
return await HttpLogic.values.http.get<FetchTextExpansionModelResponse>(
`/internal/enterprise_search/ml/models/${ELSER_MODEL_ID}`
`/internal/enterprise_search/ml/models/${modelId}`
);
};

Expand Down
15 changes: 8 additions & 7 deletions ...prise_search_content/api/ml_models/text_expansion/start_text_expansion_model_api_logic.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,22 +4,23 @@
* 2.0; you may not use this file except in compliance with the Elastic License
* 2.0.
*/
import { ELSER_MODEL_ID } from '../../../../../../common/ml_inference_pipeline';
import { Actions, createApiLogic } from '../../../../shared/api_logic/create_api_logic';
import { HttpLogic } from '../../../../shared/http';

export type StartTextExpansionModelArgs = undefined;
export interface StartTextExpansionModelArgs {
modelId: string;
}

export interface StartTextExpansionModelResponse {
deploymentState: string;
modelId: string;
}

export const startTextExpansionModel = async (): Promise<StartTextExpansionModelResponse> => {
const route = `/internal/enterprise_search/ml/models/${ELSER_MODEL_ID}/deploy`;
return await HttpLogic.values.http.post<StartTextExpansionModelResponse>(route, {
body: undefined,
});
export const startTextExpansionModel = async ({
modelId,
}: StartTextExpansionModelArgs): Promise<StartTextExpansionModelResponse> => {
const route = `/internal/enterprise_search/ml/models/${modelId}/deploy`;
return await HttpLogic.values.http.post<StartTextExpansionModelResponse>(route);
};

export const StartTextExpansionModelApiLogic = createApiLogic(
Expand Down
2 changes: 1 addition & 1 deletion ...enterprise_search_content/components/search_index/pipelines/ml_inference/deploy_model.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,7 @@ export const DeployModel = ({
data-telemetry-id={`entSearchContent-${ingestionMethod}-pipelines-textExpansionCallOut-deployModel`}
disabled={isCreateButtonDisabled}
iconType="launch"
onClick={() => createTextExpansionModel(undefined)}
onClick={() => createTextExpansionModel()}
>
{i18n.translate(
'xpack.enterpriseSearch.content.indices.pipelines.textExpansionCallOut.deployButton.label',
Expand Down
2 changes: 1 addition & 1 deletion ...terprise_search_content/components/search_index/pipelines/ml_inference/model_deployed.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -92,7 +92,7 @@ export const ModelDeployed = ({
data-telemetry-id={`entSearchContent-${ingestionMethod}-pipelines-textExpansionCallOut-startModel`}
disabled={isStartButtonDisabled}
iconType="playFilled"
onClick={() => startTextExpansionModel(undefined)}
onClick={() => startTextExpansionModel()}
>
{i18n.translate(
'xpack.enterpriseSearch.content.indices.pipelines.textExpansionCallOut.startModelButton.label',
Expand Down
1 change: 1 addition & 0 deletions ...ntent/components/search_index/pipelines/ml_inference/text_expansion_callout_logic.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@ const DEFAULT_VALUES: TextExpansionCalloutValues = {
textExpansionModel: undefined,
textExpansionModelPollTimeoutId: null,
textExpansionError: null,
elserModelId: '.elser_model_2',
};

jest.useFakeTimers();
Expand Down
50 changes: 39 additions & 11 deletions ...ch_content/components/search_index/pipelines/ml_inference/text_expansion_callout_logic.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,9 @@ import { i18n } from '@kbn/i18n';
import { HttpError, Status } from '../../../../../../../common/types/api';
import { MlModelDeploymentState } from '../../../../../../../common/types/ml';
import { getErrorsFromHttpResponse } from '../../../../../shared/flash_messages/handle_api_errors';

import { KibanaLogic } from '../../../../../shared/kibana';

import {
CreateTextExpansionModelApiLogic,
CreateTextExpansionModelApiLogicActions,
Expand All @@ -32,20 +35,24 @@ const FETCH_TEXT_EXPANSION_MODEL_POLLING_DURATION_ON_FAILURE = 30000; // 30 seco

interface TextExpansionCalloutActions {
clearTextExpansionModelPollingId: () => void;
createTextExpansionModel: CreateTextExpansionModelApiLogicActions['makeRequest'];
createTextExpansionModel: () => void;
createTextExpansionModelMakeRequest: CreateTextExpansionModelApiLogicActions['makeRequest'];
createTextExpansionModelPollingTimeout: (duration: number) => { duration: number };
createTextExpansionModelSuccess: CreateTextExpansionModelApiLogicActions['apiSuccess'];
fetchTextExpansionModel: FetchTextExpansionModelApiLogicActions['makeRequest'];
fetchTextExpansionModel: () => void;
fetchTextExpansionModelMakeRequest: FetchTextExpansionModelApiLogicActions['makeRequest'];
fetchTextExpansionModelError: FetchTextExpansionModelApiLogicActions['apiError'];
fetchTextExpansionModelSuccess: FetchTextExpansionModelApiLogicActions['apiSuccess'];
setTextExpansionModelPollingId: (pollTimeoutId: ReturnType<typeof setTimeout>) => {
pollTimeoutId: ReturnType<typeof setTimeout>;
};
startPollingTextExpansionModel: () => void;
startTextExpansionModel: StartTextExpansionModelApiLogicActions['makeRequest'];
startTextExpansionModel: () => void;
startTextExpansionModelMakeRequest: StartTextExpansionModelApiLogicActions['makeRequest'];
startTextExpansionModelSuccess: StartTextExpansionModelApiLogicActions['apiSuccess'];
stopPollingTextExpansionModel: () => void;
textExpansionModel: FetchTextExpansionModelApiLogicActions['apiSuccess'];
setElserModelId: (elserModelId: string) => { elserModelId: string };
}

export interface TextExpansionCalloutError {
Expand All @@ -54,6 +61,7 @@ export interface TextExpansionCalloutError {
}

export interface TextExpansionCalloutValues {
elserModelId: string;
createTextExpansionModelError: HttpError | undefined;
createTextExpansionModelStatus: Status;
createdTextExpansionModel: CreateTextExpansionModelResponse | undefined;
Expand Down Expand Up @@ -128,24 +136,28 @@ export const TextExpansionCalloutLogic = kea<
}),
startPollingTextExpansionModel: true,
stopPollingTextExpansionModel: true,
setElserModelId: (elserModelId) => ({ elserModelId }),
createTextExpansionModel: true,
fetchTextExpansionModel: true,
startTextExpansionModel: true,
},
connect: {
actions: [
CreateTextExpansionModelApiLogic,
[
'makeRequest as createTextExpansionModel',
'makeRequest as createTextExpansionModelMakeRequest',
'apiSuccess as createTextExpansionModelSuccess',
'apiError as createTextExpansionModelError',
],
FetchTextExpansionModelApiLogic,
[
'makeRequest as fetchTextExpansionModel',
'makeRequest as fetchTextExpansionModelMakeRequest',
'apiSuccess as fetchTextExpansionModelSuccess',
'apiError as fetchTextExpansionModelError',
],
StartTextExpansionModelApiLogic,
[
'makeRequest as startTextExpansionModel',
'makeRequest as startTextExpansionModelMakeRequest',
'apiSuccess as startTextExpansionModelSuccess',
'apiError as startTextExpansionModelError',
],
Expand All @@ -164,8 +176,12 @@ export const TextExpansionCalloutLogic = kea<
],
},
events: ({ actions, values }) => ({
afterMount: () => {
actions.fetchTextExpansionModel(undefined);
afterMount: async () => {
const elserModel = await KibanaLogic.values.ml.elasticModels?.getELSER({ version: 2 });
if (elserModel != null) {
actions.setElserModelId(elserModel.name);
actions.fetchTextExpansionModel();
}
},
beforeUnmount: () => {
if (values.textExpansionModelPollTimeoutId !== null) {
Expand All @@ -174,17 +190,23 @@ export const TextExpansionCalloutLogic = kea<
},
}),
listeners: ({ actions, values }) => ({
createTextExpansionModel: () =>
actions.createTextExpansionModelMakeRequest({ modelId: values.elserModelId }),
fetchTextExpansionModel: () =>
actions.fetchTextExpansionModelMakeRequest({ modelId: values.elserModelId }),
startTextExpansionModel: () =>
actions.startTextExpansionModelMakeRequest({ modelId: values.elserModelId }),
createTextExpansionModelPollingTimeout: ({ duration }) => {
if (values.textExpansionModelPollTimeoutId !== null) {
clearTimeout(values.textExpansionModelPollTimeoutId);
}
const timeoutId = setTimeout(() => {
actions.fetchTextExpansionModel(undefined);
actions.fetchTextExpansionModel();
}, duration);
actions.setTextExpansionModelPollingId(timeoutId);
},
createTextExpansionModelSuccess: () => {
actions.fetchTextExpansionModel(undefined);
actions.fetchTextExpansionModel();
actions.startPollingTextExpansionModel();
},
fetchTextExpansionModelError: () => {
Expand Down Expand Up @@ -217,7 +239,7 @@ export const TextExpansionCalloutLogic = kea<
actions.createTextExpansionModelPollingTimeout(FETCH_TEXT_EXPANSION_MODEL_POLLING_DURATION);
},
startTextExpansionModelSuccess: () => {
actions.fetchTextExpansionModel(undefined);
actions.fetchTextExpansionModel();
},
stopPollingTextExpansionModel: () => {
if (values.textExpansionModelPollTimeoutId !== null) {
Expand All @@ -235,6 +257,12 @@ export const TextExpansionCalloutLogic = kea<
setTextExpansionModelPollingId: (_, { pollTimeoutId }) => pollTimeoutId,
},
],
elserModelId: [
'',
{
setElserModelId: (_, { elserModelId }) => elserModelId,
},
],
},
selectors: ({ selectors }) => ({
isCreateButtonDisabled: [
Expand Down
2 changes: 2 additions & 0 deletions x-pack/plugins/enterprise_search/public/applications/index.test.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ import { dataPluginMock } from '@kbn/data-plugin/public/mocks';
import { guidedOnboardingMock } from '@kbn/guided-onboarding-plugin/public/mocks';
import { lensPluginMock } from '@kbn/lens-plugin/public/mocks';
import { licensingMock } from '@kbn/licensing-plugin/public/mocks';
import { mlPluginMock } from '@kbn/ml-plugin/public/mocks';
import { securityMock } from '@kbn/security-plugin/public/mocks';
import { sharePluginMock } from '@kbn/share-plugin/public/mocks';

Expand All @@ -38,6 +39,7 @@ describe('renderApp', () => {
licensing: licensingMock.createStart(),
security: securityMock.createStart(),
share: sharePluginMock.createStartContract(),
ml: mlPluginMock.createStartContract(),
user: {},
},
} as any;
Expand Down
3 changes: 2 additions & 1 deletion x-pack/plugins/enterprise_search/public/applications/index.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -66,7 +66,7 @@ export const renderApp = (
const { history } = params;
const { application, chrome, http, uiSettings } = core;
const { capabilities, navigateToUrl } = application;
const { charts, cloud, guidedOnboarding, lens, security, share } = plugins;
const { charts, cloud, guidedOnboarding, lens, security, share, ml } = plugins;

const entCloudHost = getCloudEnterpriseSearchHost(plugins.cloud);
externalUrl.enterpriseSearchUrl = publicUrl || entCloudHost || config.host || '';
Expand Down Expand Up @@ -107,6 +107,7 @@ export const renderApp = (
setChromeIsVisible: chrome.setIsVisible,
setDocTitle: chrome.docTitle.change,
share,
ml,
uiSettings,
});
const unmountLicensingLogic = mountLicensingLogic({
Expand Down
3 changes: 3 additions & 0 deletions x-pack/plugins/enterprise_search/public/applications/shared/kibana/kibana_logic.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ import {
import { DataPublicPluginStart } from '@kbn/data-plugin/public';
import { GuidedOnboardingPluginStart } from '@kbn/guided-onboarding-plugin/public';
import { LensPublicStart } from '@kbn/lens-plugin/public';
import { MlPluginStart } from '@kbn/ml-plugin/public';
import { SecurityPluginStart } from '@kbn/security-plugin/public';
import { SharePluginStart } from '@kbn/share-plugin/public';

Expand Down Expand Up @@ -52,6 +53,7 @@ interface KibanaLogicProps {
setChromeIsVisible(isVisible: boolean): void;
setDocTitle(title: string): void;
share: SharePluginStart;
ml: MlPluginStart;
uiSettings: IUiSettingsClient;
}

Expand Down Expand Up @@ -92,6 +94,7 @@ export const KibanaLogic = kea<MakeLogicType<KibanaValues>>({
setChromeIsVisible: [props.setChromeIsVisible, {}],
setDocTitle: [props.setDocTitle, {}],
share: [props.share, {}],
ml: [props.ml, {}],
uiSettings: [props.uiSettings, {}],
}),
selectors: ({ selectors }) => ({
Expand Down
4 changes: 3 additions & 1 deletion x-pack/plugins/enterprise_search/public/plugin.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ import { GuidedOnboardingPluginStart } from '@kbn/guided-onboarding-plugin/publi
import type { HomePublicPluginSetup } from '@kbn/home-plugin/public';
import { LensPublicStart } from '@kbn/lens-plugin/public';
import { LicensingPluginStart } from '@kbn/licensing-plugin/public';
import { MlPluginStart } from '@kbn/ml-plugin/public';
import { SecurityPluginSetup, SecurityPluginStart } from '@kbn/security-plugin/public';
import { SharePluginStart } from '@kbn/share-plugin/public';

Expand Down Expand Up @@ -65,6 +66,7 @@ export interface PluginsStart {
licensing: LicensingPluginStart;
security: SecurityPluginStart;
share: SharePluginStart;
ml: MlPluginStart;
}

export class EnterpriseSearchPlugin implements Plugin {
Expand Down Expand Up @@ -415,7 +417,7 @@ export class EnterpriseSearchPlugin implements Plugin {
}
}

public start(core: CoreStart) {
public async start(core: CoreStart) {
if (!this.config.ui?.enabled) {
return;
}
Expand Down
4 changes: 3 additions & 1 deletion x-pack/plugins/enterprise_search/server/lib/ml/ml_model_deployment_common.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,13 +5,15 @@
* 2.0.
*/

import { ELASTIC_MODEL_DEFINITIONS } from '@kbn/ml-trained-models-utils';

import {
ElasticsearchResponseError,
isNotFoundException,
isResourceNotFoundException,
} from '../../utils/identify_exceptions';

export const acceptableModelNames = ['.elser_model_1', '.elser_model_1_SNAPSHOT'];
export const acceptableModelNames = Object.keys(ELASTIC_MODEL_DEFINITIONS);

export function isNotFoundExceptionError(error: unknown): boolean {
return (
Expand Down
Loading

0 comments on commit ea8b7c0

Please sign in to comment.