Skip to content

Commit

Permalink
[ML] Fixes display of model state in trained models list with startin…
Browse files Browse the repository at this point in the history
…g and stopping deployments (elastic#188847)

## Summary

Fixes elastic#188035 and elastic#181093

<img width="1434" alt="image"
src="https://github.com/user-attachments/assets/6c14afa3-2908-45ff-a68d-88ee18f18964">

### Checklist

- [x] [Unit or functional
tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html)
were updated or added to match the most common scenarios

(cherry picked from commit 9669bfd)
  • Loading branch information
darnautov committed Jul 23, 2024
1 parent 723ff8f commit 26aba54
Show file tree
Hide file tree
Showing 4 changed files with 158 additions and 10 deletions.
1 change: 1 addition & 0 deletions x-pack/plugins/ml/common/types/trained_models.ts
Original file line number Diff line number Diff line change
Expand Up @@ -176,6 +176,7 @@ export interface TrainedModelDeploymentStatsResponse {
threads_per_allocation: number;
number_of_allocations: number;
}>;
reason?: string;
}

export interface AllocatedModel {
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,119 @@
/*
* 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 { getModelDeploymentState } from './get_model_state';
import { MODEL_STATE } from '@kbn/ml-trained-models-utils';
import type { ModelItem } from './models_list';

describe('getModelDeploymentState', () => {
it('returns STARTED if any deployment is in STARTED state', () => {
const model = {
stats: {
model_id: '.elser_model_2',
model_size_stats: {
model_size_bytes: 438123914,
required_native_memory_bytes: 2101346304,
},

deployment_stats: [
{
deployment_id: '.elser_model_2_01',
model_id: '.elser_model_2',
state: 'starting',
},
{
deployment_id: '.elser_model_2',
model_id: '.elser_model_2',
state: 'started',
allocation_status: {
allocation_count: 1,
target_allocation_count: 1,
state: 'fully_allocated',
},
},
],
},
} as unknown as ModelItem;
const result = getModelDeploymentState(model);
expect(result).toEqual(MODEL_STATE.STARTED);
});

it('returns MODEL_STATE.STARTING if any deployment is in STARTING state', () => {
const model = {
stats: {
model_id: '.elser_model_2',
model_size_stats: {
model_size_bytes: 438123914,
required_native_memory_bytes: 2101346304,
},

deployment_stats: [
{
deployment_id: '.elser_model_2',
model_id: '.elser_model_2',
state: 'stopping',
},
{
deployment_id: '.elser_model_2_01',
model_id: '.elser_model_2',
state: 'starting',
},
{
deployment_id: '.elser_model_2',
model_id: '.elser_model_2',
state: 'stopping',
},
],
},
} as unknown as ModelItem;
const result = getModelDeploymentState(model);
expect(result).toEqual(MODEL_STATE.STARTING);
});

it('returns MODEL_STATE.STOPPING if every deployment is in STOPPING state', () => {
const model = {
stats: {
model_id: '.elser_model_2',
model_size_stats: {
model_size_bytes: 438123914,
required_native_memory_bytes: 2101346304,
},

deployment_stats: [
{
deployment_id: '.elser_model_2',
model_id: '.elser_model_2',
state: 'stopping',
},
{
deployment_id: '.elser_model_2_01',
model_id: '.elser_model_2',
state: 'stopping',
},
],
},
} as unknown as ModelItem;
const result = getModelDeploymentState(model);
expect(result).toEqual(MODEL_STATE.STOPPING);
});

it('returns undefined for empty deployment stats', () => {
const model = {
stats: {
model_id: '.elser_model_2',
model_size_stats: {
model_size_bytes: 438123914,
required_native_memory_bytes: 2101346304,
},

deployment_stats: [],
},
} as unknown as ModelItem;
const result = getModelDeploymentState(model);
expect(result).toEqual(undefined);
});
});
Original file line number Diff line number Diff line change
Expand Up @@ -5,13 +5,34 @@
* 2.0.
*/

import type { ModelState } from '@kbn/ml-trained-models-utils';
import { MODEL_STATE } from '@kbn/ml-trained-models-utils';
import { DEPLOYMENT_STATE, MODEL_STATE, type ModelState } from '@kbn/ml-trained-models-utils';
import type { EuiHealthProps } from '@elastic/eui';
import { i18n } from '@kbn/i18n';
import type { ModelItem } from './models_list';

/**
* Resolves result model state based on the state of each deployment.
*
* If at least one deployment is in the STARTED state, the model state is STARTED.
* Then if none of the deployments are in the STARTED state, but at least one is in the STARTING state, the model state is STARTING.
* If all deployments are in the STOPPING state, the model state is STOPPING.
*/
export const getModelDeploymentState = (model: ModelItem): ModelState | undefined => {
if (!model.stats?.deployment_stats?.length) return;

if (model.stats?.deployment_stats?.some((v) => v.state === DEPLOYMENT_STATE.STARTED)) {
return MODEL_STATE.STARTED;
}
if (model.stats?.deployment_stats?.some((v) => v.state === DEPLOYMENT_STATE.STARTING)) {
return MODEL_STATE.STARTING;
}
if (model.stats?.deployment_stats?.every((v) => v.state === DEPLOYMENT_STATE.STOPPING)) {
return MODEL_STATE.STOPPING;
}
};

export const getModelStateColor = (
state: ModelState
state: ModelState | undefined
): { color: EuiHealthProps['color']; name: string } | null => {
switch (state) {
case MODEL_STATE.DOWNLOADED:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@ import { isDefined } from '@kbn/ml-is-defined';
import { useStorage } from '@kbn/ml-local-storage';
import { dynamic } from '@kbn/shared-ux-utility';
import useMountedState from 'react-use/lib/useMountedState';
import { getModelStateColor } from './get_model_state_color';
import { getModelStateColor, getModelDeploymentState } from './get_model_state';
import { ML_ELSER_CALLOUT_DISMISSED } from '../../../common/types/storage';
import { TechnicalPreviewBadge } from '../components/technical_preview_badge';
import { useModelActions } from './model_actions';
Expand Down Expand Up @@ -88,7 +88,11 @@ export type ModelItem = TrainedModelConfigResponse & {
origin_job_exists?: boolean;
deployment_ids: string[];
putModelConfig?: object;
state: ModelState;
state: ModelState | undefined;
/**
* Description of the current model state
*/
stateDescription?: string;
recommended?: boolean;
/**
* Model name, e.g. elser
Expand Down Expand Up @@ -374,14 +378,17 @@ export const ModelsList: FC<Props> = ({
...modelStats[0],
deployment_stats: modelStats.map((d) => d.deployment_stats).filter(isDefined),
};

// Extract deployment ids from deployment stats
model.deployment_ids = modelStats
.map((v) => v.deployment_stats?.deployment_id)
.filter(isDefined);
model.state = model.stats.deployment_stats?.some(
(v) => v.state === DEPLOYMENT_STATE.STARTED
)
? DEPLOYMENT_STATE.STARTED
: null;

model.state = getModelDeploymentState(model);
model.stateDescription = model.stats.deployment_stats.reduce((acc, c) => {
if (acc) return acc;
return c.reason ?? '';
}, '');
});

const elasticModels = models.filter((model) =>
Expand Down

0 comments on commit 26aba54

Please sign in to comment.