[ML] Trained models: Update the start model allocation response type (#210966)

Updates the type correctly, as it is currently set to `{acknowledge:
boolean}`, which is the wrong type.
This commit is contained in:
Robert Jaszczurek 2025-02-13 14:56:36 +01:00 committed by GitHub
parent 18ef744a9c
commit 26548aeaa6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 53 additions and 4 deletions

View file

@ -168,6 +168,17 @@ export type TrainedModelDeploymentStatsResponse = estypes.MlTrainedModelDeployme
};
};
export interface StartTrainedModelDeploymentResponse {
// TODO update types in elasticsearch-specification
assignment: estypes.MlStartTrainedModelDeploymentResponse['assignment'] & {
adaptive_allocations?: {
enabled: boolean;
min_number_of_allocations?: number;
max_number_of_allocations?: number;
};
};
}
export interface AllocatedModel {
key: string;
deployment_id: string;

View file

@ -12,7 +12,10 @@ import type {
TrainedModelsApiService,
} from '../services/ml_api_service/trained_models';
import { TrainedModelsService } from './trained_models_service';
import type { TrainedModelUIItem } from '../../../common/types/trained_models';
import type {
StartTrainedModelDeploymentResponse,
TrainedModelUIItem,
} from '../../../common/types/trained_models';
import { MODEL_STATE } from '@kbn/ml-trained-models-utils';
import { i18n } from '@kbn/i18n';
import type { MlTrainedModelConfig } from '@elastic/elasticsearch/lib/api/types';
@ -185,7 +188,39 @@ describe('TrainedModelsService', () => {
mockTrainedModelsApiService.getTrainedModelsList.mockResolvedValueOnce([mockModel]);
mockTrainedModelsApiService.startModelAllocation.mockReturnValueOnce(of({ acknowledge: true }));
mockTrainedModelsApiService.startModelAllocation.mockReturnValueOnce(
of({
assignment: {
task_parameters: {
model_id: 'deploy-model',
model_bytes: 1000,
allocation_id: 'test-allocation',
priority: 'normal',
number_of_allocations: 1,
threads_per_allocation: 1,
queue_capacity: 1024,
deployment_id: 'my-deployment-id',
cache_size: '1mb',
},
node_count: 1,
routing_table: {
'node-1': {
routing_state: 'started',
reason: '',
current_allocations: 1,
target_allocations: 1,
},
},
assignment_state: 'started',
start_time: 1234567890,
adaptive_allocations: {
enabled: true,
min_number_of_allocations: 1,
max_number_of_allocations: 4,
},
},
})
);
// Start deployment
trainedModelsService.startModelDeployment('deploy-model', {
@ -230,7 +265,9 @@ describe('TrainedModelsService', () => {
const deploymentError = new Error('Deployment error');
mockTrainedModelsApiService.startModelAllocation.mockReturnValueOnce(
throwError(() => deploymentError) as unknown as Observable<{ acknowledge: boolean }>
throwError(
() => deploymentError
) as unknown as Observable<StartTrainedModelDeploymentResponse>
);
trainedModelsService.startModelDeployment('error-model', {

View file

@ -27,6 +27,7 @@ import type {
ModelDownloadState,
TrainedModelUIItem,
TrainedModelConfigResponse,
StartTrainedModelDeploymentResponse,
} from '../../../../common/types/trained_models';
export interface InferenceQueryParams {
@ -238,7 +239,7 @@ export function trainedModelsApiProvider(httpService: HttpService) {
deploymentParams,
adaptiveAllocationsParams,
}: StartAllocationParams) {
return httpService.http$<{ acknowledge: boolean }>({
return httpService.http$<StartTrainedModelDeploymentResponse>({
path: `${ML_INTERNAL_BASE_PATH}/trained_models/${modelId}/deployment/_start`,
method: 'POST',
query: deploymentParams,