diff --git a/x-pack/platform/plugins/shared/ml/common/types/trained_models.ts b/x-pack/platform/plugins/shared/ml/common/types/trained_models.ts index 79d9a97c9565b..95e7c4ddc5f7e 100644 --- a/x-pack/platform/plugins/shared/ml/common/types/trained_models.ts +++ b/x-pack/platform/plugins/shared/ml/common/types/trained_models.ts @@ -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; diff --git a/x-pack/platform/plugins/shared/ml/public/application/model_management/trained_models_service.test.ts b/x-pack/platform/plugins/shared/ml/public/application/model_management/trained_models_service.test.ts index 5ab2ebbc53574..ce388d01e8046 100644 --- a/x-pack/platform/plugins/shared/ml/public/application/model_management/trained_models_service.test.ts +++ b/x-pack/platform/plugins/shared/ml/public/application/model_management/trained_models_service.test.ts @@ -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 ); trainedModelsService.startModelDeployment('error-model', { diff --git a/x-pack/platform/plugins/shared/ml/public/application/services/ml_api_service/trained_models.ts b/x-pack/platform/plugins/shared/ml/public/application/services/ml_api_service/trained_models.ts index efa491c54f248..045eac4676340 100644 --- a/x-pack/platform/plugins/shared/ml/public/application/services/ml_api_service/trained_models.ts +++ b/x-pack/platform/plugins/shared/ml/public/application/services/ml_api_service/trained_models.ts @@ -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$({ path: `${ML_INTERNAL_BASE_PATH}/trained_models/${modelId}/deployment/_start`, method: 'POST', query: deploymentParams,