-
Notifications
You must be signed in to change notification settings - Fork 8.2k
/
install.ts
112 lines (96 loc) · 3.13 KB
/
install.ts
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
/*
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
* or more contributor license agreements. Licensed under the Elastic License
* 2.0; you may not use this file except in compliance with the Elastic License
* 2.0.
*/
import type { ElasticsearchClient, Logger, SavedObjectsClientContract } from '@kbn/core/server';
import { errors } from '@elastic/elasticsearch';
import { getAssetFromAssetsMap, getPathParts } from '../../archive';
import {
ElasticsearchAssetType,
type PackageInstallContext,
} from '../../../../../common/types/models';
import type { EsAssetReference } from '../../../../../common/types/models';
import { retryTransientEsErrors } from '../retry';
import { updateEsAssetReferences } from '../../packages/es_assets_reference';
interface MlModelInstallation {
installationName: string;
content: string;
}
export const installMlModel = async (
packageInstallContext: PackageInstallContext,
esClient: ElasticsearchClient,
savedObjectsClient: SavedObjectsClientContract,
logger: Logger,
esReferences: EsAssetReference[]
) => {
const mlModelPath = packageInstallContext.paths.find((path) => isMlModel(path));
if (mlModelPath !== undefined) {
const content = getAssetFromAssetsMap(packageInstallContext.assetsMap, mlModelPath).toString(
'utf-8'
);
const pathParts = mlModelPath.split('/');
const modelId = pathParts[pathParts.length - 1].replace('.json', '');
const mlModelRef = {
id: modelId,
type: ElasticsearchAssetType.mlModel,
};
// get and save ml model refs before installing ml model
esReferences = await updateEsAssetReferences(
savedObjectsClient,
packageInstallContext.packageInfo.name,
esReferences,
{ assetsToAdd: [mlModelRef] }
);
const mlModel: MlModelInstallation = {
installationName: modelId,
content,
};
await handleMlModelInstall({ esClient, logger, mlModel });
}
return esReferences;
};
const isMlModel = (path: string) => {
const pathParts = getPathParts(path);
return !path.endsWith('/') && pathParts.type === ElasticsearchAssetType.mlModel;
};
async function handleMlModelInstall({
esClient,
logger,
mlModel,
}: {
esClient: ElasticsearchClient;
logger: Logger;
mlModel: MlModelInstallation;
}): Promise<EsAssetReference> {
try {
await retryTransientEsErrors(
() =>
esClient.ml.putTrainedModel(
{
model_id: mlModel.installationName,
defer_definition_decompression: true,
timeout: '45s',
// @ts-expect-error expects an object not a string
body: mlModel.content,
},
{
headers: {
'content-type': 'application/json',
},
}
),
{ logger }
);
} catch (err) {
// swallow the error if the ml model already exists.
const isAlreadyExistError =
err instanceof errors.ResponseError &&
err?.body?.error?.type === 'resource_already_exists_exception';
if (!isAlreadyExistError) {
throw err;
}
}
return { id: mlModel.installationName, type: ElasticsearchAssetType.mlModel };
}